Repository: polarsignals/frostdb
Pull request #635 (2023-12-16 11:35:54)

Upgrade  arrowutils.SortRecord
Followup to #628

This commit enhances `arrowutils.SortRecord` to support

- Multi column sort
- Column direction (ascending, descending)
- NullFirst (for both ascending and descending)

Also optimizes `arrowutils.ReorderRecord` to support dictionary columns.

gernest: @metalmatze @asubiotto thanks for the feedback, PTAL. I addressed all the comments.

Ignore `fix lint` commits, the linter action was generating partial recommendations, I had to fight for my life to make it happy.
gernest: @asubiotto I addressed your comments PTAL.

**NOTE** . I have left `*OptInt32Builder` api comments because changing api or api expectation is sensitive and they will need more context and input from @metalmatze . I can handle whatever feedback in followup PR.  I have also decided to leave benchmark out.
gernest: @asubiotto  thanks for the review. I addressed all comments PTAL.

Hopeful we put this into a practical use. Areas I think it will have impact

- logictest : we can now build records directly and skip parquet rows construction

- in parca

https://github.com/parca-dev/parca/blob/4f536249850d27a79b365d3d796ca8ff692896d5/pkg/parcacol/ingest.go#L63-L107
Now we have a generic builder that constructs arrow records directly, sorting was the only missing piece. 

Patch:
diff --git a/pqarrow/arrowutils/sort.go b/pqarrow/arrowutils/sort.go
index e4e3f7f..1137a95 100644
--- a/pqarrow/arrowutils/sort.go
+++ b/pqarrow/arrowutils/sort.go
@@ -2,93 +2,306 @@ package arrowutils
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"sort"
+	"strconv"
 
 	"github.com/apache/arrow/go/v14/arrow"
 	"github.com/apache/arrow/go/v14/arrow/array"
 	"github.com/apache/arrow/go/v14/arrow/compute"
 	"github.com/apache/arrow/go/v14/arrow/memory"
+	"golang.org/x/sync/errgroup"
+
+	"github.com/polarsignals/frostdb/pqarrow/builder"
+)
+
+type Direction uint
+
+const (
+	Ascending Direction = iota
+	Descending
 )
 
