Skip to content

feat(arrow/compute): scalar aggregate framework with the count and sum kernels - #1336

Open
singhpratech wants to merge 6 commits into
apache:mainfrom
singhpratech:feat-compute-scalar-aggregates
Open

singhpratech wants to merge 6 commits into
apache:mainfrom
singhpratech:feat-compute-scalar-aggregates

Conversation

@singhpratech

Copy link
Copy Markdown
Contributor

Rationale for this change

The compute package has no way to compute a summary value from an array. FuncScalarAgg exists as
a function kind but nothing implements it: exec has a kernel type for scalar and vector kernels
only, there is no aggregate function type, and execInternal returns ErrNotImplemented for that
kind. Adding a single aggregate function therefore means adding the framework first, which is what
this pull request does, together with the first two kernels.

The design was discussed on #1296 and the interface follows the changes asked for there.

What changes are included in this PR?

The framework, in arrow/compute/exec:

  • ScalarAggKernel, following the C++ ScalarAggregateKernel: an init function that creates a
    state, Consume to fold one ExecSpan into it, Merge to combine two states, Finalize to
    produce the result, plus an optional Cleanup and an Ordered flag.
  • AggregateResult, a small carrier which holds either an owned scalar.Scalar or owned
    arrow.ArrayData. exec cannot import compute and so cannot name a Datum, but an aggregate
    result is not always a scalar, so the carrier keeps the exported interface general; the compute
    executor boxes it into a Datum.
  • AggKernel, the interface the executor consumes, and MergeAll, for a caller which aggregated
    partitions of the input separately.

The executor, in arrow/compute:

  • ScalarAggregateFunction, funcImpl[exec.ScalarAggKernel] with exact dispatch, its own
    SetDefaultOptions, and AddKernel/AddNewKernel which reject a kernel missing any of the four
    lifecycle functions.
  • scalarAggExecutor, reached from execInternal for FuncScalarAgg. It iterates the input in
    spans without promoting scalars to arrays, as the C++ ScalarAggExecutor does, so a kernel sees
    a scalar weighted by the length of the span; it consumes into a single state and finalizes once.
    An empty input still reaches finalize, which is what lets min_count decide the result of an
    aggregation over no values.
  • State cleanup that runs exactly once: when Execute returns, whether it finalized or failed
    part way, and through the executor's Clear on the paths where Execute never ran, such as a
    failure in Init. The result finalize returned owns its own value and is unaffected.

The kernels, in arrow/compute/internal/kernels/aggregate_basic.go and
arrow/compute/scalar_aggregate.go:

  • count, with CountOptions and the modes CountOnlyValid, CountOnlyNull and CountAllRows,
    over any input type. Counting the valid or the null values of a run-end encoded, dictionary or
    union array needs a logical null count that an ArraySpan does not carry, so those return
    ErrNotImplemented; CountAllRows works for them because it never looks at validity.
  • sum over int8..int64 (accumulating into int64), uint8..uint64 and bool (into uint64),
    float32/float64 (into float64) and the null type (into int64), with the C++ semantics for nulls,
    skip_nulls and min_count, wrap-around on integer overflow, and the C++ pairwise summation for
    floating point input.
  • ScalarAggregateOptions, DefaultScalarAggregateOptions, DefaultCountOptions, the Count and
    Sum convenience wrappers, and registration of both option types for deserialization.

Aggregates stay out of exprs, which rejects any function whose kind is not FuncScalar
(arrow/compute/exprs/exec.go, in the *expr.ScalarFunction case), matching C++ where
aggregations run through Acero rather than through expressions.

Not in this pull request: mean, min_max, min, max, any and all, which follow in a
second one; and product, count_distinct, first/last, decimals, the statistical aggregates
and the hash_* family, which needs a grouper.

Are these changes tested?

Yes.

  • Table-driven tests per function over every input type, with and without nulls, all null, empty,
    sliced, chunked and scalar input, under the default options, skip_nulls=false, and min_count
    above and below the number of valid values. The expected values were taken from the C++
    implementation through pyarrow for the same fixtures and options; all 91 of them agree exactly.
  • A differential fuzz run of 14,033 cases across the twelve input types, plain, sliced, chunked and
    scalar, with every option combination, compared against pyarrow: 0 mismatches, 0 error
    asymmetries, 0 leaked bytes.
  • Three tests cover the framework rather than the kernels. Every aggregate runs at chunk sizes 1 to
    65 and has to agree with the single-span result, which catches a kernel that writes a cached null
    count back into the span the executor re-slices. A kernel whose result is an array, and one whose
    result is a scalar backed by an allocated buffer, check that the value finalize returned outlives
    the cleanup of the state which produced it, and that cleanup runs exactly once on success, on
    empty input, on cancellation, and on a consume, finalize or cleanup error. And the executor is
    driven directly with a batch of scalars whose logical length is greater than one, which the
    CallFunction path cannot produce.
  • The pairwise summation is pinned to 100001 doubles from a fixed generator; the test also computes
    a naive left-to-right sum of the same fixture and asserts that it differs, so the pinned value can
    only be reached with the C++ summation order.
  • go build ./..., go vet ./arrow/compute/..., go test ./arrow/compute/... with and without
    -tags assert, go test -race, and golangci-lint at 0 issues.

