Skip to content

fix(compute): preserve caller context during execution - #1154

Open
fallintoplace wants to merge 3 commits into
apache:mainfrom
fallintoplace:fix/compute-preserve-caller-context
Open

fix(compute): preserve caller context during execution#1154
fallintoplace wants to merge 3 commits into
apache:mainfrom
fallintoplace:fix/compute-preserve-caller-context

Conversation

@fallintoplace

Copy link
Copy Markdown
Contributor

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

@zeroshade

Copy link
Copy Markdown
Member

Since this is a draft, i'll hold off on further review until the conflicts are resolved and it's marked ready

@fallintoplace
fallintoplace force-pushed the fix/compute-preserve-caller-context branch from 51973d3 to 430755b Compare August 21, 2026 19:41
@fallintoplace
fallintoplace marked this pull request as ready for review August 21, 2026 19:41

@zeroshade zeroshade left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The caller-context propagation is correct in intent, but the new cancellation path exposes result ownership and error-contract problems. In particular, cancellation can return no error or an already-released result, and accumulated scalar outputs can leak. I’ve left details and reproduction cases inline.


This review was drafted by an AI-assisted tool and confirmed by an Apache Arrow Go maintainer. After you’ve addressed the points above and pushed an update, an Apache Arrow Go maintainer — a real person — will take the next look at the PR. If you think a finding is misapplied, please reply on the PR and a maintainer will weigh in.

More on how Apache Arrow Go handles contributions:
CONTRIBUTING.md.

Comment thread arrow/compute/exec.go

result = executor.WrapResults(ctx, ch, haveChunkedArray(input.Values))
if err == nil {
if ctx.Err() != nil {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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().

Comment thread arrow/compute/exec.go
ectx := GetExecCtx(ctx)

ctx, cancel := context.WithCancel(context.Background())
ctx, cancel := context.WithCancel(ctx)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread arrow/compute/exec_test.go Outdated
return ctx.Ctx.Err()
}, nil)
require.NoError(t, fn.AddKernel(kernel))
require.True(t, GetFunctionRegistry().AddFunction(fn, false))

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants