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 5aadcffb fix(compute): preserve caller context during execution (#1154)
5aadcffb is described below
commit 5aadcffb2c32dac10b70fc757433a2e785b4a15c
Author: Minh Vu <[email protected]>
AuthorDate: Tue Aug 25 18:04:47 2026 +0200
fix(compute): preserve caller context during execution (#1154)
## What
Keep the supplied context attached to compute execution so cancellation
reaches result collection. Drain pending results before returning so the
executor is not reused while work is still finishing.
## Test
- go test ./arrow/compute -count=1
- go test -race ./arrow/compute -run
^TestCallFunctionPreservesCallerCancellation -count=1
---
arrow/compute/exec.go | 22 +++-
arrow/compute/exec_test.go | 259 +++++++++++++++++++++++++++++++++++++++++
arrow/compute/executor.go | 23 +++-
arrow/compute/executor_test.go | 56 +++++++++
4 files changed, 350 insertions(+), 10 deletions(-)
diff --git a/arrow/compute/exec.go b/arrow/compute/exec.go
index 0afb3b13..710631c8 100644
--- a/arrow/compute/exec.go
+++ b/arrow/compute/exec.go
@@ -158,7 +158,7 @@ func execInternal(ctx context.Context, fn Function, opts
FunctionOptions, passed
ectx := GetExecCtx(ctx)
- ctx, cancel := context.WithCancel(context.Background())
+ ctx, cancel := context.WithCancel(ctx)
defer cancel()
ch := make(chan Datum, ectx.ExecChannelSize)
@@ -170,12 +170,22 @@ func execInternal(ctx context.Context, fn Function, opts
FunctionOptions, passed
}()
result = executor.WrapResults(ctx, ch, haveChunkedArray(input.Values))
- if err == nil {
- debug.Assert(executor.CheckResultType(result) == nil, "invalid
result type")
+ if ctx.Err() != nil {
+ for value := range ch {
+ if value != nil {
+ value.Release()
+ }
+ }
+ if err == nil {
+ err = context.Cause(ctx)
+ }
+ if result != nil {
+ result.Release()
+ result = nil
+ }
}
-
- if ctx.Err() == context.Canceled && result != nil {
- result.Release()
+ if err == nil && result != nil {
+ debug.Assert(executor.CheckResultType(result) == nil, "invalid
result type")
}
return
diff --git a/arrow/compute/exec_test.go b/arrow/compute/exec_test.go
index eb72a2bf..6f32abd3 100644
--- a/arrow/compute/exec_test.go
+++ b/arrow/compute/exec_test.go
@@ -19,15 +19,22 @@
package compute
import (
+ "context"
+ "errors"
"strings"
+ "sync"
+ "sync/atomic"
"testing"
+ "time"
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/array"
"github.com/apache/arrow-go/v18/arrow/bitutil"
"github.com/apache/arrow-go/v18/arrow/compute/exec"
"github.com/apache/arrow-go/v18/arrow/internal/debug"
+ "github.com/apache/arrow-go/v18/arrow/memory"
"github.com/apache/arrow-go/v18/arrow/scalar"
+ "github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
)
@@ -115,6 +122,258 @@ func ExecAddInt32(ctx *exec.KernelCtx, batch
*exec.ExecSpan, out *exec.ExecResul
return nil
}
+func TestCallFunctionPreservesCallerCancellation(t *testing.T) {
+ started := make(chan struct{})
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ fn := NewScalarFunction("test_preserve_caller_cancellation", Unary(),
EmptyFuncDoc)
+ kernel := exec.NewScalarKernel(
+
[]exec.InputType{exec.NewExactInput(arrow.PrimitiveTypes.Int32)},
+ exec.NewOutputType(arrow.PrimitiveTypes.Int32),
+ func(ctx *exec.KernelCtx, _ *exec.ExecSpan, _ *exec.ExecResult)
error {
+ close(started)
+ <-ctx.Ctx.Done()
+ return nil
+ }, nil)
+ require.NoError(t, fn.AddKernel(kernel))
+
+ execCtx := DefaultExecCtx()
+ execCtx.Registry = NewChildRegistry(execCtx.Registry)
+ require.True(t, execCtx.Registry.AddFunction(fn, false))
+
+ input, _, err := array.FromJSON(mem, arrow.PrimitiveTypes.Int32,
strings.NewReader(`[1]`))
+ require.NoError(t, err)
+ defer input.Release()
+
+ cancellationErr := errors.New("caller canceled")
+ ctx, cancel := context.WithCancelCause(context.Background())
+ ctx = WithAllocator(ctx, mem)
+ ctx = SetExecCtx(ctx, execCtx)
+ defer cancel(nil)
+
+ done := make(chan struct{})
+ var (
+ result Datum
+ callErr error
+ )
+ go func() {
+ result, callErr = CallFunction(ctx, fn.Name(), nil,
&ArrayDatum{Value: input.Data()})
+ close(done)
+ }()
+
+ <-started
+ cancel(cancellationErr)
+ select {
+ case <-done:
+ case <-time.After(time.Second):
+ t.Fatal("CallFunction did not stop after caller cancellation")
+ }
+ require.Nil(t, result)
+ require.ErrorIs(t, callErr, cancellationErr)
+}
+
+func TestCallFunctionReleasesPartialResultOnCallerCancellation(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ secondSpanStarted := make(chan struct{})
+ fn := NewScalarFunction("test_release_partial_result_on_cancellation",
Unary(), EmptyFuncDoc)
+ calls := 0
+ kernel := exec.NewScalarKernel(
+
[]exec.InputType{exec.NewExactInput(arrow.PrimitiveTypes.Int32)},
+ exec.NewOutputType(arrow.PrimitiveTypes.Int32),
+ func(ctx *exec.KernelCtx, _ *exec.ExecSpan, _ *exec.ExecResult)
error {
+ calls++
+ if calls == 2 {
+ close(secondSpanStarted)
+ <-ctx.Ctx.Done()
+ }
+ return nil
+ }, nil)
+ require.NoError(t, fn.AddKernel(kernel))
+
+ execCtx := DefaultExecCtx()
+ execCtx.ChunkSize = 1
+ execCtx.ExecChannelSize = 0
+ execCtx.PreallocContiguous = false
+ execCtx.Registry = NewChildRegistry(execCtx.Registry)
+ require.True(t, execCtx.Registry.AddFunction(fn, false))
+
+ input, _, err := array.FromJSON(mem, arrow.PrimitiveTypes.Int32,
strings.NewReader(`[1, 2]`))
+ require.NoError(t, err)
+ defer input.Release()
+
+ cancellationErr := errors.New("caller canceled after partial output")
+ ctx, cancel := context.WithCancelCause(context.Background())
+ ctx = WithAllocator(ctx, mem)
+ ctx = SetExecCtx(ctx, execCtx)
+ defer cancel(nil)
+
+ done := make(chan struct{})
+ var (
+ result Datum
+ callErr error
+ )
+ go func() {
+ result, callErr = CallFunction(ctx, fn.Name(), nil,
&ArrayDatum{Value: input.Data()})
+ close(done)
+ }()
+
+ <-secondSpanStarted
+ cancel(cancellationErr)
+ select {
+ case <-done:
+ case <-time.After(time.Second):
+ t.Fatal("CallFunction did not stop after caller cancellation")
+ }
+
+ require.Nil(t, result)
+ require.ErrorIs(t, callErr, cancellationErr)
+}
+
+func TestCallFunctionDrainsResultsAfterCallerCancellation(t *testing.T) {
+ const numChunks = 2
+
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ chunks := make([]arrow.Array, numChunks)
+ for i := range chunks {
+ builder := array.NewInt32Builder(mem)
+ builder.Append(int32(i))
+ chunks[i] = builder.NewInt32Array()
+ builder.Release()
+ defer chunks[i].Release()
+ }
+ chunked := arrow.NewChunked(arrow.PrimitiveTypes.Int32, chunks)
+ defer chunked.Release()
+
+ started := make(chan struct{})
+ secondStarted := make(chan struct{})
+ allowFinish := make(chan struct{})
+ var allowFinishOnce sync.Once
+ releaseSecond := func() {
+ allowFinishOnce.Do(func() { close(allowFinish) })
+ }
+ defer releaseSecond()
+
+ var calls atomic.Int32
+ fn := NewScalarFunction("test_drain_results_after_cancellation",
Unary(), EmptyFuncDoc)
+ kernel := exec.NewScalarKernel(
+
[]exec.InputType{exec.NewExactInput(arrow.PrimitiveTypes.Int32)},
+ exec.NewOutputType(arrow.PrimitiveTypes.Int32),
+ func(ctx *exec.KernelCtx, _ *exec.ExecSpan, _ *exec.ExecResult)
error {
+ switch calls.Add(1) {
+ case 1:
+ close(started)
+ <-ctx.Ctx.Done()
+ case 2:
+ close(secondStarted)
+ <-allowFinish
+ }
+ return nil
+ }, nil)
+ require.NoError(t, fn.AddKernel(kernel))
+
+ execCtx := DefaultExecCtx()
+ execCtx.Registry = NewChildRegistry(execCtx.Registry)
+ require.True(t, execCtx.Registry.AddFunction(fn, false))
+
+ cancellationErr := errors.New("caller canceled")
+ ctx, cancel := context.WithCancelCause(context.Background())
+ ctx = WithAllocator(ctx, mem)
+ ctx = SetExecCtx(ctx, execCtx)
+ defer cancel(nil)
+
+ done := make(chan struct{})
+ var (
+ result Datum
+ callErr error
+ )
+ go func() {
+ result, callErr = CallFunction(ctx, fn.Name(), nil,
&ChunkedDatum{Value: chunked})
+ close(done)
+ }()
+
+ <-started
+ cancel(cancellationErr)
+ select {
+ case <-secondStarted:
+ case <-time.After(time.Second):
+ releaseSecond()
+ <-done
+ t.Fatal("execution did not continue after caller cancellation")
+ }
+
+ select {
+ case <-done:
+ releaseSecond()
+ t.Fatal("CallFunction returned before execution finished")
+ default:
+ }
+
+ releaseSecond()
+ select {
+ case <-done:
+ case <-time.After(time.Second):
+ t.Fatal("CallFunction did not finish after execution was
released")
+ }
+ require.Nil(t, result)
+ require.ErrorIs(t, callErr, cancellationErr)
+}
+
+func TestCallFunctionPreservesCallerCancellationForVectorFunction(t
*testing.T) {
+ started := make(chan struct{})
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ fn := NewVectorFunction("test_vector_preserve_caller_cancellation",
Unary(), EmptyFuncDoc)
+ kernel := exec.NewVectorKernel(
+
[]exec.InputType{exec.NewExactInput(arrow.PrimitiveTypes.Int32)},
+ exec.NewOutputType(arrow.PrimitiveTypes.Int32),
+ func(ctx *exec.KernelCtx, _ *exec.ExecSpan, _ *exec.ExecResult)
error {
+ close(started)
+ <-ctx.Ctx.Done()
+ return nil
+ }, nil)
+ require.NoError(t, fn.AddKernel(kernel))
+
+ execCtx := DefaultExecCtx()
+ execCtx.Registry = NewChildRegistry(execCtx.Registry)
+ require.True(t, execCtx.Registry.AddFunction(fn, false))
+
+ input, _, err := array.FromJSON(mem, arrow.PrimitiveTypes.Int32,
strings.NewReader(`[1]`))
+ require.NoError(t, err)
+ defer input.Release()
+
+ cancellationErr := errors.New("caller canceled")
+ ctx, cancel := context.WithCancelCause(context.Background())
+ ctx = WithAllocator(ctx, mem)
+ ctx = SetExecCtx(ctx, execCtx)
+ defer cancel(nil)
+
+ done := make(chan struct{})
+ var (
+ result Datum
+ callErr error
+ )
+ go func() {
+ result, callErr = CallFunction(ctx, fn.Name(), nil,
&ArrayDatum{Value: input.Data()})
+ close(done)
+ }()
+
+ <-started
+ cancel(cancellationErr)
+ select {
+ case <-done:
+ case <-time.After(time.Second):
+ t.Fatal("vector CallFunction did not stop after caller
cancellation")
+ }
+ require.Nil(t, result)
+ require.ErrorIs(t, callErr, cancellationErr)
+}
+
type CallScalarFuncSuite struct {
ComputeInternalsTestSuite
}
diff --git a/arrow/compute/executor.go b/arrow/compute/executor.go
index 3468d640..febbec08 100644
--- a/arrow/compute/executor.go
+++ b/arrow/compute/executor.go
@@ -522,7 +522,14 @@ func (s *scalarExecutor) WrapResults(ctx context.Context,
out <-chan Datum, hasC
var (
output Datum
acc []arrow.Array
+ ok bool
)
+ releaseAccumulated := func() {
+ for _, c := range acc {
+ c.Release()
+ }
+ acc = nil
+ }
toChunked := func() {
acc = output.(ArrayLikeDatum).Chunks()
@@ -534,7 +541,13 @@ func (s *scalarExecutor) WrapResults(ctx context.Context,
out <-chan Datum, hasC
select {
case <-ctx.Done():
return nil
- case output = <-out:
+ case output, ok = <-out:
+ if !ok || output == nil || ctx.Err() != nil {
+ if output != nil {
+ output.Release()
+ }
+ return nil
+ }
// if the inputs contained at least one chunked array
// then we want to return chunked output
if hasChunked {
@@ -547,6 +560,10 @@ func (s *scalarExecutor) WrapResults(ctx context.Context,
out <-chan Datum, hasC
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
case o, ok := <-out:
if !ok { // channel closed, wrap it up
@@ -554,9 +571,7 @@ func (s *scalarExecutor) WrapResults(ctx context.Context,
out <-chan Datum, hasC
return output
}
- for _, c := range acc {
- defer c.Release()
- }
+ defer releaseAccumulated()
chkd := arrow.NewChunked(s.outType, acc)
defer chkd.Release()
diff --git a/arrow/compute/executor_test.go b/arrow/compute/executor_test.go
index dbc2312b..1b7bfbe6 100644
--- a/arrow/compute/executor_test.go
+++ b/arrow/compute/executor_test.go
@@ -39,6 +39,62 @@ func (d *signalChunkedDatum) Chunks() []arrow.Array {
return d.Value.Chunks()
}
+type signalArrayDatum struct {
+ *ArrayDatum
+ chunksCalled chan struct{}
+}
+
+func (d *signalArrayDatum) Chunks() []arrow.Array {
+ close(d.chunksCalled)
+ return d.ArrayDatum.Chunks()
+}
+
+func TestScalarExecutorWrapResultsReleasesAccumulatedOutputOnCancellation(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())
+ defer cancel()
+
+ executor := &scalarExecutor{
+ nonAggExecImpl: nonAggExecImpl{outType: value.DataType()},
+ }
+
+ datum := &signalArrayDatum{
+ ArrayDatum: NewDatum(value).(*ArrayDatum),
+ chunksCalled: make(chan struct{}),
+ }
+ result := make(chan Datum, 1)
+ go func() {
+ result <- executor.WrapResults(ctx, output, true)
+ }()
+
+ output <- datum
+ <-datum.chunksCalled
+ cancel()
+
+ require.Nil(t, <-result)
+ close(output)
+}
+
+func TestScalarExecutorWrapResultsHandlesClosedOutput(t *testing.T) {
+ output := make(chan Datum)
+ close(output)
+
+ executor := &scalarExecutor{
+ nonAggExecImpl: nonAggExecImpl{outType:
arrow.PrimitiveTypes.Int32},
+ }
+
+ require.Nil(t, executor.WrapResults(context.Background(), output,
false))
+}
+
func TestVectorExecutorWrapResultsReleasesEmptyArrayOutput(t *testing.T) {
mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
defer mem.AssertSize(t, 0)
