You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/16 16:58:36 UTC
(datafusion) 01/01: feat: Support lazy per-file Parquet Arrow schema derivation (#25343)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25343-140c7c5a4cbde4409a4d29a4c04032d618c3b8ce
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit f8f04b1e319cdb9a95aa246ac62b90e27791e0d9
Author: Peter Lee <pe...@gmail.com>
AuthorDate: Wed Sep 16 16:58:10 2026 +0000
feat: Support lazy per-file Parquet Arrow schema derivation (#25343)
## Which issue does this PR close?
Closes #25251.
## Rationale for this change
Readers such as Comet need each file's physical Parquet schema to choose
Arrow types for nested fields, but that schema is only available after
loading the footer.
## What changes are included in this PR?
Add an optional `ParquetFileSchemaProvider` to `ParquetSource`, invoked
after footer loading and before Arrow schema inference. It supplies a
complete file schema through Arrow's existing `with_schema` API while
retaining the original metadata. Explicit `PartitionedFile.arrow_schema`
takes precedence.
## What is the testing strategy for this PR?
Regression tests cover differing nested layouts, ENUM/BINARY handling,
cached and encrypted scans, schema precedence and validation, and
serialization. Extended workspace tests and lint checks pass.
## Are there any user-facing changes?
Adds `ParquetSource::with_schema_provider`. Sources using a provider
require a custom codec for plan serialization. Default scans are
unchanged.
---
datafusion/datasource-parquet/src/opener/mod.rs | 407 +++++++++++++++++++++++-
datafusion/datasource-parquet/src/reader.rs | 22 ++
datafusion/datasource-parquet/src/source.rs | 24 ++
datafusion/proto/tests/cases/plans/sources.rs | 31 ++
4 files changed, 473 insertions(+), 11 deletions(-)
diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs
index d430d81e44..0bd39f956d 100644
--- a/datafusion/datasource-parquet/src/opener/mod.rs
+++ b/datafusion/datasource-parquet/src/opener/mod.rs
@@ -258,6 +258,7 @@ pub(super) struct ParquetMorselizer {
pub metrics: ExecutionPlanMetricsSet,
/// Factory for instantiating parquet reader
pub parquet_file_reader_factory: Arc<dyn ParquetFileReaderFactory>,
+ pub schema_provider: Option<Arc<dyn crate::ParquetFileSchemaProvider>>,
/// Should the filters be evaluated during the parquet scan using
/// [`DatafusionArrowPredicate`](crate::row_filter::DatafusionArrowPredicate)?
pub pushdown_filters: bool,
@@ -443,6 +444,7 @@ struct PreparedParquetOpen {
metadata_size_hint: Option<usize>,
metrics: ExecutionPlanMetricsSet,
parquet_file_reader_factory: Arc<dyn ParquetFileReaderFactory>,
+ schema_provider: Option<Arc<dyn crate::ParquetFileSchemaProvider>>,
async_file_reader: Box<dyn AsyncFileReader>,
batch_size: usize,
logical_file_schema: SchemaRef,
@@ -907,6 +909,7 @@ impl ParquetMorselizer {
metadata_size_hint,
metrics: self.metrics.clone(),
parquet_file_reader_factory: Arc::clone(&self.parquet_file_reader_factory),
+ schema_provider: self.schema_provider.clone(),
async_file_reader,
batch_size: self.batch_size,
logical_file_schema: Arc::clone(&logical_file_schema),
@@ -1001,9 +1004,14 @@ impl PreparedParquetOpen {
// the returned metadata may actually include page indexes as some
// readers may return page indexes even when not requested -- for
// example when they are cached)
- let reader_metadata =
- ArrowReaderMetadata::load_async(&mut self.async_file_reader, options.clone())
- .await?;
+ let metadata = self.async_file_reader.get_metadata(Some(&options)).await?;
+ if self.partitioned_file.arrow_schema.is_none()
+ && let Some(provider) = &self.schema_provider
+ {
+ let schema = provider.schema(metadata.file_metadata().schema_descr())?;
+ options = options.with_schema(schema);
+ }
+ let reader_metadata = ArrowReaderMetadata::try_new(metadata, options.clone())?;
metadata_timer.stop();
drop(metadata_timer);
@@ -2026,9 +2034,10 @@ mod test {
use crate::metadata::DFParquetMetadata;
use crate::{
CachedParquetFileReaderFactory, DefaultParquetFileReaderFactory,
- ParquetFileReaderFactory, ParquetRowSelection, RowGroupAccess,
+ ParquetFileReaderFactory, ParquetFileSchemaProvider, ParquetRowSelection,
+ RowGroupAccess,
};
- use arrow::array::{RecordBatch, record_batch};
+ use arrow::array::{AsArray, RecordBatch, record_batch};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use bytes::{BufMut, BytesMut};
use datafusion_common::{
@@ -2053,15 +2062,17 @@ mod test {
};
use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion_pruning::MAX_IN_LIST_SIZE;
- use futures::StreamExt;
use futures::stream::BoxStream;
+ use futures::{StreamExt, TryStreamExt};
use object_store::{ObjectStore, ObjectStoreExt, memory::InMemory, path::Path};
- use parquet::arrow::{ArrowSchemaConverter, ArrowWriter};
+ use parquet::arrow::{ArrowSchemaConverter, ArrowWriter, parquet_to_arrow_schema};
+ use parquet::basic::ConvertedType;
use parquet::file::metadata::{ColumnChunkMetaData, FileMetaData, ParquetMetaData};
use parquet::file::properties::{EnabledStatistics, WriterProperties};
- use parquet::schema::types::SchemaDescPtr;
+ use parquet::schema::types::{SchemaDescPtr, SchemaDescriptor};
use std::collections::VecDeque;
use std::sync::Arc;
+ use std::sync::atomic::{AtomicUsize, Ordering};
/// Builder for creating [`ParquetMorselizer`] instances with sensible defaults for tests.
/// This helps reduce code duplication and makes it clear what differs between test cases.
@@ -2237,8 +2248,6 @@ mod test {
plan: ParquetAccessPlan,
) -> bool {
use crate::RowGroupAccessPlanFilter;
- use parquet::arrow::parquet_to_arrow_schema;
-
let arrow_schema: SchemaRef = Arc::new(
parquet_to_arrow_schema(metadata.file_metadata().schema_descr(), None)
.unwrap(),
@@ -2463,6 +2472,7 @@ mod test {
.unwrap_or_else(|| {
Arc::new(DefaultParquetFileReaderFactory::new(store)) as _
}),
+ schema_provider: None,
pushdown_filters: self.pushdown_filters,
reorder_filters: self.reorder_filters,
force_filter_selections: self.force_filter_selections,
@@ -2495,7 +2505,7 @@ mod test {
/// plans CPU work, awaits any discovered I/O futures, and feeds the planner
/// back into the ready queue until a stream morsel is ready.
async fn open_file(
- morselizer: &ParquetMorselizer,
+ morselizer: &dyn Morselizer,
file: PartitionedFile,
) -> Result<BoxStream<'static, Result<RecordBatch>>> {
let mut planners = VecDeque::from([morselizer.plan_file(file)?]);
@@ -3720,6 +3730,381 @@ mod test {
);
}
+ #[derive(Debug, Default)]
+ struct EnumSchemaProvider(AtomicUsize);
+
+ impl ParquetFileSchemaProvider for EnumSchemaProvider {
+ fn schema(&self, parquet_schema: &SchemaDescriptor) -> Result<SchemaRef> {
+ self.0.fetch_add(1, Ordering::Relaxed);
+ let schema = parquet_to_arrow_schema(parquet_schema, None)?;
+ let DataType::Struct(fields) = schema.field(1).data_type() else {
+ return internal_err!("expected a nested test schema");
+ };
+ let fields = fields
+ .iter()
+ .zip(parquet_schema.columns().iter().skip(1))
+ .map(|(field, column)| {
+ if column.converted_type() == ConvertedType::ENUM {
+ Arc::new(field.as_ref().clone().with_data_type(DataType::Utf8))
+ } else {
+ Arc::clone(field)
+ }
+ })
+ .collect::<arrow::datatypes::Fields>();
+ Ok(Arc::new(Schema::new(vec![
+ schema.field(0).clone(),
+ schema
+ .field(1)
+ .clone()
+ .with_data_type(DataType::Struct(fields)),
+ ])))
+ }
+ }
+
+ /// Write an ENUM beside raw bytes that are deliberately invalid UTF-8.
+ async fn write_enum_file(
+ store: &dyn ObjectStore,
+ name: &str,
+ extra_field: bool,
+ props: WriterProperties,
+ ) -> (PartitionedFile, SchemaRef) {
+ use parquet::data_type::{ByteArray, ByteArrayType, Int32Type};
+ use parquet::file::writer::SerializedFileWriter;
+ use parquet::schema::parser::parse_message_type;
+
+ let extra = if extra_field {
+ "required int32 extra;"
+ } else {
+ ""
+ };
+ let schema = Arc::new(
+ parse_message_type(&format!(
+ "message test {{ required int32 id; required group nested {{
+ required binary e (ENUM); required binary b; {extra}
+ }} }}"
+ ))
+ .unwrap(),
+ );
+ let arrow_schema = EnumSchemaProvider::default()
+ .schema(&SchemaDescriptor::new(Arc::clone(&schema)))
+ .unwrap();
+ let mut bytes = Vec::new();
+ let mut writer =
+ SerializedFileWriter::new(&mut bytes, schema, Arc::new(props)).unwrap();
+ let mut row_group = writer.next_row_group().unwrap();
+ let mut column = row_group.next_column().unwrap().unwrap();
+ column
+ .typed::<Int32Type>()
+ .write_batch(&[1, 2, 3], None, None)
+ .unwrap();
+ column.close().unwrap();
+ for values in [
+ vec![
+ ByteArray::from("a"),
+ ByteArray::from("b"),
+ ByteArray::from("c"),
+ ],
+ vec![ByteArray::from(vec![0xff, 0x00]); 3],
+ ] {
+ let mut column = row_group.next_column().unwrap().unwrap();
+ column
+ .typed::<ByteArrayType>()
+ .write_batch(&values, None, None)
+ .unwrap();
+ column.close().unwrap();
+ }
+ if extra_field {
+ let mut column = row_group.next_column().unwrap().unwrap();
+ column
+ .typed::<Int32Type>()
+ .write_batch(&[10, 20, 30], None, None)
+ .unwrap();
+ column.close().unwrap();
+ }
+ row_group.close().unwrap();
+ writer.close().unwrap();
+ let path = Path::from(name);
+ store.put(&path, bytes.into()).await.unwrap();
+ (
+ PartitionedFile::from(store.head(&path).await.unwrap()),
+ arrow_schema,
+ )
+ }
+
+ fn invalid_arrow_hint() -> parquet::file::properties::WriterPropertiesBuilder {
+ use parquet::file::metadata::KeyValue;
+ WriterProperties::builder()
+ .set_key_value_metadata(Some(vec![KeyValue::new(
+ "ARROW:schema".to_string(),
+ "invalid schema".to_string(),
+ )]))
+ .set_statistics_enabled(EnabledStatistics::Page)
+ .set_data_page_row_count_limit(1)
+ .set_write_batch_size(1)
+ }
+
+ #[tokio::test]
+ async fn test_schema_provider_lazy_nested_scans() {
+ use crate::source::ParquetSource;
+ use datafusion_datasource::file::FileSource;
+ use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
+ use datafusion_execution::object_store::ObjectStoreUrl;
+ use datafusion_functions::core::expr_fn::get_field;
+
+ let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
+ let (file1, schema) = write_enum_file(
+ store.as_ref(),
+ "one.parquet",
+ false,
+ invalid_arrow_hint().build(),
+ )
+ .await;
+ let (file2, _) = write_enum_file(
+ store.as_ref(),
+ "two.parquet",
+ true,
+ invalid_arrow_hint().build(),
+ )
+ .await;
+ let cache: Arc<FileMetadataCache> = Arc::new(DefaultCache::new(1024 * 1024));
+ let provider = Arc::new(EnumSchemaProvider::default());
+ let source = ParquetSource::new(Arc::clone(&schema))
+ .with_schema_provider(provider.clone())
+ .with_parquet_file_reader_factory(Arc::new(
+ CachedParquetFileReaderFactory::new(
+ Arc::clone(&store),
+ Arc::clone(&cache),
+ ),
+ ))
+ .with_predicate(logical2physical(
+ &get_field(col("nested"), "e").eq(lit("b")),
+ &schema,
+ ))
+ .with_pushdown_filters(true)
+ .with_enable_page_index(true)
+ .with_batch_size(1024);
+ let config = FileScanConfigBuilder::new(
+ ObjectStoreUrl::local_filesystem(),
+ Arc::clone(&source),
+ )
+ .build();
+ let morselizer = source
+ .create_morselizer(Arc::clone(&store), &config, 0)
+ .unwrap();
+ assert_eq!(provider.0.load(Ordering::Relaxed), 0);
+ for pass in 0..2 {
+ for file in [&file1, &file2] {
+ let batches = open_file(morselizer.as_ref(), file.clone())
+ .await
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ let batch = arrow::compute::concat_batches(&schema, &batches).unwrap();
+ assert_eq!(batch.num_rows(), 1);
+ assert_eq!(batch.schema(), schema);
+ let nested = batch.column(1).as_struct();
+ let e = nested.column(0).as_string::<i32>();
+ let b = nested.column(1).as_binary::<i32>();
+ assert_eq!(e.value(0), "b");
+ assert_eq!(b.value(0), &[0xff, 0x00]);
+ }
+ assert_eq!(provider.0.load(Ordering::Relaxed), (pass + 1) * 2);
+ }
+ }
+
+ #[tokio::test]
+ async fn test_schema_provider_preserves_metadata_and_precedence() {
+ use crate::metadata::{CachedParquetMetaData, DFParquetMetadata};
+
+ let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
+ let (file, schema) = write_enum_file(
+ store.as_ref(),
+ "test.parquet",
+ false,
+ invalid_arrow_hint().build(),
+ )
+ .await;
+ let cache: Arc<FileMetadataCache> = Arc::new(DefaultCache::new(1024 * 1024));
+ let original = DFParquetMetadata::new(store.as_ref(), &file.object_meta)
+ .with_file_metadata_cache(Some(Arc::clone(&cache)))
+ .with_page_index_policy(Some(PageIndexPolicy::Optional))
+ .fetch_metadata()
+ .await
+ .unwrap();
+ assert!(original.column_index().is_some());
+ assert!(original.offset_index().is_some());
+ let mut opener = ParquetMorselizerBuilder::new()
+ .with_store(Arc::clone(&store))
+ .with_schema(Arc::clone(&schema))
+ .with_parquet_file_reader_factory(Arc::new(
+ CachedParquetFileReaderFactory::new(
+ Arc::clone(&store),
+ Arc::clone(&cache),
+ ),
+ ))
+ .build();
+ // The default scan still rejects an invalid advisory Arrow hint.
+ assert!(
+ opener
+ .prepare_open_file(file.clone())
+ .unwrap()
+ .load()
+ .await
+ .is_err()
+ );
+ opener.schema_provider = Some(Arc::new(EnumSchemaProvider::default()));
+ for _ in 0..2 {
+ let loaded = opener
+ .prepare_open_file(file.clone())
+ .unwrap()
+ .load()
+ .await
+ .unwrap();
+ assert!(Arc::ptr_eq(&original, loaded.reader_metadata.metadata()));
+ assert_eq!(loaded.reader_metadata.schema(), &schema);
+ let prepared = loaded.prepare_filters().unwrap();
+ assert!(Arc::ptr_eq(
+ &original,
+ prepared.loaded.reader_metadata.metadata()
+ ));
+ }
+ let cached = cache.get(&file.object_meta.location).unwrap();
+ let cached = cached
+ .file_metadata
+ .as_any()
+ .downcast_ref::<CachedParquetMetaData>()
+ .unwrap();
+ assert!(Arc::ptr_eq(&original, cached.parquet_metadata()));
+
+ #[derive(Debug)]
+ struct RejectSchema;
+ impl ParquetFileSchemaProvider for RejectSchema {
+ fn schema(&self, _: &SchemaDescriptor) -> Result<SchemaRef> {
+ exec_err!("schema policy rejected file")
+ }
+ }
+ opener.schema_provider = Some(Arc::new(RejectSchema));
+ let err = opener
+ .prepare_open_file(file.clone())
+ .unwrap()
+ .load()
+ .await
+ .err()
+ .unwrap();
+ assert_contains!(err.to_string(), "schema policy rejected file");
+ // Explicit schemas skip the provider, but still undergo Arrow validation.
+ let loaded = opener
+ .prepare_open_file(file.clone().with_arrow_schema(Arc::clone(&schema)))
+ .unwrap()
+ .load()
+ .await
+ .unwrap();
+ assert_eq!(loaded.reader_metadata.schema(), &schema);
+ let incompatible = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Float64, false),
+ schema.field(1).clone(),
+ ]));
+ let err = opener
+ .prepare_open_file(file.clone().with_arrow_schema(Arc::clone(&incompatible)))
+ .unwrap()
+ .load()
+ .await
+ .err()
+ .unwrap();
+ assert_contains!(err.to_string(), "Incompatible supplied Arrow schema");
+
+ #[derive(Debug)]
+ struct FixedSchema(SchemaRef);
+ impl ParquetFileSchemaProvider for FixedSchema {
+ fn schema(&self, _: &SchemaDescriptor) -> Result<SchemaRef> {
+ Ok(Arc::clone(&self.0))
+ }
+ }
+ opener.schema_provider = Some(Arc::new(FixedSchema(incompatible)));
+ let err = opener
+ .prepare_open_file(file)
+ .unwrap()
+ .load()
+ .await
+ .err()
+ .unwrap();
+ assert_contains!(err.to_string(), "Incompatible supplied Arrow schema");
+ }
+
+ #[cfg(feature = "parquet_encryption")]
+ #[tokio::test]
+ async fn test_schema_provider_encrypted_scan() {
+ use datafusion_physical_plan::metrics::MetricValue;
+ use parquet::encryption::encrypt::FileEncryptionProperties;
+
+ let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
+ let footer_key = b"0123456789012345".to_vec();
+ let column_key = b"1234567890123450".to_vec();
+ for plaintext_footer in [false, true] {
+ let encrypt = FileEncryptionProperties::builder(footer_key.clone())
+ .with_plaintext_footer(plaintext_footer)
+ .with_column_key("nested.e", column_key.clone())
+ .build()
+ .unwrap();
+ let decrypt = FileDecryptionProperties::builder(footer_key.clone())
+ .with_column_key("nested.e", column_key.clone())
+ .build()
+ .unwrap();
+ let (file, schema) = write_enum_file(
+ store.as_ref(),
+ "encrypted.parquet",
+ false,
+ invalid_arrow_hint()
+ .with_file_encryption_properties(encrypt)
+ .build(),
+ )
+ .await;
+ let cache: Arc<FileMetadataCache> = Arc::new(DefaultCache::new(1024 * 1024));
+ let mut opener = ParquetMorselizerBuilder::new()
+ .with_store(Arc::clone(&store))
+ .with_schema(Arc::clone(&schema))
+ .with_predicate(logical2physical(&col("id").eq(lit(2i32)), &schema))
+ .with_pushdown_filters(true)
+ .with_enable_page_index(true)
+ .with_parquet_file_reader_factory(Arc::new(
+ CachedParquetFileReaderFactory::new(
+ Arc::clone(&store),
+ Arc::clone(&cache),
+ ),
+ ))
+ .build();
+ opener.file_decryption_properties = Some(decrypt);
+ opener.schema_provider = Some(Arc::new(EnumSchemaProvider::default()));
+ let batches = open_file(&opener, file.clone())
+ .await
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ let batch = arrow::compute::concat_batches(&schema, &batches).unwrap();
+ assert_eq!(batch.num_rows(), 1);
+ let nested = batch.column(1).as_struct();
+ assert_eq!(nested.column(0).as_string::<i32>().value(0), "b");
+ let page_pruning = opener
+ .metrics
+ .clone_inner()
+ .sum_by_name("page_index_rows_pruned")
+ .unwrap();
+ let MetricValue::PruningMetrics {
+ pruning_metrics, ..
+ } = page_pruning
+ else {
+ panic!("expected page pruning metrics");
+ };
+ assert_eq!(pruning_metrics.pruned(), 2);
+ assert!(
+ cache.get(&file.object_meta.location).is_none(),
+ "encrypted metadata must not enter the shared cache"
+ );
+ }
+ }
+
#[tokio::test]
async fn test_reverse_scan_row_groups() {
use parquet::file::properties::WriterProperties;
diff --git a/datafusion/datasource-parquet/src/reader.rs b/datafusion/datasource-parquet/src/reader.rs
index 71b0020f32..6224c21473 100644
--- a/datafusion/datasource-parquet/src/reader.rs
+++ b/datafusion/datasource-parquet/src/reader.rs
@@ -20,6 +20,7 @@
use crate::ParquetFileMetrics;
use crate::metadata::DFParquetMetadata;
+use arrow::datatypes::SchemaRef;
use bytes::Bytes;
use datafusion_common::HashMap;
use datafusion_datasource::PartitionedFile;
@@ -34,11 +35,32 @@ use parquet::arrow::arrow_reader::ArrowReaderOptions;
use parquet::arrow::async_reader::AsyncFileReader;
use parquet::errors::ParquetError;
use parquet::file::metadata::ParquetMetaData;
+use parquet::schema::types::SchemaDescriptor;
use std::any::Any;
use std::fmt::Debug;
use std::ops::Range;
use std::sync::Arc;
+/// Derives a file's Arrow schema from its physical Parquet schema during opening.
+///
+/// Install with [`ParquetSource::with_schema_provider`](crate::source::ParquetSource::with_schema_provider).
+/// This runs after the footer is loaded, before Arrow schema inference or filter
+/// preparation, so files may have different nested layouts without loading their
+/// footers during planning. The original Parquet metadata is retained.
+pub trait ParquetFileSchemaProvider: Debug + Send + Sync + 'static {
+ /// Return the complete Arrow schema for the physical file, excluding partition
+ /// and virtual columns. It must be compatible with the Parquet schema under
+ /// [`ArrowReaderOptions::with_schema`], which also ignores advisory `ARROW:schema`
+ /// metadata. Errors fail the file open.
+ ///
+ /// An explicit [`PartitionedFile::arrow_schema`] takes precedence and skips
+ /// this call. Existing table-schema and INT96 coercions are applied afterwards.
+ fn schema(
+ &self,
+ parquet_schema: &SchemaDescriptor,
+ ) -> datafusion_common::Result<SchemaRef>;
+}
+
/// Interface for reading Apache Parquet files.
///
/// The combined implementations of [`ParquetFileReaderFactory`] and
diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs
index 4872db9fd3..1015bdc600 100644
--- a/datafusion/datasource-parquet/src/source.rs
+++ b/datafusion/datasource-parquet/src/source.rs
@@ -22,6 +22,7 @@ use std::sync::Arc;
use crate::DefaultParquetFileReaderFactory;
use crate::ParquetFileReaderFactory;
+use crate::ParquetFileSchemaProvider;
use crate::opener::ParquetMorselizer;
use crate::opener::build_pruning_predicates;
use crate::opener::build_virtual_columns_state;
@@ -303,6 +304,8 @@ pub struct ParquetSource {
pub(crate) predicate: Option<Arc<dyn PhysicalExpr>>,
/// Optional user defined parquet file reader factory
pub(crate) parquet_file_reader_factory: Option<Arc<dyn ParquetFileReaderFactory>>,
+ /// Optional policy for deriving each file's Arrow schema after footer loading.
+ pub(crate) schema_provider: Option<Arc<dyn ParquetFileSchemaProvider>>,
/// Batch size configuration
pub(crate) batch_size: Option<usize>,
/// Optional hint for the size of the parquet metadata
@@ -338,6 +341,7 @@ impl ParquetSource {
metrics: ExecutionPlanMetricsSet::new(),
predicate: None,
parquet_file_reader_factory: None,
+ schema_provider: None,
batch_size: None,
metadata_size_hint: None,
#[cfg(feature = "parquet_encryption")]
@@ -416,6 +420,21 @@ impl ParquetSource {
self
}
+ /// Derive each file's Arrow schema from its physical Parquet schema during
+ /// lazy opening, before advisory Arrow metadata is parsed.
+ ///
+ /// Explicit [`PartitionedFile::arrow_schema`](datafusion_datasource::PartitionedFile::arrow_schema)
+ /// takes precedence. The provider is independent of the file reader factory
+ /// and metadata cache; see [`ParquetFileSchemaProvider`] for its contract.
+ /// Plans using a provider require a custom codec for serialization.
+ pub fn with_schema_provider(
+ mut self,
+ schema_provider: Arc<dyn ParquetFileSchemaProvider>,
+ ) -> Self {
+ self.schema_provider = Some(schema_provider);
+ self
+ }
+
/// If true, the predicate will be used during the parquet scan.
/// Defaults to false.
pub fn with_pushdown_filters(mut self, pushdown_filters: bool) -> Self {
@@ -645,6 +664,7 @@ impl FileSource for ParquetSource {
metadata_size_hint: self.metadata_size_hint,
metrics: self.metrics().clone(),
parquet_file_reader_factory,
+ schema_provider: self.schema_provider.clone(),
pushdown_filters: self.pushdown_filters(),
reorder_filters: self.reorder_filters(),
force_filter_selections: self.force_filter_selections(),
@@ -1089,6 +1109,10 @@ impl FileSource for ParquetSource {
use datafusion_proto_models::protobuf;
use protobuf::physical_plan_node::PhysicalPlanType;
+ if self.schema_provider.is_some() {
+ return Ok(None);
+ }
+
let predicate = self
.filter()
.map(|pred| ctx.encode_expr(&pred))
diff --git a/datafusion/proto/tests/cases/plans/sources.rs b/datafusion/proto/tests/cases/plans/sources.rs
index 7aff89eb3f..7f04a794f2 100644
--- a/datafusion/proto/tests/cases/plans/sources.rs
+++ b/datafusion/proto/tests/cases/plans/sources.rs
@@ -74,6 +74,37 @@ use std::fmt::{Display, Formatter};
use std::sync::Arc;
use std::vec;
+#[test]
+fn parquet_schema_provider_requires_extension_codec() -> Result<()> {
+ use datafusion::parquet;
+
+ #[derive(Debug)]
+ struct PhysicalSchema;
+ impl datafusion_datasource_parquet::ParquetFileSchemaProvider for PhysicalSchema {
+ fn schema(
+ &self,
+ schema: &parquet::schema::types::SchemaDescriptor,
+ ) -> Result<Arc<Schema>> {
+ Ok(Arc::new(parquet::arrow::parquet_to_arrow_schema(
+ schema, None,
+ )?))
+ }
+ }
+
+ let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
+ let source =
+ ParquetSource::new(schema).with_schema_provider(Arc::new(PhysicalSchema));
+ let config =
+ FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), Arc::new(source))
+ .build();
+ let plan = DataSourceExec::from_data_source(config);
+ let error =
+ PhysicalPlanNode::try_from_physical_plan(plan, &DefaultPhysicalExtensionCodec {})
+ .unwrap_err();
+ datafusion_common::assert_contains!(error.to_string(), "extension codec failed");
+ Ok(())
+}
+
#[test]
fn roundtrip_parquet_exec_with_pruning_predicate() -> Result<()> {
let file_schema =
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org