Merge branch 'main' into optimize_count_distinct

Author: coderfender

datafusion/functions/src/string/common.rs

@@ -19,7 +19,7 @@
 
 use std::sync::Arc;
 
-use crate::strings::make_and_append_view;
+use crate::strings::append_view;
 use arrow::array::{
     Array, ArrayRef, GenericStringArray, GenericStringBuilder, NullBufferBuilder,
     OffsetSizeTrait, StringViewArray, StringViewBuilder, new_null_array,
@@ -152,13 +152,8 @@ fn string_view_trim<Tr: Trimmer>(args: &[ArrayRef]) -> Result<ArrayRef> {
             {
                 if let Some(src_str) = src_str_opt {
                     let (trimmed, offset) = Tr::trim_ascii_char(src_str, b' ');
-                    make_and_append_view(
-                        &mut views_buf,
-                        &mut null_builder,
-                        raw_view,
-                        trimmed,
-                        offset,
-                    );
+                    append_view(&mut views_buf, raw_view, trimmed, offset);
+                    null_builder.append_non_null();
                 } else {
                     null_builder.append_null();
                     views_buf.push(0);
@@ -204,13 +199,8 @@ fn string_view_trim<Tr: Trimmer>(args: &[ArrayRef]) -> Result<ArrayRef> {
                         pattern.clear();
                         pattern.extend(characters.chars());
                         let (trimmed, offset) = Tr::trim(src_str, &pattern);
-                        make_and_append_view(
-                            &mut views_buf,
-                            &mut null_builder,
-                            raw_view,
-                            trimmed,
-                            offset,
-                        );
+                        append_view(&mut views_buf, raw_view, trimmed, offset);
+                        null_builder.append_non_null();
                     } else {
                         null_builder.append_null();
                         views_buf.push(0);
@@ -261,7 +251,8 @@ fn trim_and_append_view<Tr: Trimmer>(
 ) {
     if let Some(src_str) = src_str_opt {
         let (trimmed, offset) = Tr::trim(src_str, pattern);
-        make_and_append_view(views_buf, null_builder, original_view, trimmed, offset);
+        append_view(views_buf, original_view, trimmed, offset);
+        null_builder.append_non_null();
     } else {
         null_builder.append_null();
         views_buf.push(0);

datafusion/functions/src/strings.rs

@@ -21,7 +21,7 @@ use datafusion_common::{Result, exec_datafusion_err, internal_err};
 
 use arrow::array::{
     Array, ArrayAccessor, ArrayDataBuilder, BinaryArray, ByteView, LargeStringArray,
-    NullBufferBuilder, StringArray, StringViewArray, StringViewBuilder, make_view,
+    StringArray, StringViewArray, StringViewBuilder, make_view,
 };
 use arrow::buffer::{MutableBuffer, NullBuffer};
 use arrow::datatypes::DataType;
@@ -372,7 +372,9 @@ impl LargeStringArrayBuilder {
     }
 }
 
-/// Append a new view to the views buffer with the given substr
+/// Append a new view to the views buffer with the given substr.
+///
+/// Callers are responsible for their own null tracking.
 ///
 /// # Safety
 ///
@@ -381,13 +383,15 @@ impl LargeStringArrayBuilder {
 ///
 /// # Arguments
 /// - views_buffer: The buffer to append the new view to
-/// - null_builder: The buffer to append the null value to
 /// - original_view: The original view value
 /// - substr: The substring to append. Must be a valid substring of the original view
 /// - start_offset: The start offset of the substring in the view
-pub fn make_and_append_view(
+///
+/// LLVM is apparently overly eager to inline this function into some hot loops,
+/// which bloats them and regresses performance, so we disable inling for now.
+#[inline(never)]
+pub fn append_view(
     views_buffer: &mut Vec<u128>,
-    null_builder: &mut NullBufferBuilder,
     original_view: &u128,
     substr: &str,
     start_offset: u32,
@@ -401,11 +405,9 @@ pub fn make_and_append_view(
             view.offset + start_offset,
         )
     } else {
-        // inline value does not need block id or offset
         make_view(substr.as_bytes(), 0, 0)
     };
     views_buffer.push(sub_view);
-    null_builder.append_non_null();
 }
 
 #[derive(Debug)]

datafusion/functions/src/unicode/substr.rs

@@ -17,13 +17,13 @@
 
 use std::sync::Arc;
 
-use crate::strings::make_and_append_view;
+use crate::strings::append_view;
 use crate::utils::make_scalar_function;
 use arrow::array::{
-    Array, ArrayRef, AsArray, Int64Array, NullBufferBuilder, StringArrayType,
-    StringViewArray, StringViewBuilder,
+    Array, ArrayRef, AsArray, Int64Array, StringArrayType, StringViewArray,
+    StringViewBuilder,
 };
-use arrow::buffer::ScalarBuffer;
+use arrow::buffer::{NullBuffer, ScalarBuffer};
 use arrow::datatypes::DataType;
 use datafusion_common::cast::as_int64_array;
 use datafusion_common::types::{
@@ -278,39 +278,32 @@ fn string_view_substr(
     let enable_ascii_fast_path =
         enable_ascii_fast_path(&string_view_array, start_array, count_array_opt);
 
+    // Combine null bitmaps from all inputs in bulk.
+    let nulls = NullBuffer::union(
+        NullBuffer::union(string_view_array.nulls(), start_array.nulls()).as_ref(),
+        count_array_opt.and_then(|a| a.nulls()),
+    );
+
     let mut views_buf = Vec::with_capacity(string_view_array.len());
-    let mut null_builder = NullBufferBuilder::new(string_view_array.len());
-
-    for i in 0..string_view_array.len() {
-        if string_view_array.is_null(i)
-            || start_array.is_null(i)
-            || count_array_opt.map(|a| a.is_null(i)).unwrap_or(false)
-        {
-            null_builder.append_null();
+
+    for (i, raw_view) in string_view_array.views().iter().enumerate() {
+        if nulls.as_ref().is_some_and(|n| n.is_null(i)) {
             views_buf.push(0);
             continue;
         }
 
         let string = string_view_array.value(i);
         let start = start_array.value(i);
         let count = count_array_opt.map(|a| a.value(i));
-        let raw_view = string_view_array.views()[i];
 
         let (start, end) =
             get_true_start_end(string, start, count, enable_ascii_fast_path)?;
         let substr = &string[start..end];
 
-        make_and_append_view(
-            &mut views_buf,
-            &mut null_builder,
-            &raw_view,
-            substr,
-            start as u32,
-        );
+        append_view(&mut views_buf, raw_view, substr, start as u32);
     }
 
     let views_buf = ScalarBuffer::from(views_buf);
-    let nulls_buf = null_builder.finish();
 
     // Safety:
     // (1) The blocks of the given views are all provided
@@ -320,7 +313,7 @@ fn string_view_substr(
         let array = StringViewArray::new_unchecked(
             views_buf,
             string_view_array.data_buffers().to_vec(),
-            nulls_buf,
+            nulls,
         );
         Ok(Arc::new(array) as ArrayRef)
     }
@@ -336,13 +329,16 @@ where
     let enable_ascii_fast_path =
         enable_ascii_fast_path(&string_array, start_array, count_array_opt);
 
+    // Combine null bitmaps from all inputs in bulk.
+    let nulls = NullBuffer::union(
+        NullBuffer::union(string_array.nulls(), start_array.nulls()).as_ref(),
+        count_array_opt.and_then(|a| a.nulls()),
+    );
+
     let mut result_builder = StringViewBuilder::new();
 
     for i in 0..string_array.len() {
-        if string_array.is_null(i)
-            || start_array.is_null(i)
-            || count_array_opt.map(|a| a.is_null(i)).unwrap_or(false)
-        {
+        if nulls.as_ref().is_some_and(|n| n.is_null(i)) {
             result_builder.append_null();
             continue;
         }

datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs

@@ -202,6 +202,24 @@ impl StreamedBatch {
     }
 }
 
+/// Per-row filter outcome tracking for full outer joins.
+///
+/// In a full outer join with a filter, buffered rows that match on join
+/// keys but fail every filter evaluation must be emitted with NULLs on
+/// the streamed side. Three states are needed because a simple boolean
+/// cannot distinguish "never matched" (handled by [`BufferedBatch::null_joined`])
+/// from "matched but all filters failed" (must be emitted as null-joined).
+#[repr(u8)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(super) enum FilterState {
+    /// Row never appeared in a matched pair.
+    Unvisited = 0,
+    /// Row matched streamed rows, but all filter evaluations failed.
+    AllFailed = 1,
+    /// Row matched and at least one filter evaluation passed.
+    SomePassed = 2,
+}
+
 /// A buffered batch that contains contiguous rows with same join key
 ///
 /// `BufferedBatch` can exist as either an in-memory `RecordBatch` or a `RefCountedTempFile` on disk.
@@ -217,11 +235,9 @@ pub(super) struct BufferedBatch {
     pub null_joined: Vec<usize>,
     /// Size estimation used for reserving / releasing memory
     pub size_estimation: usize,
-    /// The indices of buffered batch that the join filter doesn't satisfy.
-    /// This is a map between right row index and a boolean value indicating whether all joined row
-    /// of the right row does not satisfy the filter .
-    /// When dequeuing the buffered batch, we need to produce null joined rows for these indices.
-    pub join_filter_not_matched_map: HashMap<u64, bool>,
+    /// Tracks filter outcomes for buffered rows in full outer joins.
+    /// Indexed by absolute row position within the batch. See [`FilterState`].
+    pub join_filter_status: Vec<FilterState>,
     /// Current buffered batch number of rows. Equal to batch.num_rows()
     /// but if batch is spilled to disk this property is preferable
     /// and less expensive
@@ -258,7 +274,7 @@ impl BufferedBatch {
             join_arrays,
             null_joined: vec![],
             size_estimation,
-            join_filter_not_matched_map: HashMap::new(),
+            join_filter_status: vec![FilterState::Unvisited; num_rows],
             num_rows,
         }
     }
@@ -1250,12 +1266,16 @@ impl MaterializingSortMergeJoinStream {
             return Ok(());
         }
 
-        // For buffered row which is joined with streamed side rows but all joined rows
-        // don't satisfy the join filter
+        // Collect buffered rows that matched on join keys but had every
+        // filter evaluation fail — these must be emitted with NULLs on
+        // the streamed side to satisfy full outer join semantics.
         let not_matched_buffered_indices = buffered_batch
-            .join_filter_not_matched_map
+            .join_filter_status
             .iter()
-            .filter_map(|(idx, failed)| if *failed { Some(*idx) } else { None })
+            .enumerate()
+            .filter_map(|(i, state)| {
+                matches!(state, FilterState::AllFailed).then_some(i as u64)
+            })
             .collect::<Vec<_>>();
 
         let buffered_indices =
@@ -1270,7 +1290,9 @@ impl MaterializingSortMergeJoinStream {
             self.joined_record_batches
                 .push_batch_with_null_metadata(record_batch, self.join_type);
         }
-        buffered_batch.join_filter_not_matched_map.clear();
+        buffered_batch
+            .join_filter_status
+            .fill(FilterState::Unvisited);
 
         Ok(())
     }
@@ -1443,15 +1465,18 @@ impl MaterializingSortMergeJoinStream {
                             if right.is_null(i) {
                                 continue;
                             }
-                            let buffered_index = right.value(i);
-                            buffered_batch.join_filter_not_matched_map.insert(
-                                buffered_index,
-                                *buffered_batch
-                                    .join_filter_not_matched_map
-                                    .get(&buffered_index)
-                                    .unwrap_or(&true)
-                                    && !pre_mask.value(offset + i),
-                            );
+                            let idx = right.value(i) as usize;
+                            match buffered_batch.join_filter_status[idx] {
+                                FilterState::SomePassed => {}
+                                _ if pre_mask.value(offset + i) => {
+                                    buffered_batch.join_filter_status[idx] =
+                                        FilterState::SomePassed;
+                                }
+                                _ => {
+                                    buffered_batch.join_filter_status[idx] =
+                                        FilterState::AllFailed;
+                                }
+                            }
                         }
                         offset += chunk_len;
                     }