Are there any user-facing changes?

Yes, all additive. arrow/compute gains the count and sum functions in the default registry,
the Count and Sum wrappers, ScalarAggregateOptions, CountOptions and their default helpers,
and ScalarAggregateFunction; arrow/compute/exec gains ScalarAggKernel, AggregateResult,
AggKernel and MergeAll. Nothing existing changes behaviour: execInternal previously fell through to its default
branch, ErrNotImplemented: direct execution of ScalarAggregate, and no function of that kind was
registered.

The compute package has no kernel type for an aggregation yet: exec only
knows the scalar and vector kernels, which produce one output value per
input value.

Add ScalarAggKernel, which follows the C++ ScalarAggregateKernel and its
init, consume, merge and finalize lifecycle, along with an optional cleanup
function and the ordered flag. Finalize returns an AggregateResult rather
than a scalar.Scalar: exec cannot import compute and so cannot name a Datum,
while an aggregate result is not always a scalar (tdigest, for instance,
finalizes to an array). The compute executor boxes the carrier into a Datum.

The ownership rule is that finalize returns a single owned result and that
cleaning up the state which produced it must not invalidate it. MergeAll
mirrors the C++ helper for a caller which aggregated partitions separately,
and cleans up every state it is given, including on an error.
Add ScalarAggregateFunction and the executor which drives its kernels, so
that a function of kind FuncScalarAgg can be registered and called through
CallFunction.

Dispatch is exact, as it is for a vector function: one kernel per input
type. The executor iterates the input in spans without promoting scalars to
arrays, as the C++ ScalarAggExecutor does, so a kernel handles a scalar
input weighted by the length of the span; it consumes into a single state
and finalizes once. An empty input still reaches finalize, which is what
lets min_count decide the result of an aggregation over no values.

The state the kernel's init function produced is cleaned up exactly once:
after finalize on the success path, and through the executor's Clear on
every other path, whether the input was empty, the context was cancelled,
or init, consume or finalize failed. The result finalize returned owns its
own value and stays valid after that cleanup.
count returns the number of values in the input, counting the non-null
values, the null values or every row depending on CountOptions.Mode, and
accepts any input type. The logical null count of a run-end encoded,
dictionary or union array is not the number of unset bits in its own
validity bitmap, so counting the valid or null values of those types
returns ErrNotImplemented rather than a wrong answer; counting every row
works for them, since it never looks at validity.

sum accepts the signed integers, which accumulate into int64, the unsigned
integers and booleans, which accumulate into uint64, the floats, which
accumulate into float64, and the null type, which sums as int64. The
semantics are the C++ ones: a null is emitted when a null was seen with
skip_nulls false, or when fewer than min_count non-null values were seen,
which is what makes an empty or all-null input null by default. An integer
sum which does not fit its accumulator wraps. Floating point input is
summed with the same pairwise summation C++ uses, so the two agree bit for
bit.

Nil options select DefaultScalarAggregateOptions, {SkipNulls: true,
MinCount: 1}, and DefaultCountOptions; an explicitly supplied zero value
keeps its own meaning. Both option types carry the C++ names and tags and
are registered for deserialization.

Aggregates stay out of the exprs package, which accepts scalar functions
only, matching C++ where aggregations run through Acero.
Table-driven tests per function over the input types, with and without
nulls, all null, empty, sliced, chunked and scalar input, under the default
options, skip_nulls false, and min_count above and below the number of
valid values. The expected values were taken from the C++ implementation
through pyarrow for the same fixtures and options, including the pairwise
summation of 100001 doubles, which a naive left-to-right sum would not
reproduce.

Three of the tests cover the framework rather than the kernels:

  - every aggregate runs at chunk sizes 1 to 65 and has to agree with the
    single-span result, which catches a kernel that wrote a cached null
    count back into the span the executor re-slices.
  - a kernel whose result is an array, and one whose result is a scalar
    backed by an allocated buffer, check that the value finalize returned
    outlives the cleanup of the state which produced it, and that the
    cleanup runs exactly once on success, on empty input, on cancellation
    and on a consume, finalize or cleanup error.
  - the executor is driven directly with a batch of scalars whose logical
    length is greater than one, which the CallFunction path cannot produce
    and which a kernel ignoring the span length for scalar input would
    otherwise pass.
singhpratech added a commit to singhpratech/ArrowMetal that referenced this pull request Sep 21, 2026
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.

1 participant