Repository: polarsignals/frostdb
Pull request #645 (2023-12-20 18:42:02)

Remove dynparquet.ValuesToBuffer
We have already removed `(*Table)Write`. There is no longer any use for this function.


Patch:
diff --git a/dynparquet/schema.go b/dynparquet/schema.go
index fa6bf37..f5d8133 100644
--- a/dynparquet/schema.go
+++ b/dynparquet/schema.go
@@ -5,21 +5,18 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"reflect"
 	"regexp"
 	"sort"
 	"strings"
 	"sync"
 	"text/tabwriter"
 
-	"github.com/google/uuid"
 	"github.com/parquet-go/parquet-go"
 	"github.com/parquet-go/parquet-go/compress"
 	"github.com/parquet-go/parquet-go/encoding"
 	"github.com/parquet-go/parquet-go/format"
 	"google.golang.org/protobuf/proto"
 
-	"github.com/polarsignals/frostdb/bufutils"
 	schemapb "github.com/polarsignals/frostdb/gen/proto/go/frostdb/schema/v1alpha1"
 	schemav2pb "github.com/polarsignals/frostdb/gen/proto/go/frostdb/schema/v1alpha2"
 )
@@ -988,178 +985,6 @@ func ToSnakeCase(str string) string {
 	return strings.ToLower(snake)
 }
 
