zeroshade commented on code in PR #1154:
URL: https://github.com/apache/arrow-go/pull/1154#discussion_r3833570208
##########
arrow/compute/exec.go:
##########
@@ -170,7 +170,14 @@ func execInternal(ctx context.Context, fn Function, opts
FunctionOptions, passed
}()
result = executor.WrapResults(ctx, ch, haveChunkedArray(input.Values))
- if err == nil {
+ if ctx.Err() != nil {
Review Comment:
**Blocking:** This cancellation branch drains the channel but does not
guarantee a cancellation error or clear a released result. A context-oblivious
kernel can finish successfully after cancellation, leaving `err == nil`. If
cancellation occurs before the first output, `CallFunction` returns `(nil,
nil)`; after the first scalar output, `WrapResults` can return a partial result
that line 185 releases but the named return still exposes. I reproduced the
latter as a non-nil `ArrayDatum` with nil data and `context.Canceled`. Please
propagate `context.Cause(ctx)` when `err` is nil, and release **and set `result
= nil`** before returning. A regression test should use a kernel that does not
itself return `ctx.Err()`.
##########
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)
Review Comment:
**Major:** Deriving this context from the caller makes the scalar
collector’s cancellation branch externally reachable. When
`scalarExecutor.WrapResults` has already moved outputs into its private `acc`,
cancellation returns without releasing those chunks. The drain below cannot
reclaim values already consumed from the channel. Please add scalar
accumulated-result cleanup equivalent to `vectorExecutor.WrapResults`’s
`releaseAccumulated()` path.
##########
arrow/compute/exec_test.go:
##########
@@ -115,6 +119,51 @@ func ExecAddInt32(ctx *exec.KernelCtx, batch
*exec.ExecSpan, out *exec.ExecResul
return nil
}
+func TestCallFunctionPreservesCallerCancellation(t *testing.T) {
+ started := make(chan struct{})
+
+ 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 ctx.Ctx.Err()
+ }, nil)
+ require.NoError(t, fn.AddKernel(kernel))
+ require.True(t, GetFunctionRegistry().AddFunction(fn, false))
Review Comment:
**Minor:** Registering this function in the process-wide registry makes the
test non-repeatable: `go test ./arrow/compute -run
'^TestCallFunctionPreservesCallerCancellation$' -count=2` fails on its second
execution. Please install a child registry through `ExecCtx` for this test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
