-
Notifications
You must be signed in to change notification settings - Fork 340
feat: native collect_list / array_agg aggregate #4720
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
c9c19e3
6d2aa2b
7496e69
9e680c4
e63a50b
589ebdd
a98c3fc
5ed285f
8d81702
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -393,6 +393,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { | |
| classOf[BitOrAgg] -> CometBitOrAgg, | ||
| classOf[BitXorAgg] -> CometBitXOrAgg, | ||
| classOf[BloomFilterAggregate] -> CometBloomFilterAggregate, | ||
| classOf[CollectList] -> CometCollectList, | ||
| classOf[CollectSet] -> CometCollectSet, | ||
| classOf[Corr] -> CometCorr, | ||
| classOf[Count] -> CometCount, | ||
|
|
@@ -439,6 +440,22 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { | |
| } | ||
| } | ||
|
|
||
| /** | ||
| * Returns true if any aggregate function produces a native intermediate buffer whose Arrow type | ||
| * (e.g. ArrayType for CollectList/CollectSet) differs from the BinaryType that Spark declares | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The comment describes this as matching any agg with a native |
||
| * for its serialized TypedImperativeAggregate buffer. Comet cannot interpret Spark's Binary | ||
| * buffer for these functions, and cannot yet represent the buffer consistently across the | ||
| * intermediate PartialMerge stages of a multi-stage aggregate (issue #4724). Such aggregates | ||
| * are therefore only safe to run natively when every stage runs in Comet and there are at most | ||
| * two stages (Partial + Final). | ||
| */ | ||
| def hasNativeArrayBufferAgg(aggExprs: Seq[AggregateExpression]): Boolean = { | ||
| aggExprs.exists(_.aggregateFunction match { | ||
| case _: CollectList | _: CollectSet => true | ||
| case _ => false | ||
| }) | ||
| } | ||
|
|
||
| // A unique id for each expression. ~used to look up QueryContext during error creation. | ||
| private val exprIdCounter = new AtomicLong(0) | ||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Could you add a comment on why this is separate from the tagging block just above? They look like duplicates. The difference is that the block above only tags when the Final can't convert, and
canAggregateBeConvertedskips the child-native check, so an all-nativecollect_listFinal converts and slips past it. This block is what catches that all-native distinct case, which is the Parquet half of the new suite test.