comphead commented on code in PR #7280:
URL: https://github.com/apache/datafusion-comet/pull/7280#discussion_r4429104967
##########
native/core/src/execution/planner/join.rs:
##########
@@ -946,6 +947,28 @@ mod tests {
.allows_runtime_filter_pushdown()
}
+ /// Wrapping a probe filter, and rebuilding it over a wrapped child, must
keep the session
+ /// batch size the planner gave it.
+ #[test]
+ fn probe_filters_keep_session_batch_size() {
Review Comment:
This test checks that `filter.clone()` and `replace_children` keep
`batch_size`, and the new public `CometFilterExec::batch_size()` exists only
for it. `runtime_filter_permission_survives_rebuilding` in
`native/operators/src/filter.rs` already walks `with_execution_input`,
`replace_children` and `with_fetch` over a wrapped filter. Building its
`filter()` fixture with `FilterExecBuilder::new(predicate,
input).with_batch_size(5).build()` and adding
`assert_eq!(rebuilt.inner().batch_size(), 5)` to that loop covers the same
rebuilds. `inner()` is visible in that crate, so the accessor and this test
could both go. I tried that change locally and it passes.
##########
native/core/src/execution/planner.rs:
##########
@@ -3812,6 +3815,73 @@ mod tests {
assert_eq!(0, filter_exec.additional_native_plans.len());
}
+ fn batch_size_planner(batch_size: usize) -> PhysicalPlanner {
+ let config = SessionConfig::new().with_batch_size(batch_size);
+
PhysicalPlanner::new(Arc::new(SessionContext::new_with_config(config)), 0)
+ }
+
+ /// The filter merges its output up to the session batch size and no
further. Without it, the
+ /// filter merges up to 8192 rows, which can exceed Arrow's 2 GiB limit
for large string values.
+ /// The NULL and multi-byte strings check that merging keeps every value.
+ #[tokio::test]
+ async fn filter_output_respects_session_batch_size() {
+ use arrow::array::AsArray;
+
+ let planner = batch_size_planner(2);
+ let scan = Operator {
+ op_struct: Some(OpStruct::Scan(spark_operator::Scan {
+ fields: vec![
+ create_proto_datatype(),
+ spark_expression::DataType {
+ type_id: 7, // STRING
+ type_info: None,
+ },
+ ],
+ source: String::new(),
+ })),
+ ..create_scan()
+ };
+ let op = create_filter(scan, 7);
+ let (mut scans, _, planned) = planner.create_plan(&op, &mut vec![],
1).unwrap();
+
+ // Five single-row batches that all pass the filter, with NULL and
multi-byte strings.
+ let strings = [Some("数据"), None, Some("é"), Some(""), Some("データ")];
+ let mut inputs = [strings
+ .iter()
+ .map(|&value| {
+ let columns = vec![
+ Arc::new(Int32Array::from(vec![7])) as ArrayRef,
+ Arc::new(StringArray::from(vec![value])) as ArrayRef,
+ ];
+ InputBatch::Batch(columns, 1)
+ })
+ .chain(std::iter::once(InputBatch::EOF))
+ .collect::<Vec<_>>()
+ .into_iter()];
+ let stream = planned
+ .native_plan
+ .execute(0, SessionContext::new().task_ctx())
+ .unwrap();
+ let batches = collect_with_inputs(stream, &mut scans, &mut
inputs).await;
+ let sizes: Vec<_> = batches.iter().map(|batch|
batch.num_rows()).collect();
+ let output: Vec<_> = batches
+ .iter()
+ .flat_map(|batch| {
+ let column = batch.column(1).as_string::<i32>();
+ column
+ .iter()
+ .map(|value| value.map(str::to_string))
+ .collect::<Vec<_>>()
+ })
+ .collect();
+ assert_eq!(sizes, vec![2, 2, 1]);
+ let expected: Vec<_> = strings
+ .iter()
+ .map(|value| value.map(str::to_string))
+ .collect();
+ assert_eq!(output, expected);
Review Comment:
`filter_varlen_values.sql` already checks NULL, empty and multi-byte strings
through the merged batches against Spark. The batch sizes are the part only
this test can see. Could it keep the single int column from `create_scan()` and
the `[2, 2, 1]` assertion? That drops the string column, the second pass over
`batches` and the `Vec` built per batch in the `flat_map`.
--
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]
