adriangb commented on code in PR #27119:
URL: https://github.com/apache/datafusion/pull/27119#discussion_r4514852742
##########
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:
Nit: skip-partial calls `select` for every batch, also when the filter keeps
all rows or no rows. `evaluate_compacted_args` skips these cases. A null
rejects its row, and `has_false` does not count nulls, so the all-rows check
also needs `null_count() == 0`.
```suggestion
// A null in `selection` rejects its row, like `false`.
if !selection.has_true() {
return Ok(batch.slice(0, 0));
}
if selection.null_count() == 0 && !selection.has_false() {
return Ok(batch.clone());
}
let schema = batch.schema_ref();
```
With an always-true filter on the #27114 data and a high-cardinality `GROUP
BY`, the summed `arguments_time` is 15 ms on main, 31 ms on this PR, and 19 ms
with this change. Wall time does not change.
##########
datafusion/physical-plan/src/aggregates/aggregate_argument.rs:
##########
@@ -7,10 +7,16 @@
use std::sync::Arc;
use crate::PhysicalExpr;
-use arrow::array::{ArrayRef, BooleanArray};
-use arrow::record_batch::RecordBatch;
-use datafusion_common::{Result, ScalarValue};
+use arrow::array::{ArrayRef, BooleanArray, NullArray, new_null_array};
Review Comment:
Needed for `null_count`.
```suggestion
use arrow::array::{Array, ArrayRef, BooleanArray, NullArray, new_null_array};
```
##########
datafusion/physical-plan/src/aggregates/aggregate_argument.rs:
##########
@@ -193,29 +343,106 @@ mod tests {
Ok(())
}
+ /// `evaluate_selection` matches [`PhysicalExpr::evaluate_selection`],
+ /// and does not evaluate the argument for rejected rows.
#[test]
- fn literal_selection_matches_evaluate_selection() -> Result<()> {
+ fn selection_matches_physical_expr_evaluate_selection() -> Result<()> {
+ // a = [0, 1, 2]
let batch = batch(3);
- let expr = lit(7i32);
- let mut argument = AggregateArgument::new(Arc::clone(&expr));
-
- for selection in [
- BooleanArray::from(vec![true, true, true]),
+ let schema = batch.schema();
+ let a = col("a", &schema)?;
+ // Every selection rejects the first row, where `10 / a` fails.
+ let selections = [
+ BooleanArray::from(vec![false, true, true]),
BooleanArray::from(vec![false, true, false]),
+ BooleanArray::from(vec![None, Some(true), Some(true)]),
BooleanArray::from(vec![Some(false), None, Some(true)]),
BooleanArray::from(vec![false, false, false]),
BooleanArray::from(vec![None, Some(false), None]),
+ ];
+ let all_rows = BooleanArray::from(vec![true, true, true]);
+
+ for (expr, fallible) in [
+ (lit(7i32), false),
+ (Arc::clone(&a), false),
+ (binary(lit(10i32), Operator::Divide, a, &schema)?, true),
] {
- let expected = expr
- .evaluate_selection(&batch, &selection)?
- .into_array(batch.num_rows())?;
- let actual = argument.evaluate_selection(&batch, &selection)?;
- assert_eq!(
- actual.as_ref(),
- expected.as_ref(),
- "selection {selection:?}"
- );
+ let mut argument = AggregateArgument::new(Arc::clone(&expr));
+ let mut filter =
+ AggregateFilter::new(lit(true),
std::slice::from_ref(&argument));
+ let all_rows = (!fallible).then_some(&all_rows);
+ for selection in selections.iter().chain(all_rows) {
+ let expected = expr
+ .evaluate_selection(&batch, selection)?
+ .into_array(batch.num_rows())?;
+ let selected = filter.select(&batch, selection)?;
+ let actual = argument.evaluate_selection(&batch, selection,
&selected)?;
+ assert_eq!(
+ actual.as_ref(),
+ expected.as_ref(),
+ "{expr} with selection {selection:?}"
+ );
+ }
}
Ok(())
}
+
+ #[test]
+ fn filter_selects_only_argument_columns() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, true),
+ Field::new("b", DataType::Utf8, true),
+ Field::new("c", DataType::Int32, false),
+ ]));
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from(vec![1, 2, 3, 4])),
+ Arc::new(StringArray::from(vec!["w", "x", "y", "z"])),
+ Arc::new(Int32Array::from(vec![10, 20, 30, 40])),
+ ],
+ )?;
+ let selection =
+ BooleanArray::from(vec![Some(true), Some(false), None,
Some(true)]);
+ let expected = filter_record_batch(&batch, &selection)?;
+
+ // `c + a` reads columns 0 and 2, so column 1 is not copied.
+ let expr = binary(
+ col("c", &schema)?,
+ Operator::Plus,
+ col("a", &schema)?,
+ &schema,
+ )?;
+ let mut argument = AggregateArgument::new(expr);
+ let mut filter = AggregateFilter::new(lit(true),
std::slice::from_ref(&argument));
+ let selected = filter.select(&batch, &selection)?;
+ assert_eq!(selected.num_rows(), 2);
+ assert_eq!(selected.column(0), expected.column(0));
+ assert_eq!(selected.column(2), expected.column(2));
+ assert_eq!(selected.column(1).data_type(), &DataType::Null);
+ assert_eq!(selected.column(1).len(), 2);
+ assert_eq!(selected.schema().field(1).name(), "b");
+ assert_eq!(
+ argument.evaluate(&selected)?.as_ref(),
+ argument.evaluate(&expected)?.as_ref()
+ );
+
+ // The schema of the selected rows is reused for later batches.
+ let next = filter.select(&batch, &BooleanArray::from(vec![true; 4]))?;
+ assert_eq!(next.num_rows(), 4);
Review Comment:
With the fast path, an all-true selection returns the input schema. A
partial selection keeps this check.
```suggestion
let next =
filter.select(&batch, &BooleanArray::from(vec![true, true,
false, true]))?;
assert_eq!(next.num_rows(), 3);
```
--
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]
