kumarUjjawal commented on code in PR #24457: URL: https://github.com/apache/datafusion/pull/24457#discussion_r3942949526
########## datafusion/physical-plan/src/joins/piecewise_merge_join/right_existence_join.rs: ########## @@ -0,0 +1,866 @@ +// 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. + +//! PiecewiseMergeJoin stream specialized for right existence joins. +//! +//! Instantiated by [`PiecewiseMergeJoinExec`] when the join type is `RightSemi` or +//! `RightAnti`. `LeftSemi`/`LeftAnti` are served by `ExistencePWMJStream` (see +//! `existence_join.rs`); the Mark joins are still rejected in +//! `PiecewiseMergeJoinExec::try_new`. +//! +//! # Algorithm +//! +//! Left and right existence joins mark opposite sides, and for a single range predicate +//! that difference is not symmetric — it collapses the work. +//! +//! `LeftSemi`/`LeftAnti` ask, for each *buffered* row, whether any streamed row matches, so +//! the answer depends on the whole streamed side and can only be emitted once it has all +//! been read. `RightSemi`/`RightAnti` ask the mirror question — for each *streamed* row, +//! does any buffered row match? — and with only `buffered_key OP streamed_key` to satisfy, +//! that is decided by a single buffered key: +//! +//! ```text +//! ∃b. b < s ⟺ min(b) < s ∃b. b > s ⟺ max(b) > s +//! ∃b. b <= s ⟺ min(b) <= s ∃b. b >= s ⟺ max(b) >= s +//! ``` +//! +//! The buffered side is therefore reduced to that one key -- `min_batch`/`max_batch`, folded +//! batch by batch as it streams in -- and each batch is dropped once folded. It is never +//! concatenated or retained, so the state this join holds is a single row however large that +//! side is. The reduction is shared, so every streamed partition reads the same key rather +//! than repeating the pass: +//! +//! ```text +//! operator `<`, buffered keys [NULL, 5, 9, 7] -> min = 5 +//! `b < s` holds for some b iff `5 < s` +//! ``` +//! +//! A min/max is `O(B)` from any order, so `required_input_ordering` returns nothing for either +//! child and no `SortExec` is planned. +//! +//! Every streamed row is then decided by comparing it against that one key, which is a +//! vectorized `cmp` kernel per batch and a filter. Nothing about a streamed row depends on +//! any other, so: +//! +//! * output is produced per batch as it arrives — no watermark, no final pass, and no +//! election among the streamed partitions, +//! * all N streamed partitions produce output, rather than one non-empty partition, +//! * the streamed side is never sorted, not even per batch. +//! +//! Rows whose join key is NULL never satisfy a comparison predicate. `min_batch`/`max_batch` +//! ignore NULLs, so the reduced key is null only when *every* buffered key is (or the buffered +//! side is empty) — and then no streamed row can match at all, which makes `RightSemi` empty +//! without reading the streamed side and `RightAnti` a passthrough of it. A NULL streamed key +//! makes its comparison NULL rather than false, which is "no match" — dropped by `RightSemi`, +//! kept by `RightAnti`. +//! +//! Picking the extreme and comparing against it must agree on ordering. `min_batch`/`max_batch` +//! and arrow's `lt`/`lt_eq`/`gt`/`gt_eq` kernels both order floats by `total_cmp` +//! (`arrow-arith`'s `MinAccumulator` compares with `ArrowNativeTypeOp::is_lt`, seeded from +//! `MAX_TOTAL_ORDER`), which puts `-0.0` strictly below `+0.0` -- but SQL comparisons treat +//! them as equal, and that is what a real `k < r` predicate actually evaluates to: any +//! `BinaryExpr` comparison, including the one the `NestedLoopJoinExec` oracle in the +//! differential fuzz test builds its filter from, normalizes `-0.0` to `+0.0` first (see +//! `apply_cmp` in `datafusion-physical-expr-common`). Both the extreme and the streamed key +//! array are normalized with [`normalize_float_zero`] before comparing, +//! after the reduction: normalizing first would not change which value the reduction picks +//! (`-0.0` and `+0.0` are numerically equal either way), and normalizing only where the +//! comparison happens keeps the reduction itself agreeing with the unnormalized `min`/`max` +//! semantics its docs above describe. `NaN` needs no such fix-up: every kernel involved orders +//! it as the maximum, so it is treated identically on both sides already. +//! +//! # Cost +//! +//! Let `B` be the buffered rows and `S` the streamed rows: `O(B)` to reduce the buffered side +//! plus `O(S)` to filter, and no sort on either side. The state retained across batches is one +//! row; while folding, each buffered partition also holds the key array of the batch it is +//! reducing, which it accounts against the memory pool for that long. +//! +//! This is why the shared state is [`BufferedExtreme`] and not `BufferedSideData`: every field +//! of the latter -- the concatenated batch, the key array, the visited-indices bitmap, the +//! final-pass counter -- would be dead here. +//! +//! [`PiecewiseMergeJoinExec`]: super::PiecewiseMergeJoinExec + +use std::sync::Arc; +use std::task::{Poll, ready}; + +use arrow::array::{Array, ArrayRef, RecordBatch, Scalar}; +use arrow::compute::filter_record_batch; +use arrow::compute::kernels::boolean::not; +use arrow::compute::kernels::cmp::{gt, gt_eq, lt, lt_eq}; +use arrow_schema::SchemaRef; +use datafusion_common::utils::normalize_float_zero; +use datafusion_common::{Result, internal_err}; +use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; +use datafusion_expr::{JoinType, Operator}; +use datafusion_physical_expr::PhysicalExprRef; +use futures::{Stream, StreamExt}; + +use crate::handle_state; +use crate::joins::piecewise_merge_join::exec::BufferedExtreme; +use crate::joins::utils::{ + BuildProbeJoinMetrics, OnceFut, StatefulStreamResult, boolean_mask_from_filter, +}; +use crate::stream::EmptyRecordBatchStream; + +pub(super) enum RightExistencePWMJStreamState { + /// Await the buffered side's reduction to a single key. + WaitBufferedExtreme, + /// Fetch streamed batches and emit the rows that do (`RightSemi`) or do not + /// (`RightAnti`) have a buffered match. + ScanStreamBatches, + Completed, +} + +pub(super) struct RightExistencePWMJStream { + /// Output schema, which for `RightSemi`/`RightAnti` is the streamed side's schema + schema: SchemaRef, + /// Physical expression evaluated on the streamed side. The buffered side's + /// equivalent is already evaluated when the buffered side is collected. + on_streamed: PhysicalExprRef, + /// `RightSemi` or `RightAnti` + join_type: JoinType, + /// Comparison operator + operator: Operator, + streamed: SendableRecordBatchStream, + /// Resolves to the whole buffered side reduced to one key. Shared with the other streamed + /// partitions, so that reduction happens exactly once. + buffered_extreme_fut: OnceFut<BufferedExtreme>, + state: RightExistencePWMJStreamState, + /// That key, held as a one-element [`Scalar`] so the comparison against a streamed batch is + /// one kernel call. `None` when the buffered side has no non-null key, i.e. nothing can ever + /// match. Only populated once `buffered_extreme_fut` has resolved. + buffered_extreme: Option<Scalar<ArrayRef>>, + join_metrics: BuildProbeJoinMetrics, +} + +impl RightExistencePWMJStream { + pub(super) fn try_new( + schema: SchemaRef, + on_streamed: PhysicalExprRef, + join_type: JoinType, + operator: Operator, + streamed: SendableRecordBatchStream, + buffered_extreme_fut: OnceFut<BufferedExtreme>, + join_metrics: BuildProbeJoinMetrics, + ) -> Self { + Self { + schema, + on_streamed, + join_type, + operator, + streamed, + buffered_extreme_fut, + state: RightExistencePWMJStreamState::WaitBufferedExtreme, + buffered_extreme: None, + join_metrics, + } + } + + fn poll_next_impl( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + loop { + return match self.state { + RightExistencePWMJStreamState::WaitBufferedExtreme => { + handle_state!(ready!(self.collect_buffered_extreme(cx))) + } + RightExistencePWMJStreamState::ScanStreamBatches => { + handle_state!(ready!(self.scan_stream_batch(cx))) + } + RightExistencePWMJStreamState::Completed => Poll::Ready(None), + }; + } + } + + /// Picks up the buffered side's reduced key, which is all this join type needs from it. + fn collect_buffered_extreme( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> { + let build_timer = self.join_metrics.build_time.timer(); + let buffered_extreme = ready!(self.buffered_extreme_fut.get_shared(cx))?; + build_timer.done(); + + // Null exactly when no buffered key is non-null, and NULLs match nothing. Cloned out by + // value -- it is one row -- so the shared state is not kept alive by this stream. + // + // Normalized because the `lt`/`lt_eq`/`gt`/`gt_eq` kernels called in + // `filter_streamed_batch` order `-0.0` strictly below `+0.0`, but SQL comparisons -- + // including `apply_cmp`, which every other `k < r` predicate in the plan goes through + // -- treat them as equal. The streamed key array is normalized the same way, right + // before that comparison. `min`/`max_batch`, which reduced this extreme, order `-0.0` + // below `+0.0` too, so normalizing after the reduction rather than before leaves the + // reduction itself agreeing with the unnormalized ordering its own docs describe. + let extreme = normalize_float_zero(buffered_extreme.extreme()); + self.buffered_extreme = (extreme.null_count() == 0).then(|| Scalar::new(extreme)); + + // With no non-null buffered key nothing matches, so `RightSemi` outputs nothing + // and does not need to read a single streamed batch. `RightAnti` still has to, + // since it outputs all of them. + self.state = match (&self.buffered_extreme, self.join_type) { + (None, JoinType::RightSemi) => { + let streamed_schema = self.streamed.schema(); + self.streamed = Box::pin(EmptyRecordBatchStream::new(streamed_schema)); + RightExistencePWMJStreamState::Completed + } + _ => RightExistencePWMJStreamState::ScanStreamBatches, + }; + + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + + /// Fetches one streamed batch and emits the rows it contributes, if any. + fn scan_stream_batch( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> { + match ready!(self.streamed.poll_next_unpin(cx)) { + None => self.state = RightExistencePWMJStreamState::Completed, + Some(Ok(batch)) => { + self.join_metrics.input_batches.add(1); + self.join_metrics.input_rows.add(batch.num_rows()); + + let output = self.filter_streamed_batch(&batch)?; + if output.num_rows() > 0 { + return Poll::Ready(Ok(StatefulStreamResult::Ready(Some(output)))); + } + // Nothing survived; take the next batch rather than yielding an empty one. + } + Some(Err(err)) => return Poll::Ready(Err(err)), + } + + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + + /// Keeps the streamed rows that have a buffered match (`RightSemi`) or that have none + /// (`RightAnti`), by comparing each against the single buffered extreme. + fn filter_streamed_batch(&self, batch: &RecordBatch) -> Result<RecordBatch> { + let columns = match &self.buffered_extreme { + // No non-null buffered key, so no streamed row matches and `RightAnti` keeps + // the batch whole. `RightSemi` never gets here: it completed without reading + // the streamed side. + None => batch.columns().to_vec(), + Some(extreme) => { + let stream_values = normalize_float_zero( + &self + .on_streamed + .evaluate(batch)? + .into_array(batch.num_rows())?, + ); + + // `extreme` is the buffered key, so it goes on the left of the operator, + // matching the `buffered OP streamed` orientation of the predicate. + let matched = match self.operator { Review Comment: These Arrow comparison kernels reject nested types such as List and Struct. The planner now routes their RightSemi/RightAnti range joins here, whereas the previous NestedLoopJoin path uses `apply_cmp`, which supports nested comparisons. For example, list keys `[1] < [2]` should produce a matching right row, but this path returns a “Nested comparison” error when PWMJ is enabled. Could we preserve the nested comparison path or retain the nested-loop fallback for unsupported types? Please add RightSemi/RightAnti regressions comparing results with PWMJ enabled and disabled. -- 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]
