Merge branch 'main' into feat_migrate_ffi_to_stabby

Author: coderfender

Cargo.lock

@@ -431,3 +431,65 @@ async fn test_select_cast_date_literal_to_timestamp_overflow() -> Result<()> {
     );
     Ok(())
 }
+
+// Regression test: a recursive CTE whose anchor aliases a computed column
+// (`upper(val) AS val`) and whose recursive term leaves the same expression
+// un-aliased must still produce batches whose schema field names come from
+// the anchor term — especially when the outer query uses ORDER BY + LIMIT
+// (the TopK path passes batch schemas through verbatim, so any drift in
+// RecursiveQueryExec's emitted batches is visible downstream).
+#[tokio::test]
+async fn test_recursive_cte_batch_schema_stable_with_order_by_limit() -> Result<()> {
+    let ctx = SessionContext::new();
+    ctx.sql(
+        "CREATE TABLE records (\
+            id VARCHAR NOT NULL, \
+            parent_id VARCHAR, \
+            ts TIMESTAMP NOT NULL, \
+            val VARCHAR\
+         )",
+    )
+    .await?
+    .collect()
+    .await?;
+    ctx.sql(
+        "INSERT INTO records VALUES \
+            ('a00', NULL,  TIMESTAMP '2025-01-01 00:00:00', 'v_span'), \
+            ('a01', 'a00', TIMESTAMP '2025-01-01 00:00:01', 'v_log_1'), \
+            ('a02', 'a01', TIMESTAMP '2025-01-01 00:00:02', 'v_log_2'), \
+            ('a03', 'a02', TIMESTAMP '2025-01-01 00:00:03', 'v_log_3'), \
+            ('a04', 'a03', TIMESTAMP '2025-01-01 00:00:04', 'v_log_4'), \
+            ('a05', 'a04', TIMESTAMP '2025-01-01 00:00:05', 'v_log_5')",
+    )
+    .await?
+    .collect()
+    .await?;
+
+    let results = ctx
+        .sql(
+            "WITH RECURSIVE descendants AS (\
+                SELECT id, parent_id, ts, upper(val) AS val \
+                  FROM records WHERE id = 'a00' \
+                UNION ALL \
+                SELECT r.id, r.parent_id, r.ts, upper(r.val) \
+                  FROM records r INNER JOIN descendants d ON r.parent_id = d.id \
+             ) \
+             SELECT id, parent_id, ts, val FROM descendants ORDER BY ts ASC LIMIT 10",
+        )
+        .await?
+        .collect()
+        .await?;
+
+    let expected_names = ["id", "parent_id", "ts", "val"];
+    assert!(!results.is_empty(), "expected at least one batch");
+    for (i, batch) in results.iter().enumerate() {
+        let schema = batch.schema();
+        let actual_names: Vec<&str> =
+            schema.fields().iter().map(|f| f.name().as_str()).collect();
+        assert_eq!(
+            actual_names, expected_names,
+            "batch {i} schema field names leaked from recursive branch"
+        );
+    }
+    Ok(())
+}

datafusion/core/tests/sql/select.rs

