use std::fmt;
use std::fmt::Debug;
use std::ops::Range;
use std::sync::Arc;
#[expect(deprecated)]
pub use crate::schema_coercion::coerce_int96_to_resolution;
pub use crate::schema_coercion::{
Int96Coercer, apply_file_schema_type_coercions, transform_binary_to_string,
transform_schema_to_view,
};
pub use crate::sink::ParquetSink;
use arrow::datatypes::{Fields, Schema, SchemaRef};
use datafusion_datasource::TableSchema;
use datafusion_datasource::file_compression_type::FileCompressionType;
use datafusion_datasource::file_sink_config::FileSinkConfig;
use datafusion_datasource::file_format::{FileFormat, FileFormatFactory};
use datafusion_common::Statistics;
use datafusion_common::config::{ConfigField, ConfigFileType, TableParquetOptions};
use datafusion_common::encryption::FileDecryptionProperties;
use datafusion_common::parsers::CompressionTypeVariant;
use datafusion_common::{
DEFAULT_PARQUET_EXTENSION, DataFusionError, GetExt, Result, internal_datafusion_err,
internal_err, not_impl_err,
};
use datafusion_datasource::file::FileSource;
use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
use datafusion_datasource::sink::DataSinkExec;
use datafusion_datasource::write::get_writer_schema;
use datafusion_expr::dml::InsertOp;
use datafusion_physical_expr_common::sort_expr::{LexOrdering, LexRequirement};
use datafusion_physical_plan::ExecutionPlan;
use datafusion_session::Session;
use crate::metadata::{DFParquetMetadata, lex_ordering_to_sorting_columns};
use crate::reader::CachedParquetFileReaderFactory;
use crate::source::{
ParquetSource, parse_coerce_int96_string, parse_coerce_int96_tz_string,
};
use async_trait::async_trait;
use bytes::Bytes;
use datafusion_datasource::source::DataSourceExec;
use datafusion_execution::cache::cache_manager::FileMetadataCache;
use futures::future::BoxFuture;
use futures::{FutureExt, StreamExt, TryStreamExt};
use object_store::path::Path;
use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt};
use parquet::arrow::async_reader::MetadataFetch;
use parquet::errors::ParquetError;
use parquet::file::metadata::ParquetMetaData;
#[derive(Default)]
pub struct ParquetFormatFactory {
pub options: Option<TableParquetOptions>,
}
impl ParquetFormatFactory {
pub fn new() -> Self {
Self { options: None }
}
pub fn new_with_options(options: TableParquetOptions) -> Self {
Self {
options: Some(options),
}
}
}
impl FileFormatFactory for ParquetFormatFactory {
fn create(
&self,
state: &dyn Session,
format_options: &std::collections::HashMap<String, String>,
) -> Result<Arc<dyn FileFormat>> {
let parquet_options = match &self.options {
None => {
let mut table_options = state.default_table_options();
table_options.set_config_format(ConfigFileType::PARQUET);
table_options.alter_with_string_hash_map(format_options)?;
table_options.parquet
}
Some(parquet_options) => {
let mut parquet_options = parquet_options.clone();
for (k, v) in format_options {
parquet_options.set(k, v)?;
}
parquet_options
}
};
Ok(Arc::new(
ParquetFormat::default().with_options(parquet_options),
))
}
fn default(&self) -> Arc<dyn FileFormat> {
Arc::new(ParquetFormat::default())
}
}
impl GetExt for ParquetFormatFactory {
fn get_ext(&self) -> String {
DEFAULT_PARQUET_EXTENSION[1..].to_string()
}
}
impl Debug for ParquetFormatFactory {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ParquetFormatFactory")
.field("ParquetFormatFactory", &self.options)
.finish()
}
}
#[derive(Debug, Default)]
pub struct ParquetFormat {
options: TableParquetOptions,
}
impl ParquetFormat {
pub fn new() -> Self {
Self::default()
}
pub fn with_enable_pruning(mut self, enable: bool) -> Self {
self.options.global.pruning = enable;
self
}
pub fn enable_pruning(&self) -> bool {
self.options.global.pruning
}
pub fn with_metadata_size_hint(mut self, size_hint: Option<usize>) -> Self {
self.options.global.metadata_size_hint = size_hint;
self
}
pub fn metadata_size_hint(&self) -> Option<usize> {
self.options.global.metadata_size_hint
}
pub fn with_skip_metadata(mut self, skip_metadata: bool) -> Self {
self.options.global.skip_metadata = skip_metadata;
self
}
pub fn skip_metadata(&self) -> bool {
self.options.global.skip_metadata
}
pub fn with_options(mut self, options: TableParquetOptions) -> Self {
self.options = options;
self
}
pub fn options(&self) -> &TableParquetOptions {
&self.options
}
pub fn force_view_types(&self) -> bool {
self.options.global.schema_force_view_types
}
pub fn with_force_view_types(mut self, use_views: bool) -> Self {
self.options.global.schema_force_view_types = use_views;
self
}
pub fn binary_as_string(&self) -> bool {
self.options.global.binary_as_string
}
pub fn with_binary_as_string(mut self, binary_as_string: bool) -> Self {
self.options.global.binary_as_string = binary_as_string;
self
}
pub fn coerce_int96(&self) -> Option<String> {
self.options.global.coerce_int96.clone()
}
pub fn with_coerce_int96(mut self, time_unit: Option<String>) -> Self {
self.options.global.coerce_int96 = time_unit;
self
}
}
fn clear_metadata(
schemas: impl IntoIterator<Item = Schema>,
) -> impl Iterator<Item = Schema> {
schemas.into_iter().map(|schema| {
let fields = schema
.fields()
.iter()
.map(|field| {
field.as_ref().clone().with_metadata(Default::default()) })
.collect::<Fields>();
Schema::new(fields)
})
}
#[cfg(feature = "parquet_encryption")]
async fn get_file_decryption_properties(
state: &dyn Session,
options: &TableParquetOptions,
file_path: &Path,
) -> Result<Option<Arc<FileDecryptionProperties>>> {
Ok(match &options.crypto.file_decryption {
Some(cfd) => Some(Arc::new(FileDecryptionProperties::try_from(cfd.clone())?)),
None => match &options.crypto.factory_id {
Some(factory_id) => {
let factory =
state.runtime_env().parquet_encryption_factory(factory_id)?;
factory
.get_file_decryption_properties(
&options.crypto.factory_options,
file_path,
)
.await?
}
None => None,
},
})
}
#[cfg(not(feature = "parquet_encryption"))]
#[expect(clippy::unused_async)]
async fn get_file_decryption_properties(
_state: &dyn Session,
_options: &TableParquetOptions,
_file_path: &Path,
) -> Result<Option<Arc<FileDecryptionProperties>>> {
Ok(None)
}
#[async_trait]
impl FileFormat for ParquetFormat {
fn get_ext(&self) -> String {
ParquetFormatFactory::new().get_ext()
}
fn get_ext_with_compression(
&self,
file_compression_type: &FileCompressionType,
) -> Result<String> {
let ext = self.get_ext();
match file_compression_type.get_variant() {
CompressionTypeVariant::UNCOMPRESSED => Ok(ext),
_ => internal_err!("Parquet FileFormat does not support compression."),
}
}
fn compression_type(&self) -> Option<FileCompressionType> {
None
}
async fn infer_schema(
&self,
state: &dyn Session,
store: &Arc<dyn ObjectStore>,
objects: &[ObjectMeta],
) -> Result<SchemaRef> {
let coerce_int96 = match self.coerce_int96() {
Some(time_unit) => Some(parse_coerce_int96_string(time_unit.as_str())?),
None => None,
};
let coerce_int96_tz = self
.options
.global
.coerce_int96_tz
.as_ref()
.map(|tz| parse_coerce_int96_tz_string(tz))
.transpose()?;
let file_metadata_cache =
state.runtime_env().cache_manager.get_file_metadata_cache();
let mut schemas: Vec<_> = futures::stream::iter(objects)
.map(|object| async {
let file_decryption_properties = get_file_decryption_properties(
state,
&self.options,
&object.location,
)
.await?;
let result = DFParquetMetadata::new(store.as_ref(), object)
.with_metadata_size_hint(self.metadata_size_hint())
.with_decryption_properties(file_decryption_properties)
.with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache)))
.with_coerce_int96(coerce_int96)
.with_coerce_int96_tz(coerce_int96_tz.clone())
.fetch_schema_with_location()
.await?;
Ok::<_, DataFusionError>(result)
})
.boxed() .buffer_unordered(
state
.config_options()
.execution
.meta_fetch_concurrency
.get(),
)
.try_collect()
.await?;
schemas
.sort_unstable_by(|(location1, _), (location2, _)| location1.cmp(location2));
let schemas = schemas.into_iter().map(|(_, schema)| schema);
let schema = if self.skip_metadata() {
Schema::try_merge(clear_metadata(schemas))
} else {
Schema::try_merge(schemas)
}?;
let schema = if self.binary_as_string() {
transform_binary_to_string(&schema)
} else {
schema
};
let schema = if self.force_view_types() {
transform_schema_to_view(&schema)
} else {
schema
};
Ok(Arc::new(schema))
}
async fn infer_stats(
&self,
state: &dyn Session,
store: &Arc<dyn ObjectStore>,
table_schema: SchemaRef,
object: &ObjectMeta,
) -> Result<Statistics> {
let file_decryption_properties =
get_file_decryption_properties(state, &self.options, &object.location)
.await?;
let file_metadata_cache =
state.runtime_env().cache_manager.get_file_metadata_cache();
DFParquetMetadata::new(store, object)
.with_metadata_size_hint(self.metadata_size_hint())
.with_decryption_properties(file_decryption_properties)
.with_file_metadata_cache(Some(file_metadata_cache))
.fetch_statistics(&table_schema)
.await
}
async fn infer_ordering(
&self,
state: &dyn Session,
store: &Arc<dyn ObjectStore>,
table_schema: SchemaRef,
object: &ObjectMeta,
) -> Result<Option<LexOrdering>> {
let file_decryption_properties =
get_file_decryption_properties(state, &self.options, &object.location)
.await?;
let file_metadata_cache =
state.runtime_env().cache_manager.get_file_metadata_cache();
let metadata = DFParquetMetadata::new(store, object)
.with_metadata_size_hint(self.metadata_size_hint())
.with_decryption_properties(file_decryption_properties)
.with_file_metadata_cache(Some(file_metadata_cache))
.fetch_metadata()
.await?;
crate::metadata::ordering_from_parquet_metadata(&metadata, &table_schema)
}
async fn infer_stats_and_ordering(
&self,
state: &dyn Session,
store: &Arc<dyn ObjectStore>,
table_schema: SchemaRef,
object: &ObjectMeta,
) -> Result<datafusion_datasource::file_format::FileMeta> {
let file_decryption_properties =
get_file_decryption_properties(state, &self.options, &object.location)
.await?;
let file_metadata_cache =
state.runtime_env().cache_manager.get_file_metadata_cache();
let metadata = DFParquetMetadata::new(store, object)
.with_metadata_size_hint(self.metadata_size_hint())
.with_decryption_properties(file_decryption_properties)
.with_file_metadata_cache(Some(file_metadata_cache))
.fetch_metadata()
.await?;
let statistics = DFParquetMetadata::statistics_from_parquet_metadata(
&metadata,
&table_schema,
)?;
let ordering =
crate::metadata::ordering_from_parquet_metadata(&metadata, &table_schema)?;
Ok(
datafusion_datasource::file_format::FileMeta::new(statistics)
.with_ordering(ordering),
)
}
async fn create_physical_plan(
&self,
state: &dyn Session,
conf: FileScanConfig,
) -> Result<Arc<dyn ExecutionPlan>> {
let mut metadata_size_hint = None;
if let Some(metadata) = self.metadata_size_hint() {
metadata_size_hint = Some(metadata);
}
let mut source = conf
.file_source()
.downcast_ref::<ParquetSource>()
.cloned()
.ok_or_else(|| internal_datafusion_err!("Expected ParquetSource"))?;
source = source.with_table_parquet_options(self.options.clone());
let metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache();
let store = state
.runtime_env()
.object_store(conf.object_store_url.clone())?;
let cached_parquet_read_factory =
Arc::new(CachedParquetFileReaderFactory::new(store, metadata_cache));
source = source.with_parquet_file_reader_factory(cached_parquet_read_factory);
if let Some(metadata_size_hint) = metadata_size_hint {
source = source.with_metadata_size_hint(metadata_size_hint)
}
source = self.set_source_encryption_factory(source, state)?;
let conf = FileScanConfigBuilder::from(conf)
.with_source(Arc::new(source))
.build();
Ok(DataSourceExec::from_data_source(conf))
}
async fn create_writer_physical_plan(
&self,
input: Arc<dyn ExecutionPlan>,
_state: &dyn Session,
conf: FileSinkConfig,
order_requirements: Option<LexRequirement>,
) -> Result<Arc<dyn ExecutionPlan>> {
if conf.insert_op != InsertOp::Append {
return not_impl_err!("Overwrites are not implemented yet for Parquet");
}
let sorting_columns = if let Some(ref requirements) = order_requirements {
let ordering: LexOrdering = requirements.clone().into();
let writer_schema = get_writer_schema(&conf);
lex_ordering_to_sorting_columns(
&ordering,
conf.output_schema(),
&writer_schema,
)
.ok()
.filter(|columns| !columns.is_empty())
} else {
None
};
let sink = Arc::new(
ParquetSink::new(conf, self.options.clone())
.with_sorting_columns(sorting_columns),
);
Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _)
}
fn file_source(&self, table_schema: TableSchema) -> Arc<dyn FileSource> {
Arc::new(
ParquetSource::new(table_schema)
.with_table_parquet_options(self.options.clone()),
)
}
}
#[cfg(feature = "parquet_encryption")]
impl ParquetFormat {
fn set_source_encryption_factory(
&self,
source: ParquetSource,
state: &dyn Session,
) -> Result<ParquetSource> {
if let Some(encryption_factory_id) = &self.options.crypto.factory_id {
Ok(source.with_encryption_factory(
state
.runtime_env()
.parquet_encryption_factory(encryption_factory_id)?,
))
} else {
Ok(source)
}
}
}
#[cfg(not(feature = "parquet_encryption"))]
impl ParquetFormat {
fn set_source_encryption_factory(
&self,
source: ParquetSource,
_state: &dyn Session,
) -> Result<ParquetSource> {
if let Some(encryption_factory_id) = &self.options.crypto.factory_id {
Err(DataFusionError::Configuration(format!(
"Parquet encryption factory id is set to '{encryption_factory_id}' but the parquet_encryption feature is disabled"
)))
} else {
Ok(source)
}
}
}
pub struct ObjectStoreFetch<'a> {
store: &'a dyn ObjectStore,
meta: &'a ObjectMeta,
}
impl<'a> ObjectStoreFetch<'a> {
pub fn new(store: &'a dyn ObjectStore, meta: &'a ObjectMeta) -> Self {
Self { store, meta }
}
}
impl MetadataFetch for ObjectStoreFetch<'_> {
fn fetch(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes, ParquetError>> {
async {
self.store
.get_range(&self.meta.location, range)
.await
.map_err(ParquetError::from)
}
.boxed()
}
}
#[deprecated(
since = "50.0.0",
note = "Use `DFParquetMetadata::fetch_metadata` instead"
)]
pub async fn fetch_parquet_metadata(
store: &dyn ObjectStore,
object_meta: &ObjectMeta,
size_hint: Option<usize>,
decryption_properties: Option<&FileDecryptionProperties>,
file_metadata_cache: Option<Arc<FileMetadataCache>>,
) -> Result<Arc<ParquetMetaData>> {
let decryption_properties = decryption_properties.cloned().map(Arc::new);
DFParquetMetadata::new(store, object_meta)
.with_metadata_size_hint(size_hint)
.with_decryption_properties(decryption_properties)
.with_file_metadata_cache(file_metadata_cache)
.fetch_metadata()
.await
}
#[deprecated(
since = "50.0.0",
note = "Use `DFParquetMetadata::fetch_statistics` instead"
)]
pub async fn fetch_statistics(
store: &dyn ObjectStore,
table_schema: SchemaRef,
file: &ObjectMeta,
metadata_size_hint: Option<usize>,
decryption_properties: Option<&FileDecryptionProperties>,
file_metadata_cache: Option<Arc<FileMetadataCache>>,
) -> Result<Statistics> {
let decryption_properties = decryption_properties.cloned().map(Arc::new);
DFParquetMetadata::new(store, file)
.with_metadata_size_hint(metadata_size_hint)
.with_decryption_properties(decryption_properties)
.with_file_metadata_cache(file_metadata_cache)
.fetch_statistics(&table_schema)
.await
}
#[deprecated(
since = "50.0.0",
note = "Use `DFParquetMetadata::statistics_from_parquet_metadata` instead"
)]
#[expect(clippy::needless_pass_by_value)]
pub fn statistics_from_parquet_meta_calc(
metadata: &ParquetMetaData,
table_schema: SchemaRef,
) -> Result<Statistics> {
DFParquetMetadata::statistics_from_parquet_metadata(metadata, &table_schema)
}
#[cfg(feature = "proto")]
use datafusion_proto_models::protobuf::{self, parquet_column_options, parquet_options};
#[cfg(feature = "proto")]
impl From<&ParquetFormatFactory> for protobuf::TableParquetOptions {
fn from(factory: &ParquetFormatFactory) -> Self {
let global_options = if let Some(ref options) = factory.options {
options.clone()
} else {
return protobuf::TableParquetOptions::default();
};
let column_specific_options = global_options.column_specific_options;
protobuf::TableParquetOptions {
global: Some(protobuf::ParquetOptions {
enable_page_index: global_options.global.enable_page_index,
pruning: global_options.global.pruning,
skip_metadata: global_options.global.skip_metadata,
metadata_size_hint_opt: global_options.global.metadata_size_hint.map(|size| {
parquet_options::MetadataSizeHintOpt::MetadataSizeHint(size as u64)
}),
pushdown_filters: global_options.global.pushdown_filters,
reorder_filters: global_options.global.reorder_filters,
force_filter_selections: global_options.global.force_filter_selections,
data_pagesize_limit: global_options.global.data_pagesize_limit as u64,
write_batch_size: global_options.global.write_batch_size as u64,
writer_version: global_options.global.writer_version.to_string(),
compression_opt: global_options.global.compression.map(|compression| {
parquet_options::CompressionOpt::Compression(compression)
}),
dictionary_enabled_opt: global_options.global.dictionary_enabled.map(|enabled| {
parquet_options::DictionaryEnabledOpt::DictionaryEnabled(enabled)
}),
dictionary_page_size_limit: global_options.global.dictionary_page_size_limit as u64,
statistics_enabled_opt: global_options.global.statistics_enabled.map(|enabled| {
parquet_options::StatisticsEnabledOpt::StatisticsEnabled(enabled)
}),
max_row_group_size: global_options.global.max_row_group_size as u64,
max_in_list_size: global_options.global.max_in_list_size as u64,
created_by: global_options.global.created_by.clone(),
column_index_truncate_length_opt: global_options.global.column_index_truncate_length.map(|length| {
parquet_options::ColumnIndexTruncateLengthOpt::ColumnIndexTruncateLength(length as u64)
}),
statistics_truncate_length_opt: global_options.global.statistics_truncate_length.map(|length| {
parquet_options::StatisticsTruncateLengthOpt::StatisticsTruncateLength(length as u64)
}),
data_page_row_count_limit: global_options.global.data_page_row_count_limit as u64,
encoding_opt: global_options.global.encoding.map(|encoding| {
parquet_options::EncodingOpt::Encoding(encoding)
}),
bloom_filter_on_read: global_options.global.bloom_filter_on_read,
bloom_filter_on_write: global_options.global.bloom_filter_on_write,
bloom_filter_fpp_opt: global_options.global.bloom_filter_fpp.map(|fpp| {
parquet_options::BloomFilterFppOpt::BloomFilterFpp(fpp)
}),
bloom_filter_ndv_opt: global_options.global.bloom_filter_ndv.map(|ndv| {
parquet_options::BloomFilterNdvOpt::BloomFilterNdv(ndv)
}),
allow_single_file_parallelism: global_options.global.allow_single_file_parallelism,
maximum_parallel_row_group_writers: global_options.global.maximum_parallel_row_group_writers as u64,
maximum_buffered_record_batches_per_stream: global_options.global.maximum_buffered_record_batches_per_stream as u64,
schema_force_view_types: global_options.global.schema_force_view_types,
binary_as_string: global_options.global.binary_as_string,
skip_arrow_metadata: global_options.global.skip_arrow_metadata,
coerce_int96_opt: global_options.global.coerce_int96.map(|compression| {
parquet_options::CoerceInt96Opt::CoerceInt96(compression)
}),
coerce_int96_tz_opt: global_options.global.coerce_int96_tz.map(|tz| {
parquet_options::CoerceInt96TzOpt::CoerceInt96Tz(tz)
}),
max_predicate_cache_size_opt: global_options.global.max_predicate_cache_size.map(|size| {
parquet_options::MaxPredicateCacheSizeOpt::MaxPredicateCacheSize(size as u64)
}),
max_row_group_bytes_opt: global_options.global.max_row_group_bytes.map(|size| {
parquet_options::MaxRowGroupBytesOpt::MaxRowGroupBytes(size.get() as u64)
}),
content_defined_chunking: Some(protobuf::ParquetCdcOptions {
enabled: global_options.global.content_defined_chunking.enabled,
min_chunk_size: global_options.global.content_defined_chunking.min_chunk_size as u64,
max_chunk_size: global_options.global.content_defined_chunking.max_chunk_size as u64,
norm_level: global_options.global.content_defined_chunking.norm_level,
}),
}),
column_specific_options: column_specific_options.into_iter().map(|(column_name, options)| {
protobuf::ParquetColumnSpecificOptions {
column_name,
options: Some(protobuf::ParquetColumnOptions {
bloom_filter_enabled_opt: options.bloom_filter_enabled.map(|enabled| {
parquet_column_options::BloomFilterEnabledOpt::BloomFilterEnabled(enabled)
}),
encoding_opt: options.encoding.map(|encoding| {
parquet_column_options::EncodingOpt::Encoding(encoding)
}),
dictionary_enabled_opt: options.dictionary_enabled.map(|enabled| {
parquet_column_options::DictionaryEnabledOpt::DictionaryEnabled(enabled)
}),
compression_opt: options.compression.map(|compression| {
parquet_column_options::CompressionOpt::Compression(compression)
}),
statistics_enabled_opt: options.statistics_enabled.map(|enabled| {
parquet_column_options::StatisticsEnabledOpt::StatisticsEnabled(enabled)
}),
bloom_filter_fpp_opt: options.bloom_filter_fpp.map(|fpp| {
parquet_column_options::BloomFilterFppOpt::BloomFilterFpp(fpp)
}),
bloom_filter_ndv_opt: options.bloom_filter_ndv.map(|ndv| {
parquet_column_options::BloomFilterNdvOpt::BloomFilterNdv(ndv)
}),
})
}
}).collect(),
key_value_metadata: global_options.key_value_metadata
.iter()
.filter_map(|(key, value)| {
value.as_ref().map(|v| (key.clone(), v.clone()))
})
.collect(),
}
}
}