xudong963 commented on code in PR #27119:
URL: https://github.com/apache/datafusion/pull/27119#discussion_r4520235389
##########
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:
Updated the cache check to use two partial selections. I also added separate
assertions that an all-selected result shares the original schema and each
input array.
--
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]
