From 19a0fcaa276c86beda544c6e01c75f6e0639767e Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Wed, 4 Mar 2026 13:58:50 -0500 Subject: [PATCH 1/8] [branch-52] SortMergeJoin don't wait for all input before emitting (#20699) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? Backport of #20482 to branch-52. ## Rationale for this change Cherry-pick fix and prerequisites so that SortMergeJoin emits output incrementally instead of waiting for all input to complete. This resolves OOM issues Comet is seeing with DataFusion 52. ## What changes are included in this PR? Cherry-picks of the following commits from `main`: 1. #19614 — Extract sort-merge join filter logic into separate module 2. #20463 — Use zero-copy slice instead of take kernel in sort merge join 3. #20482 — Fix SortMergeJoin to not wait for all input before emitting ## Are these changes tested? Yes, covered by existing and new tests included in #20482. ## Are there any user-facing changes? No. --------- Co-authored-by: Liang-Chi Hsieh Co-authored-by: Claude Sonnet 4.5 Co-authored-by: Andy Grove Co-authored-by: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> --- .../src/joins/sort_merge_join/filter.rs | 595 +++++++++++++++++ .../src/joins/sort_merge_join/mod.rs | 1 + .../src/joins/sort_merge_join/stream.rs | 607 ++++-------------- .../src/joins/sort_merge_join/tests.rs | 514 +++++++++++++-- datafusion/physical-plan/src/test/exec.rs | 111 +++- 5 files changed, 1307 insertions(+), 521 deletions(-) create mode 100644 datafusion/physical-plan/src/joins/sort_merge_join/filter.rs diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/filter.rs b/datafusion/physical-plan/src/joins/sort_merge_join/filter.rs new file mode 100644 index 0000000000000..d598442b653eb --- /dev/null +++ b/datafusion/physical-plan/src/joins/sort_merge_join/filter.rs @@ -0,0 +1,595 @@ +// 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. + +//! Filter handling for Sort-Merge Join +//! +//! This module encapsulates the complexity of join filter evaluation, including: +//! - Immediate filtering for INNER joins +//! - Deferred filtering for outer/semi/anti/mark joins +//! - Metadata tracking for grouping output rows by input row +//! - Correcting filter masks to handle multiple matches per input row + +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayBuilder, ArrayRef, BooleanArray, BooleanBuilder, RecordBatch, + UInt64Array, UInt64Builder, +}; +use arrow::compute::{self, concat_batches, filter_record_batch}; +use arrow::datatypes::SchemaRef; +use datafusion_common::{JoinSide, JoinType, Result}; + +use crate::joins::utils::JoinFilter; + +/// Metadata for tracking filter results during deferred filtering +/// +/// When a join filter is present and we need to ensure each input row produces +/// at least one output (outer joins) or exactly one output (semi joins), we can't +/// filter immediately. Instead, we accumulate all joined rows with metadata, +/// then post-process to determine which rows to output. +#[derive(Debug)] +pub struct FilterMetadata { + /// Did each output row pass the join filter? + /// Used to detect if an input row found ANY match + pub filter_mask: BooleanBuilder, + + /// Which input row (within batch) produced each output row? + /// Used for grouping output rows by input row + pub row_indices: UInt64Builder, + + /// Which input batch did each output row come from? + /// Used to disambiguate row_indices across multiple batches + pub batch_ids: Vec, +} + +impl FilterMetadata { + /// Create new empty filter metadata + pub fn new() -> Self { + Self { + filter_mask: BooleanBuilder::new(), + row_indices: UInt64Builder::new(), + batch_ids: vec![], + } + } + + /// Returns (row_indices, filter_mask, batch_ids_ref) and clears builders + pub fn finish_metadata(&mut self) -> (UInt64Array, BooleanArray, &[usize]) { + let row_indices = self.row_indices.finish(); + let filter_mask = self.filter_mask.finish(); + (row_indices, filter_mask, &self.batch_ids) + } + + /// Add metadata for null-joined rows (no filter applied) + pub fn append_nulls(&mut self, num_rows: usize) { + self.filter_mask.append_nulls(num_rows); + self.row_indices.append_nulls(num_rows); + self.batch_ids.resize( + self.batch_ids.len() + num_rows, + 0, // batch_id = 0 for null-joined rows + ); + } + + /// Add metadata for filtered rows + pub fn append_filter_metadata( + &mut self, + row_indices: &UInt64Array, + filter_mask: &BooleanArray, + batch_id: usize, + ) { + debug_assert_eq!( + row_indices.len(), + filter_mask.len(), + "row_indices and filter_mask must have same length" + ); + + for i in 0..row_indices.len() { + if filter_mask.is_null(i) { + self.filter_mask.append_null(); + } else if filter_mask.value(i) { + self.filter_mask.append_value(true); + } else { + self.filter_mask.append_value(false); + } + + if row_indices.is_null(i) { + self.row_indices.append_null(); + } else { + self.row_indices.append_value(row_indices.value(i)); + } + + self.batch_ids.push(batch_id); + } + } + + /// Verify that metadata arrays are aligned (same length) + pub fn debug_assert_metadata_aligned(&self) { + if self.filter_mask.len() > 0 { + debug_assert_eq!( + self.filter_mask.len(), + self.row_indices.len(), + "filter_mask and row_indices must have same length when metadata is used" + ); + debug_assert_eq!( + self.filter_mask.len(), + self.batch_ids.len(), + "filter_mask and batch_ids must have same length when metadata is used" + ); + } else { + debug_assert_eq!( + self.filter_mask.len(), + 0, + "filter_mask should be empty when batches is empty" + ); + } + } +} + +impl Default for FilterMetadata { + fn default() -> Self { + Self::new() + } +} + +/// Determines if a join type needs deferred filtering +/// +/// Deferred filtering is required when: +/// - A filter exists AND +/// - The join type requires ensuring each input row produces at least one output +/// (or exactly one for semi joins) +pub fn needs_deferred_filtering( + filter: &Option, + join_type: JoinType, +) -> bool { + filter.is_some() + && matches!( + join_type, + JoinType::Left + | JoinType::LeftSemi + | JoinType::LeftMark + | JoinType::Right + | JoinType::RightSemi + | JoinType::RightMark + | JoinType::LeftAnti + | JoinType::RightAnti + | JoinType::Full + ) +} + +/// Gets the arrays which join filters are applied on +/// +/// Extracts the columns needed for filter evaluation from left and right batch columns +pub fn get_filter_columns( + join_filter: &Option, + left_columns: &[ArrayRef], + right_columns: &[ArrayRef], +) -> Vec { + let mut filter_columns = vec![]; + + if let Some(f) = join_filter { + let left_columns: Vec = f + .column_indices() + .iter() + .filter(|col_index| col_index.side == JoinSide::Left) + .map(|i| Arc::clone(&left_columns[i.index])) + .collect(); + let right_columns: Vec = f + .column_indices() + .iter() + .filter(|col_index| col_index.side == JoinSide::Right) + .map(|i| Arc::clone(&right_columns[i.index])) + .collect(); + + filter_columns.extend(left_columns); + filter_columns.extend(right_columns); + } + + filter_columns +} + +/// Determines if current index is the last occurrence of a row +/// +/// Used during filter mask correction to detect row boundaries when grouping +/// output rows by input row. +fn last_index_for_row( + row_index: usize, + indices: &UInt64Array, + batch_ids: &[usize], + indices_len: usize, +) -> bool { + debug_assert_eq!( + indices.len(), + indices_len, + "indices.len() should match indices_len parameter" + ); + debug_assert_eq!( + batch_ids.len(), + indices_len, + "batch_ids.len() should match indices_len" + ); + debug_assert!( + row_index < indices_len, + "row_index {row_index} should be < indices_len {indices_len}", + ); + + // If this is the last index overall, it's definitely the last for this row + if row_index == indices_len - 1 { + return true; + } + + // Check if next row has different (batch_id, index) pair + let current_batch_id = batch_ids[row_index]; + let next_batch_id = batch_ids[row_index + 1]; + + if current_batch_id != next_batch_id { + return true; + } + + // Same batch_id, check if row index is different + // Both current and next should be non-null (already joined rows) + if indices.is_null(row_index) || indices.is_null(row_index + 1) { + return true; + } + + indices.value(row_index) != indices.value(row_index + 1) +} + +/// Corrects the filter mask for joins with deferred filtering +/// +/// When an input row joins with multiple buffered rows, we get multiple output rows. +/// This function groups them by input row and applies join-type-specific logic: +/// +/// - **Outer joins**: Keep first matching row, convert rest to nulls, add null-joined for unmatched +/// - **Semi joins**: Keep first matching row, discard rest +/// - **Anti joins**: Keep row only if NO matches passed filter +/// - **Mark joins**: Like semi but first match only +/// +/// # Arguments +/// * `join_type` - The type of join being performed +/// * `row_indices` - Which input row produced each output row +/// * `batch_ids` - Which batch each output row came from +/// * `filter_mask` - Whether each output row passed the filter +/// * `expected_size` - Total number of input rows (for adding unmatched) +/// +/// # Returns +/// Corrected mask indicating which rows to include in final output: +/// - `true`: Include this row +/// - `false`: Convert to null-joined row (outer joins) or include as unmatched (anti joins) +/// - `null`: Discard this row +pub fn get_corrected_filter_mask( + join_type: JoinType, + row_indices: &UInt64Array, + batch_ids: &[usize], + filter_mask: &BooleanArray, + expected_size: usize, +) -> Option { + let row_indices_length = row_indices.len(); + let mut corrected_mask: BooleanBuilder = + BooleanBuilder::with_capacity(row_indices_length); + let mut seen_true = false; + + match join_type { + JoinType::Left | JoinType::Right => { + // For outer joins: Keep first matching row per input row, + // convert rest to nulls, add null-joined rows for unmatched + for i in 0..row_indices_length { + let last_index = + last_index_for_row(i, row_indices, batch_ids, row_indices_length); + if filter_mask.value(i) { + seen_true = true; + corrected_mask.append_value(true); + } else if seen_true || !filter_mask.value(i) && !last_index { + corrected_mask.append_null(); // to be ignored and not set to output + } else { + corrected_mask.append_value(false); // to be converted to null joined row + } + + if last_index { + seen_true = false; + } + } + + // Generate null joined rows for records which have no matching join key + corrected_mask.append_n(expected_size - corrected_mask.len(), false); + Some(corrected_mask.finish()) + } + JoinType::LeftMark | JoinType::RightMark => { + // For mark joins: Like outer but only keep first match, mark with boolean + for i in 0..row_indices_length { + let last_index = + last_index_for_row(i, row_indices, batch_ids, row_indices_length); + if filter_mask.value(i) && !seen_true { + seen_true = true; + corrected_mask.append_value(true); + } else if seen_true || !filter_mask.value(i) && !last_index { + corrected_mask.append_null(); // to be ignored and not set to output + } else { + corrected_mask.append_value(false); // to be converted to null joined row + } + + if last_index { + seen_true = false; + } + } + + // Generate null joined rows for records which have no matching join key + corrected_mask.append_n(expected_size - corrected_mask.len(), false); + Some(corrected_mask.finish()) + } + JoinType::LeftSemi | JoinType::RightSemi => { + // For semi joins: Keep only first matching row per input row, discard rest + for i in 0..row_indices_length { + let last_index = + last_index_for_row(i, row_indices, batch_ids, row_indices_length); + if filter_mask.value(i) && !seen_true { + seen_true = true; + corrected_mask.append_value(true); + } else { + corrected_mask.append_null(); // to be ignored and not set to output + } + + if last_index { + seen_true = false; + } + } + + Some(corrected_mask.finish()) + } + JoinType::LeftAnti | JoinType::RightAnti => { + // For anti joins: Keep row only if NO matches passed the filter + for i in 0..row_indices_length { + let last_index = + last_index_for_row(i, row_indices, batch_ids, row_indices_length); + + if filter_mask.value(i) { + seen_true = true; + } + + if last_index { + if !seen_true { + corrected_mask.append_value(true); + } else { + corrected_mask.append_null(); + } + + seen_true = false; + } else { + corrected_mask.append_null(); + } + } + // Generate null joined rows for records which have no matching join key, + // for LeftAnti non-matched considered as true + corrected_mask.append_n(expected_size - corrected_mask.len(), true); + Some(corrected_mask.finish()) + } + JoinType::Full => { + // For full joins: Similar to outer but handle both sides + for i in 0..row_indices_length { + let last_index = + last_index_for_row(i, row_indices, batch_ids, row_indices_length); + + if filter_mask.is_null(i) { + // null joined + corrected_mask.append_value(true); + } else if filter_mask.value(i) { + seen_true = true; + corrected_mask.append_value(true); + } else if seen_true || !filter_mask.value(i) && !last_index { + corrected_mask.append_null(); // to be ignored and not set to output + } else { + corrected_mask.append_value(false); // to be converted to null joined row + } + + if last_index { + seen_true = false; + } + } + // Generate null joined rows for records which have no matching join key + corrected_mask.append_n(expected_size - corrected_mask.len(), false); + Some(corrected_mask.finish()) + } + JoinType::Inner => { + // Inner joins don't need deferred filtering + None + } + } +} + +/// Applies corrected filter mask to record batch based on join type +/// +/// Different join types require different handling of filtered results: +/// - Outer joins: Add null-joined rows for false mask values +/// - Semi/Anti joins: May need projection to remove right columns +/// - Full joins: Add null-joined rows for both sides +pub fn filter_record_batch_by_join_type( + record_batch: &RecordBatch, + corrected_mask: &BooleanArray, + join_type: JoinType, + schema: &SchemaRef, + streamed_schema: &SchemaRef, + buffered_schema: &SchemaRef, +) -> Result { + let filtered_record_batch = filter_record_batch(record_batch, corrected_mask)?; + + match join_type { + JoinType::Left | JoinType::LeftMark => { + // For left joins, add null-joined rows where mask is false + let null_mask = compute::not(corrected_mask)?; + let null_joined_batch = filter_record_batch(record_batch, &null_mask)?; + + if null_joined_batch.num_rows() == 0 { + return Ok(filtered_record_batch); + } + + // Create null columns for right side + let null_joined_streamed_batch = create_null_joined_batch( + &null_joined_batch, + buffered_schema, + JoinSide::Left, + join_type, + schema, + )?; + + Ok(concat_batches( + schema, + &[filtered_record_batch, null_joined_streamed_batch], + )?) + } + JoinType::LeftSemi + | JoinType::LeftAnti + | JoinType::RightSemi + | JoinType::RightAnti => { + // For semi/anti joins, project to only include the outer side columns + // Both Left and Right semi/anti use streamed_schema.len() because: + // - For Left: columns are [left, right], so we take first streamed_schema.len() + // - For Right: columns are [right, left], and streamed side is right, so we take first streamed_schema.len() + let output_column_indices: Vec = + (0..streamed_schema.fields().len()).collect(); + Ok(filtered_record_batch.project(&output_column_indices)?) + } + JoinType::Right | JoinType::RightMark => { + // For right joins, add null-joined rows where mask is false + let null_mask = compute::not(corrected_mask)?; + let null_joined_batch = filter_record_batch(record_batch, &null_mask)?; + + if null_joined_batch.num_rows() == 0 { + return Ok(filtered_record_batch); + } + + // Create null columns for left side (buffered side for RIGHT join) + let null_joined_buffered_batch = create_null_joined_batch( + &null_joined_batch, + buffered_schema, // Pass buffered (left) schema to create nulls for it + JoinSide::Right, + join_type, + schema, + )?; + + Ok(concat_batches( + schema, + &[filtered_record_batch, null_joined_buffered_batch], + )?) + } + JoinType::Full => { + // For full joins, add null-joined rows for both sides + let joined_filter_not_matched_mask = compute::not(corrected_mask)?; + let joined_filter_not_matched_batch = + filter_record_batch(record_batch, &joined_filter_not_matched_mask)?; + + if joined_filter_not_matched_batch.num_rows() == 0 { + return Ok(filtered_record_batch); + } + + // Create null-joined batches for both sides + let left_null_joined_batch = create_null_joined_batch( + &joined_filter_not_matched_batch, + buffered_schema, + JoinSide::Left, + join_type, + schema, + )?; + + Ok(concat_batches( + schema, + &[filtered_record_batch, left_null_joined_batch], + )?) + } + JoinType::Inner => Ok(filtered_record_batch), + } +} + +/// Creates a batch with null columns for the non-joined side +/// +/// Note: The input `batch` is assumed to be a fully-joined batch that already contains +/// columns from both sides. We need to extract the data side columns and replace the +/// null side columns with actual nulls. +fn create_null_joined_batch( + batch: &RecordBatch, + null_schema: &SchemaRef, + join_side: JoinSide, + join_type: JoinType, + output_schema: &SchemaRef, +) -> Result { + let num_rows = batch.num_rows(); + + // The input batch is a fully-joined batch [left_cols..., right_cols...] + // We need to extract the appropriate side and replace the other with nulls (or mark column) + let columns = match (join_side, join_type) { + (JoinSide::Left, JoinType::LeftMark) => { + // For LEFT mark: output is [left_cols..., mark_col] + // Batch is [left_cols..., right_cols...], extract left from beginning + // Number of left columns = output columns - 1 (mark column) + let left_col_count = output_schema.fields().len() - 1; + let mut result: Vec = batch.columns()[..left_col_count].to_vec(); + result.push(Arc::new(BooleanArray::from(vec![false; num_rows])) as ArrayRef); + result + } + (JoinSide::Right, JoinType::RightMark) => { + // For RIGHT mark: output is [right_cols..., mark_col] + // For RIGHT joins, batch is [right_cols..., left_cols...] (right comes first!) + // Extract right columns from the beginning + let right_col_count = output_schema.fields().len() - 1; // -1 for mark column + let mut result: Vec = batch.columns()[..right_col_count].to_vec(); + result.push(Arc::new(BooleanArray::from(vec![false; num_rows])) as ArrayRef); + result + } + (JoinSide::Left, _) => { + // For LEFT join: output is [left_cols..., right_cols...] + // Extract left columns, then add null right columns + let null_columns: Vec = null_schema + .fields() + .iter() + .map(|field| arrow::array::new_null_array(field.data_type(), num_rows)) + .collect(); + let left_col_count = output_schema.fields().len() - null_columns.len(); + let mut result: Vec = batch.columns()[..left_col_count].to_vec(); + result.extend(null_columns); + result + } + (JoinSide::Right, _) => { + // For RIGHT join: batch is [left_cols..., right_cols...] (same as schema) + // We want: [null_left..., actual_right...] + // Extract left columns from beginning, replace with nulls, keep right columns + let null_columns: Vec = null_schema + .fields() + .iter() + .map(|field| arrow::array::new_null_array(field.data_type(), num_rows)) + .collect(); + let left_col_count = null_columns.len(); + let mut result = null_columns; + // Extract right columns starting after left columns + result.extend_from_slice(&batch.columns()[left_col_count..]); + result + } + (JoinSide::None, _) => { + // This should not happen in normal join operations + unreachable!( + "JoinSide::None should not be used in null-joined batch creation" + ) + } + }; + + // Create the batch - don't validate nullability since outer joins can have + // null values in columns that were originally non-nullable + use arrow::array::RecordBatchOptions; + let mut options = RecordBatchOptions::new(); + options = options.with_row_count(Some(num_rows)); + Ok(RecordBatch::try_new_with_options( + Arc::clone(output_schema), + columns, + &options, + )?) +} diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/mod.rs b/datafusion/physical-plan/src/joins/sort_merge_join/mod.rs index 82f18e7414095..06290ec4d0908 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/mod.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/mod.rs @@ -20,6 +20,7 @@ pub use exec::SortMergeJoinExec; mod exec; +mod filter; mod metrics; mod stream; diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/stream.rs index b36992caf4b45..3a57dc6b41a57 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/stream.rs @@ -33,6 +33,10 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering::Relaxed; use std::task::{Context, Poll}; +use crate::joins::sort_merge_join::filter::{ + FilterMetadata, filter_record_batch_by_join_type, get_corrected_filter_mask, + get_filter_columns, needs_deferred_filtering, +}; use crate::joins::sort_merge_join::metrics::SortMergeJoinMetrics; use crate::joins::utils::{JoinFilter, compare_join_arrays}; use crate::metrics::RecordOutput; @@ -42,15 +46,13 @@ use crate::{PhysicalExpr, RecordBatchStream, SendableRecordBatchStream}; use arrow::array::{types::UInt64Type, *}; use arrow::compute::{ self, BatchCoalescer, SortOptions, concat_batches, filter_record_batch, is_not_null, - take, + take, take_arrays, }; use arrow::datatypes::{DataType, SchemaRef, TimeUnit}; -use arrow::error::ArrowError; use arrow::ipc::reader::StreamReader; use datafusion_common::config::SpillCompression; use datafusion_common::{ - DataFusionError, HashSet, JoinSide, JoinType, NullEquality, Result, exec_err, - internal_err, not_impl_err, + HashSet, JoinType, NullEquality, Result, exec_err, internal_err, not_impl_err, }; use datafusion_execution::disk_manager::RefCountedTempFile; use datafusion_execution::memory_pool::MemoryReservation; @@ -68,6 +70,8 @@ pub(super) enum SortMergeJoinState { Polling, /// Joining polled data and making output JoinOutput, + /// Emit ready data if have any and then go back to [`Self::Init`] state + EmitReadyThenInit, /// No more output Exhausted, } @@ -370,12 +374,8 @@ pub(super) struct SortMergeJoinStream { pub(super) struct JoinedRecordBatches { /// Joined batches. Each batch is already joined columns from left and right sources pub(super) joined_batches: BatchCoalescer, - /// Did each output row pass the join filter? (detect if input row found any match) - pub(super) filter_mask: BooleanBuilder, - /// Which input row (within batch) produced each output row? (for grouping by input row) - pub(super) row_indices: UInt64Builder, - /// Which input batch did each output row come from? (disambiguate row_indices) - pub(super) batch_ids: Vec, + /// Filter metadata for deferred filtering + pub(super) filter_metadata: FilterMetadata, } impl JoinedRecordBatches { @@ -398,61 +398,28 @@ impl JoinedRecordBatches { } } - /// Finishes and returns the metadata arrays, clearing the builders - /// - /// Returns (row_indices, filter_mask, batch_ids_ref) - /// Note: batch_ids is returned as a reference since it's still needed in the struct - fn finish_metadata(&mut self) -> (UInt64Array, BooleanArray, &[usize]) { - let row_indices = self.row_indices.finish(); - let filter_mask = self.filter_mask.finish(); - (row_indices, filter_mask, &self.batch_ids) - } - /// Clears batches without touching metadata (for early return when no filtering needed) fn clear_batches(&mut self, schema: &SchemaRef, batch_size: usize) { self.joined_batches = BatchCoalescer::new(Arc::clone(schema), batch_size) .with_biggest_coalesce_batch_size(Option::from(batch_size / 2)); } - /// Asserts that internal metadata arrays are consistent with each other - /// Only checks if metadata is actually being used (i.e., not all empty) - #[inline] - fn debug_assert_metadata_aligned(&self) { - // Metadata arrays should be aligned IF they're being used - // (For non-filtered joins, they may all be empty) - if self.filter_mask.len() > 0 - || self.row_indices.len() > 0 - || !self.batch_ids.is_empty() - { - debug_assert_eq!( - self.filter_mask.len(), - self.row_indices.len(), - "filter_mask and row_indices must have same length when metadata is used" - ); - debug_assert_eq!( - self.filter_mask.len(), - self.batch_ids.len(), - "filter_mask and batch_ids must have same length when metadata is used" - ); - } - } - /// Asserts that if batches is empty, metadata is also empty #[inline] fn debug_assert_empty_consistency(&self) { if self.joined_batches.is_empty() { debug_assert_eq!( - self.filter_mask.len(), + self.filter_metadata.filter_mask.len(), 0, "filter_mask should be empty when batches is empty" ); debug_assert_eq!( - self.row_indices.len(), + self.filter_metadata.row_indices.len(), 0, "row_indices should be empty when batches is empty" ); debug_assert_eq!( - self.batch_ids.len(), + self.filter_metadata.batch_ids.len(), 0, "batch_ids should be empty when batches is empty" ); @@ -473,14 +440,9 @@ impl JoinedRecordBatches { let num_rows = batch.num_rows(); - self.filter_mask.append_nulls(num_rows); - self.row_indices.append_nulls(num_rows); - self.batch_ids.resize( - self.batch_ids.len() + num_rows, - 0, // batch_id = 0 for null-joined rows - ); + self.filter_metadata.append_nulls(num_rows); - self.debug_assert_metadata_aligned(); + self.filter_metadata.debug_assert_metadata_aligned(); self.joined_batches .push_batch(batch) .expect("Failed to push batch to BatchCoalescer"); @@ -525,13 +487,13 @@ impl JoinedRecordBatches { "row_indices and filter_mask must have same length" ); - // For Full joins, we keep the pre_mask (with nulls), for others we keep the cleaned mask - self.filter_mask.extend(filter_mask); - self.row_indices.extend(row_indices); - self.batch_ids - .resize(self.batch_ids.len() + row_indices.len(), streamed_batch_id); + self.filter_metadata.append_filter_metadata( + row_indices, + filter_mask, + streamed_batch_id, + ); - self.debug_assert_metadata_aligned(); + self.filter_metadata.debug_assert_metadata_aligned(); self.joined_batches .push_batch(batch) .expect("Failed to push batch to BatchCoalescer"); @@ -551,9 +513,7 @@ impl JoinedRecordBatches { fn clear(&mut self, schema: &SchemaRef, batch_size: usize) { self.joined_batches = BatchCoalescer::new(Arc::clone(schema), batch_size) .with_biggest_coalesce_batch_size(Option::from(batch_size / 2)); - self.batch_ids.clear(); - self.filter_mask = BooleanBuilder::new(); - self.row_indices = UInt64Builder::new(); + self.filter_metadata = FilterMetadata::new(); self.debug_assert_empty_consistency(); } } @@ -563,199 +523,6 @@ impl RecordBatchStream for SortMergeJoinStream { } } -/// True if next index refers to either: -/// - another batch id -/// - another row index within same batch id -/// - end of row indices -#[inline(always)] -fn last_index_for_row( - row_index: usize, - indices: &UInt64Array, - batch_ids: &[usize], - indices_len: usize, -) -> bool { - debug_assert_eq!( - indices.len(), - indices_len, - "indices.len() should match indices_len parameter" - ); - debug_assert_eq!( - batch_ids.len(), - indices_len, - "batch_ids.len() should match indices_len" - ); - debug_assert!( - row_index < indices_len, - "row_index {row_index} should be < indices_len {indices_len}", - ); - - row_index == indices_len - 1 - || batch_ids[row_index] != batch_ids[row_index + 1] - || indices.value(row_index) != indices.value(row_index + 1) -} - -// Returns a corrected boolean bitmask for the given join type -// Values in the corrected bitmask can be: true, false, null -// `true` - the row found its match and sent to the output -// `null` - the row ignored, no output -// `false` - the row sent as NULL joined row -pub(super) fn get_corrected_filter_mask( - join_type: JoinType, - row_indices: &UInt64Array, - batch_ids: &[usize], - filter_mask: &BooleanArray, - expected_size: usize, -) -> Option { - let row_indices_length = row_indices.len(); - let mut corrected_mask: BooleanBuilder = - BooleanBuilder::with_capacity(row_indices_length); - let mut seen_true = false; - - match join_type { - JoinType::Left | JoinType::Right => { - for i in 0..row_indices_length { - let last_index = - last_index_for_row(i, row_indices, batch_ids, row_indices_length); - if filter_mask.value(i) { - seen_true = true; - corrected_mask.append_value(true); - } else if seen_true || !filter_mask.value(i) && !last_index { - corrected_mask.append_null(); // to be ignored and not set to output - } else { - corrected_mask.append_value(false); // to be converted to null joined row - } - - if last_index { - seen_true = false; - } - } - - // Generate null joined rows for records which have no matching join key - corrected_mask.append_n(expected_size - corrected_mask.len(), false); - Some(corrected_mask.finish()) - } - JoinType::LeftMark | JoinType::RightMark => { - for i in 0..row_indices_length { - let last_index = - last_index_for_row(i, row_indices, batch_ids, row_indices_length); - if filter_mask.value(i) && !seen_true { - seen_true = true; - corrected_mask.append_value(true); - } else if seen_true || !filter_mask.value(i) && !last_index { - corrected_mask.append_null(); // to be ignored and not set to output - } else { - corrected_mask.append_value(false); // to be converted to null joined row - } - - if last_index { - seen_true = false; - } - } - - // Generate null joined rows for records which have no matching join key - corrected_mask.append_n(expected_size - corrected_mask.len(), false); - Some(corrected_mask.finish()) - } - JoinType::LeftSemi | JoinType::RightSemi => { - for i in 0..row_indices_length { - let last_index = - last_index_for_row(i, row_indices, batch_ids, row_indices_length); - if filter_mask.value(i) && !seen_true { - seen_true = true; - corrected_mask.append_value(true); - } else { - corrected_mask.append_null(); // to be ignored and not set to output - } - - if last_index { - seen_true = false; - } - } - - Some(corrected_mask.finish()) - } - JoinType::LeftAnti | JoinType::RightAnti => { - for i in 0..row_indices_length { - let last_index = - last_index_for_row(i, row_indices, batch_ids, row_indices_length); - - if filter_mask.value(i) { - seen_true = true; - } - - if last_index { - if !seen_true { - corrected_mask.append_value(true); - } else { - corrected_mask.append_null(); - } - - seen_true = false; - } else { - corrected_mask.append_null(); - } - } - // Generate null joined rows for records which have no matching join key, - // for LeftAnti non-matched considered as true - corrected_mask.append_n(expected_size - corrected_mask.len(), true); - Some(corrected_mask.finish()) - } - JoinType::Full => { - let mut mask: Vec> = vec![Some(true); row_indices_length]; - let mut last_true_idx = 0; - let mut first_row_idx = 0; - let mut seen_false = false; - - for i in 0..row_indices_length { - let last_index = - last_index_for_row(i, row_indices, batch_ids, row_indices_length); - let val = filter_mask.value(i); - let is_null = filter_mask.is_null(i); - - if val { - // memoize the first seen matched row - if !seen_true { - last_true_idx = i; - } - seen_true = true; - } - - if is_null || val { - mask[i] = Some(true); - } else if !is_null && !val && (seen_true || seen_false) { - mask[i] = None; - } else { - mask[i] = Some(false); - } - - if !is_null && !val { - seen_false = true; - } - - if last_index { - // If the left row seen as true its needed to output it once - // To do that we mark all other matches for same row as null to avoid the output - if seen_true { - #[expect(clippy::needless_range_loop)] - for j in first_row_idx..last_true_idx { - mask[j] = None; - } - } - - seen_true = false; - seen_false = false; - last_true_idx = 0; - first_row_idx = i + 1; - } - } - - Some(BooleanArray::from(mask)) - } - // Only outer joins needs to keep track of processed rows and apply corrected filter mask - _ => None, - } -} - impl Stream for SortMergeJoinStream { type Item = Result; @@ -778,7 +545,10 @@ impl Stream for SortMergeJoinStream { match self.current_ordering { Ordering::Less | Ordering::Equal => { if !streamed_exhausted { - if self.needs_deferred_filtering() { + if needs_deferred_filtering( + &self.filter, + self.join_type, + ) { match self.process_filtered_batches()? { Poll::Ready(Some(batch)) => { return Poll::Ready(Some(Ok(batch))); @@ -830,22 +600,56 @@ impl Stream for SortMergeJoinStream { self.current_ordering = self.compare_streamed_buffered()?; self.state = SortMergeJoinState::JoinOutput; } + SortMergeJoinState::EmitReadyThenInit => { + // If have data to emit, emit it and if no more, change to next + + // Verify metadata alignment before checking if we have batches to output + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); + + // For filtered joins, skip output and let Init state handle it + if needs_deferred_filtering(&self.filter, self.join_type) { + self.state = SortMergeJoinState::Init; + continue; + } + + // For non-filtered joins, only output if we have a completed batch + // (opportunistic output when target batch size is reached) + if self + .joined_record_batches + .joined_batches + .has_completed_batch() + { + let record_batch = self + .joined_record_batches + .joined_batches + .next_completed_batch() + .expect("has_completed_batch was true"); + (&record_batch) + .record_output(&self.join_metrics.baseline_metrics()); + return Poll::Ready(Some(Ok(record_batch))); + } + self.state = SortMergeJoinState::Init; + } SortMergeJoinState::JoinOutput => { self.join_partial()?; if self.num_unfrozen_pairs() < self.batch_size { if self.buffered_data.scanning_finished() { self.buffered_data.scanning_reset(); - self.state = SortMergeJoinState::Init; + self.state = SortMergeJoinState::EmitReadyThenInit; } } else { self.freeze_all()?; // Verify metadata alignment before checking if we have batches to output - self.joined_record_batches.debug_assert_metadata_aligned(); + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); // For filtered joins, skip output and let Init state handle it - if self.needs_deferred_filtering() { + if needs_deferred_filtering(&self.filter, self.join_type) { continue; } @@ -872,10 +676,12 @@ impl Stream for SortMergeJoinStream { self.freeze_all()?; // Verify metadata alignment before final output - self.joined_record_batches.debug_assert_metadata_aligned(); + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); // For filtered joins, must concat and filter ALL data at once - if self.needs_deferred_filtering() + if needs_deferred_filtering(&self.filter, self.join_type) && !self.joined_record_batches.joined_batches.is_empty() { let record_batch = self.filter_joined_batch()?; @@ -975,9 +781,7 @@ impl SortMergeJoinStream { joined_record_batches: JoinedRecordBatches { joined_batches: BatchCoalescer::new(Arc::clone(&schema), batch_size) .with_biggest_coalesce_batch_size(Option::from(batch_size / 2)), - filter_mask: BooleanBuilder::new(), - row_indices: UInt64Builder::new(), - batch_ids: vec![], + filter_metadata: FilterMetadata::new(), }, output: BatchCoalescer::new(schema, batch_size) .with_biggest_coalesce_batch_size(Option::from(batch_size / 2)), @@ -996,26 +800,6 @@ impl SortMergeJoinStream { self.streamed_batch.num_output_rows() } - /// Returns true if this join needs deferred filtering - /// - /// Deferred filtering is needed when a filter exists and the join type requires - /// ensuring each input row produces at least one output row (or exactly one for semi). - fn needs_deferred_filtering(&self) -> bool { - self.filter.is_some() - && matches!( - self.join_type, - JoinType::Left - | JoinType::LeftSemi - | JoinType::LeftMark - | JoinType::Right - | JoinType::RightSemi - | JoinType::RightMark - | JoinType::LeftAnti - | JoinType::RightAnti - | JoinType::Full - ) - } - /// Process accumulated batches for filtered joins /// /// Freezes unfrozen pairs, applies deferred filtering, and outputs if ready. @@ -1023,7 +807,9 @@ impl SortMergeJoinStream { fn process_filtered_batches(&mut self) -> Poll>> { self.freeze_all()?; - self.joined_record_batches.debug_assert_metadata_aligned(); + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); if !self.joined_record_batches.joined_batches.is_empty() { let out_filtered_batch = self.filter_joined_batch()?; @@ -1399,7 +1185,9 @@ impl SortMergeJoinStream { self.freeze_streamed()?; // After freezing, metadata should be aligned - self.joined_record_batches.debug_assert_metadata_aligned(); + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); Ok(()) } @@ -1414,7 +1202,9 @@ impl SortMergeJoinStream { self.freeze_buffered(1)?; // After freezing, metadata should be aligned - self.joined_record_batches.debug_assert_metadata_aligned(); + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); Ok(()) } @@ -1490,13 +1280,19 @@ impl SortMergeJoinStream { continue; } - let mut left_columns = self - .streamed_batch - .batch - .columns() - .iter() - .map(|column| take(column, &left_indices, None)) - .collect::, ArrowError>>()?; + let mut left_columns = if let Some(range) = is_contiguous_range(&left_indices) + { + // When indices form a contiguous range (common for the streamed + // side which advances sequentially), use zero-copy slice instead + // of the O(n) take kernel. + self.streamed_batch + .batch + .slice(range.start, range.len()) + .columns() + .to_vec() + } else { + take_arrays(self.streamed_batch.batch.columns(), &left_indices, None)? + }; // The row indices of joined buffered batch let right_indices: UInt64Array = chunk.buffered_indices.finish(); @@ -1541,7 +1337,7 @@ impl SortMergeJoinStream { &right_indices, )?; - get_filter_column(&self.filter, &left_columns, &right_cols) + get_filter_columns(&self.filter, &left_columns, &right_cols) } else if matches!( self.join_type, JoinType::RightAnti | JoinType::RightSemi | JoinType::RightMark @@ -1552,12 +1348,12 @@ impl SortMergeJoinStream { &right_indices, )?; - get_filter_column(&self.filter, &right_cols, &left_columns) + get_filter_columns(&self.filter, &right_cols, &left_columns) } else { - get_filter_column(&self.filter, &left_columns, &right_columns) + get_filter_columns(&self.filter, &left_columns, &right_columns) } } else { - get_filter_column(&self.filter, &right_columns, &left_columns) + get_filter_columns(&self.filter, &right_columns, &left_columns) } } else { // This chunk is totally for null joined rows (outer join), we don't need to apply join filter. @@ -1679,11 +1475,13 @@ impl SortMergeJoinStream { fn filter_joined_batch(&mut self) -> Result { // Metadata should be aligned before processing - self.joined_record_batches.debug_assert_metadata_aligned(); + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); let record_batch = self.joined_record_batches.concat_batches(&self.schema)?; let (mut out_indices, mut out_mask, mut batch_ids) = - self.joined_record_batches.finish_metadata(); + self.joined_record_batches.filter_metadata.finish_metadata(); let default_batch_ids = vec![0; record_batch.num_rows()]; // If only nulls come in and indices sizes doesn't match with expected record batch count @@ -1754,139 +1552,14 @@ impl SortMergeJoinStream { record_batch: &RecordBatch, corrected_mask: &BooleanArray, ) -> Result { - // Corrected mask should have length matching or exceeding record_batch rows - // (for outer joins it may be longer to include null-joined rows) - debug_assert!( - corrected_mask.len() >= record_batch.num_rows(), - "corrected_mask length ({}) should be >= record_batch rows ({})", - corrected_mask.len(), - record_batch.num_rows() - ); - - let mut filtered_record_batch = - filter_record_batch(record_batch, corrected_mask)?; - let left_columns_length = self.streamed_schema.fields.len(); - let right_columns_length = self.buffered_schema.fields.len(); - - if matches!( + let filtered_record_batch = filter_record_batch_by_join_type( + record_batch, + corrected_mask, self.join_type, - JoinType::Left | JoinType::LeftMark | JoinType::Right | JoinType::RightMark - ) { - let null_mask = compute::not(corrected_mask)?; - let null_joined_batch = filter_record_batch(record_batch, &null_mask)?; - - let mut right_columns = create_unmatched_columns( - self.join_type, - &self.buffered_schema, - null_joined_batch.num_rows(), - ); - - let columns = match self.join_type { - JoinType::Right => { - // The first columns are the right columns. - let left_columns = null_joined_batch - .columns() - .iter() - .skip(right_columns_length) - .cloned() - .collect::>(); - - right_columns.extend(left_columns); - right_columns - } - JoinType::Left | JoinType::LeftMark | JoinType::RightMark => { - // The first columns are the left columns. - let mut left_columns = null_joined_batch - .columns() - .iter() - .take(left_columns_length) - .cloned() - .collect::>(); - - left_columns.extend(right_columns); - left_columns - } - _ => exec_err!("Did not expect join type {}", self.join_type)?, - }; - - // Push the streamed/buffered batch joined nulls to the output - let null_joined_streamed_batch = - RecordBatch::try_new(Arc::clone(&self.schema), columns)?; - - filtered_record_batch = concat_batches( - &self.schema, - &[filtered_record_batch, null_joined_streamed_batch], - )?; - } else if matches!( - self.join_type, - JoinType::LeftSemi - | JoinType::LeftAnti - | JoinType::RightAnti - | JoinType::RightSemi - ) { - let output_column_indices = (0..left_columns_length).collect::>(); - filtered_record_batch = - filtered_record_batch.project(&output_column_indices)?; - } else if matches!(self.join_type, JoinType::Full) - && corrected_mask.false_count() > 0 - { - // Find rows which joined by key but Filter predicate evaluated as false - let joined_filter_not_matched_mask = compute::not(corrected_mask)?; - let joined_filter_not_matched_batch = - filter_record_batch(record_batch, &joined_filter_not_matched_mask)?; - - // Add left unmatched rows adding the right side as nulls - let right_null_columns = self - .buffered_schema - .fields() - .iter() - .map(|f| { - new_null_array( - f.data_type(), - joined_filter_not_matched_batch.num_rows(), - ) - }) - .collect::>(); - - let mut result_joined = joined_filter_not_matched_batch - .columns() - .iter() - .take(left_columns_length) - .cloned() - .collect::>(); - - result_joined.extend(right_null_columns); - - let left_null_joined_batch = - RecordBatch::try_new(Arc::clone(&self.schema), result_joined)?; - - // Add right unmatched rows adding the left side as nulls - let mut result_joined = self - .streamed_schema - .fields() - .iter() - .map(|f| { - new_null_array( - f.data_type(), - joined_filter_not_matched_batch.num_rows(), - ) - }) - .collect::>(); - - let right_data = joined_filter_not_matched_batch - .columns() - .iter() - .skip(left_columns_length) - .cloned() - .collect::>(); - - result_joined.extend(right_data); - - filtered_record_batch = concat_batches( - &self.schema, - &[filtered_record_batch, left_null_joined_batch], - )?; - } + &self.schema, + &self.streamed_schema, + &self.buffered_schema, + )?; self.joined_record_batches .clear(&self.schema, self.batch_size); @@ -1911,36 +1584,6 @@ fn create_unmatched_columns( } } -/// Gets the arrays which join filters are applied on. -fn get_filter_column( - join_filter: &Option, - streamed_columns: &[ArrayRef], - buffered_columns: &[ArrayRef], -) -> Vec { - let mut filter_columns = vec![]; - - if let Some(f) = join_filter { - let left_columns = f - .column_indices() - .iter() - .filter(|col_index| col_index.side == JoinSide::Left) - .map(|i| Arc::clone(&streamed_columns[i.index])) - .collect::>(); - - let right_columns = f - .column_indices() - .iter() - .filter(|col_index| col_index.side == JoinSide::Right) - .map(|i| Arc::clone(&buffered_columns[i.index])) - .collect::>(); - - filter_columns.extend(left_columns); - filter_columns.extend(right_columns); - } - - filter_columns -} - fn produce_buffered_null_batch( schema: &SchemaRef, streamed_schema: &SchemaRef, @@ -1970,6 +1613,30 @@ fn produce_buffered_null_batch( )?)) } +/// Checks if a `UInt64Array` contains a contiguous ascending range (e.g. \[3,4,5,6\]). +/// Returns `Some(start..start+len)` if so, `None` otherwise. +/// This allows replacing an O(n) `take` with an O(1) `slice`. +#[inline] +fn is_contiguous_range(indices: &UInt64Array) -> Option> { + if indices.is_empty() || indices.null_count() > 0 { + return None; + } + let values = indices.values(); + let start = values[0]; + let len = values.len() as u64; + // Quick rejection: if last element doesn't match expected, not contiguous + if values[values.len() - 1] != start + len - 1 { + return None; + } + // Verify every element is sequential (handles duplicates and gaps) + for i in 1..values.len() { + if values[i] != start + i as u64 { + return None; + } + } + Some(start as usize..(start + len) as usize) +} + /// Get `buffered_indices` rows for `buffered_data[buffered_batch_idx]` by specific column indices #[inline(always)] fn fetch_right_columns_by_idxs( @@ -1990,12 +1657,16 @@ fn fetch_right_columns_from_batch_by_idxs( ) -> Result> { match &buffered_batch.batch { // In memory batch - BufferedBatchState::InMemory(batch) => Ok(batch - .columns() - .iter() - .map(|column| take(column, &buffered_indices, None)) - .collect::, ArrowError>>() - .map_err(Into::::into)?), + // In memory batch + BufferedBatchState::InMemory(batch) => { + // When indices form a contiguous range (common in SMJ since the + // buffered side is scanned sequentially), use zero-copy slice. + if let Some(range) = is_contiguous_range(buffered_indices) { + Ok(batch.slice(range.start, range.len()).columns().to_vec()) + } else { + Ok(take_arrays(batch.columns(), buffered_indices, None)?) + } + } // If the batch was spilled to disk, less likely BufferedBatchState::Spilled(spill_file) => { let mut buffered_cols: Vec = diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs index d0bcc79636f75..4329abdd522dd 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -24,42 +24,44 @@ //! //! Add relevant tests under the specified sections. -use std::sync::Arc; - +use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn}; +use crate::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec}; +use crate::test::TestMemoryExec; +use crate::test::exec::BarrierExec; +use crate::test::{build_table_i32, build_table_i32_two_cols}; +use crate::{ExecutionPlan, common}; +use crate::{ + expressions::Column, joins::sort_merge_join::filter::get_corrected_filter_mask, + joins::sort_merge_join::stream::JoinedRecordBatches, +}; use arrow::array::{ BinaryArray, BooleanArray, Date32Array, Date64Array, FixedSizeBinaryArray, Int32Array, RecordBatch, UInt64Array, - builder::{BooleanBuilder, UInt64Builder}, }; use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch}; use arrow::datatypes::{DataType, Field, Schema}; - +use arrow_ord::sort::SortColumn; +use arrow_schema::SchemaRef; use datafusion_common::JoinType::*; use datafusion_common::{ - JoinSide, + JoinSide, internal_err, test_util::{batches_to_sort_string, batches_to_string}, }; use datafusion_common::{ JoinType, NullEquality, Result, assert_batches_eq, assert_contains, }; -use datafusion_execution::TaskContext; +use datafusion_common_runtime::JoinSet; use datafusion_execution::config::SessionConfig; use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::BinaryExpr; +use futures::StreamExt; use insta::{allow_duplicates, assert_snapshot}; - -use crate::{ - expressions::Column, - joins::sort_merge_join::stream::{JoinedRecordBatches, get_corrected_filter_mask}, -}; - -use crate::joins::SortMergeJoinExec; -use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn}; -use crate::test::TestMemoryExec; -use crate::test::{build_table_i32, build_table_i32_two_cols}; -use crate::{ExecutionPlan, common}; +use itertools::Itertools; +use std::sync::Arc; +use std::task::Poll; fn build_table( a: (&str, &Vec), @@ -2375,9 +2377,7 @@ fn build_joined_record_batches() -> Result { let mut batches = JoinedRecordBatches { joined_batches: BatchCoalescer::new(Arc::clone(&schema), 8192), - filter_mask: BooleanBuilder::new(), - row_indices: UInt64Builder::new(), - batch_ids: vec![], + filter_metadata: crate::joins::sort_merge_join::filter::FilterMetadata::new(), }; // Insert already prejoined non-filtered rows @@ -2432,44 +2432,73 @@ fn build_joined_record_batches() -> Result { )?)?; let streamed_indices = vec![0, 0]; - batches.batch_ids.extend(vec![0; streamed_indices.len()]); batches + .filter_metadata + .batch_ids + .extend(vec![0; streamed_indices.len()]); + batches + .filter_metadata .row_indices .extend(&UInt64Array::from(streamed_indices)); let streamed_indices = vec![1]; - batches.batch_ids.extend(vec![0; streamed_indices.len()]); batches + .filter_metadata + .batch_ids + .extend(vec![0; streamed_indices.len()]); + batches + .filter_metadata .row_indices .extend(&UInt64Array::from(streamed_indices)); let streamed_indices = vec![0, 0]; - batches.batch_ids.extend(vec![1; streamed_indices.len()]); batches + .filter_metadata + .batch_ids + .extend(vec![1; streamed_indices.len()]); + batches + .filter_metadata .row_indices .extend(&UInt64Array::from(streamed_indices)); let streamed_indices = vec![0]; - batches.batch_ids.extend(vec![2; streamed_indices.len()]); batches + .filter_metadata + .batch_ids + .extend(vec![2; streamed_indices.len()]); + batches + .filter_metadata .row_indices .extend(&UInt64Array::from(streamed_indices)); let streamed_indices = vec![0, 0]; - batches.batch_ids.extend(vec![3; streamed_indices.len()]); batches + .filter_metadata + .batch_ids + .extend(vec![3; streamed_indices.len()]); + batches + .filter_metadata .row_indices .extend(&UInt64Array::from(streamed_indices)); batches + .filter_metadata .filter_mask .extend(&BooleanArray::from(vec![true, false])); - batches.filter_mask.extend(&BooleanArray::from(vec![true])); batches + .filter_metadata + .filter_mask + .extend(&BooleanArray::from(vec![true])); + batches + .filter_metadata .filter_mask .extend(&BooleanArray::from(vec![false, true])); - batches.filter_mask.extend(&BooleanArray::from(vec![false])); batches + .filter_metadata + .filter_mask + .extend(&BooleanArray::from(vec![false])); + batches + .filter_metadata .filter_mask .extend(&BooleanArray::from(vec![false, false])); @@ -2482,8 +2511,8 @@ async fn test_left_outer_join_filtered_mask() -> Result<()> { let schema = joined_batches.joined_batches.schema(); let output = joined_batches.concat_batches(&schema)?; - let out_mask = joined_batches.filter_mask.finish(); - let out_indices = joined_batches.row_indices.finish(); + let out_mask = joined_batches.filter_metadata.filter_mask.finish(); + let out_indices = joined_batches.filter_metadata.row_indices.finish(); assert_eq!( get_corrected_filter_mask( @@ -2620,7 +2649,7 @@ async fn test_left_outer_join_filtered_mask() -> Result<()> { let corrected_mask = get_corrected_filter_mask( Left, &out_indices, - &joined_batches.batch_ids, + &joined_batches.filter_metadata.batch_ids, &out_mask, output.num_rows(), ) @@ -2689,8 +2718,8 @@ async fn test_semi_join_filtered_mask() -> Result<()> { let schema = joined_batches.joined_batches.schema(); let output = joined_batches.concat_batches(&schema)?; - let out_mask = joined_batches.filter_mask.finish(); - let out_indices = joined_batches.row_indices.finish(); + let out_mask = joined_batches.filter_metadata.filter_mask.finish(); + let out_indices = joined_batches.filter_metadata.row_indices.finish(); assert_eq!( get_corrected_filter_mask( @@ -2791,7 +2820,7 @@ async fn test_semi_join_filtered_mask() -> Result<()> { let corrected_mask = get_corrected_filter_mask( join_type, &out_indices, - &joined_batches.batch_ids, + &joined_batches.filter_metadata.batch_ids, &out_mask, output.num_rows(), ) @@ -2864,8 +2893,8 @@ async fn test_anti_join_filtered_mask() -> Result<()> { let schema = joined_batches.joined_batches.schema(); let output = joined_batches.concat_batches(&schema)?; - let out_mask = joined_batches.filter_mask.finish(); - let out_indices = joined_batches.row_indices.finish(); + let out_mask = joined_batches.filter_metadata.filter_mask.finish(); + let out_indices = joined_batches.filter_metadata.row_indices.finish(); assert_eq!( get_corrected_filter_mask( @@ -2966,7 +2995,7 @@ async fn test_anti_join_filtered_mask() -> Result<()> { let corrected_mask = get_corrected_filter_mask( join_type, &out_indices, - &joined_batches.batch_ids, + &joined_batches.filter_metadata.batch_ids, &out_mask, output.num_rows(), ) @@ -3104,6 +3133,419 @@ fn test_partition_statistics() -> Result<()> { Ok(()) } +fn build_batches( + a: (&str, &[Vec]), + b: (&str, &[Vec]), + c: (&str, &[Vec]), +) -> (Vec, SchemaRef) { + assert_eq!(a.1.len(), b.1.len()); + let mut batches = vec![]; + + let schema = Arc::new(Schema::new(vec![ + Field::new(a.0, DataType::Boolean, false), + Field::new(b.0, DataType::Int32, false), + Field::new(c.0, DataType::Int32, false), + ])); + + for i in 0..a.1.len() { + batches.push( + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(BooleanArray::from(a.1[i].clone())), + Arc::new(Int32Array::from(b.1[i].clone())), + Arc::new(Int32Array::from(c.1[i].clone())), + ], + ) + .unwrap(), + ); + } + let schema = batches[0].schema(); + (batches, schema) +} + +fn build_batched_finish_barrier_table( + a: (&str, &[Vec]), + b: (&str, &[Vec]), + c: (&str, &[Vec]), +) -> (Arc, Arc) { + let (batches, schema) = build_batches(a, b, c); + + let memory_exec = TestMemoryExec::try_new_exec( + std::slice::from_ref(&batches), + Arc::clone(&schema), + None, + ) + .unwrap(); + + let barrier_exec = Arc::new( + BarrierExec::new(vec![batches], schema) + .with_log(false) + .without_start_barrier() + .with_finish_barrier(), + ); + + (barrier_exec, memory_exec) +} + +/// Concat and sort batches by all the columns to make sure we can compare them with different join +fn prepare_record_batches_for_cmp(output: Vec) -> RecordBatch { + let output_batch = arrow::compute::concat_batches(output[0].schema_ref(), &output) + .expect("failed to concat batches"); + + // Sort on all columns to make sure we have a deterministic order for the assertion + let sort_columns = output_batch + .columns() + .iter() + .map(|c| SortColumn { + values: Arc::clone(c), + options: None, + }) + .collect::>(); + + let sorted_columns = + arrow::compute::lexsort(&sort_columns, None).expect("failed to sort"); + + RecordBatch::try_new(output_batch.schema(), sorted_columns) + .expect("failed to create batch") +} + +#[expect(clippy::too_many_arguments)] +async fn join_get_stream_and_get_expected( + left: Arc, + right: Arc, + oracle_left: Arc, + oracle_right: Arc, + on: JoinOn, + join_type: JoinType, + filter: Option, + batch_size: usize, +) -> Result<(SendableRecordBatchStream, RecordBatch)> { + let sort_options = vec![SortOptions::default(); on.len()]; + let null_equality = NullEquality::NullEqualsNothing; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::default().with_batch_size(batch_size)), + ); + + let expected_output = { + let oracle = HashJoinExec::try_new( + oracle_left, + oracle_right, + on.clone(), + filter.clone(), + &join_type, + None, + PartitionMode::Partitioned, + null_equality, + )?; + + let stream = oracle.execute(0, Arc::clone(&task_ctx))?; + + let batches = common::collect(stream).await?; + + prepare_record_batches_for_cmp(batches) + }; + + let join = SortMergeJoinExec::try_new( + left, + right, + on, + filter, + join_type, + sort_options, + null_equality, + )?; + + let stream = join.execute(0, task_ctx)?; + + Ok((stream, expected_output)) +} + +fn generate_data_for_emit_early_test( + batch_size: usize, + number_of_batches: usize, + join_type: JoinType, +) -> ( + Arc, + Arc, + Arc, + Arc, +) { + let number_of_rows_per_batch = number_of_batches * batch_size; + // Prepare data + let left_a1 = (0..number_of_rows_per_batch as i32) + .chunks(batch_size) + .into_iter() + .map(|chunk| chunk.collect::>()) + .collect::>(); + let left_b1 = (0..1000000) + .filter(|item| { + match join_type { + LeftAnti | RightAnti => { + let remainder = item % (batch_size as i32); + + // Make sure to have one that match and one that don't + remainder == 0 || remainder == 1 + } + // Have at least 1 that is not matching + _ => item % batch_size as i32 != 0, + } + }) + .take(number_of_rows_per_batch) + .chunks(batch_size) + .into_iter() + .map(|chunk| chunk.collect::>()) + .collect::>(); + + let left_bool_col1 = left_a1 + .clone() + .into_iter() + .map(|b| { + b.into_iter() + // Mostly true but have some false that not overlap with the right column + .map(|a| a % (batch_size as i32) != (batch_size as i32) - 2) + .collect::>() + }) + .collect::>(); + + let (left, left_memory) = build_batched_finish_barrier_table( + ("bool_col1", left_bool_col1.as_slice()), + ("b1", left_b1.as_slice()), + ("a1", left_a1.as_slice()), + ); + + let right_a2 = (0..number_of_rows_per_batch as i32) + .map(|item| item * 11) + .chunks(batch_size) + .into_iter() + .map(|chunk| chunk.collect::>()) + .collect::>(); + let right_b1 = (0..1000000) + .filter(|item| { + match join_type { + LeftAnti | RightAnti => { + let remainder = item % (batch_size as i32); + + // Make sure to have one that match and one that don't + remainder == 1 || remainder == 2 + } + // Have at least 1 that is not matching + _ => item % batch_size as i32 != 1, + } + }) + .take(number_of_rows_per_batch) + .chunks(batch_size) + .into_iter() + .map(|chunk| chunk.collect::>()) + .collect::>(); + let right_bool_col2 = right_a2 + .clone() + .into_iter() + .map(|b| { + b.into_iter() + // Mostly true but have some false that not overlap with the left column + .map(|a| a % (batch_size as i32) != (batch_size as i32) - 1) + .collect::>() + }) + .collect::>(); + + let (right, right_memory) = build_batched_finish_barrier_table( + ("bool_col2", right_bool_col2.as_slice()), + ("b1", right_b1.as_slice()), + ("a2", right_a2.as_slice()), + ); + + (left, right, left_memory, right_memory) +} + +#[tokio::test] +async fn test_should_emit_early_when_have_enough_data_to_emit() -> Result<()> { + for with_filtering in [false, true] { + let join_types = vec![ + Inner, Left, Right, RightSemi, Full, LeftSemi, LeftAnti, LeftMark, RightMark, + ]; + const BATCH_SIZE: usize = 10; + for join_type in join_types { + for output_batch_size in [ + BATCH_SIZE / 3, + BATCH_SIZE / 2, + BATCH_SIZE, + BATCH_SIZE * 2, + BATCH_SIZE * 3, + ] { + // Make sure the number of batches is enough for all join type to emit some output + let number_of_batches = if output_batch_size <= BATCH_SIZE { + 100 + } else { + // Have enough batches + (output_batch_size * 100) / BATCH_SIZE + }; + + let (left, right, left_memory, right_memory) = + generate_data_for_emit_early_test( + BATCH_SIZE, + number_of_batches, + join_type, + ); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let join_filter = if with_filtering { + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("bool_col1", 0)), + Operator::And, + Arc::new(Column::new("bool_col2", 1)), + )), + vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("bool_col1", DataType::Boolean, true), + Field::new("bool_col2", DataType::Boolean, true), + ])), + ); + Some(filter) + } else { + None + }; + + // select * + // from t1 + // right join t2 on t1.b1 = t2.b1 and t1.bool_col1 AND t2.bool_col2 + let (mut output_stream, expected) = join_get_stream_and_get_expected( + Arc::clone(&left) as Arc, + Arc::clone(&right) as Arc, + left_memory as Arc, + right_memory as Arc, + on, + join_type, + join_filter, + output_batch_size, + ) + .await?; + + let (output_batched, output_batches_after_finish) = + consume_stream_until_finish_barrier_reached(left, right, &mut output_stream).await.unwrap_or_else(|e| panic!("Failed to consume stream for join type: '{join_type}' and with filtering '{with_filtering}': {e:?}")); + + // It should emit more than that, but we are being generous + // and to make sure the test pass for all + const MINIMUM_OUTPUT_BATCHES: usize = 5; + assert!( + MINIMUM_OUTPUT_BATCHES <= number_of_batches / 5, + "Make sure that the minimum output batches is realistic" + ); + // Test to make sure that we are not waiting for input to be fully consumed to emit some output + assert!( + output_batched.len() >= MINIMUM_OUTPUT_BATCHES, + "[Sort Merge Join {join_type}] Stream must have at least emit {} batches, but only got {} batches", + MINIMUM_OUTPUT_BATCHES, + output_batched.len() + ); + + // Just sanity test to make sure we are still producing valid output + { + let output = [output_batched, output_batches_after_finish].concat(); + let actual_prepared = prepare_record_batches_for_cmp(output); + + assert_eq!(actual_prepared.columns(), expected.columns()); + } + } + } + } + Ok(()) +} + +/// Polls the stream until both barriers are reached, +/// collecting the emitted batches along the way. +/// +/// If the stream is pending for too long (5s) without emitting any batches, +/// it panics to avoid hanging the test indefinitely. +/// +/// Note: The left and right BarrierExec might be the input of the output stream +async fn consume_stream_until_finish_barrier_reached( + left: Arc, + right: Arc, + output_stream: &mut SendableRecordBatchStream, +) -> Result<(Vec, Vec)> { + let mut switch_to_finish_barrier = false; + let mut output_batched = vec![]; + let mut after_finish_barrier_reached = vec![]; + let mut background_task = JoinSet::new(); + + let mut start_time_since_last_ready = datafusion_common::instant::Instant::now(); + loop { + let next_item = output_stream.next(); + + // Manual polling + let poll_output = futures::poll!(next_item); + + // Wake up the stream to make sure it makes progress + tokio::task::yield_now().await; + + match poll_output { + Poll::Ready(Some(Ok(batch))) => { + if batch.num_rows() == 0 { + return internal_err!("join stream should not emit empty batch"); + } + if switch_to_finish_barrier { + after_finish_barrier_reached.push(batch); + } else { + output_batched.push(batch); + } + start_time_since_last_ready = datafusion_common::instant::Instant::now(); + } + Poll::Ready(Some(Err(e))) => return Err(e), + Poll::Ready(None) if !switch_to_finish_barrier => { + unreachable!("Stream should not end before manually finishing it") + } + Poll::Ready(None) => { + break; + } + Poll::Pending => { + if right.is_finish_barrier_reached() + && left.is_finish_barrier_reached() + && !switch_to_finish_barrier + { + switch_to_finish_barrier = true; + + let right = Arc::clone(&right); + background_task.spawn(async move { + right.wait_finish().await; + }); + let left = Arc::clone(&left); + background_task.spawn(async move { + left.wait_finish().await; + }); + } + + // Make sure the test doesn't run forever + if start_time_since_last_ready.elapsed() + > std::time::Duration::from_secs(5) + { + return internal_err!( + "Stream should have emitted data by now, but it's still pending. Output batches so far: {}", + output_batched.len() + ); + } + } + } + } + + Ok((output_batched, after_finish_barrier_reached)) +} + /// Returns the column names on the schema fn columns(schema: &Schema) -> Vec { schema.fields().iter().map(|f| f.name().clone()).collect() diff --git a/datafusion/physical-plan/src/test/exec.rs b/datafusion/physical-plan/src/test/exec.rs index 4507cccba05a9..c2293184151fd 100644 --- a/datafusion/physical-plan/src/test/exec.rs +++ b/datafusion/physical-plan/src/test/exec.rs @@ -17,13 +17,6 @@ //! Simple iterator over batches for use in testing -use std::{ - any::Any, - pin::Pin, - sync::{Arc, Weak}, - task::{Context, Poll}, -}; - use crate::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, common, @@ -33,6 +26,13 @@ use crate::{ execution_plan::EmissionType, stream::{RecordBatchReceiverStream, RecordBatchStreamAdapter}, }; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::{ + any::Any, + pin::Pin, + sync::{Arc, Weak}, + task::{Context, Poll}, +}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; @@ -298,29 +298,91 @@ pub struct BarrierExec { schema: SchemaRef, /// all streams wait on this barrier to produce - barrier: Arc, + start_data_barrier: Option>, + + /// the stream wait for this to return Poll::Ready(None) + finish_barrier: Option>, + cache: PlanProperties, + + log: bool, } impl BarrierExec { /// Create a new exec with some number of partitions. pub fn new(data: Vec>, schema: SchemaRef) -> Self { // wait for all streams and the input - let barrier = Arc::new(Barrier::new(data.len() + 1)); + let barrier = Some(Arc::new(Barrier::new(data.len() + 1))); let cache = Self::compute_properties(Arc::clone(&schema), &data); Self { data, schema, - barrier, + start_data_barrier: barrier, cache, + finish_barrier: None, + log: true, } } + pub fn with_log(mut self, log: bool) -> Self { + self.log = log; + self + } + + pub fn without_start_barrier(mut self) -> Self { + self.start_data_barrier = None; + self + } + + pub fn with_finish_barrier(mut self) -> Self { + let barrier = Arc::new(( + // wait for all streams and the input + Barrier::new(self.data.len() + 1), + AtomicUsize::new(0), + )); + + self.finish_barrier = Some(barrier); + self + } + /// wait until all the input streams and this function is ready pub async fn wait(&self) { - println!("BarrierExec::wait waiting on barrier"); - self.barrier.wait().await; - println!("BarrierExec::wait done waiting"); + let barrier = &self + .start_data_barrier + .as_ref() + .expect("Must only be called when having a start barrier"); + if self.log { + println!("BarrierExec::wait waiting on barrier"); + } + barrier.wait().await; + if self.log { + println!("BarrierExec::wait done waiting"); + } + } + + pub async fn wait_finish(&self) { + let (barrier, _) = &self + .finish_barrier + .as_deref() + .expect("Must only be called when having a finish barrier"); + + if self.log { + println!("BarrierExec::wait_finish waiting on barrier"); + } + barrier.wait().await; + if self.log { + println!("BarrierExec::wait_finish done waiting"); + } + } + + /// Return true if the finish barrier has been reached in all partitions + pub fn is_finish_barrier_reached(&self) -> bool { + let (_, reached_finish) = self + .finish_barrier + .as_deref() + .expect("Must only be called when having finish barrier"); + + reached_finish.load(Ordering::Relaxed) == self.data.len() } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -391,17 +453,32 @@ impl ExecutionPlan for BarrierExec { // task simply sends data in order after barrier is reached let data = self.data[partition].clone(); - let b = Arc::clone(&self.barrier); + let start_barrier = self.start_data_barrier.as_ref().map(Arc::clone); + let finish_barrier = self.finish_barrier.as_ref().map(Arc::clone); + let log = self.log; let tx = builder.tx(); builder.spawn(async move { - println!("Partition {partition} waiting on barrier"); - b.wait().await; + if let Some(barrier) = start_barrier { + if log { + println!("Partition {partition} waiting on barrier"); + } + barrier.wait().await; + } for batch in data { - println!("Partition {partition} sending batch"); + if log { + println!("Partition {partition} sending batch"); + } if let Err(e) = tx.send(Ok(batch)).await { println!("ERROR batch via barrier stream stream: {e}"); } } + if let Some((barrier, reached_finish)) = finish_barrier.as_deref() { + if log { + println!("Partition {partition} waiting on finish barrier"); + } + reached_finish.fetch_add(1, Ordering::Relaxed); + barrier.wait().await; + } Ok(()) }); From 9a67de58c027e6057aa37327ae4d0192d5c45fc5 Mon Sep 17 00:00:00 2001 From: Haresh Khanna Date: Wed, 4 Mar 2026 20:13:36 +0000 Subject: [PATCH 2/8] [branch-52] Fix Arrow Spill Underrun (#20159) (#20684) ## Which issue does this PR close? - Related to #20681 - Backport of https://github.com/apache/datafusion/pull/20159 ## Rationale for this change This adjusts the way that the spill channel works. Currently we have a spill writer & reader pairing which uses a mutex to coordindate when a file is ready to be read. What happens is, that because we were using a `spawn_buffered` call, the read task would race ahead trying to read a file which is yet to be written out completely. Alongside this, we need to flush each write to the file, as there is a chance that another thread may see stale data. ## What changes are included in this PR? Adds a flush on write, and converts the read task to not buffer reads. ## Are these changes tested? I haven't written a test, but I have been running the example in the attached issue. While it now fails with allocation errors, the original error goes away. ## Are there any user-facing changes? Nope Co-authored-by: Peter L --- .../src/spill/in_progress_spill_file.rs | 7 +++++++ datafusion/physical-plan/src/spill/mod.rs | 5 +++++ datafusion/physical-plan/src/spill/spill_manager.rs | 13 +++++++++++++ datafusion/physical-plan/src/spill/spill_pool.rs | 8 +++++++- 4 files changed, 32 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs index d2acf4993b857..b9ff6b2f3b655 100644 --- a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs +++ b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs @@ -88,6 +88,13 @@ impl InProgressSpillFile { Ok(()) } + pub fn flush(&mut self) -> Result<()> { + if let Some(writer) = &mut self.writer { + writer.flush()?; + } + Ok(()) + } + /// Returns a reference to the in-progress file, if it exists. /// This can be used to get the file path for creating readers before the file is finished. pub fn file(&self) -> Option<&RefCountedTempFile> { diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index 78dea99ac820c..3c4ee065c3151 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -310,6 +310,11 @@ impl IPCStreamWriter { Ok((delta_num_rows, delta_num_bytes)) } + pub fn flush(&mut self) -> Result<()> { + self.writer.flush()?; + Ok(()) + } + /// Finish the writer pub fn finish(&mut self) -> Result<()> { self.writer.finish().map_err(Into::into) diff --git a/datafusion/physical-plan/src/spill/spill_manager.rs b/datafusion/physical-plan/src/spill/spill_manager.rs index 89b0276206774..6d931112ad888 100644 --- a/datafusion/physical-plan/src/spill/spill_manager.rs +++ b/datafusion/physical-plan/src/spill/spill_manager.rs @@ -188,6 +188,19 @@ impl SpillManager { Ok(spawn_buffered(stream, self.batch_read_buffer_capacity)) } + + /// Same as `read_spill_as_stream`, but without buffering. + pub fn read_spill_as_stream_unbuffered( + &self, + spill_file_path: RefCountedTempFile, + max_record_batch_memory: Option, + ) -> Result { + Ok(Box::pin(cooperative(SpillReaderStream::new( + Arc::clone(&self.schema), + spill_file_path, + max_record_batch_memory, + )))) + } } pub(crate) trait GetSlicedSize { diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index e3b547b5731f3..e8eea360da8cc 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -194,6 +194,8 @@ impl SpillPoolWriter { // Append the batch if let Some(ref mut writer) = file_shared.writer { writer.append_batch(batch)?; + // make sure we flush the writer for readers + writer.flush()?; file_shared.batches_written += 1; file_shared.estimated_size += batch_size; } @@ -535,7 +537,11 @@ impl Stream for SpillFile { // Step 2: Lazy-create reader stream if needed if self.reader.is_none() && should_read { if let Some(file) = file { - match self.spill_manager.read_spill_as_stream(file, None) { + // we want this unbuffered because files are actively being written to + match self + .spill_manager + .read_spill_as_stream_unbuffered(file, None) + { Ok(stream) => { self.reader = Some(SpillFileReader { stream, From 72ea8ec086e59220f6b255ea565e710990ad7967 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 4 Mar 2026 17:43:21 -0500 Subject: [PATCH 3/8] [branch-52] Fix constant value from stats (#20042) (#20709) - Part of https://github.com/apache/datafusion/issues/20681 - Closes https://github.com/apache/datafusion/issues/20041 on branch-52 This PR: - Backports https://github.com/apache/datafusion/pull/20042 from @gabotechs to the `branch-52` line Co-authored-by: Gabriel <45515538+gabotechs@users.noreply.github.com> --- .../src/datasource/physical_plan/parquet.rs | 48 +++++++++++++++++-- datafusion/datasource-parquet/src/opener.rs | 4 ++ 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/datafusion/core/src/datasource/physical_plan/parquet.rs b/datafusion/core/src/datasource/physical_plan/parquet.rs index 4703b55ecc0de..dde40cc0603b4 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet.rs +++ b/datafusion/core/src/datasource/physical_plan/parquet.rs @@ -38,10 +38,10 @@ mod tests { use crate::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; use crate::test::object_store::local_unpartitioned_file; use arrow::array::{ - ArrayRef, AsArray, Date64Array, Int8Array, Int32Array, Int64Array, StringArray, - StringViewArray, StructArray, TimestampNanosecondArray, + ArrayRef, AsArray, Date64Array, DictionaryArray, Int8Array, Int32Array, + Int64Array, StringArray, StringViewArray, StructArray, TimestampNanosecondArray, }; - use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaBuilder}; + use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaBuilder, UInt16Type}; use arrow::record_batch::RecordBatch; use arrow::util::pretty::pretty_format_batches; use arrow_schema::{SchemaRef, TimeUnit}; @@ -2249,6 +2249,48 @@ mod tests { Ok(()) } + /// Tests that constant dictionary columns (where min == max in statistics) + /// are correctly handled. This reproduced a bug where the constant value + /// from statistics had type Utf8 but the schema expected Dictionary. + #[tokio::test] + async fn test_constant_dictionary_column_parquet() -> Result<()> { + let tmp_dir = TempDir::new()?; + let path = tmp_dir.path().to_str().unwrap().to_string() + "/test.parquet"; + + // Write parquet with dictionary column where all values are the same + let schema = Arc::new(Schema::new(vec![Field::new( + "status", + DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)), + false, + )])); + let status: DictionaryArray = + vec!["active", "active"].into_iter().collect(); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(status)])?; + let file = File::create(&path)?; + let props = WriterProperties::builder() + .set_statistics_enabled(parquet::file::properties::EnabledStatistics::Page) + .build(); + let mut writer = ArrowWriter::try_new(file, schema, Some(props))?; + writer.write(&batch)?; + writer.close()?; + + // Query the constant dictionary column + let ctx = SessionContext::new(); + ctx.register_parquet("t", &path, ParquetReadOptions::default()) + .await?; + let result = ctx.sql("SELECT status FROM t").await?.collect().await?; + + insta::assert_snapshot!(batches_to_string(&result),@r" + +--------+ + | status | + +--------+ + | active | + | active | + +--------+ + "); + Ok(()) + } + fn write_file(file: &String) { let struct_fields = Fields::from(vec![ Field::new("id", DataType::Int64, false), diff --git a/datafusion/datasource-parquet/src/opener.rs b/datafusion/datasource-parquet/src/opener.rs index 83bdf79c8fcc0..719a3afc764f8 100644 --- a/datafusion/datasource-parquet/src/opener.rs +++ b/datafusion/datasource-parquet/src/opener.rs @@ -696,6 +696,10 @@ fn constant_value_from_stats( && !min.is_null() && matches!(column_stats.null_count, Precision::Exact(0)) { + // Cast to the expected data type if needed (e.g., Utf8 -> Dictionary) + if min.data_type() != *data_type { + return min.cast_to(data_type).ok(); + } return Some(min.clone()); } From d317d00b886bbf11cb489e4c4bdc2280b3ca9e07 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 4 Mar 2026 19:01:19 -0500 Subject: [PATCH 4/8] [branch-52] fix: `HashJoin` panic with String dictionary keys (don't flatten keys) (#20505) (#20708) - Part of https://github.com/apache/datafusion/issues/20681 - Closes https://github.com/apache/datafusion/issues/20696 on branch-52 This PR: - Backports https://github.com/apache/datafusion/pull/20505 from to the [branch-52] line --- .../src/joins/hash_join/inlist_builder.rs | 39 +++---- .../test_files/parquet_filter_pushdown.slt | 100 ++++++++++++++++++ 2 files changed, 113 insertions(+), 26 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/inlist_builder.rs b/datafusion/physical-plan/src/joins/hash_join/inlist_builder.rs index 9bf59d9e333d6..0ca338265ecc6 100644 --- a/datafusion/physical-plan/src/joins/hash_join/inlist_builder.rs +++ b/datafusion/physical-plan/src/joins/hash_join/inlist_builder.rs @@ -20,7 +20,6 @@ use std::sync::Arc; use arrow::array::{ArrayRef, StructArray}; -use arrow::compute::cast; use arrow::datatypes::{Field, FieldRef, Fields}; use arrow_schema::DataType; use datafusion_common::Result; @@ -33,19 +32,6 @@ pub(super) fn build_struct_fields(data_types: &[DataType]) -> Result { .collect() } -/// Casts dictionary-encoded arrays to their underlying value type, preserving row count. -/// Non-dictionary arrays are returned as-is. -fn flatten_dictionary_array(array: &ArrayRef) -> Result { - match array.data_type() { - DataType::Dictionary(_, value_type) => { - let casted = cast(array, value_type)?; - // Recursively flatten in case of nested dictionaries - flatten_dictionary_array(&casted) - } - _ => Ok(Arc::clone(array)), - } -} - /// Builds InList values from join key column arrays. /// /// If `join_key_arrays` is: @@ -65,20 +51,14 @@ fn flatten_dictionary_array(array: &ArrayRef) -> Result { pub(super) fn build_struct_inlist_values( join_key_arrays: &[ArrayRef], ) -> Result> { - // Flatten any dictionary-encoded arrays - let flattened_arrays: Vec = join_key_arrays - .iter() - .map(flatten_dictionary_array) - .collect::>>()?; - // Build the source array/struct - let source_array: ArrayRef = if flattened_arrays.len() == 1 { + let source_array: ArrayRef = if join_key_arrays.len() == 1 { // Single column: use directly - Arc::clone(&flattened_arrays[0]) + Arc::clone(&join_key_arrays[0]) } else { // Multi-column: build StructArray once from all columns let fields = build_struct_fields( - &flattened_arrays + &join_key_arrays .iter() .map(|arr| arr.data_type().clone()) .collect::>(), @@ -88,7 +68,7 @@ pub(super) fn build_struct_inlist_values( let arrays_with_fields: Vec<(FieldRef, ArrayRef)> = fields .iter() .cloned() - .zip(flattened_arrays.iter().cloned()) + .zip(join_key_arrays.iter().cloned()) .collect(); Arc::new(StructArray::from(arrays_with_fields)) @@ -152,7 +132,14 @@ mod tests { assert_eq!( *result.data_type(), DataType::Struct( - build_struct_fields(&[DataType::Utf8, DataType::Int32]).unwrap() + build_struct_fields(&[ + DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Utf8) + ), + DataType::Int32 + ]) + .unwrap() ) ); } @@ -168,6 +155,6 @@ mod tests { .unwrap(); assert_eq!(result.len(), 3); - assert_eq!(*result.data_type(), DataType::Utf8); + assert_eq!(result.data_type(), dict_array.data_type()); } } diff --git a/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt b/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt index 8bb79d576990e..5e643273baed4 100644 --- a/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt +++ b/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt @@ -563,3 +563,103 @@ ORDER BY start_timestamp, trace_id LIMIT 1; ---- 2024-10-01T00:00:00 + + +statement ok +set datafusion.execution.parquet.pushdown_filters = false; + +# Regression test for https://github.com/apache/datafusion/issues/20696 +# Multi-column INNER JOIN with dictionary fails +# when parquet pushdown filters are enabled. + +statement ok +COPY ( + SELECT + to_timestamp_nanos(time_ns) AS time, + arrow_cast(state, 'Dictionary(Int32, Utf8)') AS state, + arrow_cast(city, 'Dictionary(Int32, Utf8)') AS city, + temp + FROM ( + VALUES + (200, 'CA', 'LA', 90.0), + (250, 'MA', 'Boston', 72.4), + (100, 'MA', 'Boston', 70.4), + (350, 'CA', 'LA', 90.0) + ) AS t(time_ns, state, city, temp) +) +TO 'test_files/scratch/parquet_filter_pushdown/issue_20696/h2o/data.parquet'; + +statement ok +COPY ( + SELECT + to_timestamp_nanos(time_ns) AS time, + arrow_cast(state, 'Dictionary(Int32, Utf8)') AS state, + arrow_cast(city, 'Dictionary(Int32, Utf8)') AS city, + temp, + reading + FROM ( + VALUES + (250, 'MA', 'Boston', 53.4, 51.0), + (100, 'MA', 'Boston', 50.4, 50.0) + ) AS t(time_ns, state, city, temp, reading) +) +TO 'test_files/scratch/parquet_filter_pushdown/issue_20696/o2/data.parquet'; + +statement ok +CREATE EXTERNAL TABLE h2o_parquet_20696 STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_filter_pushdown/issue_20696/h2o/'; + +statement ok +CREATE EXTERNAL TABLE o2_parquet_20696 STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_filter_pushdown/issue_20696/o2/'; + +# Query should work both with and without filters +statement ok +set datafusion.execution.parquet.pushdown_filters = false; + +query RRR +SELECT + h2o_parquet_20696.temp AS h2o_temp, + o2_parquet_20696.temp AS o2_temp, + o2_parquet_20696.reading +FROM h2o_parquet_20696 +INNER JOIN o2_parquet_20696 + ON h2o_parquet_20696.time = o2_parquet_20696.time + AND h2o_parquet_20696.state = o2_parquet_20696.state + AND h2o_parquet_20696.city = o2_parquet_20696.city +WHERE h2o_parquet_20696.time >= '1970-01-01T00:00:00.000000050Z' + AND h2o_parquet_20696.time <= '1970-01-01T00:00:00.000000300Z'; +---- +72.4 53.4 51 +70.4 50.4 50 + + +statement ok +set datafusion.execution.parquet.pushdown_filters = true; + +query RRR +SELECT + h2o_parquet_20696.temp AS h2o_temp, + o2_parquet_20696.temp AS o2_temp, + o2_parquet_20696.reading +FROM h2o_parquet_20696 +INNER JOIN o2_parquet_20696 + ON h2o_parquet_20696.time = o2_parquet_20696.time + AND h2o_parquet_20696.state = o2_parquet_20696.state + AND h2o_parquet_20696.city = o2_parquet_20696.city +WHERE h2o_parquet_20696.time >= '1970-01-01T00:00:00.000000050Z' + AND h2o_parquet_20696.time <= '1970-01-01T00:00:00.000000300Z'; +---- +72.4 53.4 51 +70.4 50.4 50 + +# Cleanup +statement ok +DROP TABLE h2o_parquet_20696; + +statement ok +DROP TABLE o2_parquet_20696; + +# Cleanup settings +statement ok +set datafusion.execution.parquet.pushdown_filters = false; From afc1c72a15bdd31e15a7e354e86a505be7882f08 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 5 Mar 2026 12:22:47 -0500 Subject: [PATCH 5/8] [branch-52] FFI_TableOptions are using default values only (#20705) ## Which issue does this PR close? - Addresses part of https://github.com/apache/datafusion/issues/20704 ## Rationale for this change FFI_TableOptions fails with a warning that is getting swallowed in the unit tests. ## What changes are included in this PR? Correctly check format for table options. ## Are these changes tested? Unit test now passes with patch described in topic. ## Are there any user-facing changes? None, internal only. --- datafusion/ffi/src/session/mod.rs | 107 ++++++++++++++++++++++++++---- 1 file changed, 95 insertions(+), 12 deletions(-) diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index aa910abb9149a..6b8664a437495 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -26,7 +26,7 @@ use arrow_schema::SchemaRef; use arrow_schema::ffi::FFI_ArrowSchema; use async_ffi::{FfiFuture, FutureExt}; use async_trait::async_trait; -use datafusion_common::config::{ConfigOptions, TableOptions}; +use datafusion_common::config::{ConfigFileType, ConfigOptions, TableOptions}; use datafusion_common::{DFSchema, DataFusionError}; use datafusion_execution::TaskContext; use datafusion_execution::config::SessionConfig; @@ -240,12 +240,30 @@ unsafe extern "C" fn window_functions_fn_wrapper( .collect() } -fn table_options_to_rhash(options: &TableOptions) -> RHashMap { - options +fn table_options_to_rhash(mut options: TableOptions) -> RHashMap { + // It is important that we mutate options here and set current format + // to None so that when we call `entries()` we get ALL format entries. + // We will pass current_format as a special case and strip it on the + // other side of the boundary. + let current_format = options.current_format.take(); + let mut options: HashMap = options .entries() .into_iter() .filter_map(|entry| entry.value.map(|v| (entry.key.into(), v.into()))) - .collect() + .collect(); + if let Some(current_format) = current_format { + options.insert( + "datafusion_ffi.table_current_format".into(), + match current_format { + ConfigFileType::JSON => "json", + ConfigFileType::PARQUET => "parquet", + ConfigFileType::CSV => "csv", + } + .into(), + ); + } + + options.into() } unsafe extern "C" fn table_options_fn_wrapper( @@ -253,7 +271,7 @@ unsafe extern "C" fn table_options_fn_wrapper( ) -> RHashMap { let session = session.inner(); let table_options = session.table_options(); - table_options_to_rhash(table_options) + table_options_to_rhash(table_options.clone()) } unsafe extern "C" fn default_table_options_fn_wrapper( @@ -262,7 +280,7 @@ unsafe extern "C" fn default_table_options_fn_wrapper( let session = session.inner(); let table_options = session.default_table_options(); - table_options_to_rhash(&table_options) + table_options_to_rhash(table_options) } unsafe extern "C" fn task_ctx_fn_wrapper(session: &FFI_SessionRef) -> FFI_TaskContext { @@ -438,15 +456,70 @@ impl Clone for FFI_SessionRef { } fn table_options_from_rhashmap(options: RHashMap) -> TableOptions { - let options = options + let mut options: HashMap = options .into_iter() .map(|kv_pair| (kv_pair.0.into_string(), kv_pair.1.into_string())) .collect(); + let current_format = options.remove("datafusion_ffi.table_current_format"); + + let mut table_options = TableOptions::default(); + let formats = [ + ConfigFileType::CSV, + ConfigFileType::JSON, + ConfigFileType::PARQUET, + ]; + for format in formats { + // It is imperative that if new enum variants are added below that they be + // included in the formats list above and in the extension check below. + let format_name = match &format { + ConfigFileType::CSV => "csv", + ConfigFileType::PARQUET => "parquet", + ConfigFileType::JSON => "json", + }; + let format_options: HashMap = options + .iter() + .filter_map(|(k, v)| { + let (prefix, key) = k.split_once(".")?; + if prefix == format_name { + Some((format!("format.{key}"), v.to_owned())) + } else { + None + } + }) + .collect(); + if !format_options.is_empty() { + table_options.current_format = Some(format.clone()); + table_options + .alter_with_string_hash_map(&format_options) + .unwrap_or_else(|err| log::warn!("Error parsing table options: {err}")); + } + } + + let extension_options: HashMap = options + .iter() + .filter_map(|(k, v)| { + let (prefix, _) = k.split_once(".")?; + if !["json", "parquet", "csv"].contains(&prefix) { + Some((k.to_owned(), v.to_owned())) + } else { + None + } + }) + .collect(); + if !extension_options.is_empty() { + table_options + .alter_with_string_hash_map(&extension_options) + .unwrap_or_else(|err| log::warn!("Error parsing table options: {err}")); + } - TableOptions::from_string_hash_map(&options).unwrap_or_else(|err| { - log::warn!("Error parsing default table options: {err}"); - TableOptions::default() - }) + table_options.current_format = + current_format.and_then(|format| match format.as_str() { + "csv" => Some(ConfigFileType::CSV), + "parquet" => Some(ConfigFileType::PARQUET), + "json" => Some(ConfigFileType::JSON), + _ => None, + }); + table_options } #[async_trait] @@ -556,6 +629,7 @@ mod tests { use std::sync::Arc; use arrow_schema::{DataType, Field, Schema}; + use datafusion::execution::SessionStateBuilder; use datafusion_common::DataFusionError; use datafusion_expr::col; use datafusion_expr::registry::FunctionRegistry; @@ -566,7 +640,16 @@ mod tests { #[tokio::test] async fn test_ffi_session() -> Result<(), DataFusionError> { let (ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); - let state = ctx.state(); + let mut table_options = TableOptions::default(); + table_options.csv.has_header = Some(true); + table_options.json.schema_infer_max_rec = Some(10); + table_options.parquet.global.coerce_int96 = Some("123456789".into()); + table_options.current_format = Some(ConfigFileType::JSON); + + let state = SessionStateBuilder::new_from_existing(ctx.state()) + .with_table_options(table_options) + .build(); + let logical_codec = FFI_LogicalExtensionCodec::new( Arc::new(DefaultLogicalExtensionCodec {}), None, From 9797095e152749721bec07c0944fe664acaa0849 Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Thu, 5 Mar 2026 17:04:12 -0500 Subject: [PATCH 6/8] [branch-52] perf: sort replace free()->try_grow() pattern with try_resize() to reduce memory pool interactions (#20732) Backport #20729 to `branch-52`. --- datafusion/physical-plan/src/sorts/sort.rs | 50 +++++++++------------- 1 file changed, 20 insertions(+), 30 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 3e8fdf1f3ed7e..475738cca3f05 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -728,37 +728,27 @@ impl ExternalSorter { let sorted_batches = sort_batch_chunked(&batch, &expressions, batch_size)?; drop(batch); - // Free the old reservation and grow it to match the actual sorted output size - reservation.free(); + // Resize the reservation to match the actual sorted output size. + // Using try_resize avoids a release-then-reacquire cycle, which + // matters for MemoryPool implementations where grow/shrink have + // non-trivial cost (e.g. JNI calls in Comet). + let total_sorted_size: usize = sorted_batches + .iter() + .map(get_record_batch_memory_size) + .sum(); + reservation + .try_resize(total_sorted_size) + .map_err(Self::err_with_oom_context)?; - Result::<_, DataFusionError>::Ok((schema, sorted_batches, reservation)) - }) - .then({ - move |batches| async move { - match batches { - Ok((schema, sorted_batches, mut reservation)) => { - // Calculate the total size of sorted batches - let total_sorted_size: usize = sorted_batches - .iter() - .map(get_record_batch_memory_size) - .sum(); - reservation - .try_grow(total_sorted_size) - .map_err(Self::err_with_oom_context)?; - - // Wrap in ReservationStream to hold the reservation - Ok(Box::pin(ReservationStream::new( - Arc::clone(&schema), - Box::pin(RecordBatchStreamAdapter::new( - schema, - futures::stream::iter(sorted_batches.into_iter().map(Ok)), - )), - reservation, - )) as SendableRecordBatchStream) - } - Err(e) => Err(e), - } - } + // Wrap in ReservationStream to hold the reservation + Result::<_, DataFusionError>::Ok(Box::pin(ReservationStream::new( + Arc::clone(&schema), + Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(sorted_batches.into_iter().map(Ok)), + )), + reservation, + )) as SendableRecordBatchStream) }) .try_flatten() .map(move |batch| match batch { From 1bd7082b798d0d55c1e90c7be1d7e3dba057c288 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sun, 8 Mar 2026 11:14:32 -0400 Subject: [PATCH 7/8] [branch-52] Fix repartition from dropping data when spilling (#20672) (#20777) - Part of https://github.com/apache/datafusion/issues/20681 - Closes https://github.com/apache/datafusion/issues/20683 on branch-52 This PR: - Backports https://github.com/apache/datafusion/pull/20672 from @xanderbailey to the `branch-52` line Co-authored-by: Xander --- .../physical-plan/src/spill/spill_pool.rs | 102 +++++++++++++++++- 1 file changed, 101 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index e8eea360da8cc..36920c7fd0601 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -61,6 +61,10 @@ struct SpillPoolShared { /// Writer's reference to the current file (shared by all cloned writers). /// Has its own lock to allow I/O without blocking queue access. current_write_file: Option>>, + /// Number of active writer clones. Only when this reaches zero should + /// `writer_dropped` be set to true. This prevents premature EOF signaling + /// when one writer clone is dropped while others are still active. + active_writer_count: usize, } impl SpillPoolShared { @@ -72,6 +76,7 @@ impl SpillPoolShared { waker: None, writer_dropped: false, current_write_file: None, + active_writer_count: 1, } } @@ -97,7 +102,6 @@ impl SpillPoolShared { /// The writer automatically manages file rotation based on the `max_file_size_bytes` /// configured in [`channel`]. When the last writer clone is dropped, it finalizes the /// current file so readers can access all written data. -#[derive(Clone)] pub struct SpillPoolWriter { /// Maximum size in bytes before rotating to a new file. /// Typically set from configuration `datafusion.execution.max_spill_file_size_bytes`. @@ -106,6 +110,18 @@ pub struct SpillPoolWriter { shared: Arc>, } +impl Clone for SpillPoolWriter { + fn clone(&self) -> Self { + // Increment the active writer count so that `writer_dropped` is only + // set to true when the *last* clone is dropped. + self.shared.lock().active_writer_count += 1; + Self { + max_file_size_bytes: self.max_file_size_bytes, + shared: Arc::clone(&self.shared), + } + } +} + impl SpillPoolWriter { /// Spills a batch to the pool, rotating files when necessary. /// @@ -233,6 +249,15 @@ impl Drop for SpillPoolWriter { fn drop(&mut self) { let mut shared = self.shared.lock(); + shared.active_writer_count -= 1; + let is_last_writer = shared.active_writer_count == 0; + + if !is_last_writer { + // Other writer clones are still active; do not finalize or + // signal EOF to readers. + return; + } + // Finalize the current file when the last writer is dropped if let Some(current_file) = shared.current_write_file.take() { // Release shared lock before locking file @@ -1343,6 +1368,81 @@ mod tests { Ok(()) } + /// Verifies that the reader stays alive as long as any writer clone exists. + /// + /// `SpillPoolWriter` is `Clone`, and in non-preserve-order repartitioning + /// mode multiple input partition tasks share clones of the same writer. + /// The reader must not see EOF until **all** clones have been dropped, + /// even if the queue is temporarily empty between writes from different + /// clones. + /// + /// The test sequence is: + /// + /// 1. writer1 writes a batch, then is dropped. + /// 2. The reader consumes that batch (queue is now empty). + /// 3. writer2 (still alive) writes a batch. + /// 4. The reader must see that batch. + /// 5. EOF is only signalled after writer2 is also dropped. + #[tokio::test] + async fn test_clone_drop_does_not_signal_eof_prematurely() -> Result<()> { + let (writer1, mut reader) = create_spill_channel(1024 * 1024); + let writer2 = writer1.clone(); + + // Synchronization: tell writer2 when it may proceed. + let (proceed_tx, proceed_rx) = tokio::sync::oneshot::channel::<()>(); + + // Spawn writer2 — it waits for the signal before writing. + let writer2_handle = SpawnedTask::spawn(async move { + proceed_rx.await.unwrap(); + writer2.push_batch(&create_test_batch(10, 10)).unwrap(); + // writer2 is dropped here (last clone → true EOF) + }); + + // Writer1 writes one batch, then drops. + writer1.push_batch(&create_test_batch(0, 10))?; + drop(writer1); + + // Read writer1's batch. + let batch1 = reader.next().await.unwrap()?; + assert_eq!(batch1.num_rows(), 10); + let col = batch1 + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(col.value(0), 0); + + // Signal writer2 to write its batch. It will execute when the + // current task yields (i.e. when reader.next() returns Pending). + proceed_tx.send(()).unwrap(); + + // The reader should wait (Pending) for writer2's data, not EOF. + let batch2 = + tokio::time::timeout(std::time::Duration::from_secs(5), reader.next()) + .await + .expect("Reader timed out — should not hang"); + + assert!( + batch2.is_some(), + "Reader must not return EOF while a writer clone is still alive" + ); + let batch2 = batch2.unwrap()?; + assert_eq!(batch2.num_rows(), 10); + let col = batch2 + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(col.value(0), 10); + + writer2_handle.await.unwrap(); + + // All writers dropped — reader should see real EOF now. + assert!(reader.next().await.is_none()); + + Ok(()) + } + #[tokio::test] async fn test_disk_usage_decreases_as_files_consumed() -> Result<()> { use datafusion_execution::runtime_env::RuntimeEnvBuilder; From 28d012a41a3017b5f682ef6b01468a7ff9a48fb7 Mon Sep 17 00:00:00 2001 From: Haresh Khanna Date: Sun, 8 Mar 2026 15:48:23 +0000 Subject: [PATCH 8/8] [branch-52] Bump to 52.3.0 and changelog (#20790) ## Which issue does this PR close? - Related to #20681 ## Rationale for this change ## What changes are included in this PR? Bumps to 52.3.0 and changelog ## Are these changes tested? ## Are there any user-facing changes? --------- Co-authored-by: Haresh Khanna Co-authored-by: Tim Saucer --- Cargo.lock | 84 +++++++++++++++---------------- Cargo.toml | 76 ++++++++++++++-------------- dev/changelog/52.3.0.md | 50 ++++++++++++++++++ docs/source/user-guide/configs.md | 2 +- 4 files changed, 131 insertions(+), 81 deletions(-) create mode 100644 dev/changelog/52.3.0.md diff --git a/Cargo.lock b/Cargo.lock index b48f4a777aca5..dc938853f3aad 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1767,7 +1767,7 @@ dependencies = [ [[package]] name = "datafusion" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "arrow-schema", @@ -1839,7 +1839,7 @@ dependencies = [ [[package]] name = "datafusion-benchmarks" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "clap", @@ -1864,7 +1864,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "async-trait", @@ -1887,7 +1887,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "async-trait", @@ -1909,7 +1909,7 @@ dependencies = [ [[package]] name = "datafusion-cli" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "async-trait", @@ -1940,7 +1940,7 @@ dependencies = [ [[package]] name = "datafusion-common" -version = "52.2.0" +version = "52.3.0" dependencies = [ "ahash", "apache-avro", @@ -1967,7 +1967,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "52.2.0" +version = "52.3.0" dependencies = [ "futures", "log", @@ -1976,7 +1976,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "async-compression", @@ -2011,7 +2011,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "arrow-ipc", @@ -2034,7 +2034,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-avro" -version = "52.2.0" +version = "52.3.0" dependencies = [ "apache-avro", "arrow", @@ -2053,7 +2053,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "async-trait", @@ -2074,7 +2074,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "async-trait", @@ -2094,7 +2094,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "async-trait", @@ -2123,11 +2123,11 @@ dependencies = [ [[package]] name = "datafusion-doc" -version = "52.2.0" +version = "52.3.0" [[package]] name = "datafusion-examples" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "arrow-flight", @@ -2166,7 +2166,7 @@ dependencies = [ [[package]] name = "datafusion-execution" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "async-trait", @@ -2187,7 +2187,7 @@ dependencies = [ [[package]] name = "datafusion-expr" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "async-trait", @@ -2211,7 +2211,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "datafusion-common", @@ -2222,7 +2222,7 @@ dependencies = [ [[package]] name = "datafusion-ffi" -version = "52.2.0" +version = "52.3.0" dependencies = [ "abi_stable", "arrow", @@ -2256,7 +2256,7 @@ dependencies = [ [[package]] name = "datafusion-functions" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "arrow-buffer", @@ -2289,7 +2289,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "52.2.0" +version = "52.3.0" dependencies = [ "ahash", "arrow", @@ -2310,7 +2310,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" -version = "52.2.0" +version = "52.3.0" dependencies = [ "ahash", "arrow", @@ -2323,7 +2323,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "arrow-ord", @@ -2346,7 +2346,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "async-trait", @@ -2360,7 +2360,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "datafusion-common", @@ -2376,7 +2376,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "52.2.0" +version = "52.3.0" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -2384,7 +2384,7 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "52.2.0" +version = "52.3.0" dependencies = [ "datafusion-doc", "quote", @@ -2393,7 +2393,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "async-trait", @@ -2420,7 +2420,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "52.2.0" +version = "52.3.0" dependencies = [ "ahash", "arrow", @@ -2447,7 +2447,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "datafusion-common", @@ -2460,7 +2460,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" -version = "52.2.0" +version = "52.3.0" dependencies = [ "ahash", "arrow", @@ -2475,7 +2475,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "datafusion-common", @@ -2495,7 +2495,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" -version = "52.2.0" +version = "52.3.0" dependencies = [ "ahash", "arrow", @@ -2531,7 +2531,7 @@ dependencies = [ [[package]] name = "datafusion-proto" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "async-trait", @@ -2568,7 +2568,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "datafusion-common", @@ -2580,7 +2580,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "datafusion-common", @@ -2598,7 +2598,7 @@ dependencies = [ [[package]] name = "datafusion-session" -version = "52.2.0" +version = "52.3.0" dependencies = [ "async-trait", "datafusion-common", @@ -2610,7 +2610,7 @@ dependencies = [ [[package]] name = "datafusion-spark" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "bigdecimal", @@ -2632,7 +2632,7 @@ dependencies = [ [[package]] name = "datafusion-sql" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "bigdecimal", @@ -2658,7 +2658,7 @@ dependencies = [ [[package]] name = "datafusion-sqllogictest" -version = "52.2.0" +version = "52.3.0" dependencies = [ "arrow", "async-trait", @@ -2689,7 +2689,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" -version = "52.2.0" +version = "52.3.0" dependencies = [ "async-recursion", "async-trait", @@ -2710,7 +2710,7 @@ dependencies = [ [[package]] name = "datafusion-wasmtest" -version = "52.2.0" +version = "52.3.0" dependencies = [ "chrono", "console_error_panic_hook", diff --git a/Cargo.toml b/Cargo.toml index 2c0764c06b6bf..9a0766efcb364 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -79,7 +79,7 @@ repository = "https://github.com/apache/datafusion" # Define Minimum Supported Rust Version (MSRV) rust-version = "1.88.0" # Define DataFusion version -version = "52.2.0" +version = "52.3.0" [workspace.dependencies] # We turn off default-features for some dependencies here so the workspaces which inherit them can @@ -112,43 +112,43 @@ chrono = { version = "0.4.42", default-features = false } criterion = "0.8" ctor = "0.6.3" dashmap = "6.0.1" -datafusion = { path = "datafusion/core", version = "52.2.0", default-features = false } -datafusion-catalog = { path = "datafusion/catalog", version = "52.2.0" } -datafusion-catalog-listing = { path = "datafusion/catalog-listing", version = "52.2.0" } -datafusion-common = { path = "datafusion/common", version = "52.2.0", default-features = false } -datafusion-common-runtime = { path = "datafusion/common-runtime", version = "52.2.0" } -datafusion-datasource = { path = "datafusion/datasource", version = "52.2.0", default-features = false } -datafusion-datasource-arrow = { path = "datafusion/datasource-arrow", version = "52.2.0", default-features = false } -datafusion-datasource-avro = { path = "datafusion/datasource-avro", version = "52.2.0", default-features = false } -datafusion-datasource-csv = { path = "datafusion/datasource-csv", version = "52.2.0", default-features = false } -datafusion-datasource-json = { path = "datafusion/datasource-json", version = "52.2.0", default-features = false } -datafusion-datasource-parquet = { path = "datafusion/datasource-parquet", version = "52.2.0", default-features = false } -datafusion-doc = { path = "datafusion/doc", version = "52.2.0" } -datafusion-execution = { path = "datafusion/execution", version = "52.2.0", default-features = false } -datafusion-expr = { path = "datafusion/expr", version = "52.2.0", default-features = false } -datafusion-expr-common = { path = "datafusion/expr-common", version = "52.2.0" } -datafusion-ffi = { path = "datafusion/ffi", version = "52.2.0" } -datafusion-functions = { path = "datafusion/functions", version = "52.2.0" } -datafusion-functions-aggregate = { path = "datafusion/functions-aggregate", version = "52.2.0" } -datafusion-functions-aggregate-common = { path = "datafusion/functions-aggregate-common", version = "52.2.0" } -datafusion-functions-nested = { path = "datafusion/functions-nested", version = "52.2.0", default-features = false } -datafusion-functions-table = { path = "datafusion/functions-table", version = "52.2.0" } -datafusion-functions-window = { path = "datafusion/functions-window", version = "52.2.0" } -datafusion-functions-window-common = { path = "datafusion/functions-window-common", version = "52.2.0" } -datafusion-macros = { path = "datafusion/macros", version = "52.2.0" } -datafusion-optimizer = { path = "datafusion/optimizer", version = "52.2.0", default-features = false } -datafusion-physical-expr = { path = "datafusion/physical-expr", version = "52.2.0", default-features = false } -datafusion-physical-expr-adapter = { path = "datafusion/physical-expr-adapter", version = "52.2.0", default-features = false } -datafusion-physical-expr-common = { path = "datafusion/physical-expr-common", version = "52.2.0", default-features = false } -datafusion-physical-optimizer = { path = "datafusion/physical-optimizer", version = "52.2.0" } -datafusion-physical-plan = { path = "datafusion/physical-plan", version = "52.2.0" } -datafusion-proto = { path = "datafusion/proto", version = "52.2.0" } -datafusion-proto-common = { path = "datafusion/proto-common", version = "52.2.0" } -datafusion-pruning = { path = "datafusion/pruning", version = "52.2.0" } -datafusion-session = { path = "datafusion/session", version = "52.2.0" } -datafusion-spark = { path = "datafusion/spark", version = "52.2.0" } -datafusion-sql = { path = "datafusion/sql", version = "52.2.0" } -datafusion-substrait = { path = "datafusion/substrait", version = "52.2.0" } +datafusion = { path = "datafusion/core", version = "52.3.0", default-features = false } +datafusion-catalog = { path = "datafusion/catalog", version = "52.3.0" } +datafusion-catalog-listing = { path = "datafusion/catalog-listing", version = "52.3.0" } +datafusion-common = { path = "datafusion/common", version = "52.3.0", default-features = false } +datafusion-common-runtime = { path = "datafusion/common-runtime", version = "52.3.0" } +datafusion-datasource = { path = "datafusion/datasource", version = "52.3.0", default-features = false } +datafusion-datasource-arrow = { path = "datafusion/datasource-arrow", version = "52.3.0", default-features = false } +datafusion-datasource-avro = { path = "datafusion/datasource-avro", version = "52.3.0", default-features = false } +datafusion-datasource-csv = { path = "datafusion/datasource-csv", version = "52.3.0", default-features = false } +datafusion-datasource-json = { path = "datafusion/datasource-json", version = "52.3.0", default-features = false } +datafusion-datasource-parquet = { path = "datafusion/datasource-parquet", version = "52.3.0", default-features = false } +datafusion-doc = { path = "datafusion/doc", version = "52.3.0" } +datafusion-execution = { path = "datafusion/execution", version = "52.3.0", default-features = false } +datafusion-expr = { path = "datafusion/expr", version = "52.3.0", default-features = false } +datafusion-expr-common = { path = "datafusion/expr-common", version = "52.3.0" } +datafusion-ffi = { path = "datafusion/ffi", version = "52.3.0" } +datafusion-functions = { path = "datafusion/functions", version = "52.3.0" } +datafusion-functions-aggregate = { path = "datafusion/functions-aggregate", version = "52.3.0" } +datafusion-functions-aggregate-common = { path = "datafusion/functions-aggregate-common", version = "52.3.0" } +datafusion-functions-nested = { path = "datafusion/functions-nested", version = "52.3.0", default-features = false } +datafusion-functions-table = { path = "datafusion/functions-table", version = "52.3.0" } +datafusion-functions-window = { path = "datafusion/functions-window", version = "52.3.0" } +datafusion-functions-window-common = { path = "datafusion/functions-window-common", version = "52.3.0" } +datafusion-macros = { path = "datafusion/macros", version = "52.3.0" } +datafusion-optimizer = { path = "datafusion/optimizer", version = "52.3.0", default-features = false } +datafusion-physical-expr = { path = "datafusion/physical-expr", version = "52.3.0", default-features = false } +datafusion-physical-expr-adapter = { path = "datafusion/physical-expr-adapter", version = "52.3.0", default-features = false } +datafusion-physical-expr-common = { path = "datafusion/physical-expr-common", version = "52.3.0", default-features = false } +datafusion-physical-optimizer = { path = "datafusion/physical-optimizer", version = "52.3.0" } +datafusion-physical-plan = { path = "datafusion/physical-plan", version = "52.3.0" } +datafusion-proto = { path = "datafusion/proto", version = "52.3.0" } +datafusion-proto-common = { path = "datafusion/proto-common", version = "52.3.0" } +datafusion-pruning = { path = "datafusion/pruning", version = "52.3.0" } +datafusion-session = { path = "datafusion/session", version = "52.3.0" } +datafusion-spark = { path = "datafusion/spark", version = "52.3.0" } +datafusion-sql = { path = "datafusion/sql", version = "52.3.0" } +datafusion-substrait = { path = "datafusion/substrait", version = "52.3.0" } doc-comment = "0.3" env_logger = "0.11" diff --git a/dev/changelog/52.3.0.md b/dev/changelog/52.3.0.md new file mode 100644 index 0000000000000..ed505b7fc2d0a --- /dev/null +++ b/dev/changelog/52.3.0.md @@ -0,0 +1,50 @@ + + +# Apache DataFusion 52.3.0 Changelog + +This release consists of 7 commits from 4 contributors. See credits at the end of this changelog for more information. + +See the [upgrade guide](https://datafusion.apache.org/library-user-guide/upgrading.html) for information on how to upgrade from previous versions. + +**Performance related:** + +- [branch-52] perf: sort replace free()->try_grow() pattern with try_resize() to reduce memory pool interactions [#20732](https://github.com/apache/datafusion/pull/20732) (mbutrovich) + +**Other:** + +- [branch-52] Backport fix: SortMergeJoin don't wait for all input before emitting #20482 [#20699](https://github.com/apache/datafusion/pull/20699) (mbutrovich) +- [branch-52] Fix Arrow Spill Underrun (#20159) [#20684](https://github.com/apache/datafusion/pull/20684) (hareshkh) +- [branch-52] Fix constant value from stats (#20042) [#20709](https://github.com/apache/datafusion/pull/20709) (alamb) +- [branch-52] fix: `HashJoin` panic with String dictionary keys (don't flatten keys) (#20505) [#20708](https://github.com/apache/datafusion/pull/20708) (alamb) +- [branch-52] FFI_TableOptions are using default values only [#20705](https://github.com/apache/datafusion/pull/20705) (timsaucer) +- [branch-52] Fix repartition from dropping data when spilling (#20672) [#20777](https://github.com/apache/datafusion/pull/20777) (alamb) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 3 Andrew Lamb + 2 Matt Butrovich + 1 Haresh Khanna + 1 Tim Saucer +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index b862d65401363..8423480d53305 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -99,7 +99,7 @@ The following configuration settings are available: | datafusion.execution.parquet.dictionary_page_size_limit | 1048576 | (writing) Sets best effort maximum dictionary page size, in bytes | | datafusion.execution.parquet.statistics_enabled | page | (writing) Sets if statistics are enabled for any column Valid values are: "none", "chunk", and "page" These values are not case sensitive. If NULL, uses default parquet writer setting | | datafusion.execution.parquet.max_row_group_size | 1048576 | (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. | -| datafusion.execution.parquet.created_by | datafusion version 52.2.0 | (writing) Sets "created by" property | +| datafusion.execution.parquet.created_by | datafusion version 52.3.0 | (writing) Sets "created by" property | | datafusion.execution.parquet.column_index_truncate_length | 64 | (writing) Sets column index truncate length | | datafusion.execution.parquet.statistics_truncate_length | 64 | (writing) Sets statistics truncate length. If NULL, uses default parquet writer setting | | datafusion.execution.parquet.data_page_row_count_limit | 20000 | (writing) Sets best effort maximum number of rows in data page |