Merge branch 'spiceai-52' into viktor/duckdb-deletes-stack-overflow

Author: krinart

Cargo.lock

@@ -41,7 +41,7 @@ arrow-json = "57.1.0"
 arrow-odbc = { version = "23.1" }
 datafusion = { version = "52", default-features = false }
 datafusion-expr = { version = "52" }
-datafusion-federation = { git = "https://github.com/spiceai/datafusion-federation.git", rev = "a4fce79433e5c2fe779427c0bfce6a599193e300" }
+datafusion-federation = { git = "https://github.com/spiceai/datafusion-federation.git", rev = "6b6bfb0d30da8e5c2eb851094e366d98fa839575" }
 datafusion-ffi = { version = "52" }
 datafusion-proto = { version = "52" }
 datafusion-physical-expr = { version = "52" }

Cargo.toml

@@ -121,6 +121,7 @@ duckdb = [
   "dep:arrow-schema",
   "dep:byte-unit",
   "dep:datafusion-physical-expr",
+  "federation",
 ]
 duckdb-federation = ["duckdb", "federation"]
 federation = ["dep:datafusion-federation"]

core/Cargo.toml

@@ -453,7 +453,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 +501,20 @@ 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 datafusion_federation::schema_cast::record_convert::try_cast_to(batch, Arc::clone(target_schema)) {
+                            Ok(casted) => yield Ok(casted),
+                            Err(e) => yield Err(DataFusionError::Execution(format!(
+                                "Failed to cast DuckDB result batch to projected schema: {e}"
+                            ))),
+                        }
+                    } else {
+                        yield Ok(batch);
+                    }
                 }
 
                 match join_handle.await {
@@ -522,7 +533,7 @@ impl SyncDbConnection<r2d2::PooledConnection<DuckdbConnectionManager>, DuckDBPar
             };
 
             Ok(Box::pin(RecordBatchStreamAdapter::new(
-                schema,
+                stream_schema,
                 output_stream,
             )))
         };
@@ -945,4 +956,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/sql/db_connection_pool/dbconnection/duckdbconn.rs

@@ -8,15 +8,14 @@ use datafusion::execution::context::SessionContext;
 use datafusion::logical_expr::dml::InsertOp;
 use datafusion::logical_expr::CreateExternalTable;
 use datafusion::physical_plan::collect;
-use datafusion_federation::schema_cast::record_convert::try_cast_to;
 use datafusion_table_providers::duckdb::DuckDBTableProviderFactory;
 use rstest::rstest;
 use std::collections::HashMap;
 use std::sync::Arc;
 
 async fn arrow_duckdb_round_trip(
     arrow_record: RecordBatch,
-    source_schema: SchemaRef,
+    _source_schema: SchemaRef,
     table_name: &str,
 ) {
     let factory = DuckDBTableProviderFactory::new(duckdb::AccessMode::ReadWrite);
@@ -68,7 +67,6 @@ async fn arrow_duckdb_round_trip(
         .expect("DataFrame should be created from query");
 
     let record_batch = df.collect().await.expect("RecordBatch should be collected");
-    let casted_record = try_cast_to(record_batch[0].clone(), source_schema).unwrap();
 
     tracing::debug!("Original Arrow Record Batch: {:?}", arrow_record.columns());
     tracing::debug!(
@@ -80,7 +78,7 @@ async fn arrow_duckdb_round_trip(
     assert_eq!(record_batch.len(), 1);
     assert_eq!(record_batch[0].num_rows(), arrow_record.num_rows());
     assert_eq!(record_batch[0].num_columns(), arrow_record.num_columns());
-    assert_eq!(casted_record, arrow_record);
+    assert_eq!(record_batch[0], arrow_record);
 }
 
 #[rstest]