diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index d430d81e44bae..0bd39f956dbad 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, + pub schema_provider: Option>, /// 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, metrics: ExecutionPlanMetricsSet, parquet_file_reader_factory: Arc, + schema_provider: Option>, async_file_reader: Box, 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>> { 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 { + 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::(); + 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::() + .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::() + .write_batch(&values, None, None) + .unwrap(); + column.close().unwrap(); + } + if extra_field { + let mut column = row_group.next_column().unwrap().unwrap(); + column + .typed::() + .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 = 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 = 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::>() + .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::(); + let b = nested.column(1).as_binary::(); + 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 = Arc::new(InMemory::new()); + let (file, schema) = write_enum_file( + store.as_ref(), + "test.parquet", + false, + invalid_arrow_hint().build(), + ) + .await; + let cache: Arc = 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::() + .unwrap(); + assert!(Arc::ptr_eq(&original, cached.parquet_metadata())); + + #[derive(Debug)] + struct RejectSchema; + impl ParquetFileSchemaProvider for RejectSchema { + fn schema(&self, _: &SchemaDescriptor) -> Result { + 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 { + 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 = 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 = 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::>() + .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::().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 71b0020f32f64..6224c21473139 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; +} + /// 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 4872db9fd3329..1015bdc600bdd 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>, /// Optional user defined parquet file reader factory pub(crate) parquet_file_reader_factory: Option>, + /// Optional policy for deriving each file's Arrow schema after footer loading. + pub(crate) schema_provider: Option>, /// Batch size configuration pub(crate) batch_size: Option, /// 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, + ) -> 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 7aff89eb3fdf7..7f04a794f2803 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> { + 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 =