fix(duckdb): cast query_arrow results to projected_schema

DuckDB's query_arrow ignored the projected_schema parameter, returning
batches with DuckDB's native types (e.g. Timestamp(µs)) even when the
caller expected different types (e.g. Timestamp(ns)). This caused schema
mismatches for downstream operators pushed below SchemaCastScanExec.

Cast result batches to projected_schema in the output stream when types
differ. Add shared cast_batch_to_schema utility in util/arrow.rs for
reuse by other Arrow-native connectors (ADBC, ODBC).

Author: ewgenius

core/src/sql/db_connection_pool/dbconnection/duckdbconn.rs

@@ -23,6 +23,7 @@ use snafu::{prelude::*, ResultExt};
 use tokio::sync::mpsc::Sender;
 
 use crate::sql::db_connection_pool::runtime::run_sync_with_tokio;
+use crate::util::arrow::cast_batch_to_schema;
 use crate::util::schema::SchemaValidator;
 use crate::UnsupportedTypeAction;
 
@@ -453,7 +454,7 @@ impl SyncDbConnection<r2d2::PooledConnection<DuckdbConnectionManager>, DuckDBPar
         &self,
         sql: &str,
         params: &[DuckDBParameter],
-        _projected_schema: Option<SchemaRef>,
+        projected_schema: Option<SchemaRef>,
     ) -> Result<SendableRecordBatchStream> {
         let (batch_tx, mut batch_rx) = tokio::sync::mpsc::channel::<RecordBatch>(4);
 
@@ -501,9 +502,18 @@ impl SyncDbConnection<r2d2::PooledConnection<DuckdbConnectionManager>, DuckDBPar
                 Ok::<_, Box<dyn std::error::Error + Send + Sync>>(())
             });
 
+            let stream_schema = projected_schema.clone().unwrap_or(schema);
+
             let output_stream = stream! {
                 while let Some(batch) = batch_rx.recv().await {
-                    yield Ok(batch);
+                    if let Some(ref target_schema) = projected_schema {
+                        match cast_batch_to_schema(&batch, target_schema) {
+                            Ok(casted) => yield Ok(casted),
+                            Err(e) => yield Err(e),
+                        }
+                    } else {
+                        yield Ok(batch);
+                    }
                 }
 
                 match join_handle.await {
@@ -522,7 +532,7 @@ impl SyncDbConnection<r2d2::PooledConnection<DuckdbConnectionManager>, DuckDBPar
             };
 
             Ok(Box::pin(RecordBatchStreamAdapter::new(
-                schema,
+                stream_schema,
                 output_stream,
             )))
         };
@@ -945,4 +955,70 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_query_arrow_casts_to_projected_schema() {
+        use arrow::datatypes::{Schema, TimeUnit};
+        use futures::StreamExt;
+
+        use crate::sql::db_connection_pool::duckdbpool::DuckDbConnectionPool;
+        use crate::sql::db_connection_pool::DbConnectionPool;
+
+        let rt = tokio::runtime::Runtime::new().expect("runtime");
+        rt.block_on(async {
+            let pool = DuckDbConnectionPool::new_memory().expect("pool created");
+
+            // Create table with TIMESTAMPTZ (DuckDB stores as Microsecond)
+            let conn = pool.connect().await.expect("connection");
+            let conn = conn.as_sync().expect("sync connection");
+            conn.execute(
+                "CREATE TABLE test_ts (id INTEGER, created_at TIMESTAMPTZ)",
+                &[],
+            )
+            .expect("table created");
+            conn.execute(
+                "INSERT INTO test_ts VALUES (1, '2023-01-01T00:00:00Z')",
+                &[],
+            )
+            .expect("data inserted");
+
+            // Request Nanosecond via projected_schema
+            let projected_schema = Arc::new(Schema::new(vec![
+                Field::new("id", DataType::Int32, true),
+                Field::new(
+                    "created_at",
+                    DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
+                    true,
+                ),
+            ]));
+
+            let stream = conn
+                .query_arrow(
+                    "SELECT id, created_at FROM test_ts",
+                    &[],
+                    Some(projected_schema.clone()),
+                )
+                .expect("query_arrow should succeed");
+
+            // Verify stream schema matches projected_schema
+            assert_eq!(stream.schema(), projected_schema);
+
+            let mut batches = vec![];
+            let mut stream = stream;
+            while let Some(batch) = stream.next().await {
+                batches.push(batch.expect("batch should be Ok"));
+            }
+
+            assert_eq!(batches.len(), 1);
+            let batch = &batches[0];
+            assert_eq!(batch.schema(), projected_schema);
+
+            // Verify the timestamp column is Nanosecond
+            let ts_col = batch.column(1);
+            assert_eq!(
+                ts_col.data_type(),
+                &DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into()))
+            );
+        });
+    }
 }

