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 0702f072 fix(compute): discard results after cancellation (#1278)
0702f072 is described below
commit 0702f0724e01d24b6f6db3b8eba80b9e529dc454
Author: Matt Topol <[email protected]>
AuthorDate: Thu Sep 3 11:09:05 2026 -0400
fix(compute): discard results after cancellation (#1278)
### Rationale for this change
`WrapResults` selects between result delivery and context cancellation.
When cancellation and channel closure become ready together, the channel
branch can win and return a partial result from a canceled
`CallFunction`. This has caused intermittent failures of the
caller-cancellation tests on Windows, macOS, and release verification.
### What changes are included in this PR?
- Discard and release scalar and vector partial outputs on cancellation
- Recheck cancellation after receiving output or channel closure
- Add direct coverage for releasing a single scalar output
- Release unexpected results in cancellation tests so regressions report
the primary assertion rather than a secondary allocator leak
### Are these changes tested?
- `go test ./arrow/compute -run
'Test(ScalarExecutorWrapResults|CallFunctionPreservesCallerCancellation)'
-count=100`
- `go test ./arrow/compute`
- pre-commit hooks
### Are there any user-facing changes?
Canceled compute calls now deterministically return no result instead of
occasionally returning a partial result.
Signed-off-by: Matt Topol <[email protected]>
---
arrow/compute/exec_test.go | 12 +++++++++
arrow/compute/executor.go | 58 ++++++++++++++++++++++++++++++++----------
arrow/compute/executor_test.go | 27 ++++++++++++++++++++
3 files changed, 83 insertions(+), 14 deletions(-)
diff --git a/arrow/compute/exec_test.go b/arrow/compute/exec_test.go
index 6f32abd3..65007dd5 100644
--- a/arrow/compute/exec_test.go
+++ b/arrow/compute/exec_test.go
@@ -169,6 +169,9 @@ func TestCallFunctionPreservesCallerCancellation(t
*testing.T) {
case <-time.After(time.Second):
t.Fatal("CallFunction did not stop after caller cancellation")
}
+ if result != nil {
+ defer result.Release()
+ }
require.Nil(t, result)
require.ErrorIs(t, callErr, cancellationErr)
}
@@ -228,6 +231,9 @@ func
TestCallFunctionReleasesPartialResultOnCallerCancellation(t *testing.T) {
t.Fatal("CallFunction did not stop after caller cancellation")
}
+ if result != nil {
+ defer result.Release()
+ }
require.Nil(t, result)
require.ErrorIs(t, callErr, cancellationErr)
}
@@ -319,6 +325,9 @@ func TestCallFunctionDrainsResultsAfterCallerCancellation(t
*testing.T) {
case <-time.After(time.Second):
t.Fatal("CallFunction did not finish after execution was
released")
}
+ if result != nil {
+ defer result.Release()
+ }
require.Nil(t, result)
require.ErrorIs(t, callErr, cancellationErr)
}
@@ -370,6 +379,9 @@ func
TestCallFunctionPreservesCallerCancellationForVectorFunction(t *testing.T)
case <-time.After(time.Second):
t.Fatal("vector CallFunction did not stop after caller
cancellation")
}
+ if result != nil {
+ defer result.Release()
+ }
require.Nil(t, result)
require.ErrorIs(t, callErr, cancellationErr)
}
diff --git a/arrow/compute/executor.go b/arrow/compute/executor.go
index c9e4c2ef..a804aae9 100644
--- a/arrow/compute/executor.go
+++ b/arrow/compute/executor.go
@@ -534,6 +534,13 @@ func (s *scalarExecutor) WrapResults(ctx context.Context,
out <-chan Datum, hasC
}
acc = nil
}
+ releaseOutput := func() {
+ if output != nil {
+ output.Release()
+ output = nil
+ }
+ releaseAccumulated()
+ }
toChunked := func() {
acc = output.(ArrayLikeDatum).Chunks()
@@ -562,15 +569,18 @@ func (s *scalarExecutor) WrapResults(ctx context.Context,
out <-chan Datum, hasC
for {
select {
case <-ctx.Done():
- // context is done, either cancelled or a timeout.
- // either way, we end early and return what we've got
so far.
- if output == nil {
- releaseAccumulated()
- return nil
- }
- return output
+ // A canceled call must not return a partial result.
+ releaseOutput()
+ return nil
case o, ok := <-out:
if !ok { // channel closed, wrap it up
+ // Cancellation and channel closure can become
ready together.
+ // Recheck the context so selecting the channel
cannot make a
+ // canceled call return a result
nondeterministically.
+ if ctx.Err() != nil {
+ releaseOutput()
+ return nil
+ }
if output != nil {
return output
}
@@ -1022,6 +1032,12 @@ func (v *vectorExecutor) WrapResults(ctx
context.Context, out <-chan Datum, hasC
case <-ctx.Done():
return nil
case output = <-out:
+ if output == nil || ctx.Err() != nil {
+ if output != nil {
+ output.Release()
+ }
+ return nil
+ }
}
// we got an output datum, but let's wait for the channel to
@@ -1031,6 +1047,10 @@ func (v *vectorExecutor) WrapResults(ctx
context.Context, out <-chan Datum, hasC
output.Release()
return nil
case <-out:
+ if ctx.Err() != nil {
+ output.Release()
+ return nil
+ }
return output
}
}
@@ -1046,6 +1066,13 @@ func (v *vectorExecutor) WrapResults(ctx
context.Context, out <-chan Datum, hasC
}
acc = nil
}
+ releaseOutput := func() {
+ if output != nil {
+ output.Release()
+ output = nil
+ }
+ releaseAccumulated()
+ }
toChunked := func() {
out := output.(ArrayLikeDatum).Chunks()
@@ -1089,15 +1116,18 @@ func (v *vectorExecutor) WrapResults(ctx
context.Context, out <-chan Datum, hasC
for {
select {
case <-ctx.Done():
- // context is done, either cancelled or a timeout.
- // either way, we end early and return what we've got
so far.
- if output == nil {
- releaseAccumulated()
- return nil
- }
- return output
+ // A canceled call must not return a partial result.
+ releaseOutput()
+ return nil
case o, ok := <-out:
if !ok { // channel closed, wrap it up
+ // Cancellation and channel closure can become
ready together.
+ // Recheck the context so selecting the channel
cannot make a
+ // canceled call return a result
nondeterministically.
+ if ctx.Err() != nil {
+ releaseOutput()
+ return nil
+ }
if output != nil {
return output
}
diff --git a/arrow/compute/executor_test.go b/arrow/compute/executor_test.go
index 1b7bfbe6..8de080c2 100644
--- a/arrow/compute/executor_test.go
+++ b/arrow/compute/executor_test.go
@@ -84,6 +84,33 @@ func
TestScalarExecutorWrapResultsReleasesAccumulatedOutputOnCancellation(t *tes
close(output)
}
+func TestScalarExecutorWrapResultsReleasesSingleOutputOnCancellation(t
*testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ builder := array.NewInt32Builder(mem)
+ builder.Append(42)
+ value := builder.NewInt32Array()
+ builder.Release()
+ defer value.Release()
+
+ output := make(chan Datum)
+ ctx, cancel := context.WithCancel(context.Background())
+ executor := &scalarExecutor{
+ nonAggExecImpl: nonAggExecImpl{outType: value.DataType()},
+ }
+
+ result := make(chan Datum, 1)
+ go func() {
+ result <- executor.WrapResults(ctx, output, false)
+ }()
+
+ output <- NewDatum(value)
+ cancel()
+ require.Nil(t, <-result)
+ close(output)
+}
+
func TestScalarExecutorWrapResultsHandlesClosedOutput(t *testing.T) {
output := make(chan Datum)
close(output)
