Repository: polarsignals/frostdb
Pull request #766 (2024-03-12 15:44:59)

pqarrow/arrowutils: Use SortingColumns in MergeRecords
This allows to merge records in arbitrary orders and columns.


Patch:
diff --git a/pqarrow/arrowutils/groupranges.go b/pqarrow/arrowutils/groupranges.go
index afdf25e..72b6ab5 100644
--- a/pqarrow/arrowutils/groupranges.go
+++ b/pqarrow/arrowutils/groupranges.go
@@ -212,7 +212,7 @@ func GetGroupsAndOrderedSetRanges(
 // If the returned boolean is false, the comparison should be disregarded.
 func nullComparison(leftNull, rightNull bool) (int, bool) {
 	if !leftNull && !rightNull {
-		// Both are null, this implies that the null comparison should be
+		// Both are not null, this implies that the null comparison should be
 		// disregarded.
 		return 0, false
 	}
diff --git a/pqarrow/arrowutils/merge.go b/pqarrow/arrowutils/merge.go
index 8ff1831..011eadb 100644
--- a/pqarrow/arrowutils/merge.go
+++ b/pqarrow/arrowutils/merge.go
@@ -5,6 +5,7 @@ import (
 	"container/heap"
 	"fmt"
 	"math"
+	"strings"
 
 	"github.com/apache/arrow/go/v15/arrow"
 	"github.com/apache/arrow/go/v15/arrow/array"
@@ -22,7 +23,7 @@ import (
 func MergeRecords(
 	mem memory.Allocator,
 	records []arrow.Record,
-	orderByCols []int,
+	orderByCols []SortingColumn,
 	limit uint64,
 ) (arrow.Record, error) {
 	h := cursorHeap{
@@ -73,7 +74,7 @@ type cursor struct {
 
 type cursorHeap struct {
 	cursors     []cursor
-	orderByCols []int
+	orderByCols []SortingColumn
 }
 
 func (h cursorHeap) Len() int {
@@ -81,67 +82,93 @@ func (h cursorHeap) Len() int {
 }
 
 func (h cursorHeap) Less(i, j int) bool {
-	c1 := h.cursors[i]
-	c2 := h.cursors[j]
-	for _, i := range h.orderByCols {
-		col1 := c1.r.Column(i)
-		col2 := c2.r.Column(i)
+	for idx := range h.orderByCols {
+		c1 := h.cursors[i]
+		c2 := h.cursors[j]
+		sc := h.orderByCols[idx]
+		col1 := c1.r.Column(sc.Index)
+		col2 := c2.r.Column(sc.Index)
 		if cmp, ok := nullComparison(col1.IsNull(c1.curIdx), col2.IsNull(c2.curIdx)); ok {
-			if cmp == 0 {
-				continue
+			if h.orderByCols[idx].NullsFirst {
+				return cmp == -1
+			}
+			if !h.orderByCols[idx].NullsFirst {
+				return cmp == 1
 			}
-			return cmp < 0
 		}
-		switch arr1 := c1.r.Column(i).(type) {
+
+		cmp := h.compare(idx, i, j)
+		if cmp != 0 {
+			// Use direction to reorder the comparison. Direction determines if the list
+			// is in ascending or descending.
+			//
+			// For instance if comparison between i,j value is -1 and direction is -1
+			// this will resolve to true hence the list will be in ascending order. Same
+			// principle applies for descending.
+			return cmp == h.orderByCols[idx].Direction.comparison()
+		}
+		// Try comparing the next column
+	}
+
+	return false
+}
+
+func (h cursorHeap) compare(idx, i, j int) int {
+	c1 := h.cursors[i]
+	c2 := h.cursors[j]
+	sc := h.orderByCols[idx]
+	switch arr1 := c1.r.Column(sc.Index).(type) {
+	case *array.Binary:
+		arr2 := c2.r.Column(sc.Index).(*array.Binary)
+		return bytes.Compare(arr1.Value(c1.curIdx), arr2.Value(c2.curIdx))
+	case *array.String:
+		arr2 := c2.r.Column(sc.Index).(*array.String)
+		return strings.Compare(arr1.Value(c1.curIdx), arr2.Value(c2.curIdx))
+	case *array.Int64:
+		arr2 := c2.r.Column(sc.Index).(*array.Int64)
+		v1 := arr1.Value(c1.curIdx)
+		v2 := arr2.Value(c2.curIdx)
+		if v1 == v2 {
+			return 0
+		}
+		if v1 < v2 {
+			return -1
+		}
+		return 1
+	case *array.Int32:
+		arr2 := c2.r.Column(sc.Index).(*array.Int32)
+		v1 := arr1.Value(c1.curIdx)
+		v2 := arr2.Value(c2.curIdx)
+		if v1 == v2 {
+			return 0
+		}
+		if v1 < v2 {
+			return -1
+		}
+		return 1
+	case *array.Uint64:
+		arr2 := c2.r.Column(sc.Index).(*array.Uint64)
+		v1 := arr1.Value(c1.curIdx)
+		v2 := arr2.Value(c2.curIdx)
+		if v1 == v2 {
+			return 0
+		}
+		if v1 < v2 {
+			return -1
+		}
+		return 1
+	case *array.Dictionary:
+		switch dict := arr1.Dictionary().(type) {
 		case *array.Binary:
-			arr2 := c2.r.Column(i).(*array.Binary)
-			cmp := bytes.Compare(arr1.Value(c1.curIdx), arr2.Value(c2.curIdx))
-			if cmp == 0 {
-				continue
-			}
-			return cmp < 0
-		case *array.Int64:
-			arr2 := c2.r.Column(i).(*array.Int64)
-			v1 := arr1.Value(c1.curIdx)
-			v2 := arr2.Value(c2.curIdx)
-			if v1 == v2 {
-				continue
-			}
-			return v1 < v2
-		case *array.Int32:
-			arr2 := c2.r.Column(i).(*array.Int32)
-			v1 := arr1.Value(c1.curIdx)
-			v2 := arr2.Value(c2.curIdx)
-			if v1 == v2 {
-				continue
-			}
-			return v1 < v2
-		case *array.Uint64:
-			arr2 := c2.r.Column(i).(*array.Uint64)
-			v1 := arr1.Value(c1.curIdx)
-			v2 := arr2.Value(c2.curIdx)
-			if v1 == v2 {
-				continue
-			}
-			return v1 < v2
-		case *array.Dictionary:
-			switch dict := arr1.Dictionary().(type) {
-			case *array.Binary:
-				arr2 := c2.r.Column(i).(*array.Dictionary)
-				dict2 := arr2.Dictionary().(*array.Binary)
-				cmp := bytes.Compare(dict.Value(arr1.GetValueIndex(c1.curIdx)), dict2.Value(arr2.GetValueIndex(c2.curIdx)))
-				if cmp == 0 {
-					continue
-				}
-				return cmp < 0
-			default:
-				panic(fmt.Sprintf("unsupported dictionary type for record merging %T", dict))
-			}
+			arr2 := c2.r.Column(sc.Index).(*array.Dictionary)
+			dict2 := arr2.Dictionary().(*array.Binary)
+			return bytes.Compare(dict.Value(arr1.GetValueIndex(c1.curIdx)), dict2.Value(arr2.GetValueIndex(c2.curIdx)))
 		default:
-			panic(fmt.Sprintf("unsupported type for record merging %T", arr1))
+			panic(fmt.Sprintf("unsupported dictionary type for record merging %T", dict))
 		}
+	default:
+		panic(fmt.Sprintf("unsupported type for record merging %T", arr1))
 	}
-	return false
 }
 
 func (h cursorHeap) Swap(i, j int) {
diff --git a/query/physicalplan/ordered_aggregate.go b/query/physicalplan/ordered_aggregate.go
index 1ef5130..1c85655 100644
--- a/query/physicalplan/ordered_aggregate.go
+++ b/query/physicalplan/ordered_aggregate.go
@@ -482,9 +482,9 @@ func (a *OrderedAggregate) Finish(ctx context.Context) error {
 		}
 	} else {
 		// The aggregation results must be merged.
-		orderByCols := make([]int, len(a.groupColOrdering))
+		orderByCols := make([]arrowutils.SortingColumn, len(a.groupColOrdering))
 		for i := range orderByCols {
-			orderByCols[i] = i
+			orderByCols[i] = arrowutils.SortingColumn{Index: i}
 		}
 		mergedRecord, err := arrowutils.MergeRecords(a.pool, records, orderByCols, 0)
 		if err != nil {
diff --git a/query/physicalplan/ordered_synchronizer.go b/query/physicalplan/ordered_synchronizer.go
index b7fc5c6..3e3de22 100644
--- a/query/physicalplan/ordered_synchronizer.go
+++ b/query/physicalplan/ordered_synchronizer.go
@@ -24,7 +24,8 @@ import (
 type OrderedSynchronizer struct {
 	pool         memory.Allocator
 	orderByExprs []logicalplan.Expr
-	orderByCols  []int
+	// orderByCols  []int
+	orderByCols []arrowutils.SortingColumn
 
 	sync struct {
 		mtx        sync.Mutex
@@ -197,7 +198,7 @@ func (o *OrderedSynchronizer) ensureSameSchema(records []arrow.Record) error {
 
 	o.orderByCols = o.orderByCols[:0]
 	for i := range newFields {
-		o.orderByCols = append(o.orderByCols, i)
+		o.orderByCols = append(o.orderByCols, arrowutils.SortingColumn{Index: i})
 	}
 
 	for _, field := range leftoverCols {


Test patch:
diff --git a/pqarrow/arrowutils/merge_test.go b/pqarrow/arrowutils/merge_test.go
index 5f2a474..723fc93 100644
--- a/pqarrow/arrowutils/merge_test.go
+++ b/pqarrow/arrowutils/merge_test.go
@@ -1,6 +1,7 @@
 package arrowutils_test
 
 import (
+	"context"
 	"testing"
 
 	"github.com/apache/arrow/go/v15/arrow"
@@ -8,50 +9,325 @@ import (
 	"github.com/apache/arrow/go/v15/arrow/memory"
 	"github.com/stretchr/testify/require"
 
+	"github.com/polarsignals/frostdb/internal/records"
 	"github.com/polarsignals/frostdb/pqarrow/arrowutils"
 )
 
 func TestMerge(t *testing.T) {
-	schema := arrow.NewSchema(
-		[]arrow.Field{{Name: "test", Type: arrow.PrimitiveTypes.Int64}},
-		nil,
-	)
-
-	b := array.NewInt64Builder(memory.DefaultAllocator)
-	b.Append(0)
-	b.Append(1)
-	b.Append(3)
-	b.Append(5)
-	a := b.NewArray()
-	record1 := array.NewRecord(schema, []arrow.Array{a}, int64(a.Len()))
-	b.Append(2)
-	b.Append(4)
-	b.Append(5)
-	b.Append(6)
-	a = b.NewArray()
-	record2 := array.NewRecord(schema, []arrow.Array{a}, int64(a.Len()))
-	b.AppendNull()
-	a = b.NewArray()
-	record3 := array.NewRecord(schema, []arrow.Array{a}, int64(a.Len()))
-
-	res, err := arrowutils.MergeRecords(
-		memory.DefaultAllocator, []arrow.Record{record1, record2, record3}, []int{0}, 0,
-	)
-	require.NoError(t, err)
-	require.Equal(t, int64(1), res.NumCols())
-	col := res.Column(0).(*array.Int64)
-	require.Equal(t, 9, col.Len())
-	// Nulls sort first.
-	require.True(t, col.IsNull(0))
-	expected := []int64{0, 1, 2, 3, 4, 5, 5, 6}
-	for i := 1; i < col.Len(); i++ {
-		require.Equal(t, expected[i-1], col.Value(i))
+	type row struct {
+		Number *int64  `frostdb:",asc(0)"`
+		Text   *string `frostdb:",asc(1)"`
 	}
+	int64Ptr := func(i int64) *int64 { return &i }
+	stringPtr := func(s string) *string { return &s }
+
+	// mem := memory.NewGoAllocator()
+	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+	defer mem.AssertSize(t, 0)
 
-	// check that we can merge with a limit
-	res, err = arrowutils.MergeRecords(
-		memory.DefaultAllocator, []arrow.Record{record1, record2, record3}, []int{0}, 3,
-	)
+	t.Parallel()
+
+	build := records.NewBuild[row](mem)
+	defer build.Release()
+	err := build.Append([]row{
+		{Number: nil, Text: nil},
+		{Number: int64Ptr(0), Text: stringPtr("a")},
+		{Number: int64Ptr(1), Text: stringPtr("b")},
+		{Number: int64Ptr(3), Text: stringPtr("d")},
+		{Number: int64Ptr(5), Text: stringPtr("d")},
+	}...)
 	require.NoError(t, err)
-	require.Equal(t, int64(3), res.NumRows())
+	record1 := build.NewRecord()
+	defer record1.Release()
+
+	for _, tc := range []struct {
+		name string
+		// inputs
+		records [][]row
+		orderBy []arrowutils.SortingColumn
+		limit   uint64
+		// expected output
+		expected []row
+	}{{
+		name: "merge-ascending",
+		records: [][]row{{
+			{Number: nil},         // 0
+			{Number: int64Ptr(0)}, // 1
+			{Number: int64Ptr(2)}, // 3
+			{Number: int64Ptr(4)}, // 5
+			{Number: int64Ptr(5)}, // 7
+		}, {
+			{Number: int64Ptr(1)}, // 2
+			{Number: int64Ptr(3)}, // 4
+			{Number: int64Ptr(5)}, // 6
+			{Number: int64Ptr(6)}, // 8
+			{Number: int64Ptr(7)}, // 9
+		}},
+		orderBy: []arrowutils.SortingColumn{
+			{Index: 0, Direction: arrowutils.Ascending, NullsFirst: true},
+		},
+
+		expected: []row{
+			{Number: nil, Text: nil},
+			{Number: int64Ptr(0)},
+			{Number: int64Ptr(1)},
+			{Number: int64Ptr(2)},
+			{Number: int64Ptr(3)},
+			{Number: int64Ptr(4)},
+			{Number: int64Ptr(5)},
+			{Number: int64Ptr(5)},
+			{Number: int64Ptr(6)},
+			{Number: int64Ptr(7)},
+		},
+	}, {
+		name: "merge-ascending-limit",
+		records: [][]row{{
+			{Number: nil},         // 0
+			{Number: int64Ptr(0)}, // 1
+			{Number: int64Ptr(2)}, // 3
+			{Number: int64Ptr(4)}, // 5
+			{Number: int64Ptr(5)}, // 7
+		}, {
+			{Number: int64Ptr(1)}, // 2
+			{Number: int64Ptr(3)}, // 4
+			{Number: int64Ptr(5)}, // 6
+			{Number: int64Ptr(6)}, // 8
+			{Number: int64Ptr(7)}, // 9
+		}},
+		orderBy: []arrowutils.SortingColumn{
+			{Index: 0, Direction: arrowutils.Ascending, NullsFirst: true},
+		},
+		limit: 3,
+
+		expected: []row{
+			{Number: nil, Text: nil},
+			{Number: int64Ptr(0)},
+			{Number: int64Ptr(1)},
+		},
+	}, {
+		name: "merge-descending",
+		records: [][]row{{
+			{Number: nil},         // 0
+			{Number: int64Ptr(5)}, // 3
+			{Number: int64Ptr(3)}, // 6
+			{Number: int64Ptr(1)}, // 7
+			{Number: int64Ptr(0)}, // 8
+		}, {
+			{Number: int64Ptr(7)}, // 1
+			{Number: int64Ptr(6)}, // 2
+			{Number: int64Ptr(5)}, // 4
+			{Number: int64Ptr(4)}, // 5
+		}},
+		orderBy: []arrowutils.SortingColumn{
+			{Index: 0, Direction: arrowutils.Descending, NullsFirst: true},
+		},
+
+		expected: []row{
+			{Number: nil},
+			{Number: int64Ptr(7)},
+			{Number: int64Ptr(6)},
+			{Number: int64Ptr(5)},
+			{Number: int64Ptr(5)},
+			{Number: int64Ptr(4)},
+			{Number: int64Ptr(3)},
+			{Number: int64Ptr(1)},
+			{Number: int64Ptr(0)},
+		},
+	}, {
+		name: "merge-descending-limit",
+		records: [][]row{{
+			{Number: nil},         // 0
+			{Number: int64Ptr(5)}, // 4
+			{Number: int64Ptr(3)}, // limited
+			{Number: int64Ptr(1)}, // limited
+			{Number: int64Ptr(0)}, // limited
+		}, {
+			{Number: int64Ptr(7)}, // 1
+			{Number: int64Ptr(6)}, // 2
+			{Number: int64Ptr(5)}, // 3
+			{Number: int64Ptr(4)}, // 5
+		}},
+		orderBy: []arrowutils.SortingColumn{
+			{Index: 0, Direction: arrowutils.Descending, NullsFirst: true},
+		},
+		limit: 6,
+
+		expected: []row{
+			{Number: nil},
+			{Number: int64Ptr(7)},
+			{Number: int64Ptr(6)},
+			{Number: int64Ptr(5)},
+			{Number: int64Ptr(5)},
+			{Number: int64Ptr(4)},
+		},
+	}, {
+		name: "multiple-ascending",
+		records: [][]row{{
+			{Number: int64Ptr(0), Text: stringPtr("a")}, // 1
+			{Number: int64Ptr(0), Text: stringPtr("c")}, // 3
+			{Number: int64Ptr(1), Text: stringPtr("e")}, // 4
+		}, {
+			{Number: nil, Text: nil},                    // 0
+			{Number: int64Ptr(0), Text: stringPtr("b")}, // 2
+			{Number: int64Ptr(1), Text: stringPtr("d")}, // 3
+			{Number: int64Ptr(2), Text: stringPtr("f")}, // 5
+		}},
+		orderBy: []arrowutils.SortingColumn{
+			{Index: 0, Direction: arrowutils.Ascending, NullsFirst: true},
+			{Index: 1, Direction: arrowutils.Ascending, NullsFirst: true},
+		},
+
+		expected: []row{
+			{Number: nil, Text: nil},
+			{Number: int64Ptr(0), Text: stringPtr("a")},
+			{Number: int64Ptr(0), Text: stringPtr("b")},
+			{Number: int64Ptr(0), Text: stringPtr("c")},
+			{Number: int64Ptr(1), Text: stringPtr("d")},
+			{Number: int64Ptr(1), Text: stringPtr("e")},
+			{Number: int64Ptr(2), Text: stringPtr("f")},
+		},
+	}, {
+		name: "multiple-descending",
+		records: [][]row{{
+			{Number: int64Ptr(1), Text: stringPtr("e")},
+			{Number: int64Ptr(0), Text: stringPtr("c")},
+			{Number: int64Ptr(0), Text: stringPtr("a")},
+		}, {
+			{Number: nil, Text: nil},                    // 0
+			{Number: int64Ptr(2), Text: stringPtr("f")}, //
+			{Number: int64Ptr(1), Text: stringPtr("d")},
+			{Number: int64Ptr(0), Text: stringPtr("b")},
+		}},
+		orderBy: []arrowutils.SortingColumn{
+			{Index: 0, Direction: arrowutils.Descending, NullsFirst: true},
+			{Index: 1, Direction: arrowutils.Descending, NullsFirst: true},
+		},
+
+		expected: []row{
+			{Number: nil, Text: nil},
+			{Number: int64Ptr(2), Text: stringPtr("f")},
+			{Number: int64Ptr(1), Text: stringPtr("e")},
+			{Number: int64Ptr(1), Text: stringPtr("d")},
+			{Number: int64Ptr(0), Text: stringPtr("c")},
+			{Number: int64Ptr(0), Text: stringPtr("b")},
+			{Number: int64Ptr(0), Text: stringPtr("a")},
+		},
+	}, {
+		name: "multiple-mixed",
+		records: [][]row{{
+			{Number: int64Ptr(1), Text: stringPtr("e")},
+			{Number: int64Ptr(0), Text: stringPtr("a")},
+			{Number: int64Ptr(0), Text: stringPtr("c")},
+		}, {
+			{Number: nil, Text: nil},
+			{Number: int64Ptr(2), Text: stringPtr("f")},
+			{Number: int64Ptr(1), Text: stringPtr("d")},
+			{Number: int64Ptr(0), Text: stringPtr("b")},
+		}},
+		orderBy: []arrowutils.SortingColumn{
+			{Index: 0, Direction: arrowutils.Descending, NullsFirst: true},
+			{Index: 1, Direction: arrowutils.Ascending, NullsFirst: true},
+		},
+
+		expected: []row{
+			{Number: nil, Text: nil},
+			{Number: int64Ptr(2), Text: stringPtr("f")},
+			{Number: int64Ptr(1), Text: stringPtr("d")},
+			{Number: int64Ptr(1), Text: stringPtr("e")},
+			{Number: int64Ptr(0), Text: stringPtr("a")},
+			{Number: int64Ptr(0), Text: stringPtr("b")},
+			{Number: int64Ptr(0), Text: stringPtr("c")},
+		},
+	}} {
+		t.Run(tc.name, func(t *testing.T) {
+			builder := records.NewBuild[row](mem)
+			defer builder.Release()
+
+			records := make([]arrow.Record, 0, len(tc.records))
+			defer func() {
+				for _, record := range records {
+					record.Release()
+				}
+			}()
+
+			for _, rows := range tc.records {
+				err := builder.Append(rows...)
+				require.NoError(t, err)
+				records = append(records, builder.NewRecord())
+			}
+
+			res, err := arrowutils.MergeRecords(mem, records, tc.orderBy, tc.limit)
+			require.NoError(t, err)
+			defer res.Release()
+
+			numbers := res.Column(0).(*array.Int64)
+			texts := res.Column(1).(*array.String)
+
+			// TODO: Create an equivalent to the records.NewBuild for reading back into the struct.
+			result := make([]row, res.NumRows())
+			for i := 0; i < int(res.NumRows()); i++ {
+				if numbers.IsValid(i) {
+					result[i].Number = int64Ptr(numbers.Value(i))
+				}
+				if texts.IsValid(i) {
+					result[i].Text = stringPtr(texts.Value(i))
+				}
+			}
+
+			require.Equal(t, tc.expected, result)
+		})
+	}
+}
+
+func BenchmarkMergeRecords(b *testing.B) {
+	ctx := context.Background()
+	mem := memory.NewGoAllocator()
+
+	var r1, r2 arrow.Record
+	{
+		build := records.NewBuild[Model](mem)
+		err := build.Append(makeModels(40)...)
+		require.NoError(b, err)
+		r := build.NewRecord()
+		defer r.Release()
+
+		indices, err := arrowutils.SortRecord(r.Record, r.SortingColumns)
+		require.NoError(b, err)
+		defer indices.Release()
+
+		r1, err = arrowutils.Take(ctx, r.Record, indices)
+		require.NoError(b, err)
+		defer r1.Release()
+	}
+	{
+		build := records.NewBuild[Model](mem)
+		err := build.Append(makeModels(60)...)
+		require.NoError(b, err)
+		r := build.NewRecord()
+		defer r.Release()
+
+		indices, err := arrowutils.SortRecord(r.Record, r.SortingColumns)
+		require.NoError(b, err)
+		defer indices.Release()
+
+		r2, err = arrowutils.Take(ctx, r.Record, indices)
+		require.NoError(b, err)
+		defer r2.Release()
+	}
+
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		_, err := arrowutils.MergeRecords(
+			mem,
+			[]arrow.Record{r1, r2},
+			[]arrowutils.SortingColumn{
+				{Index: 0, Direction: arrowutils.Ascending, NullsFirst: true},
+			},
+			0,
+		)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
 }