viirya commented on code in PR #26213:
URL: https://github.com/apache/datafusion/pull/26213#discussion_r4320244434
##########
datafusion/physical-plan/src/joins/hash_join/stream.rs:
##########
@@ -930,31 +940,55 @@ impl HashJoinStream {
// Must precede the empty-build-map return below: an all-NULL-key
// build side has an empty full-key map but still needs UNKNOWN
marks.
if mode.is_correlated() && state.is_first_chunk() {
+ // `RightAnti` marks the rows of this probe batch, the others
+ // the build rows shared by all probe partitions.
+ let probe_marks = (!mode.marks_build_rows()).then(|| {
Review Comment:
Done in 785332e66: a correlated `RightAnti` now allocates the bitmap only
when the build side has a NULL-valued row or the probe batch has a NULL in a
value key. `SharedBitmapBuilder` is a plain `Mutex<BooleanBufferBuilder>` (no
`Arc`); I kept it since `UnknownMarks` shares the pairing helpers with the
build-side bitmap, which needs the mutex, and it is uncontended here.
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -11307,6 +11301,141 @@ mod tests {
Ok(())
}
+ /// Swaps a null-aware `LeftAnti` into the `RightAnti` that `JoinSelection`
+ /// plans, with its own build state so that both can be executed.
+ fn swap_for_separate_execution(
+ left_anti: &HashJoinExec,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let swapped = left_anti.swap_inputs(PartitionMode::CollectLeft)?;
+ let swapped = swapped.downcast_ref::<HashJoinExec>().unwrap();
+ Ok(Arc::new(swapped.builder().reset_state().build()?))
+ }
+
+ /// Uncorrelated two-column `(a, b) NOT IN (SELECT a, b ...)`, run both as
+ /// `LeftAnti` (built on the outer side) and as the swapped `RightAnti`
+ /// (built on the subquery side), which decides each outer row per probe
+ /// batch.
+ #[apply(hash_join_exec_configs)]
+ #[tokio::test]
+ async fn test_null_aware_multi_column_right_anti(batch_size: usize) ->
Result<()> {
Review Comment:
Done: dropped `test_null_aware_multi_column_right_anti` and added the
all-NULL-valued subquery (`mc_inner_all_null`) to
`null_aware_multi_column.slt`, which runs with join reordering on and off. The
two construction tests are now one parametrized
`test_null_aware_correlated_anti_join_accepted`.
##########
datafusion/core/tests/physical_optimizer/join_selection.rs:
##########
@@ -641,16 +642,70 @@ async fn
test_null_aware_left_anti_with_filter_does_not_swap() -> Result<()> {
let optimized_join =
JoinSelection::new().optimize(Arc::new(join), &ConfigOptions::new())?;
- let unswapped_join = optimized_join
+ let swapped_join = optimized_join
.downcast_ref::<HashJoinExec>()
- .expect("filtered null-aware anti join should not swap");
+ .expect("anti join swap should not require a projection");
- assert_eq!(*unswapped_join.join_type(), JoinType::LeftAnti);
- assert_eq!(*unswapped_join.partition_mode(), PartitionMode::CollectLeft);
- assert!(unswapped_join.null_aware);
- assert!(unswapped_join.filter().is_some());
- assert_eq!(unswapped_join.left().schema().field(0).name(), "big_col");
- assert_eq!(unswapped_join.right().schema().field(0).name(), "small_col");
+ assert_eq!(*swapped_join.join_type(), JoinType::RightAnti);
+ assert_eq!(*swapped_join.partition_mode(), PartitionMode::CollectLeft);
+ assert!(swapped_join.null_aware);
+ let filter = swapped_join.filter().expect("filter should be kept");
+ assert_eq!(
+ filter
+ .column_indices()
+ .iter()
+ .map(|c| c.side)
+ .collect::<Vec<_>>(),
+ vec![JoinSide::Right, JoinSide::Left]
+ );
+ assert_eq!(swapped_join.left().schema().field(0).name(), "small_col");
+ assert_eq!(swapped_join.right().schema().field(0).name(), "big_col");
+
+ Ok(())
+}
+
+/// A multi-key null-aware `LeftAnti` (a multi-column `NOT IN`, correlation
+/// scope keys, or both) swaps to a `RightAnti` that keeps the number of
+/// `NOT IN` value keys and the key order.
+#[rstest]
+#[case(1)]
+#[case(2)]
+#[tokio::test]
+async fn test_null_aware_multi_key_left_anti_swaps_to_right_anti(
Review Comment:
Done: folded into `test_null_aware_left_anti_swaps_to_right_anti` with
filter, scope-key and multi-column cases; it still asserts the swapped filter
sides.
##########
datafusion/core/tests/fuzz_cases/null_aware_join_fuzz.rs:
##########
@@ -68,6 +78,43 @@ async fn null_aware_join_matches_reference() {
}
}
}
+ // A fuzz run that never reaches one of the plan shapes would be vacuous.
+ assert!(plans.left_anti > 0, "{plans:?}");
+ assert!(plans.right_anti > 0, "{plans:?}");
+ assert!(plans.correlated_right_anti > 0, "{plans:?}");
+}
+
+/// Counts the null-aware anti join plans the fuzz queries ran.
+#[derive(Debug, Default)]
+struct NullAwarePlans {
+ left_anti: usize,
+ right_anti: usize,
+ /// `RightAnti` joins with more than one key or a join filter, which
+ /// decide `NOT IN` per probe row.
+ correlated_right_anti: usize,
+}
+
+impl NullAwarePlans {
+ async fn record(&mut self, ctx: &SessionContext, sql: &str) {
+ let plan = ctx
+ .sql(sql)
+ .await
+ .unwrap()
+ .create_physical_plan()
+ .await
+ .unwrap();
+ let plan = displayable(plan.as_ref()).indent(true).to_string();
+ for line in plan.lines().filter(|line| line.contains("null_aware")) {
+ if line.contains("join_type=LeftAnti") {
+ self.left_anti += 1;
+ } else if line.contains("join_type=RightAnti") {
+ self.right_anti += 1;
+ if line.contains("), (") || line.contains("filter=") {
Review Comment:
Done: `record` walks the plan for `HashJoinExec` and checks `join_type()`,
`on().len()` and `filter()`, and reuses the physical plan that `run` executes,
so each query is planned once. (Only the `NOT IN` query was recorded before,
not the reference.) The counts are unchanged.
##########
datafusion/core/tests/memory_limit/null_aware_not_in.rs:
##########
@@ -0,0 +1,103 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! A multi-column `NOT IN` over a small subquery must build its hash join on
+//! the subquery, so a large outer table runs under a memory limit that only
+//! fits the subquery side.
+//!
+//! The null-aware join of a multi-column `NOT IN` used to stay a `LeftAnti`
+//! that builds on the outer side, so it ran out of memory where the scalar
+//! `NOT IN`, swapped to a `RightAnti` that builds on the subquery, ran.
+
+use std::sync::Arc;
+
+use arrow::array::{Int64Array, RecordBatch};
+use arrow::datatypes::{DataType, Field, Schema};
+use datafusion::datasource::MemTable;
+use datafusion::prelude::*;
+use datafusion_common::{Result, assert_contains};
+use datafusion_execution::runtime_env::RuntimeEnvBuilder;
+
+const OUTER_ROWS: i64 = 200_000;
+/// Far below the outer side's build state, far above the subquery's.
+const MEMORY_LIMIT: usize = 2 * 1024 * 1024;
+
+/// A table `(a, b)` of `a = v, b = v % 1000` for `v` in `1..=rows`, followed
by
+/// `extra` rows.
+fn table(rows: i64, extra: &[(Option<i64>, Option<i64>)]) -> Arc<MemTable> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, true),
+ Field::new("b", DataType::Int64, true),
+ ]));
+ let a = (1..=rows).map(Some).chain(extra.iter().map(|row| row.0));
+ let b = (1..=rows)
+ .map(|v| Some(v % 1000))
+ .chain(extra.iter().map(|row| row.1));
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(a.collect::<Int64Array>()),
+ Arc::new(b.collect::<Int64Array>()),
+ ],
+ )
+ .unwrap();
+ // Several batches, so the outer side streams through the probe.
+ let batches = (0..batch.num_rows())
+ .step_by(8192)
+ .map(|offset| batch.slice(offset, 8192.min(batch.num_rows() - offset)))
+ .collect();
+ Arc::new(MemTable::try_new(schema, vec![batches]).unwrap())
+}
+
+async fn count_not_in(sql: &str, join_reordering: bool) -> Result<i64> {
+ let config = SessionConfig::new()
+ .with_target_partitions(1)
+ .set_bool("datafusion.optimizer.join_reordering", join_reordering);
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_limit(MEMORY_LIMIT, 1.0)
+ .build_arc()?;
+ let ctx = SessionContext::new_with_config_rt(config, runtime);
+ ctx.register_table("big", table(OUTER_ROWS, &[]))?;
+ // A NULL-valued subquery row makes `NOT IN` UNKNOWN for every outer row
+ // with `b = 7` (and no definite mismatch in `a`).
+ ctx.register_table("small", table(100, &[(None, Some(7))]))?;
+
+ let batches = ctx.sql(sql).await?.collect().await?;
+ Ok(batches[0]
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap()
+ .value(0))
+}
+
+#[tokio::test]
+async fn multi_column_not_in_builds_on_small_subquery() -> Result<()> {
Review Comment:
Done: moved to `null_aware_memory_limit.slt`; 200k rows run fine there. It
needs a `RESET datafusion.catalog.create_default_catalog_and_schema` at the
end, since setting a runtime option rebuilds the session state. It still fails
without the `JoinSelection` change.
--
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]
