Repository: polarsignals/frostdb
Pull request #658 (2023-12-26 06:44:24)

Support repeated fields of all base types on recordToRows
Previously we were operating under assumptions that repeated fields were dictionary of `*array.Binary`.

This commit implements conversion on all arrow arrays of repeated base types `int64|float64|bool|string`

I have also added `*array.Binary` in the mix although it is not categorized as base type (Based on the use, It will make sense to introduce binary type in our schema proto)


Patch:
diff --git a/pqarrow/parquet.go b/pqarrow/parquet.go
index 6c82484..96c2894 100644
--- a/pqarrow/parquet.go
+++ b/pqarrow/parquet.go
@@ -13,7 +13,6 @@ import (
 	"github.com/parquet-go/parquet-go"
 
 	"github.com/polarsignals/frostdb/dynparquet"
-	"github.com/polarsignals/frostdb/pqarrow/arrowutils"
 )
 
 func ArrowScalarToParquetValue(sc scalar.Scalar) (parquet.Value, error) {
@@ -112,6 +111,10 @@ func recordToRows(w dynparquet.ParquetWriter, dcv dynamicColumnVerifier, record
 			writers[i] = writeInt32(def, i, recordStart, a)
 		case *array.Int64:
 			writers[i] = writeInt64(def, i, recordStart, a)
+		case *array.String:
+			writers[i] = writeString(def, i, recordStart, a)
+		case *array.Binary:
+			writers[i] = writeBinary(def, i, recordStart, a)
 		default:
 			writers[i] = writeGeneral(def, i, recordStart, a)
 		}
@@ -170,6 +173,32 @@ func writeInt32(def, column, startIdx int, a *array.Int32) arrowToParquet {
 	}
 }
 
