fix default schema adapter for arrow

Author: adriangb

datafusion/core/tests/physical_optimizer/filter_pushdown/util.rs

@@ -18,7 +18,7 @@
 use arrow::datatypes::SchemaRef;
 use arrow::{array::RecordBatch, compute::concat_batches};
 use datafusion::{datasource::object_store::ObjectStoreUrl, physical_plan::PhysicalExpr};
-use datafusion_common::{config::ConfigOptions, internal_err, Result, Statistics};
+use datafusion_common::{config::ConfigOptions, internal_err, Result};
 use datafusion_datasource::{
     file::FileSource, file_scan_config::FileScanConfig,
     file_scan_config::FileScanConfigBuilder, file_stream::FileOpenFuture,

datafusion/datasource-arrow/src/file_format.rs

@@ -208,14 +208,21 @@ impl FileFormat for ArrowFormat {
             conf.table_partition_cols().clone(),
         );
 
-        let source: Arc<dyn FileSource> =
+        let mut source: Arc<dyn FileSource> =
             match is_object_in_arrow_ipc_file_format(object_store, object_location).await
             {
                 Ok(true) => Arc::new(ArrowSource::new_file_source(table_schema)),
                 Ok(false) => Arc::new(ArrowSource::new_stream_file_source(table_schema)),
                 Err(e) => Err(e)?,
             };
 
+        // Preserve projection from the original file source
+        if let Some(projection) = conf.file_source.projection() {
+            if let Some(new_source) = source.try_pushdown_projection(projection)? {
+                source = new_source;
+            }
+        }
+
         let config = FileScanConfigBuilder::from(conf)
             .with_source(source)
             .build();

datafusion/datasource-arrow/src/source.rs

@@ -34,10 +34,13 @@
 use std::sync::Arc;
 use std::{any::Any, io::Cursor};
 
-use datafusion_datasource::schema_adapter::SchemaAdapterFactory;
+use datafusion_datasource::schema_adapter::{
+    DefaultSchemaAdapterFactory, SchemaAdapterFactory,
+};
 use datafusion_datasource::{as_file_source, TableSchema};
 
 use arrow::buffer::Buffer;
+use arrow::datatypes::SchemaRef;
 use arrow::ipc::reader::{FileDecoder, FileReader, StreamReader};
 use datafusion_common::error::Result;
 use datafusion_common::exec_datafusion_err;
@@ -51,7 +54,7 @@ use datafusion_physical_plan::projection::ProjectionExprs;
 
 use datafusion_datasource::file_stream::FileOpenFuture;
 use datafusion_datasource::file_stream::FileOpener;
-use futures::StreamExt;
+use futures::{StreamExt, TryStreamExt};
 use itertools::Itertools;
 use object_store::{GetOptions, GetRange, GetResultPayload, ObjectStore};
 
@@ -68,6 +71,8 @@ enum ArrowFormat {
 pub(crate) struct ArrowStreamFileOpener {
     object_store: Arc<dyn ObjectStore>,
     projection: Option<Vec<usize>>,
+    projected_schema: Option<SchemaRef>,
+    schema_adapter_factory: Option<Arc<dyn SchemaAdapterFactory>>,
 }
 
 impl FileOpener for ArrowStreamFileOpener {
@@ -79,27 +84,52 @@ impl FileOpener for ArrowStreamFileOpener {
         }
         let object_store = Arc::clone(&self.object_store);
         let projection = self.projection.clone();
+        let projected_schema = self.projected_schema.clone();
+        let schema_adapter_factory = self.schema_adapter_factory.clone();
+
         Ok(Box::pin(async move {
             let r = object_store
                 .get(&partitioned_file.object_meta.location)
                 .await?;
-            match r.payload {
+
+            let stream = match r.payload {
                 #[cfg(not(target_arch = "wasm32"))]
-                GetResultPayload::File(file, _) => Ok(futures::stream::iter(
+                GetResultPayload::File(file, _) => futures::stream::iter(
                     StreamReader::try_new(file.try_clone()?, projection.clone())?,
                 )
                 .map(|r| r.map_err(Into::into))
-                .boxed()),
+                .boxed(),
                 GetResultPayload::Stream(_) => {
                     let bytes = r.bytes().await?;
                     let cursor = Cursor::new(bytes);
-                    Ok(futures::stream::iter(StreamReader::try_new(
+                    futures::stream::iter(StreamReader::try_new(
                         cursor,
                         projection.clone(),
                     )?)
                     .map(|r| r.map_err(Into::into))
-                    .boxed())
+                    .boxed()
                 }
+            };
+
+            // If we have a schema adapter factory and projected schema, use them to normalize the schema
+            if let (Some(factory), Some(proj_schema)) =
+                (schema_adapter_factory, projected_schema)
+            {
+                Ok(stream
+                    .and_then(move |batch| {
+                        let factory = factory.clone();
+                        let proj_schema = proj_schema.clone();
+                        async move {
+                            let schema_adapter =
+                                factory.create_with_projected_schema(proj_schema);
+                            let (schema_mapper, _) =
+                                schema_adapter.map_schema(batch.schema().as_ref())?;
+                            schema_mapper.map_batch(batch)
+                        }
+                    })
+                    .boxed())
+            } else {
+                Ok(stream)
             }
         }))
     }
@@ -109,36 +139,62 @@ impl FileOpener for ArrowStreamFileOpener {
 pub(crate) struct ArrowFileOpener {
     object_store: Arc<dyn ObjectStore>,
     projection: Option<Vec<usize>>,
+    projected_schema: Option<SchemaRef>,
+    schema_adapter_factory: Option<Arc<dyn SchemaAdapterFactory>>,
 }
 
 impl FileOpener for ArrowFileOpener {
     fn open(&self, partitioned_file: PartitionedFile) -> Result<FileOpenFuture> {
         let object_store = Arc::clone(&self.object_store);
         let projection = self.projection.clone();
+        let projected_schema = self.projected_schema.clone();
+        let schema_adapter_factory = self.schema_adapter_factory.clone();
+
         Ok(Box::pin(async move {
             let range = partitioned_file.range.clone();
             match range {
                 None => {
                     let r = object_store
                         .get(&partitioned_file.object_meta.location)
                         .await?;
-                    match r.payload {
+                    let stream = match r.payload {
                         #[cfg(not(target_arch = "wasm32"))]
-                        GetResultPayload::File(file, _) => Ok(futures::stream::iter(
+                        GetResultPayload::File(file, _) => futures::stream::iter(
                             FileReader::try_new(file.try_clone()?, projection.clone())?,
                         )
                         .map(|r| r.map_err(Into::into))
-                        .boxed()),
+                        .boxed(),
                         GetResultPayload::Stream(_) => {
                             let bytes = r.bytes().await?;
                             let cursor = Cursor::new(bytes);
-                            Ok(futures::stream::iter(FileReader::try_new(
+                            futures::stream::iter(FileReader::try_new(
                                 cursor,
                                 projection.clone(),
                             )?)
                             .map(|r| r.map_err(Into::into))
-                            .boxed())
+                            .boxed()
                         }
+                    };
+
+                    // Apply schema adaptation if available
+                    if let (Some(factory), Some(proj_schema)) =
+                        (schema_adapter_factory, projected_schema)
+                    {
+                        Ok(stream
+                            .and_then(move |batch| {
+                                let factory = factory.clone();
+                                let proj_schema = proj_schema.clone();
+                                async move {
+                                    let schema_adapter =
+                                        factory.create_with_projected_schema(proj_schema);
+                                    let (schema_mapper, _) = schema_adapter
+                                        .map_schema(batch.schema().as_ref())?;
+                                    schema_mapper.map_batch(batch)
+                                }
+                            })
+                            .boxed())
+                    } else {
+                        Ok(stream)
                     }
                 }
                 Some(range) => {
@@ -228,7 +284,7 @@ impl FileOpener for ArrowFileOpener {
                         )
                         .await?;
 
-                    Ok(futures::stream::iter(
+                    let stream = futures::stream::iter(
                         recordbatches
                             .into_iter()
                             .zip(recordbatch_results)
@@ -239,7 +295,29 @@ impl FileOpener for ArrowFileOpener {
                             }),
                     )
                     .map(|r| r.map_err(Into::into))
-                    .boxed())
+                    .boxed();
+
+                    // Apply schema adaptation if available
+                    if let (Some(factory), Some(proj_schema)) =
+                        (schema_adapter_factory, projected_schema)
+                    {
+                        Ok(stream
+                            .and_then(move |batch| {
+                                let factory = factory.clone();
+                                let proj_schema = proj_schema.clone();
+                                async move {
+                                    let schema_adapter =
+                                        factory.create_with_projected_schema(proj_schema);
+                                    let (schema_mapper, projection) = schema_adapter
+                                        .map_schema(batch.schema().as_ref())?;
+                                    let batch = batch.project(&projection)?;
+                                    schema_mapper.map_batch(batch)
+                                }
+                            })
+                            .boxed())
+                    } else {
+                        Ok(stream)
+                    }
                 }
             }
         }))
@@ -290,14 +368,32 @@ impl FileSource for ArrowSource {
         _partition: usize,
     ) -> Result<Arc<dyn FileOpener>> {
         let split_projection = self.projection.clone();
+        // For schema adaptation, we only use the file schema (not partition columns)
+        let projected_file_schema = SchemaRef::from(
+            self.table_schema
+                .file_schema()
+                .project(&split_projection.file_indices)?,
+        );
+
+        // Use provided schema adapter factory, or default to DefaultSchemaAdapterFactory
+        // This ensures schema normalization (removing metadata differences) happens during execution
+        let schema_adapter_factory = self
+            .schema_adapter_factory
+            .clone()
+            .unwrap_or_else(|| Arc::new(DefaultSchemaAdapterFactory));
+
         let opener: Arc<dyn FileOpener> = match self.format {
             ArrowFormat::File => Arc::new(ArrowFileOpener {
                 object_store,
                 projection: Some(split_projection.file_indices.clone()),
+                projected_schema: Some(projected_file_schema.clone()),
+                schema_adapter_factory: Some(schema_adapter_factory),
             }),
             ArrowFormat::Stream => Arc::new(ArrowStreamFileOpener {
                 object_store,
                 projection: Some(split_projection.file_indices.clone()),
+                projected_schema: Some(projected_file_schema),
+                schema_adapter_factory: Some(schema_adapter_factory),
             }),
         };
         Ok(ProjectionOpener::new(split_projection, opener)?)
@@ -429,6 +525,8 @@ impl ArrowOpener {
             inner: Arc::new(ArrowFileOpener {
                 object_store,
                 projection,
+                projected_schema: None,
+                schema_adapter_factory: None,
             }),
         }
     }
@@ -441,6 +539,8 @@ impl ArrowOpener {
             inner: Arc::new(ArrowStreamFileOpener {
                 object_store,
                 projection,
+                projected_schema: None,
+                schema_adapter_factory: None,
             }),
         }
     }
@@ -636,6 +736,8 @@ mod tests {
         let opener = ArrowStreamFileOpener {
             object_store,
             projection: Some(vec![0]), // just the first column
+            projected_schema: None,
+            schema_adapter_factory: None,
         };
 
         let mut stream = opener.open(partitioned_file)?.await?;