xudong963 commented on code in PR #27119:
URL: https://github.com/apache/datafusion/pull/27119#discussion_r4520242815
##########
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())?;
Review Comment:
Agreed, a scalar predicate could avoid materializing a BooleanArray. I’ve
left that for a follow-up, as suggested; it would need to handle scalar
true/false/NULL consistently across the compacted and skip-partial paths.
--
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]