@@ -22,7 +22,8 @@ use std::sync::Arc;
 
 use crate::string::concat;
 use crate::strings::{
-    ColumnarValueRef, LargeStringArrayBuilder, StringArrayBuilder, StringViewArrayBuilder,
+    ColumnarValueRef, ConcatLargeStringBuilder, ConcatStringBuilder,
+    ConcatStringViewBuilder,
 };
 use datafusion_common::cast::{as_binary_array, as_string_array, as_string_view_array};
 use datafusion_common::{
@@ -242,7 +243,7 @@ impl ScalarUDFImpl for ConcatFunc {
 
         match return_datatype {
             DataType::Utf8 => {
-                let mut builder = StringArrayBuilder::with_capacity(len, data_size);
+                let mut builder = ConcatStringBuilder::with_capacity(len, data_size);
                 for i in 0..len {
                     columns
                         .iter()
@@ -254,7 +255,7 @@ impl ScalarUDFImpl for ConcatFunc {
                 Ok(ColumnarValue::Array(Arc::new(string_array)))
             }
             DataType::Utf8View => {
-                let mut builder = StringViewArrayBuilder::with_capacity(len, data_size);
+                let mut builder = ConcatStringViewBuilder::with_capacity(len, data_size);
                 for i in 0..len {
                     columns
                         .iter()
@@ -266,7 +267,7 @@ impl ScalarUDFImpl for ConcatFunc {
                 Ok(ColumnarValue::Array(Arc::new(string_array)))
             }
             DataType::LargeUtf8 => {
-                let mut builder = LargeStringArrayBuilder::with_capacity(len, data_size);
+                let mut builder = ConcatLargeStringBuilder::with_capacity(len, data_size);
                 for i in 0..len {
                     columns
                         .iter()

datafusion/functions/src/string/concat.rs

@@ -24,7 +24,8 @@ use crate::string::concat;
 use crate::string::concat::simplify_concat;
 use crate::string::concat_ws;
 use crate::strings::{
-    ColumnarValueRef, LargeStringArrayBuilder, StringArrayBuilder, StringViewArrayBuilder,
+    ColumnarValueRef, ConcatLargeStringBuilder, ConcatStringBuilder,
+    ConcatStringViewBuilder,
 };
 use datafusion_common::cast::{
     as_large_string_array, as_string_array, as_string_view_array,
@@ -311,7 +312,7 @@ impl ScalarUDFImpl for ConcatWsFunc {
 
         match return_datatype {
             DataType::Utf8View => {
-                let mut builder = StringViewArrayBuilder::with_capacity(len, data_size);
+                let mut builder = ConcatStringViewBuilder::with_capacity(len, data_size);
                 for i in 0..len {
                     if !sep.is_valid(i) {
                         builder.append_offset()?;
@@ -332,7 +333,7 @@ impl ScalarUDFImpl for ConcatWsFunc {
                 Ok(ColumnarValue::Array(Arc::new(builder.finish(sep.nulls())?)))
             }
             DataType::LargeUtf8 => {
-                let mut builder = LargeStringArrayBuilder::with_capacity(len, data_size);
+                let mut builder = ConcatLargeStringBuilder::with_capacity(len, data_size);
                 for i in 0..len {
                     if !sep.is_valid(i) {
                         builder.append_offset()?;
@@ -353,7 +354,7 @@ impl ScalarUDFImpl for ConcatWsFunc {
                 Ok(ColumnarValue::Array(Arc::new(builder.finish(sep.nulls())?)))
             }
             _ => {
-                let mut builder = StringArrayBuilder::with_capacity(len, data_size);
+                let mut builder = ConcatStringBuilder::with_capacity(len, data_size);
                 for i in 0..len {
                     if !sep.is_valid(i) {
                         builder.append_offset()?;

datafusion/functions/src/string/concat_ws.rs

@@ -26,17 +26,20 @@ use arrow::array::{
 use arrow::buffer::{Buffer, MutableBuffer, NullBuffer, ScalarBuffer};
 use arrow::datatypes::DataType;
 
-/// Optimized version of the StringBuilder in Arrow that:
-/// 1. Precalculating the expected length of the result, avoiding reallocations.
-/// 2. Avoids creating / incrementally creating a `NullBufferBuilder`
-pub struct StringArrayBuilder {
+/// Builder used by `concat`/`concat_ws` to assemble a [`StringArray`] one row
+/// at a time from multiple input columns.
+///
+/// Each row is written via repeated `write` calls, followed by a single
+/// `append_offset` call to commit the row. The output null buffer is supplied
+/// by the caller at `finish` time.
+pub(crate) struct ConcatStringBuilder {
     offsets_buffer: MutableBuffer,
     value_buffer: MutableBuffer,
     /// If true, a safety check is required during the `finish` call
     tainted: bool,
 }
 
-impl StringArrayBuilder {
+impl ConcatStringBuilder {
     pub fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
         let capacity = item_capacity
             .checked_add(1)
@@ -151,21 +154,21 @@ impl StringArrayBuilder {
     }
 }
 
-/// Optimized version of Arrow's [`StringViewBuilder`]. Rather than adding NULLs
-/// on a row-by-row basis, the caller should provide nulls when calling
-/// [`finish`](Self::finish). This allows callers to compute nulls more
-/// efficiently (e.g., via bulk bitmap operations).
+/// Builder used by `concat`/`concat_ws` to assemble a [`StringViewArray`] one
+/// row at a time from multiple input columns.
 ///
-/// [`StringViewBuilder`]: arrow::array::StringViewBuilder
-pub struct StringViewArrayBuilder {
+/// Each row is written via repeated `write` calls, followed by a single
+/// `append_offset` call to commit the row as a single string view. The output
+/// null buffer is supplied by the caller at `finish` time.
+pub(crate) struct ConcatStringViewBuilder {
     views: Vec<u128>,
     data: Vec<u8>,
     block: Vec<u8>,
     /// If true, a safety check is required during the `append_offset` call
     tainted: bool,
 }
 
-impl StringViewArrayBuilder {
+impl ConcatStringViewBuilder {
     pub fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
         Self {
             views: Vec::with_capacity(item_capacity),
@@ -286,14 +289,17 @@ impl StringViewArrayBuilder {
     }
 }
 
-pub struct LargeStringArrayBuilder {
+/// Builder used by `concat`/`concat_ws` to assemble a [`LargeStringArray`] one
+/// row at a time from multiple input columns. See [`ConcatStringBuilder`] for
+/// details on the row-composition contract.
+pub(crate) struct ConcatLargeStringBuilder {
     offsets_buffer: MutableBuffer,
     value_buffer: MutableBuffer,
     /// If true, a safety check is required during the `finish` call
     tainted: bool,
 }
 
-impl LargeStringArrayBuilder {
+impl ConcatLargeStringBuilder {
     pub fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
         let capacity = item_capacity
             .checked_add(1)
@@ -426,7 +432,7 @@ impl LargeStringArrayBuilder {
 /// LLVM is apparently overly eager to inline this function into some hot loops,
 /// which bloats them and regresses performance, so we disable inlining for now.
 #[inline(never)]
-pub fn append_view(
+pub(crate) fn append_view(
     views_buffer: &mut Vec<u128>,
     original_view: &u128,
     substr: &str,
@@ -447,7 +453,7 @@ pub fn append_view(
 }
 
 #[derive(Debug)]
-pub enum ColumnarValueRef<'a> {
+pub(crate) enum ColumnarValueRef<'a> {
     Scalar(&'a [u8]),
     NullableArray(&'a StringArray),
     NonNullableArray(&'a StringArray),
@@ -497,13 +503,13 @@ mod tests {
 
     #[test]
     #[should_panic(expected = "capacity integer overflow")]
-    fn test_overflow_string_array_builder() {
-        let _builder = StringArrayBuilder::with_capacity(usize::MAX, usize::MAX);
+    fn test_overflow_concat_string_builder() {
+        let _builder = ConcatStringBuilder::with_capacity(usize::MAX, usize::MAX);
     }
 
     #[test]
     #[should_panic(expected = "capacity integer overflow")]
-    fn test_overflow_large_string_array_builder() {
-        let _builder = LargeStringArrayBuilder::with_capacity(usize::MAX, usize::MAX);
+    fn test_overflow_concat_large_string_builder() {
+        let _builder = ConcatLargeStringBuilder::with_capacity(usize::MAX, usize::MAX);
     }
 }

datafusion/functions/src/strings.rs

@@ -317,6 +317,20 @@ impl RecursiveQueryStream {
         mut batch: RecordBatch,
     ) -> Poll<Option<Result<RecordBatch>>> {
         let baseline_metrics = self.baseline_metrics.clone();
+
+        // Rebind to the declared output schema. The recursive term is planned
+        // independently from the static term and its projection may leave
+        // columns un-aliased (e.g. `upper(r.val)` vs the anchor's
+        // `upper(val) AS val`); downstream consumers that key on
+        // `batch.schema().field(i).name()` (TopK, CSV/JSON writers, custom
+        // collectors) would otherwise see the recursive branch's names leak
+        // through. Logical-plan coercion guarantees matching types, so this
+        // is a zero-copy field rebind.
+        if batch.schema() != self.schema {
+            batch =
+                RecordBatch::try_new(Arc::clone(&self.schema), batch.columns().to_vec())?;
+        }
+
         if let Some(deduplicator) = &mut self.distinct_deduplicator {
             let _timer_guard = baseline_metrics.elapsed_compute().timer();
             batch = deduplicator.deduplicate(&batch)?;

datafusion/physical-plan/src/recursive_query.rs

@@ -1223,6 +1223,83 @@ physical_plan
 07)------RepartitionExec: partitioning=Hash([start@0], 4), input_partitions=1
 08)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/recursive_cte/closure.csv]]}, projection=[start, end], file_type=csv, has_header=true
 
+# Regression test: recursive CTE column-name stability under ORDER BY + LIMIT.
+#
+# The recursive CTE below has an anchor that aliases a computed column
+# (`upper(val) AS val`) while the recursive term leaves the same expression
+# un-aliased (`upper(r.val)`). RecursiveQueryExec declares its output schema
+# from the anchor, so the plan-level column is `val`; but if it forwards the
+# recursive term's batches without rebinding their schema to the anchor's,
+# downstream consumers that read `batch.schema().field(i).name()` observe
+# the leaked `upper(r.val)` name. The outer `ORDER BY ts LIMIT 10` routes
+# through the TopK path in SortExec, which passes batch schemas through
+# verbatim — so the leak surfaces there.
+#
+# sqllogictest compares row values and can't assert column names directly.
+# To surface the leak inside SLT, we `COPY` the result to a CSV file with a
+# header row (the CSV writer names header columns from each batch's own
+# schema) and then read the file back as headerless CSV so the header line
+# appears as a data row we can compare against the expected column names.
+statement ok
+CREATE TABLE cte_schema_records (
+    id VARCHAR NOT NULL,
+    parent_id VARCHAR,
+    ts TIMESTAMP NOT NULL,
+    val VARCHAR
+);
+
+statement ok
+INSERT INTO cte_schema_records VALUES
+  ('a00', NULL,  TIMESTAMP '2025-01-01 00:00:00', 'v_span'),
+  ('a01', 'a00', TIMESTAMP '2025-01-01 00:00:01', 'v_log_1'),
+  ('a02', 'a01', TIMESTAMP '2025-01-01 00:00:02', 'v_log_2'),
+  ('a03', 'a02', TIMESTAMP '2025-01-01 00:00:03', 'v_log_3'),
+  ('a04', 'a03', TIMESTAMP '2025-01-01 00:00:04', 'v_log_4'),
+  ('a05', 'a04', TIMESTAMP '2025-01-01 00:00:05', 'v_log_5');
+
+query I
+COPY (WITH RECURSIVE descendants AS (
+    SELECT id, parent_id, ts, upper(val) AS val
+      FROM cte_schema_records WHERE id = 'a00'
+    UNION ALL
+    SELECT r.id, r.parent_id, r.ts, upper(r.val)
+      FROM cte_schema_records r
+      INNER JOIN descendants d ON r.parent_id = d.id
+  )
+  SELECT id, parent_id, ts, val
+    FROM descendants
+    ORDER BY ts ASC
+    LIMIT 10)
+TO 'test_files/scratch/cte/recursive_order_limit_schema/'
+STORED AS CSV OPTIONS ('format.has_header' 'true');
+----
+6
+
+statement ok
+CREATE EXTERNAL TABLE cte_schema_reread (
+    id VARCHAR, parent_id VARCHAR, ts VARCHAR, val VARCHAR
+)
+STORED AS CSV
+LOCATION 'test_files/scratch/cte/recursive_order_limit_schema/'
+OPTIONS ('format.has_header' 'false');
+
+# The CSV header row appears as a data row here because we re-read with
+# has_header='false'. It must match the CTE's declared column names; if
+# RecursiveQueryExec leaked the recursive branch's schema, the last field
+# would read `upper(r.val)` instead of `val`.
+query TTTT
+SELECT id, parent_id, ts, val
+  FROM cte_schema_reread
+  WHERE id = 'id';
+----
+id parent_id ts val
+
+statement ok
+DROP TABLE cte_schema_reread;
+
+statement ok
+DROP TABLE cte_schema_records;
+
 statement count 0
 set datafusion.execution.enable_recursive_ctes = false;
 

datafusion/sqllogictest/test_files/cte.slt

@@ -320,6 +320,14 @@ The difference is only observable for strings containing combining characters
 clusters (e.g., ZWJ emoji sequences). For ASCII and most common Unicode text,
 behavior is unchanged.
 
+### Items in `datafusion_functions::strings` are no longer public
+
+`StringArrayBuilder`, `LargeStringArrayBuilder`, `StringViewArrayBuilder`,
+`ColumnarValueRef`, and `append_view` have been reduced to `pub(crate)`. They
+were only ever used to implement `concat` and `concat_ws` inside the crate. If
+you were importing them externally, use Arrow's corresponding builders with a
+caller-computed `NullBuffer`.
+
 [#17861]: https://github.com/apache/datafusion/pull/17861
 
 ### `approx_percentile_cont`, `approx_percentile_cont_with_weight`, `approx_median` now coerce to floats