xudong963 commented on code in PR #27119:
URL: https://github.com/apache/datafusion/pull/27119#discussion_r4520251581
##########
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();
+ // As in `filter_record_batch`, optimizing the filter pays off only
+ // when it is applied to several arrays.
+ let optimize = match self.argument_columns.as_slice() {
Review Comment:
The current implementation already uses` Arrow’s FilterBuilder `and
`FilterPredicate::filter`. The complication with `RecordBatch::project` is that
it changes column positions, while the argument expressions still reference the
original indices.
I tried projecting, filtering, and restoring the original positions locally,
but that was slightly slower in the cases tested. I’ve kept the current
approach here; remapping the expressions once to use a compact projected batch
could be worth exploring separately.
--
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]
