fix(DuckDB): replace cast_batch_to_schema with federation's try_cast_to (#653)

Author: sgrebnov

core/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/src/sql/db_connection_pool/dbconnection/duckdbconn.rs

@@ -23,7 +23,6 @@ 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;
 
@@ -507,9 +506,11 @@ impl SyncDbConnection<r2d2::PooledConnection<DuckdbConnectionManager>, DuckDBPar
             let output_stream = stream! {
                 while let Some(batch) = batch_rx.recv().await {
                     if let Some(ref target_schema) = projected_schema {
-                        match cast_batch_to_schema(&batch, target_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(e),
+                            Err(e) => yield Err(DataFusionError::Execution(format!(
+                                "Failed to cast DuckDB result batch to projected schema: {e}"
+                            ))),
                         }
                     } else {
                         yield Ok(batch);

core/src/util/arrow.rs

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

core/src/util/mod.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]