-// SortRecord sorts the given record's rows by the given column. Currently only supports int64, string and binary columns.
-func SortRecord(mem memory.Allocator, r arrow.Record, cols []int) (*array.Int64, error) {
-	if len(cols) > 1 {
-		return nil, fmt.Errorf("sorting by multiple columns isn't implemented yet")
+func (d Direction) comparison() int {
+	switch d {
+	case Ascending:
+		return -1
+	case Descending:
+		return 1
+	default:
+		panic("unexpected direction value " + strconv.Itoa(int(d)) + " only -1 and 1 are allowed")
+	}
+}
+
+// SortingColumn describes a sorting column on a arrow.Record.
+type SortingColumn struct {
+	Index      int
+	Direction  Direction
+	NullsFirst bool
+}
+
+// SortRecord sorts given arrow.Record by columns. Returns *array.Int32 of
+// indices to sorted rows or record r.
+//
+// Comparison is made sequentially by each column. When rows are equal in the
+// first column we compare the rows om the second column and so on and so forth
+// until rows that are not equal are found.
+func SortRecord(r arrow.Record, columns []SortingColumn) (*array.Int32, error) {
+	if len(columns) == 0 {
+		return nil, errors.New("pqarrow/arrowutils: at least one column is needed for sorting")
 	}
-	indicesBuilder := array.NewInt64Builder(mem)
+	indicesBuilder := builder.NewOptInt32Builder(arrow.PrimitiveTypes.Int32)
+	defer indicesBuilder.Release()
 
 	if r.NumRows() == 0 {
-		return indicesBuilder.NewInt64Array(), nil
+		return indicesBuilder.NewArray().(*array.Int32), nil
 	}
 	if r.NumRows() == 1 {
 		indicesBuilder.Append(0)
-		return indicesBuilder.NewInt64Array(), nil
+		return indicesBuilder.NewArray().(*array.Int32), nil
+	}
+	indicesBuilder.Reserve(int(r.NumRows()))
+	for i := 0; i < int(r.NumRows()); i++ {
+		indicesBuilder.UnsafeSet(i, int32(i))
+	}
+	ms, err := newMultiColSorter(r, indicesBuilder, columns)
+	if err != nil {
+		return nil, err
 	}
+	sort.Sort(ms)
+	return indicesBuilder.NewArray().(*array.Int32), nil
+}
 
-	indices := make([]int64, r.NumRows())
-	// populate indices
-	for i := range indices {
-		indices[i] = int64(i)
+// ReorderRecord uses indices to create a new record that is sorted according to
+// the indices array.
+func ReorderRecord(ctx context.Context, mem memory.Allocator, r arrow.Record, indices *array.Int32) (arrow.Record, error) {
+	ctx = compute.WithAllocator(ctx, mem)
+	// compute.Take doesn't support dictionaries. Use take on r when r does not have
+	// dictionary column.
+	var hasDictionary bool
+	for i := 0; i < int(r.NumCols()); i++ {
+		if r.Column(i).DataType().ID() == arrow.DICTIONARY {
+			hasDictionary = true
+			break
+		}
 	}
+	if !hasDictionary {
+		res, err := compute.Take(
+			ctx,
+			compute.TakeOptions{BoundsCheck: true},
+			compute.NewDatum(r),
+			compute.NewDatum(indices),
+		)
+		if err != nil {
+			return nil, err
+		}
+		return res.(*compute.RecordDatum).Value, nil
+	}
+	resArr := make([]arrow.Array, r.NumCols())
 
-	switch c := r.Column(cols[0]).(type) {
-	case *array.Int64:
-		sort.Sort(orderedSorter[int64]{array: c, indices: indices})
-	case *array.String:
-		sort.Sort(orderedSorter[string]{array: c, indices: indices})
-	default:
-		return nil, fmt.Errorf("unsupported column type for sorting %T", c)
+	defer func() {
+		for _, a := range resArr {
+			if a != nil {
+				a.Release()
+			}
+		}
+	}()
+	var g errgroup.Group
+	for i := 0; i < int(r.NumCols()); i++ {
+		i := i
+		col := r.Column(i)
+		if d, ok := col.(*array.Dictionary); ok {
+			g.Go(func() error {
+				return takeDictColumn(ctx, d, i, resArr, indices)
+			})
+		} else {
+			g.Go(func() error {
+				return takeColumn(ctx, col, i, resArr, indices)
+			})
+		}
+	}
+	err := g.Wait()
+	if err != nil {
+		return nil, err
+	}
+	return array.NewRecord(r.Schema(), resArr, r.NumRows()), nil
+}
+
+func takeColumn(ctx context.Context, a arrow.Array, idx int, arr []arrow.Array, indices *array.Int32) error {
+	r, err := compute.TakeArray(ctx, a, indices)
+	if err != nil {
+		return err
 	}
+	arr[idx] = r
+	return nil
+}
 
-	indicesBuilder.Reserve(len(indices))
-	for _, i := range indices {
-		indicesBuilder.Append(i)
+func takeDictColumn(ctx context.Context, a *array.Dictionary, idx int, arr []arrow.Array, indices *array.Int32) error {
+	base := a.Dictionary().(*array.String)
+	// Only string dictionaries are supported.
+	r := array.NewBuilder(compute.GetAllocator(ctx), a.DataType()).(*array.BinaryDictionaryBuilder)
+	defer r.Release()
+	r.Reserve(indices.Len())
+	for _, i := range indices.Int32Values() {
+		if a.IsNull(int(i)) {
+			r.AppendNull()
+			continue
+		}
+		err := r.AppendString(base.Value(a.GetValueIndex(int(i))))
+		if err != nil {
+			return err
+		}
 	}
+	arr[idx] = r.NewArray()
+	return nil
+}
 
-	return indicesBuilder.NewInt64Array(), nil
+type multiColSorter struct {
+	indices     *builder.OptInt32Builder
+	comparisons []comparator
+	directions  []int
+	nullsFirst  []bool
 }
 
-// ReorderRecord reorders the given record's rows by the given indices.
-// This is a wrapper around compute.Take which handles the type castings.
-func ReorderRecord(ctx context.Context, r arrow.Record, indices arrow.Array) (arrow.Record, error) {
-	res, err := compute.Take(
-		ctx,
-		*compute.DefaultTakeOptions(),
-		compute.NewDatum(r),
-		compute.NewDatum(indices),
-	)
-	if err != nil {
-		return nil, err
+func newMultiColSorter(
+	r arrow.Record,
+	indices *builder.OptInt32Builder,
+	columns []SortingColumn,
+) (*multiColSorter, error) {
+	ms := &multiColSorter{
+		indices:     indices,
+		directions:  make([]int, len(columns)),
+		nullsFirst:  make([]bool, len(columns)),
+		comparisons: make([]comparator, len(columns)),
+	}
+	for i := range columns {
+		ms.directions[i] = int(columns[i].Direction.comparison())
+		ms.nullsFirst[i] = columns[i].NullsFirst
+	}
+	for i, col := range columns {
+		switch e := r.Column(col.Index).(type) {
+		case *array.Int64:
+			ms.comparisons[i] = &orderedSorter[int64]{array: e}
+		case *array.Float64:
+			ms.comparisons[i] = &orderedSorter[float64]{array: e}
+		case *array.String:
+			ms.comparisons[i] = &orderedSorter[string]{array: e}
+		case *array.Dictionary:
+			switch elem := e.Dictionary().(type) {
+			case *array.String:
+				ms.comparisons[i] = &orderedSorter[string]{
+					array: &stringDictionary{
+						dict: e,
+						elem: elem,
+					},
+				}
+			default:
+				return nil, fmt.Errorf("unsupported dictionary column type for sorting %T", e)
+			}
+		default:
+			return nil, fmt.Errorf("unsupported column type for sorting %T", e)
+		}
+	}
+	return ms, nil
+}
+
+var _ sort.Interface = (*multiColSorter)(nil)
+
+func (m *multiColSorter) Len() int { return m.indices.Len() }
+
+func (m *multiColSorter) Less(i, j int) bool {
+	for idx := range m.comparisons {
+		cmp := m.compare(idx, int(m.indices.UnsafeValue(i)), int(m.indices.UnsafeValue(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 == m.directions[idx]
+		}
+		// Try comparing the next column
+	}
+	return false
+}
+
+func (m *multiColSorter) compare(idx, i, j int) int {
+	x := m.comparisons[idx]
+	if x.IsNull(i) {
+		if x.IsNull(j) {
+			return 0
+		}
+		if m.directions[idx] == 1 {
+			if m.nullsFirst[idx] {
+				return 1
+			}
+			return -1
+		}
+		if m.nullsFirst[idx] {
+			return -1
+		}
+		return 1
+	}
+	if x.IsNull(j) {
+		if m.directions[idx] == 1 {
+			if m.nullsFirst[idx] {
+				return -1
+			}
+			return 1
+		}
+		if m.nullsFirst[idx] {
+			return 1
+		}
+		return -1
 	}
-	return res.(*compute.RecordDatum).Value, nil
+	return x.Compare(i, j)
+}
+
+func (m *multiColSorter) Swap(i, j int) {
+	m.indices.UnsafeSwap(i, j)
+}
+
+type comparator interface {
+	Compare(i, j int) int
+	IsNull(int) bool
 }
 
 type orderedArray[T int64 | float64 | string] interface {
 	Value(int) T
 	IsNull(int) bool
-	Len() int
 }
 
 type orderedSorter[T int64 | float64 | string] struct {
-	array   orderedArray[T]
-	indices []int64
+	array orderedArray[T]
 }
 
-func (s orderedSorter[T]) Len() int {
-	return s.array.Len()
+func (s *orderedSorter[T]) IsNull(i int) bool {
+	return s.array.IsNull(i)
 }
 
-func (s orderedSorter[T]) Less(i, j int) bool {
-	if s.array.IsNull(int(s.indices[i])) {
-		return false
+func (s *orderedSorter[T]) Compare(i, j int) int {
+	x := s.array.Value(i)
+	y := s.array.Value(j)
+	if x < y {
+		return -1
 	}
-	if s.array.IsNull(int(s.indices[j])) {
-		return true
+	if x > y {
+		return 1
 	}
-	return s.array.Value(int(s.indices[i])) < s.array.Value(int(s.indices[j]))
+	return 0
+}
+
+type stringDictionary struct {
+	dict *array.Dictionary
+	elem *array.String
+}
+
+func (s *stringDictionary) IsNull(i int) bool {
+	return s.dict.IsNull(i)
 }
 
-func (s orderedSorter[T]) Swap(i, j int) {
-	s.indices[i], s.indices[j] = s.indices[j], s.indices[i]
+func (s *stringDictionary) Value(i int) string {
+	return s.elem.Value(s.dict.GetValueIndex(i))
 }
diff --git a/pqarrow/builder/optbuilders.go b/pqarrow/builder/optbuilders.go
index a5dbaaf..a84172e 100644
--- a/pqarrow/builder/optbuilders.go
+++ b/pqarrow/builder/optbuilders.go
@@ -3,6 +3,7 @@ package builder
 import (
 	"fmt"
 	"math"
+	"slices"
 	"sync/atomic"
 	"unsafe"
 
@@ -637,6 +638,22 @@ func (b *OptInt32Builder) Set(i int, v int32) {
 	b.data[i] = v
 }
 
+// UnsafeSwap swaps without bounds check.
+func (b *OptInt32Builder) UnsafeSwap(i, j int) {
+	b.data[i], b.data[j] = b.data[j], b.data[i]
+}
+
+// UnsafeSet sets v at index i without bounds check.
+func (b *OptInt32Builder) UnsafeSet(i int, v int32) {
+	b.data[i] = v
+	bitutil.SetBit(b.validityBitmap, i)
+}
+
+// UnsafeValue returns value at index i without bounds check.
+func (b *OptInt32Builder) UnsafeValue(i int) int32 {
+	return b.data[i]
+}
+
 func (b *OptInt32Builder) Add(i int, v int32) {
 	if i < 0 || i >= len(b.data) {
 		panic("arrow/array: index out of range")
@@ -688,3 +705,9 @@ func (b *OptInt32Builder) ResetToLength(n int) {
 	b.data = b.data[:n]
 	b.validityBitmap = resizeBitmap(b.validityBitmap, n)
 }
+
+func (b *OptInt32Builder) Reserve(n int) {
+	b.length = n
+	b.data = slices.Grow(b.data, n)[:n]
+	b.validityBitmap = resizeBitmap(b.validityBitmap, n)
+}


Test patch:
diff --git a/pqarrow/arrowutils/sort_test.go b/pqarrow/arrowutils/sort_test.go
index 8f94586..d6855fd 100644
--- a/pqarrow/arrowutils/sort_test.go
+++ b/pqarrow/arrowutils/sort_test.go
@@ -11,73 +11,341 @@ import (
 )
 
 func TestSortRecord(t *testing.T) {
-	ctx := context.Background()
-	schema := arrow.NewSchema(
-		[]arrow.Field{
-			{Name: "int", Type: arrow.PrimitiveTypes.Int64},
-			{Name: "string", Type: arrow.BinaryTypes.String},
+	null := func(v int64) *int64 {
+		return &v
+	}
+
+	cases := []SortCase{
+		{
+			Name: "must provide at least one column",
+			Samples: Samples{
+				{},
+			},
+			Error: "expected missing column error",
+		},
+
+		{
+			Name:    "No Nows",
+			Samples: Samples{},
+			Columns: []SortingColumn{{Index: 0}},
+		},
+		{
+			Name: "One Row",
+			Samples: Samples{
+				{},
+			},
+			Columns: []SortingColumn{
+				{
+					Index: 0,
+				},
+			},
+			Indices: []int32{0},
+		},
+		{
+			Name: "By Integer column ascending",
+			Samples: Samples{
+				{Int: 3},
+				{Int: 2},
+				{Int: 1},
+			},
+			Columns: []SortingColumn{
+				{Index: 0},
+			},
+			Indices: []int32{2, 1, 0},
+		},
+		{
+			Name: "By Integer column descending",
+			Samples: Samples{
+				{Int: 1},
+				{Int: 2},
+				{Int: 3},
+			},
+
+			Columns: []SortingColumn{
+				{Index: 0, Direction: Descending},
+			},
+			Indices: []int32{2, 1, 0},
+		},
+		{
+			Name: "By Double column ascending",
+			Samples: Samples{
+				{Double: 3},
+				{Double: 2},
+				{Double: 1},
+			},
+			Columns: []SortingColumn{{Index: 1}},
+			Indices: []int32{2, 1, 0},
+		},
+		{
+			Name: "By Double column descending",
+			Samples: Samples{
+				{Double: 1},
+				{Double: 2},
+				{Double: 3},
+			},
+			Columns: []SortingColumn{{Index: 1, Direction: Descending}},
+			Indices: []int32{2, 1, 0},
+		},
+		{
+			Name: "By String column ascending",
+			Samples: Samples{
+				{String: "3"},
+				{String: "2"},
+				{String: "1"},
+			},
+			Columns: []SortingColumn{{Index: 2}},
+			Indices: []int32{2, 1, 0},
 		},
-		nil,
+		{
+			Name: "By String column descending",
+			Samples: Samples{
+				{String: "1"},
+				{String: "2"},
+				{String: "3"},
+			},
+			Columns: []SortingColumn{{Index: 2, Direction: Descending}},
+			Indices: []int32{2, 1, 0},
+		},
+		{
+			Name: "By Dict column ascending",
+			Samples: Samples{
+				{Dict: "3"},
+				{Dict: "2"},
+				{Dict: "1"},
+			},
+			Columns: []SortingColumn{{Index: 3}},
+			Indices: []int32{2, 1, 0},
+		},
+		{
+			Name: "By Dict column descending",
+			Samples: Samples{
+				{Dict: "1"},
+				{Dict: "2"},
+				{Dict: "3"},
+			},
+			Columns: []SortingColumn{{Index: 3, Direction: Descending}},
+			Indices: []int32{2, 1, 0},
+		},
+		{
+			Name: "By Null column ascending",
+			Samples: Samples{
+				{},
+				{},
+				{Nullable: null(1)},
+			},
+			Columns: []SortingColumn{{Index: 4}},
+			Indices: []int32{2, 0, 1},
+		},
+		{
+			Name: "By Null column ascending nullsFirst",
+			Samples: Samples{
+				{},
+				{},
+				{Nullable: null(1)},
+			},
+			Columns: []SortingColumn{{Index: 4, NullsFirst: true}},
+			Indices: []int32{0, 1, 2},
+		},
+		{
+			Name: "By Null column descending",
+			Samples: Samples{
+				{},
+				{},
+				{Nullable: null(1)},
+			},
+			Columns: []SortingColumn{{Index: 4, Direction: Descending}},
+			Indices: []int32{2, 0, 1},
+		},
+		{
+			Name: "By Null column descending nullsFirst",
+			Samples: Samples{
+				{},
+				{},
+				{Nullable: null(1)},
+			},
+			Columns: []SortingColumn{{Index: 4, Direction: Descending, NullsFirst: true}},
+			Indices: []int32{0, 1, 2},
+		},
+		{
+			Name: "Multiple columns same direction",
+			Samples: Samples{
+				{String: "1", Int: 3},
+				{String: "2", Int: 2},
+				{String: "3", Int: 2},
+				{String: "4", Int: 1},
+			},
+			Columns: []SortingColumn{
+				{Index: 0},
+				{Index: 2},
+			},
+			Indices: []int32{3, 1, 2, 0},
+		},
+		{
+			Name: "Multiple columns different direction",
+			Samples: Samples{
+				{String: "1", Int: 3},
+				{String: "2", Int: 2},
+				{String: "3", Int: 2},
+				{String: "4", Int: 1},
+			},
+			Columns: []SortingColumn{
+				{Index: 0, Direction: Ascending},
+				{Index: 2, Direction: Descending},
+			},
+			Indices: []int32{3, 2, 1, 0},
+		},
+	}
+
+	for _, kase := range cases {
+		t.Run(kase.Name, func(t *testing.T) {
+			sortAndCompare(t, kase)
+		})
+	}
+}
+
+func TestReorderRecord(t *testing.T) {
+	t.Run("Without dictionary field", func(t *testing.T) {
+		mem := memory.NewGoAllocator()
+		b := array.NewRecordBuilder(mem, arrow.NewSchema(
+			[]arrow.Field{
+				{
+					Name: "int",
+					Type: arrow.PrimitiveTypes.Int64,
+				},
+			}, nil,
+		))
+		defer b.Release()
+		b.Field(0).(*array.Int64Builder).AppendValues([]int64{3, 2, 1}, nil)
+		r := b.NewRecord()
+		defer r.Release()
+
+		indices := array.NewInt32Builder(mem)
+		indices.AppendValues([]int32{2, 1, 0}, nil)
+		by := indices.NewInt32Array()
+		result, err := ReorderRecord(context.Background(), mem, r, by)
+		require.Nil(t, err)
+		defer result.Release()
+
+		want := []int64{1, 2, 3}
+		require.Equal(t, want, result.Column(0).(*array.Int64).Int64Values())
+	})
+	t.Run("With dictionary field", func(t *testing.T) {
+		mem := memory.NewGoAllocator()
+		b := array.NewRecordBuilder(mem, arrow.NewSchema(
+			[]arrow.Field{
+				{
+					Name: "dict",
+					Type: &arrow.DictionaryType{
+						IndexType: arrow.PrimitiveTypes.Int32,
+						ValueType: arrow.BinaryTypes.String,
+					},
+				},
+			}, nil,
+		))
+		defer b.Release()
+		d := b.Field(0).(*array.BinaryDictionaryBuilder)
+		d.Reserve(3)
+		require.Nil(t, d.AppendString("3"))
+		require.Nil(t, d.AppendString("2"))
+		require.Nil(t, d.AppendString("1"))
+		r := b.NewRecord()
+		defer r.Release()
+
+		indices := array.NewInt32Builder(mem)
+		indices.AppendValues([]int32{2, 1, 0}, nil)
+		by := indices.NewInt32Array()
+		result, err := ReorderRecord(context.Background(), mem, r, by)
+		require.Nil(t, err)
+		defer result.Release()
+
+		want := []string{"1", "2", "3"}
+		got := result.Column(0).(*array.Dictionary)
+		str := got.Dictionary().(*array.String)
+		require.Equal(t, len(want), got.Len())
+		for i := range want {
+			require.Equal(t, want[i], str.Value(got.GetValueIndex(i)))
+		}
+	})
+}
+
+// Use all supported sort field.
+type Sample struct {
+	Int      int64
+	Double   float64
+	String   string
+	Dict     string
+	Nullable *int64
+}
+
+type Samples []Sample
+
+func (s Samples) Record() arrow.Record {
+	b := array.NewRecordBuilder(memory.NewGoAllocator(),
+		arrow.NewSchema([]arrow.Field{
+			{
+				Name: "int",
+				Type: arrow.PrimitiveTypes.Int64,
+			},
+			{
+				Name: "double",
+				Type: arrow.PrimitiveTypes.Float64,
+			},
+			{
+				Name: "string",
+				Type: arrow.BinaryTypes.String,
+			},
+			{
+				Name: "dict",
+				Type: &arrow.DictionaryType{
+					IndexType: arrow.PrimitiveTypes.Int32,
+					ValueType: arrow.BinaryTypes.String,
+				},
+			},
+			{
+				Name:     "nullable",
+				Type:     arrow.PrimitiveTypes.Int64,
+				Nullable: true,
+			},
+		}, nil),
 	)
 
-	mem := memory.DefaultAllocator
-	ib := array.NewInt64Builder(mem)
-	ib.Append(0)
-	ib.AppendNull()
-	ib.Append(3)
-	ib.Append(5)
-	ib.Append(1)
-
-	sb := array.NewStringBuilder(mem)
-	sb.Append("d")
-	sb.Append("c")
-	sb.Append("b")
-	sb.AppendNull()
-	sb.Append("a")
-
-	record := array.NewRecord(schema, []arrow.Array{ib.NewArray(), sb.NewArray()}, int64(5))
-
-	// Sort the record by the first column - int64
-	{
-		sortedIndices, err := SortRecord(mem, record, []int{record.Schema().FieldIndices("int")[0]})
-		require.NoError(t, err)
-		require.Equal(t, []int64{0, 4, 2, 3, 1}, sortedIndices.Int64Values())
-
-		sortedByInts, err := ReorderRecord(ctx, record, sortedIndices)
-		require.NoError(t, err)
-
-		// check that the column got sortedIndices
-		intCol := sortedByInts.Column(0).(*array.Int64)
-		require.Equal(t, []int64{0, 1, 3, 5, 0}, intCol.Int64Values())
-		require.True(t, intCol.IsNull(intCol.Len()-1)) // last is NULL
-		// make sure the other column got updated too
-		strings := make([]string, sortedByInts.NumRows())
-		stringCol := sortedByInts.Column(1).(*array.String)
-		for i := 0; i < int(sortedByInts.NumRows()); i++ {
-			strings[i] = stringCol.Value(i)
+	fInt := b.Field(0).(*array.Int64Builder)
+	fDouble := b.Field(1).(*array.Float64Builder)
+	fString := b.Field(2).(*array.StringBuilder)
+	fDict := b.Field(3).(*array.BinaryDictionaryBuilder)
+	fNullable := b.Field(4).(*array.Int64Builder)
+
+	for _, v := range s {
+		fInt.Append(v.Int)
+		fDouble.Append(v.Double)
+		fString.Append(v.String)
+		_ = fDict.AppendString(v.Dict)
+		if v.Nullable != nil {
+			fNullable.Append(*v.Nullable)
+		} else {
+			fNullable.AppendNull()
 		}
-		require.Equal(t, []string{"d", "a", "b", "", "c"}, strings)
 	}
+	return b.NewRecord()
+}
 
-	// Sort the record by the second column - string
-	{
-		sortedIndices, err := SortRecord(mem, record, []int{record.Schema().FieldIndices("string")[0]})
-		require.NoError(t, err)
-		require.Equal(t, []int64{4, 2, 1, 0, 3}, sortedIndices.Int64Values())
-
-		sortedByStrings, err := ReorderRecord(ctx, record, sortedIndices)
-		require.NoError(t, err)
-
-		// check that the column got sortedByInts
-		intCol := sortedByStrings.Column(0).(*array.Int64)
-		require.Equal(t, []int64{1, 3, 0, 0, 5}, intCol.Int64Values())
-		// make sure the other column got updated too
-		strings := make([]string, sortedByStrings.NumRows())
-		stringCol := sortedByStrings.Column(1).(*array.String)
-		for i := 0; i < int(sortedByStrings.NumRows()); i++ {
-			strings[i] = stringCol.Value(i)
-		}
-		require.Equal(t, []string{"a", "b", "c", "d", ""}, strings)
-		require.True(t, stringCol.IsNull(stringCol.Len()-1)) // last is NULL
+type SortCase struct {
+	Name    string
+	Samples Samples
+	Columns []SortingColumn
+	Indices []int32
+	Error   string
+}
+
+func sortAndCompare(t *testing.T, kase SortCase) {
+	t.Helper()
+
+	got, err := SortRecord(kase.Samples.Record(), kase.Columns)
+	if kase.Error != "" {
+		require.NotNil(t, err, kase.Error)
+		return
 	}
+	defer got.Release()
+
+	require.Equal(t, kase.Indices, got.Int32Values())
 }