This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new bea104d0 fix(arrow/compute): release sort scratch buffers (#1431)
bea104d0 is described below
commit bea104d01941d9672629bf700f6e7a099e7a6399
Author: Matt Topol <[email protected]>
AuthorDate: Wed Oct 28 19:41:10 2026 +0000
fix(arrow/compute): release sort scratch buffers (#1431)
### Rationale for this change
SortIndices allocates temporary index buffers for single-column chunked
sorting, null/NaN partitioning, and multi-key sorting without releasing
them. A checked allocator reports outstanding memory even after callers
release the returned indices.
### What changes are included in this PR?
- Defer Release immediately after each of the three previously
unreleased scratch allocations, matching the existing integer radix-sort
pattern.
- Preserve ownership of the output buffer, which is transferred to the
returned result.
- Add checked-allocator regression coverage for an array with nulls, a
float64 array with NaN, a multi-chunk array, a multi-key record batch,
and a multi-key multi-chunk table. Each case also checks the resulting
permutation.
- Document temporary-buffer lifetime and result ownership on
SortIndices.
### Are these changes tested?
Yes. Before the fix, `go test ./arrow/compute -run
'^TestSortIndicesReleasesScratchBuffers$' -count=1` failed on
upstream/main with the new test added: all five subtests reported a
64-byte leak from `exec.(*KernelCtx).Allocate`.
After the fix, `go test ./arrow/compute/... -count=1` passed for all
four packages, including the new regression test.
### Are there any user-facing changes?
Sort operations no longer leak temporary allocator-backed scratch
buffers. Sort order and API signatures are unchanged; callers still own
and must release returned results.
---
arrow/compute/internal/kernels/vector_sort.go | 3 ++
arrow/compute/vector_sort.go | 20 ++++++++
arrow/compute/vector_sort_test.go | 68 +++++++++++++++++++++++++++
3 files changed, 91 insertions(+)
diff --git a/arrow/compute/internal/kernels/vector_sort.go
b/arrow/compute/internal/kernels/vector_sort.go
index fc09c5ba..025c3dd3 100644
--- a/arrow/compute/internal/kernels/vector_sort.go
+++ b/arrow/compute/internal/kernels/vector_sort.go
@@ -448,6 +448,7 @@ func SortIndices(ctx *exec.KernelCtx, columns
[]*arrow.Chunked, keys []SortKey)
chunks := columns[0].Chunks()
if len(chunks) > 1 {
tmpBuf := ctx.Allocate(nRows * arrow.Uint64SizeBytes)
+ defer tmpBuf.Release()
tmp := arrow.GetData[uint64](tmpBuf.Buf())[:nRows]
spanScratch := make([]chunkIndexSpan, len(chunks))
sortIndicesSingleColumnChunked(indices, chunks,
comparators[0], keys[0], tmp, spanScratch)
@@ -463,6 +464,7 @@ func SortIndices(ctx *exec.KernelCtx, columns
[]*arrow.Chunked, keys []SortKey)
slices.SortStableFunc(indices, func(a, b
uint64) int { return c0.compareRowsForKey(a, b, k0) })
} else {
tmpBuf := ctx.Allocate(nRows *
arrow.Uint64SizeBytes)
+ defer tmpBuf.Release()
tmp :=
arrow.GetData[uint64](tmpBuf.Buf())[:nRows]
arraySortOneColumnRange(indices, tmp, c0, k0,
0, nRows)
}
@@ -477,6 +479,7 @@ func SortIndices(ctx *exec.KernelCtx, columns
[]*arrow.Chunked, keys []SortKey)
multiChunkMerge := aligned && nSeg > 1
tmpBuf := ctx.Allocate(nRows * arrow.Uint64SizeBytes)
+ defer tmpBuf.Release()
tmp := arrow.GetData[uint64](tmpBuf.Buf())[:nRows]
var spanScratch []chunkIndexSpan
diff --git a/arrow/compute/vector_sort.go b/arrow/compute/vector_sort.go
index 2baaa95e..bad14e4f 100644
--- a/arrow/compute/vector_sort.go
+++ b/arrow/compute/vector_sort.go
@@ -206,11 +206,15 @@ func sortIndicesImpl(ctx context.Context, opts
FunctionOptions, input Datum) (Da
// SortIndices computes the indices that would sort the input.
// For arrays and chunked arrays, pass a single SortKey (ColumnIndex is
ignored).
// For record batches and tables, pass []SortKey to specify multi-column sort
order.
+// Temporary sorting buffers use the context's allocator and are released
before returning.
+// The caller must release the returned datum.
func SortIndices(ctx context.Context, input Datum, keys SortOptions) (Datum,
error) {
return CallFunction(ctx, "sort_indices", keys, input)
}
// SortIndicesArray computes the indices that would sort the input array.
+// As with [SortIndices], temporary sorting buffers use the context's allocator
+// and are released before returning. The caller must release the returned
array.
func SortIndicesArray(ctx context.Context, input arrow.Array, key SortKey)
(arrow.Array, error) {
v := NewDatumWithoutOwning(input)
@@ -224,6 +228,8 @@ func SortIndicesArray(ctx context.Context, input
arrow.Array, key SortKey) (arro
}
// SortIndicesChunked computes the indices that would sort the input chunked
array.
+// As with [SortIndices], temporary sorting buffers use the context's allocator
+// and are released before returning. The caller must release the returned
array.
func SortIndicesChunked(ctx context.Context, input *arrow.Chunked, key
SortKey) (arrow.Array, error) {
v := NewDatumWithoutOwning(input)
@@ -237,6 +243,8 @@ func SortIndicesChunked(ctx context.Context, input
*arrow.Chunked, key SortKey)
}
// SortIndicesRecordBatch computes the indices that would sort the record
batch (stable, lexicographic by keys).
+// As with [SortIndices], temporary sorting buffers use the context's allocator
+// and are released before returning. The caller must release the returned
array.
func SortIndicesRecordBatch(ctx context.Context, batch arrow.RecordBatch, keys
[]SortKey) (arrow.Array, error) {
if len(keys) == 0 {
return nil, fmt.Errorf("%w: at least one sort key is required",
arrow.ErrInvalid)
@@ -254,6 +262,8 @@ func SortIndicesRecordBatch(ctx context.Context, batch
arrow.RecordBatch, keys [
}
// SortIndicesTable computes the indices that would sort the table (stable,
lexicographic by keys).
+// As with [SortIndices], temporary sorting buffers use the context's allocator
+// and are released before returning. The caller must release the returned
array.
func SortIndicesTable(ctx context.Context, tbl arrow.Table, keys []SortKey)
(arrow.Array, error) {
if len(keys) == 0 {
return nil, fmt.Errorf("%w: at least one sort key is required",
arrow.ErrInvalid)
@@ -272,11 +282,15 @@ func SortIndicesTable(ctx context.Context, tbl
arrow.Table, keys []SortKey) (arr
// Sort returns a sorted copy of the input datum by calling the registered
"sort" function
// ([SortIndices] then [Take]). Supported kinds are [KindArray],
[KindChunked], [KindRecord], and [KindTable].
+// As with [SortIndices], temporary sorting buffers use the context's allocator
+// and are released before returning. The caller must release the returned
datum.
func Sort(ctx context.Context, input Datum, keys SortOptions) (Datum, error) {
return CallFunction(ctx, "sort", keys, input)
}
// SortArray returns a sorted copy of the input array.
+// As with [SortIndices], temporary sorting buffers use the context's allocator
+// and are released before returning. The caller must release the returned
array.
func SortArray(ctx context.Context, input arrow.Array, key SortKey)
(arrow.Array, error) {
indicesArr, err := SortIndicesArray(ctx, input, key)
if err != nil {
@@ -288,6 +302,8 @@ func SortArray(ctx context.Context, input arrow.Array, key
SortKey) (arrow.Array
}
// SortChunked returns a sorted copy of the input chunked array.
+// As with [SortIndices], temporary sorting buffers use the context's allocator
+// and are released before returning. The caller must release the returned
chunked array.
func SortChunked(ctx context.Context, input *arrow.Chunked, key SortKey)
(*arrow.Chunked, error) {
inputDatum := NewDatumWithoutOwning(input)
@@ -311,6 +327,8 @@ func SortChunked(ctx context.Context, input *arrow.Chunked,
key SortKey) (*arrow
// SortRecordBatch returns a sorted copy of the record batch using
lexicographic ordering across the specified columns.
// Each SortKey specifies a column index and its sort order and null placement.
// When multiple keys are provided, ties in earlier columns are broken by
later columns.
+// As with [SortIndices], temporary sorting buffers use the context's allocator
+// and are released before returning. The caller must release the returned
record batch.
//
// Example:
//
@@ -347,6 +365,8 @@ func SortRecordBatch(ctx context.Context, batch
arrow.RecordBatch, keys []kernel
// SortTable returns a sorted copy of the table using lexicographic ordering
across the specified columns.
// Each SortKey specifies a column index and its sort order and null placement.
// When multiple keys are provided, ties in earlier columns are broken by
later columns.
+// As with [SortIndices], temporary sorting buffers use the context's allocator
+// and are released before returning. The caller must release the returned
table.
//
// Example:
//
diff --git a/arrow/compute/vector_sort_test.go
b/arrow/compute/vector_sort_test.go
index 39bf5e95..a4922baa 100644
--- a/arrow/compute/vector_sort_test.go
+++ b/arrow/compute/vector_sort_test.go
@@ -323,6 +323,74 @@ func TestSortIndices(t *testing.T) {
}
}
+func TestSortIndicesReleasesScratchBuffers(t *testing.T) {
+ for _, input := range []string{"ArrayNulls", "ArrayNaN", "Chunked",
"RecordBatch", "Table"} {
+ t.Run(input, func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ bldr := array.NewFloat64Builder(mem)
+ defer bldr.Release()
+ valid := []bool{true, true, true, false}
+ if input == "ArrayNaN" {
+ valid = nil
+ }
+ bldr.AppendValues([]float64{2, 1, 1, math.NaN()}, valid)
+ arr := bldr.NewArray()
+ defer arr.Release()
+
+ key := compute.DefaultSortKey()
+ want := []uint64{1, 2, 0, 3}
+ var (
+ indices arrow.Array
+ err error
+ )
+ switch input {
+ case "ArrayNulls", "ArrayNaN":
+ indices, err = compute.SortIndicesArray(ctx,
arr, key)
+ case "Chunked":
+ first := array.NewSlice(arr, 0, 2)
+ defer first.Release()
+ second := array.NewSlice(arr, 2, 4)
+ defer second.Release()
+ chunked := arrow.NewChunked(arr.DataType(),
[]arrow.Array{first, second})
+ defer chunked.Release()
+ indices, err = compute.SortIndicesChunked(ctx,
chunked, key)
+ case "RecordBatch", "Table":
+ secondary, _, buildErr := array.FromJSON(mem,
arrow.PrimitiveTypes.Int32,
+ strings.NewReader("[0, 2, 1, 0]"))
+ require.NoError(t, buildErr)
+ defer secondary.Release()
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "primary", Type: arr.DataType(),
Nullable: true},
+ {Name: "secondary", Type:
secondary.DataType()},
+ }, nil)
+ batch := array.NewRecordBatch(schema,
[]arrow.Array{arr, secondary}, 4)
+ defer batch.Release()
+ secondKey := compute.DefaultSortKey()
+ secondKey.ColumnIndex = 1
+ keys := []compute.SortKey{key, secondKey}
+ want = []uint64{2, 1, 0, 3}
+ if input == "RecordBatch" {
+ indices, err =
compute.SortIndicesRecordBatch(ctx, batch, keys)
+ } else {
+ first := batch.NewSlice(0, 2)
+ defer first.Release()
+ second := batch.NewSlice(2, 4)
+ defer second.Release()
+ tbl :=
array.NewTableFromRecords(schema, []arrow.RecordBatch{first, second})
+ defer tbl.Release()
+ indices, err =
compute.SortIndicesTable(ctx, tbl, keys)
+ }
+ }
+ require.NoError(t, err)
+ defer indices.Release()
+ assert.Equal(t, want,
indices.(*array.Uint64).Uint64Values())
+ })
+ }
+}
+
func TestSortArray(t *testing.T) {
mem := memory.NewGoAllocator()
ctx := context.Background()
