parthchandra commented on code in PR #7280:
URL: https://github.com/apache/datafusion-comet/pull/7280#discussion_r4429410468
##########
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:
Done. The rebuild check now lives in
runtime_filter_permission_survives_rebuilding, and I removed the accessor and
the join test. Thanks
##########
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:
Done. Back to one int column with the exact [2, 2, 1] check.
--
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]
