feat(arrow/compute): scalar aggregate framework with the count and sum kernels - #1336
Open
singhpratech wants to merge 6 commits into
Open
singhpratech wants to merge 6 commits into
singhpratech wants to merge 6 commits into
Conversation
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.
…egistration comment
singhpratech
added a commit
to singhpratech/ArrowMetal
that referenced
this pull request
Sep 21, 2026
…ramework with count and sum
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Rationale for this change
The compute package has no way to compute a summary value from an array.
FuncScalarAggexists asa function kind but nothing implements it:
exechas a kernel type for scalar and vector kernelsonly, there is no aggregate function type, and
execInternalreturnsErrNotImplementedfor thatkind. 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 astate,
Consumeto fold oneExecSpaninto it,Mergeto combine two states,Finalizetoproduce the result, plus an optional
Cleanupand anOrderedflag.AggregateResult, a small carrier which holds either an ownedscalar.Scalaror ownedarrow.ArrayData.execcannot importcomputeand so cannot name aDatum, but an aggregateresult 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, andMergeAll, for a caller which aggregatedpartitions of the input separately.
The executor, in
arrow/compute:ScalarAggregateFunction,funcImpl[exec.ScalarAggKernel]with exact dispatch, its ownSetDefaultOptions, andAddKernel/AddNewKernelwhich reject a kernel missing any of the fourlifecycle functions.
scalarAggExecutor, reached fromexecInternalforFuncScalarAgg. It iterates the input inspans without promoting scalars to arrays, as the C++
ScalarAggExecutordoes, so a kernel seesa 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_countdecide the result of anaggregation over no values.
Executereturns, whether it finalized or failedpart way, and through the executor's
Clearon the paths whereExecutenever ran, such as afailure in
Init. The result finalize returned owns its own value and is unaffected.The kernels, in
arrow/compute/internal/kernels/aggregate_basic.goandarrow/compute/scalar_aggregate.go:count, withCountOptionsand the modesCountOnlyValid,CountOnlyNullandCountAllRows,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
ArraySpandoes not carry, so those returnErrNotImplemented;CountAllRowsworks for them because it never looks at validity.sumover 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_nullsandmin_count, wrap-around on integer overflow, and the C++ pairwise summation forfloating point input.
ScalarAggregateOptions,DefaultScalarAggregateOptions,DefaultCountOptions, theCountandSumconvenience wrappers, and registration of both option types for deserialization.Aggregates stay out of
exprs, which rejects any function whose kind is notFuncScalar(
arrow/compute/exprs/exec.go, in the*expr.ScalarFunctioncase), matching C++ whereaggregations run through Acero rather than through expressions.
Not in this pull request:
mean,min_max,min,max,anyandall, which follow in asecond one; and
product,count_distinct,first/last, decimals, the statistical aggregatesand the
hash_*family, which needs a grouper.Are these changes tested?
Yes.
sliced, chunked and scalar input, under the default options,
skip_nulls=false, andmin_countabove 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.
scalar, with every option combination, compared against pyarrow: 0 mismatches, 0 error
asymmetries, 0 leaked bytes.
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
CallFunctionpath cannot produce.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/computegains thecountandsumfunctions in the default registry,the
CountandSumwrappers,ScalarAggregateOptions,CountOptionsand their default helpers,and
ScalarAggregateFunction;arrow/compute/execgainsScalarAggKernel,AggregateResult,AggKernelandMergeAll. Nothing existing changes behaviour:execInternalpreviously fell through to its defaultbranch,
ErrNotImplemented: direct execution of ScalarAggregate, and no function of that kind wasregistered.