+func writeBinary(def, column, startIdx int, a *array.Binary) arrowToParquet {
+	return func(w parquet.Row, row int) parquet.Row {
+		if a.IsNull(row + startIdx) {
+			return append(w,
+				parquet.Value{}.Level(0, 0, column),
+			)
+		}
+		return append(w,
+			parquet.ByteArrayValue(a.Value(row+startIdx)).Level(0, def, column),
+		)
+	}
+}
+
+func writeString(def, column, startIdx int, a *array.String) arrowToParquet {
+	return func(w parquet.Row, row int) parquet.Row {
+		if a.IsNull(row + startIdx) {
+			return append(w,
+				parquet.Value{}.Level(0, 0, column),
+			)
+		}
+		return append(w,
+			parquet.ByteArrayValue([]byte(a.Value(row+startIdx))).Level(0, def, column),
+		)
+	}
+}
+
 func writeDictionary(def, column, startIdx int, a *array.Dictionary) arrowToParquet {
 	value := func(row int) parquet.Value {
 		return parquet.ValueOf(a.GetOneForMarshal(row))
@@ -194,16 +223,52 @@ func writeDictionary(def, column, startIdx int, a *array.Dictionary) arrowToParq
 }
 
 func writeList(def, column, startIdx int, a *array.List) (arrowToParquet, error) {
-	// We are just replicating current behavior so that this change is not a
-	// breaking change, however the assumption here is wrong.
-	//
-	// We support lists of all base types, and there is no requirement for the list
-	// to be a dictionary
-	//
-	// TODO:(gernest) properly cast to base types + their dictionary counterparts ?
-	dictionaryList, binaryDictionaryList, err := arrowutils.ToConcreteList(a)
-	if err != nil {
-		return nil, err
+	var lw arrowToParquet
+	switch e := a.ListValues().(type) {
+	case *array.Int32:
+		// WHile this is not base type. To avoid breaking things I have left it here.
+		lw = writeListOf(def, column, startIdx, a, func(idx int) parquet.Value {
+			return parquet.Int32Value(e.Value(idx))
+		})
+	case *array.Int64:
+		lw = writeListOf(def, column, startIdx, a, func(idx int) parquet.Value {
+			return parquet.Int64Value(e.Value(idx))
+		})
+	case *array.Boolean:
+		lw = writeListOf(def, column, startIdx, a, func(idx int) parquet.Value {
+			return parquet.BooleanValue(e.Value(idx))
+		})
+	case *array.Float64:
+		lw = writeListOf(def, column, startIdx, a, func(idx int) parquet.Value {
+			return parquet.DoubleValue(e.Value(idx))
+		})
+	case *array.String:
+		lw = writeListOf(def, column, startIdx, a, func(idx int) parquet.Value {
+			return parquet.ByteArrayValue([]byte(e.Value(idx)))
+		})
+	case *array.Binary:
+		lw = writeListOf(def, column, startIdx, a, func(idx int) parquet.Value {
+			return parquet.ByteArrayValue([]byte(e.Value(idx)))
+		})
+	case *array.Dictionary:
+		switch d := e.Dictionary().(type) {
+		case *array.Binary:
+			lw = writeListOf(def, column, startIdx, a, func(idx int) parquet.Value {
+				return parquet.ByteArrayValue(
+					d.Value(e.GetValueIndex(idx)),
+				)
+			})
+		case *array.String:
+			lw = writeListOf(def, column, startIdx, a, func(idx int) parquet.Value {
+				return parquet.ByteArrayValue(
+					[]byte(d.Value(e.GetValueIndex(idx))),
+				)
+			})
+		default:
+			return nil, fmt.Errorf("list dictionary not of expected type: %T", d)
+		}
+	default:
+		return nil, fmt.Errorf("list not of expected type: %T", e)
 	}
 	return func(w parquet.Row, row int) parquet.Row {
 		if a.IsNull(row + startIdx) {
@@ -211,17 +276,22 @@ func writeList(def, column, startIdx int, a *array.List) (arrowToParquet, error)
 				parquet.Value{}.Level(0, 0, column),
 			)
 		}
+		return lw(w, row)
+	}, nil
+}
+
+func writeListOf(def, column, startIdx int, a *array.List, value func(idx int) parquet.Value) arrowToParquet {
+	return func(w parquet.Row, row int) parquet.Row {
 		start, end := a.ValueOffsets(row + startIdx)
 		for k := start; k < end; k++ {
-			v := binaryDictionaryList.Value(dictionaryList.GetValueIndex(int(k)))
 			rep := 0
 			if k != start {
 				rep = 1
 			}
-			w = append(w, parquet.ByteArrayValue(v).Level(rep, def+1, column))
+			w = append(w, value(int(k)).Level(rep, def+1, column))
 		}
 		return w
-	}, nil
+	}
 }
 
 func writeNull(column int) arrowToParquet {


Test patch:
diff --git a/pqarrow/parquet_test.go b/pqarrow/parquet_test.go
index c64966a..804bbdc 100644
--- a/pqarrow/parquet_test.go
+++ b/pqarrow/parquet_test.go
@@ -87,6 +87,149 @@ func BenchmarkRecordsToFile(b *testing.B) {
 	}
 }
 