-func ValuesToBuffer(schema *Schema, vals ...any) (*Buffer, error) {
-	dynamicColumns := map[string][]string{}
-	rows := make([]parquet.Row, 0, len(vals))
-
-	findColumn := func(val reflect.Value, col string, v any) any {
-		for i := 0; i < val.NumField(); i++ {
-			if ToSnakeCase(val.Type().Field(i).Name) == col {
-				return val.Field(i).Interface()
-			}
-		}
-		return nil
-	}
-
-	// Collect dynamic columns
-	for _, v := range vals {
-		val := deref(reflect.ValueOf(v))
-		for _, col := range schema.Columns() {
-			cv := findColumn(val, col.Name, v)
-			if cv == nil {
-				continue
-			}
-			switch col.Dynamic {
-			case true:
-				switch derefType(reflect.TypeOf(cv)).Kind() {
-				case reflect.Struct:
-					dynVals := deref(reflect.ValueOf(cv))
-					for j := 0; j < dynVals.NumField(); j++ {
-						dynamicColumns[col.Name] = append(dynamicColumns[col.Name], ToSnakeCase(dynVals.Type().Field(j).Name))
-					}
-				case reflect.Slice:
-					dynVals := reflect.ValueOf(cv)
-					for j := 0; j < dynVals.Len(); j++ {
-						pair := deref(reflect.ValueOf(dynVals.Index(j).Interface()))
-						dynamicColumns[col.Name] = append(dynamicColumns[col.Name], ToSnakeCase(fmt.Sprintf("%v", pair.Field(0))))
-					}
-				case reflect.Map:
-					dynVals := reflect.ValueOf(cv)
-					for _, key := range dynVals.MapKeys() {
-						switch key.Kind() {
-						case reflect.String:
-							dynamicColumns[col.Name] = append(dynamicColumns[col.Name], ToSnakeCase(fmt.Sprintf("%v", key.String())))
-						default:
-							return nil, fmt.Errorf("unsupported dynamic map key type for column for column '%s'", col.Name)
-						}
-					}
-				default:
-					return nil, fmt.Errorf("unsupported dynamic type for column '%s'", col.Name)
-				}
-			}
-		}
-	}
-
-	dynamicColumns = bufutils.Dedupe(dynamicColumns)
-
-	// Create all rows
-	for _, v := range vals {
-		row := []parquet.Value{}
-		val := deref(reflect.ValueOf(v))
-
-		colIdx := 0
-		for _, col := range schema.Columns() {
-			cv := findColumn(val, col.Name, v)
-			if cv == nil {
-				continue
-			}
-			switch col.Dynamic {
-			case true:
-				switch derefType(reflect.TypeOf(cv)).Kind() {
-				case reflect.Struct:
-					dynVals := deref(reflect.ValueOf(cv))
-					for _, dyncol := range dynamicColumns[col.Name] {
-						found := false
-						for j := 0; j < dynVals.NumField(); j++ {
-							if ToSnakeCase(dynVals.Type().Field(j).Name) == dyncol {
-								row = append(row, parquet.ValueOf(dynVals.Field(j).Interface()).Level(0, 1, colIdx))
-								colIdx++
-								found = true
-								break
-							}
-						}
-						if !found {
-							row = append(row, parquet.ValueOf(nil).Level(0, 0, colIdx))
-							colIdx++
-						}
-					}
-				case reflect.Slice:
-					dynVals := reflect.ValueOf(cv)
-					for _, dyncol := range dynamicColumns[col.Name] {
-						found := false
-						for j := 0; j < dynVals.Len(); j++ {
-							pair := reflect.ValueOf(dynVals.Index(j).Interface())
-							if ToSnakeCase(fmt.Sprintf("%v", pair.Field(0).Interface())) == dyncol {
-								row = append(row, parquet.ValueOf(pair.Field(1).Interface()).Level(0, 1, colIdx))
-								colIdx++
-								found = true
-								break
-							}
-						}
-						if !found {
-							row = append(row, parquet.ValueOf(nil).Level(0, 0, colIdx))
-							colIdx++
-						}
-					}
-				case reflect.Map:
-					dynVals := reflect.ValueOf(cv)
-					for _, key := range dynVals.MapKeys() {
-						switch key.Kind() {
-						case reflect.String:
-							for _, dyncol := range dynamicColumns[col.Name] {
-								found := false
-								for _, key := range dynVals.MapKeys() {
-									if ToSnakeCase(fmt.Sprintf("%v", key.String())) == dyncol {
-										row = append(row, parquet.ValueOf(dynVals.MapIndex(key).Interface()).Level(0, 1, colIdx))
-										colIdx++
-										found = true
-										break
-									}
-								}
-
-								if !found {
-									row = append(row, parquet.ValueOf(nil).Level(0, 0, colIdx))
-									colIdx++
-								}
-							}
-						default:
-							return nil, fmt.Errorf("unsupported dynamic map key type for column '%s'", col.Name)
-						}
-					}
-				default:
-					return nil, fmt.Errorf("unsupported dynamic type for column '%s'", col.Name)
-				}
-			default:
-				switch t := cv.(type) {
-				case []uuid.UUID: // Special handling for this type
-					row = append(row, parquet.ValueOf(ExtractLocationIDs(t)).Level(0, 0, colIdx))
-				default:
-					row = append(row, parquet.ValueOf(cv).Level(0, 0, colIdx))
-				}
-				colIdx++
-			}
-		}
-
-		rows = append(rows, row)
-	}
-
-	pb, err := schema.NewBuffer(dynamicColumns)
-	if err != nil {
-		return nil, err
-	}
-
-	_, err = pb.WriteRows(rows)
-	if err != nil {
-		return nil, err
-	}
-
-	return pb, nil
-}
-
-func deref(v reflect.Value) reflect.Value {
-	for v.Kind() == reflect.Ptr {
-		v = v.Elem()
-	}
-	return v
-}
-
-func derefType(v reflect.Type) reflect.Type {
-	for v.Kind() == reflect.Ptr {
-		v = v.Elem()
-	}
-	return v
-}
-
 // NewBuffer returns a new buffer with a concrete parquet schema generated
 // using the given concrete dynamic column names.
 func (s *Schema) NewBuffer(dynamicColumns map[string][]string) (*Buffer, error) {
diff --git a/parts/arrow.go b/parts/arrow.go
index 16ff96a..c94ce1d 100644
--- a/parts/arrow.go
+++ b/parts/arrow.go
@@ -1,10 +1,7 @@
 package parts
 
 import (
-	"bytes"
-
 	"github.com/apache/arrow/go/v14/arrow"
-	"github.com/parquet-go/parquet-go"
 
 	"github.com/polarsignals/frostdb/dynparquet"
 	"github.com/polarsignals/frostdb/pqarrow"
@@ -48,29 +45,7 @@ func (p *arrowPart) SerializeBuffer(schema *dynparquet.Schema, w dynparquet.Parq
 }
 
 func (p *arrowPart) AsSerializedBuffer(schema *dynparquet.Schema) (*dynparquet.SerializedBuffer, error) {
-	// If this is a Arrow record part, convert the record into a serialized buffer
-	b := &bytes.Buffer{}
-
-	w, err := schema.GetWriter(b, pqarrow.RecordDynamicCols(p.record), false)
-	if err != nil {
-		return nil, err
-	}
-	defer schema.PutWriter(w)
-	if err := p.SerializeBuffer(schema, w.ParquetWriter); err != nil {
-		return nil, err
-	}
-
-	f, err := parquet.OpenFile(bytes.NewReader(b.Bytes()), int64(b.Len()))
-	if err != nil {
-		return nil, err
-	}
-
-	buf, err := dynparquet.NewSerializedBuffer(f)
-	if err != nil {
-		return nil, err
-	}
-
-	return buf, nil
+	return pqarrow.SerializeRecord(p.record, schema)
 }
 
 func (p *arrowPart) NumRows() int64 {
diff --git a/pqarrow/parquet.go b/pqarrow/parquet.go
index ac571e9..d06332b 100644
--- a/pqarrow/parquet.go
+++ b/pqarrow/parquet.go
@@ -1,6 +1,7 @@
 package pqarrow
 
 import (
+	"bytes"
 	"fmt"
 	"io"
 	"strings"
@@ -180,6 +181,27 @@ func RecordToDynamicRow(dynSchema *dynparquet.Schema, pqSchema *parquet.Schema,
 	return dynparquet.NewDynamicRow(row, pqSchema, dyncols, pqSchema.Fields()), nil
 }
 
+func SerializeRecord(r arrow.Record, schema *dynparquet.Schema) (*dynparquet.SerializedBuffer, error) {
+	b := &bytes.Buffer{}
+	w, err := schema.GetWriter(b, RecordDynamicCols(r), false)
+	if err != nil {
+		return nil, err
+	}
+	defer schema.PutWriter(w)
+	if err := RecordToFile(schema, w.ParquetWriter, r); err != nil {
+		return nil, err
+	}
+	f, err := parquet.OpenFile(bytes.NewReader(b.Bytes()), int64(b.Len()))
+	if err != nil {
+		return nil, err
+	}
+	buf, err := dynparquet.NewSerializedBuffer(f)
+	if err != nil {
+		return nil, err
+	}
+	return buf, nil
+}
+
 func RecordToFile(schema *dynparquet.Schema, w dynparquet.ParquetWriter, r arrow.Record) error {
 	return RecordsToFile(schema, w, []arrow.Record{r})
 }


Test patch:
diff --git a/parts/part_test.go b/parts/part_test.go
index 53bda21..e25644c 100644
--- a/parts/part_test.go
+++ b/parts/part_test.go
@@ -1,13 +1,14 @@
 package parts
 
 import (
-	"bytes"
 	"testing"
 
+	"github.com/apache/arrow/go/v14/arrow/memory"
 	"github.com/stretchr/testify/require"
 
 	"github.com/polarsignals/frostdb/dynparquet"
 	schemapb "github.com/polarsignals/frostdb/gen/proto/go/frostdb/schema/v1alpha1"
+	"github.com/polarsignals/frostdb/pqarrow"
 )
 
 func TestFindMaximumNonOverlappingSet(t *testing.T) {
@@ -29,7 +30,7 @@ func TestFindMaximumNonOverlappingSet(t *testing.T) {
 		end   int64
 	}
 	type dataModel struct {
-		Ints int64
+		Ints int64 `frostdb:",asc,plain"`
 	}
 	for _, tc := range []struct {
 		name                   string
@@ -83,14 +84,18 @@ func TestFindMaximumNonOverlappingSet(t *testing.T) {
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			parts := make([]Part, len(tc.ranges))
+			b := dynparquet.NewBuild[dataModel](memory.NewGoAllocator())
+			schema, err := dynparquet.SchemaFromDefinition(b.Schema("test_schema"))
+			require.NoError(t, err)
+			defer b.Release()
 			for i := range parts {
 				start := dataModel{Ints: tc.ranges[i].start}
 				end := dataModel{Ints: tc.ranges[i].end}
-				buf, err := dynparquet.ValuesToBuffer(testSchema, start, end)
+				require.Nil(t, b.Append(start, end))
+				r := b.NewRecord()
+				serBuf, err := pqarrow.SerializeRecord(r, schema)
 				require.NoError(t, err)
-				var b bytes.Buffer
-				require.NoError(t, testSchema.SerializeBuffer(&b, buf))
-				serBuf, err := dynparquet.ReaderFromBytes(b.Bytes())
+				r.Release()
 				require.NoError(t, err)
 				parts[i] = NewParquetPart(0, serBuf)
 			}