core/src/util/arrow.rs

@@ -0,0 +1,125 @@
+use std::sync::Arc;
+
+use arrow::array::RecordBatch;
+use arrow::compute::cast;
+use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::error::DataFusionError;
+
+/// Cast a `RecordBatch` to match the target schema, casting columns whose types differ.
+/// Columns that already match are passed through unchanged.
+pub fn cast_batch_to_schema(
+    batch: &RecordBatch,
+    target_schema: &SchemaRef,
+) -> Result<RecordBatch, DataFusionError> {
+    let columns: Vec<_> = batch
+        .columns()
+        .iter()
+        .zip(target_schema.fields())
+        .map(|(col, target_field)| {
+            if col.data_type() == target_field.data_type() {
+                Ok(Arc::clone(col))
+            } else {
+                cast(col, target_field.data_type()).map_err(|e| {
+                    DataFusionError::Execution(format!(
+                        "Failed to cast column '{}' from {:?} to {:?}: {e}",
+                        target_field.name(),
+                        col.data_type(),
+                        target_field.data_type(),
+                    ))
+                })
+            }
+        })
+        .collect::<Result<_, _>>()?;
+
+    RecordBatch::try_new(Arc::clone(target_schema), columns).map_err(|e| {
+        DataFusionError::Execution(format!("Failed to create RecordBatch after cast: {e}"))
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::array::{Int32Array, TimestampMicrosecondArray, TimestampNanosecondArray};
+    use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
+
+    #[test]
+    fn test_cast_timestamp_us_to_ns() {
+        let source_schema = Arc::new(Schema::new(vec![
+            Field::new("id", DataType::Int32, false),
+            Field::new(
+                "ts",
+                DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
+                false,
+            ),
+        ]));
+
+        let target_schema = Arc::new(Schema::new(vec![
+            Field::new("id", DataType::Int32, false),
+            Field::new(
+                "ts",
+                DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
+                false,
+            ),
+        ]));
+
+        let batch = RecordBatch::try_new(
+            source_schema,
+            vec![
+                Arc::new(Int32Array::from(vec![1, 2, 3])),
+                Arc::new(
+                    TimestampMicrosecondArray::from(vec![1_000_000, 2_000_000, 3_000_000])
+                        .with_timezone("UTC"),
+                ),
+            ],
+        )
+        .unwrap();
+
+        let result = cast_batch_to_schema(&batch, &target_schema).unwrap();
+        assert_eq!(result.schema(), target_schema);
+        assert_eq!(
+            result.column(1).data_type(),
+            &DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into()))
+        );
+
+        // Values should be multiplied by 1000
+        let ts_col = result
+            .column(1)
+            .as_any()
+            .downcast_ref::<TimestampNanosecondArray>()
+            .unwrap();
+        assert_eq!(ts_col.value(0), 1_000_000_000);
+        assert_eq!(ts_col.value(1), 2_000_000_000);
+        assert_eq!(ts_col.value(2), 3_000_000_000);
+    }
+
+    #[test]
+    fn test_no_cast_when_types_match() {
+        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
+
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
+        )
+        .unwrap();
+
+        let result = cast_batch_to_schema(&batch, &schema).unwrap();
+        assert_eq!(result.schema(), schema);
+    }
+
+    #[test]
+    fn test_cast_incompatible_types_returns_error() {
+        use arrow::array::StringArray;
+
+        let source_schema = Arc::new(Schema::new(vec![Field::new("val", DataType::Utf8, false)]));
+        let target_schema = Arc::new(Schema::new(vec![Field::new("val", DataType::Int32, false)]));
+
+        let batch = RecordBatch::try_new(
+            source_schema,
+            vec![Arc::new(StringArray::from(vec!["not_a_number"]))],
+        )
+        .unwrap();
+
+        let result = cast_batch_to_schema(&batch, &target_schema);
+        assert!(result.is_err());
+    }
+}

core/src/util/mod.rs

@@ -6,6 +6,7 @@ use std::collections::HashMap;
 
 use crate::UnsupportedTypeAction;
 
+pub mod arrow;
 pub mod column_reference;
 pub mod constraints;
 pub mod count_exec;