xudong963 commented on code in PR #27119:
URL: https://github.com/apache/datafusion/pull/27119#discussion_r4520230240
##########
datafusion/physical-plan/src/aggregates/aggregate_argument.rs:
##########
@@ -77,16 +87,154 @@ impl AggregateArgument {
&mut self,
batch: &RecordBatch,
selection: &BooleanArray,
+ selected: &RecordBatch,
) -> Result<ArrayRef> {
+ let num_rows = batch.num_rows();
+ if selected.num_rows() == num_rows {
+ return self.evaluate(batch);
+ }
+ if selected.num_rows() == 0 {
+ // A fallible expression may fail even for an empty batch, so it
+ // is not evaluated.
+ let data_type = self.expr.data_type(batch.schema_ref())?;
+ return Ok(new_null_array(&data_type, num_rows));
+ }
match &mut self.literal {
- Some(literal) if selection.has_true() => {
- literal.array_of_size(batch.num_rows())
- }
- _ => self
- .expr
- .evaluate_selection(batch, selection)?
- .into_array_of_size(batch.num_rows()),
+ Some(literal) => literal.array_of_size(num_rows),
+ None => scatter(selection, self.evaluate(selected)?.as_ref()),
+ }
+ }
+}
+
+/// The `FILTER` clause of one aggregate function, such as `x > 0` in
+/// `SUM(y) FILTER (WHERE x > 0)`.
+///
+/// The arguments of the aggregate are evaluated only for the rows the filter
+/// selects, so those rows are copied out of each input batch first: evaluating
+/// the arguments for every row could raise errors for rejected rows, such as
+/// division by zero in `SUM(1 / x) FILTER (WHERE x <> 0)`.
+///
+/// Only the input columns that the arguments read are copied. A query may have
+/// many filtered aggregates over a wide input, and copying every column for
+/// every aggregate would dominate the cost of the aggregation.
+pub(in crate::aggregates) struct AggregateFilter {
+ predicate: Arc<dyn PhysicalExpr>,
+ /// Sorted indices of the input columns that the arguments read.
+ argument_columns: Vec<usize>,
+ /// The input schema last passed to [`Self::select`], and the schema of the
+ /// batches it returns for that input.
+ selected_schema: Option<(SchemaRef, SchemaRef)>,
+}
+
+impl AggregateFilter {
+ pub(in crate::aggregates) fn new(
+ predicate: Arc<dyn PhysicalExpr>,
+ arguments: &[AggregateArgument],
+ ) -> Self {
+ let argument_columns = arguments
+ .iter()
+ .flat_map(|argument| collect_columns(argument.expr()))
+ .map(|column| column.index())
+ .sorted_unstable()
+ .dedup()
+ .collect();
+ Self {
+ predicate,
+ argument_columns,
+ selected_schema: None,
+ }
+ }
+
+ pub(in crate::aggregates) fn predicate(&self) -> &Arc<dyn PhysicalExpr> {
+ &self.predicate
+ }
+
+ /// Evaluates the filter predicate for every row of `batch`.
+ pub(in crate::aggregates) fn evaluate(
+ &self,
+ batch: &RecordBatch,
+ ) -> Result<BooleanArray> {
+ let selection = self
+ .predicate
+ .evaluate(batch)?
+ .into_array(batch.num_rows())?;
+ Ok(as_boolean_array(&selection)?.clone())
+ }
+
+ /// Returns the rows of `batch` where `selection` is true, for evaluating
+ /// the arguments of the aggregate.
+ ///
+ /// Only the columns that the arguments read are copied. The other columns
+ /// are replaced by [`NullArray`]s, which hold no data, so the column
+ /// indices in the arguments remain valid.
+ pub(in crate::aggregates) fn select(
+ &mut self,
+ batch: &RecordBatch,
+ selection: &BooleanArray,
+ ) -> Result<RecordBatch> {
+ let schema = batch.schema_ref();
Review Comment:
Added the all-selected fast path with the `null_count() == 0 `guard. For
empty selections, the skip-partial caller now directly creates typed null
arguments and retains the original filter, avoiding selected-batch construction
entirely. I avoided `batch.slice(0, 0) `because it still slices every input
column, which was slower for wide batches in local testing. In a local
all-selected/high-cardinality run, summed partial `arguments_time` dropped from
about 12.6 ms to 6.4 ms, with wall time essentially unchanged.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
