alamb commented on code in PR #27119:
URL: https://github.com/apache/datafusion/pull/27119#discussion_r4517730897
##########
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())?;
Review Comment:
The group by should already know its output -- calling `data_type` once per
batch is non trivial in some cases -- I woudl recommend computing this once and
then just copying the type per batch
##########
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:
why turn it into an array? If it was a ColumnarValue::Scalar we could
probably have a special path (can do as a follow on)
##########
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]
Review Comment:
can we test this with SLT rather than with a new unit test?
##########
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:
how about just using
https://docs.rs/arrow/latest/arrow/array/struct.RecordBatch.html#method.project
and then calling the actual filter kernel?
I think that would be a lot faster and you can take advantage of all the
crazy filter ing tricks in arrow
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -7220,6 +7220,247 @@ mod tests {
Ok(())
}
+ /// Aggregate argument that reads `column` and records the data types of
+ /// the columns of every batch it is evaluated against.
+ #[derive(Debug)]
Review Comment:
this seems like a lot of test code, but it does test a reasonable property
so it is probably ok
##########
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() {
+ [] => false,
+ [index] => schema.fields().get(*index).is_some_and(|field| {
+ FilterBuilder::is_optimize_beneficial(field.data_type())
+ }),
+ _ => true,
+ };
+ let mut builder = FilterBuilder::new(selection);
+ if optimize {
+ builder = builder.optimize();
+ }
+ let predicate = builder.build();
+
+ let num_rows = predicate.count();
+ // Every column the arguments do not read shares one placeholder.
+ let placeholder: ArrayRef = Arc::new(NullArray::new(num_rows));
+ let mut columns = vec![placeholder; batch.num_columns()];
+ for &index in &self.argument_columns {
+ let Some(column) = batch.columns().get(index) else {
+ return internal_err!(
+ "Aggregate argument references column {index}, but the
input has {} columns",
+ batch.num_columns()
+ );
+ };
+ columns[index] = predicate.filter(column.as_ref())?;
+ }
+
+ let options = RecordBatchOptions::new().with_row_count(Some(num_rows));
+ Ok(RecordBatch::try_new_with_options(
+ self.selected_schema(schema),
+ columns,
+ &options,
+ )?)
+ }
+
+ /// Returns `schema` with the fields of columns that the arguments do not
+ /// read changed to [`DataType::Null`].
+ fn selected_schema(&mut self, schema: &SchemaRef) -> SchemaRef {
+ if let Some((input, selected)) = &self.selected_schema
+ && input == schema
+ {
+ return Arc::clone(selected);
}
+
+ let fields = schema
Review Comment:
This looks like maybe it is a manual reimplementation of
https://docs.rs/arrow/latest/arrow/datatypes/struct.Schema.html#method.project ?
--
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]
