feat(arrow/compute): sort support

[hamilton-earthscope]

Summary

Implements stable sort_indices (and sort via take) for arrays, chunked arrays, record batches, and tables using logical row indices over Chunked data without concatenating chunks. The control flow and ordering rules are modeled on Apache Arrow C++ vector_sort.cc / vector_sort_internal.h, with a few Go- and performance-driven differences called out below.

Parity with Arrow C++ (vector_sort.cc / vector_sort_internal.h)

Same overall structure

• Single sort key, one column

  • Multiple chunks: per-chunk sort then pairwise merge of sorted spans (C++ ChunkedArraySorter / ChunkedMergeImpl idea).
  • Single chunk, no validity nulls and no null-likes: direct stable sort on indices (C++ skips null partitioning when null_count == 0 and there are no null-likes).
  • Otherwise: partition validity nulls, partition float null-likes (NaN), stable sort of finite values, then VisitConstantRanges-style handling of ties (vector_sort_internal.go).

• Multiple sort keys

  • len(keys) <= kMaxRadixSortKeys (8): MSD radix path per record-batch range (radixRecordBatchSortRange ↔ ConcreteRecordBatchColumnSorter::SortRange).
  • More than 8 keys: MultipleKeyRecordBatchSorter-style global stable sort with lexicographic compare across keys (multipleKeyRecordBatchSortRange).
  • Aligned chunk boundaries across all keyed columns (typical table): sort each chunk slice with the same strategy, then merge spans like C++ TableSorter batch merge.

Same ordering semantics (intended match to C++)

• Per-key ascending / descending and null placement (including NaN as null-like for floats).
• Stable ordering: merge and slices.SortStableFunc are used so tie-breaking matches the C++ “left before right” stable merge behavior where documented in code.

Same “column comparator” role

• Go columnComparator interface ↔ C++ ColumnComparator: compareRowsForKey, null / null-like metadata, columnHasValidityNulls (skip PartitionNullsOnly when there are no validity nulls).

Physical types

• One monomorphic comparator type per supported physical pattern in vector_sort_physical.go, analogous to C++ ConcreteColumnComparator<T> (concrete *array.T + direct Value / Cmp / special cases for bool and intervals).

Intentional differences and rationale

Area	C++	This Go port
Resolving logical row → (chunk, offset)	Chunk / resolver machinery in C++	Dense logicalRowMap: one rowMapCell{chunk, local} per logical row when len(chunks) > 1; pair(i,j) resolves two rows in one shot. Why: random compares during sort/merge need O(1) resolution; a flat table + co-located fields beats repeated resolver work and improves locality vs separate chunk/local slices.
physicalColumnBase methods	N/A (different language)	Pointer receivers on pair / isNullAtGlobal / cell. Why: value receivers would copy slice headers (and map state) on every compare.
Stable sort primitive	std::stable_sort	slices.SortStableFunc (Go 1.21+). Why: library primitive; semantics aligned with stable weak ordering used elsewhere in the port.
Column dispatch at runtime	Templates + virtuals	columnComparator interface for “which column” in multi-key and merge loops. Why: idiomatic Go; per-type work stays in concrete compareRowsForKey implementations.
Chunked merge with null-likes (e.g. float)	C++ can split merge for null-like vs non-null-like regions (ChunkedMergeImpl)	Single less over full row order after per-chunk partitioning/sort. Why: simpler merge while preserving order as long as per-chunk phases match C++; documented in vector_sort.go comments.
Generics for physical columns	Templates instantiate fully	Explicit monomorphs only for the hot compare path. Why: measured regression vs Go generics on this hot path (inlining / assertions); verbosity traded for performance.

File Layout

• arrow/compute/vector_sort.go — sort_indices / sort registration and datum dispatch.
• arrow/compute/vector_sort_test.go — functional tests.
• arrow/compute/internal/kernels/vector_sort.go — orchestration, merge, SortIndices kernel.
• arrow/compute/internal/kernels/vector_sort_internal.go — null partitions, radix / multi-key batch sort.
• arrow/compute/internal/kernels/vector_sort_support.go — logicalRowMap and ordering helpers.
• arrow/compute/internal/kernels/vector_sort_physical.go — per-type column comparators.
• arrow/compute/internal/kernels/vector_sort_bench_test.go — benchmarks.

Testing

• go test ./arrow/compute -run TestSort -count=1
• Benchmarks: go test ./arrow/compute/internal/kernels -bench=BenchmarkSortIndices -benchmem .

References

• Arrow C++: cpp/src/arrow/compute/kernels/vector_sort.cc and vector_sort_internal.h (and related comparators).
  • https://github.com/apache/arrow/blob/main/cpp/src/arrow/compute/kernels/vector_sort.cc
  • https://github.com/apache/arrow/blob/main/cpp/src/arrow/compute/kernels/vector_sort_internal.h

Related Issues

• Closes [Go][Table] Implement sort function #66

[zeroshade]

can't we use go generics here and just embed the base?

i.e.

type physicalSortColumn[T arrow.ValueTypes] struct { physicalColumnBase }

...

func (c *physicalSortColumn[T]) compareRowsForKey(i, j uint64, key SortKey) int {
	ai, aj, li, lj := c.pair(i, j)
	a := ai.(arrow.TypedArray[T])
	b := aj.(arrow.TypedArray[T])
	if c.validityNulls {
		if v, stop := compareKeyedNulls(a.IsNull(li), b.IsNull(lj), key); stop {
			return v
		}
	}
	return compareOrdered(key.Order, a.Value(li), b.Value(lj))
}

func (c *physicalSortColumn[T]) isNullAt(row uint64) bool { return c.isNullAtGlobal(row) }

etc.

[hamilton-earthscope]

We certainly can and I made several attempts at a generic solution to avoid code duplication. Correctness was straightforward given the test suite, but I failed to match the performance of these verbose single-use classes via go's generics. The closest I got was ~25% slower. I opted for the performance over maintenance thinking this would be in the hot path for some critical operations (sorted iceberg writes/maintenance activities).

When I say performance, I mean the following benchmark test suite. It is rather limited in that it's only testing int64 and string, but the results were consistent across the types.

go test ./arrow/compute/internal/kernels -bench=BenchmarkSortIndices -benchmem

[zeroshade]

Interesting, I'm surprised that there was that much of a slowdown using generics. I'll take a look at this but for now I agree that the performance is definitely a higher priority. Maybe we set this up like elsewhere that we do codegen? Just so that we can prevent future bugs if we change this logic

[zeroshade]

if we're pre-ordering by the sort keys ColumnIndex fields, we should at least document that in the internal kernel implementation that this is assumed to have happened already.

[hamilton-earthscope]

done. let me know if the docstring on the kernel method is insufficient.