+func TestRecordToRows_list(t *testing.T) {
+	b := array.NewRecordBuilder(memory.NewGoAllocator(),
+		arrow.NewSchema(
+			[]arrow.Field{
+				{
+					Name:     "int",
+					Type:     arrow.ListOf(arrow.PrimitiveTypes.Int64),
+					Nullable: true,
+				},
+				{
+					Name:     "double",
+					Type:     arrow.ListOf(arrow.PrimitiveTypes.Float64),
+					Nullable: true,
+				},
+				{
+					Name:     "string",
+					Type:     arrow.ListOf(arrow.BinaryTypes.String),
+					Nullable: true,
+				},
+				{
+					Name:     "binary",
+					Type:     arrow.ListOf(arrow.BinaryTypes.Binary),
+					Nullable: true,
+				},
+				{
+					Name:     "bool",
+					Type:     arrow.ListOf(arrow.FixedWidthTypes.Boolean),
+					Nullable: true,
+				},
+			}, nil,
+		),
+	)
+	defer b.Release()
+
+	type Sample struct {
+		Int    []int64
+		Double []float64
+		String []string
+		Binary [][]byte
+		Bool   []bool
+	}
+	samples := []Sample{
+		{}, // handle nulls
+		{
+			Int:    []int64{1},
+			Double: []float64{1},
+			String: []string{"1"},
+			Binary: [][]byte{
+				[]byte("1"),
+			},
+			Bool: []bool{true},
+		},
+	}
+	ints := b.Field(0).(*array.ListBuilder)
+	intsBuild := ints.ValueBuilder().(*array.Int64Builder)
+	double := b.Field(1).(*array.ListBuilder)
+	doubleBuild := double.ValueBuilder().(*array.Float64Builder)
+	str := b.Field(2).(*array.ListBuilder)
+	strBuild := str.ValueBuilder().(*array.StringBuilder)
+	bin := b.Field(3).(*array.ListBuilder)
+	binBuild := bin.ValueBuilder().(*array.BinaryBuilder)
+	boolean := b.Field(4).(*array.ListBuilder)
+	booleanBuild := boolean.ValueBuilder().(*array.BooleanBuilder)
+
+	for _, s := range samples {
+		appendList[int64](ints, s.Int, intsBuild.Append)
+		appendList[float64](double, s.Double, doubleBuild.Append)
+		appendList[string](str, s.String, strBuild.Append)
+		appendList[[]byte](bin, s.Binary, binBuild.Append)
+		appendList[bool](boolean, s.Bool, booleanBuild.Append)
+	}
+	r := b.NewRecord()
+	defer r.Release()
+
+	parquetFields := parquet.Group{}
+	for _, f := range r.Schema().Fields() {
+		parquetFields[f.Name] = parquet.Node(nil)
+	}
+	clone := &cloneWriter{}
+	if err := recordToRows(
+		clone, func(string) bool { return false }, r, 0, int(r.NumRows()), parquetFields.Fields(),
+	); err != nil {
+		t.Fatal(err)
+	}
+	// parquetFields.Fields() changes the order of the rows
+	// From
+	//  int, double, string, binary, bool
+	// To
+	//  binary, bool, double, int, string
+	want := []parquet.Row{
+		{
+			parquet.Value{}.Level(0, 0, 0),
+			parquet.Value{}.Level(0, 0, 1),
+			parquet.Value{}.Level(0, 0, 2),
+			parquet.Value{}.Level(0, 0, 3),
+			parquet.Value{}.Level(0, 0, 4),
+		},
+		{
+			parquet.ByteArrayValue([]byte("1")).Level(0, 1, 0),
+			parquet.BooleanValue(true).Level(0, 1, 1),
+			parquet.DoubleValue(1).Level(0, 1, 2),
+			parquet.Int64Value(1).Level(0, 1, 3),
+			parquet.ByteArrayValue([]byte("1")).Level(0, 1, 4),
+		},
+	}
+	require.Equal(t, len(want), len(clone.rows))
+	for i := range want {
+		require.True(t, want[i].Equal(clone.rows[i]))
+	}
+}
+
+type cloneWriter struct {
+	rows []parquet.Row
+}
+
+func (w cloneWriter) Schema() *parquet.Schema { return nil }
+
+func (w cloneWriter) Write(r []any) (int, error) { return len(r), nil }
+
+func (w *cloneWriter) WriteRows(r []parquet.Row) (int, error) {
+	for i := range r {
+		w.rows = append(w.rows, r[i].Clone())
+	}
+	return len(r), nil
+}
+
+func (w cloneWriter) Flush() error { return nil }
+
+func (w cloneWriter) Close() error { return nil }
+
+func (w cloneWriter) Reset(_ io.Writer) {}
+
+func appendList[T any](ls *array.ListBuilder, values []T, add func(T)) {
+	if values == nil {
+		ls.AppendNull()
+		return
+	}
+	ls.Append(true)
+	for i := range values {
+		add(values[i])
+	}
+}
+
 func TestRecordDynamicCols(t *testing.T) {
 	build := array.NewRecordBuilder(memory.NewGoAllocator(),
 		arrow.NewSchema([]arrow.Field{