Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions datafusion/common/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1381,6 +1381,17 @@ config_namespace! {
/// Defaults to 20.
pub max_in_list_size: usize, default = 20

/// (reading) If true, top-level string and binary Parquet columns with
/// dictionary pages are inferred and scanned as
/// `Dictionary<Int32, Utf8>` / `Dictionary<Int32, Binary>` instead of
/// their plain value type.
///
/// This applies only when DataFusion infers the table schema. Tables with
/// a user-supplied schema are not promoted because the Parquet footer is
/// not read at DDL time, so dictionary pages cannot be detected per column.
/// See <https://github.com/apache/datafusion/issues/24112>
pub enable_rle_to_dictionary: bool, default = false

// The following options affect writing to parquet files
// and map to parquet::file::properties::WriterProperties

Expand Down
4 changes: 4 additions & 0 deletions datafusion/common/src/file_options/parquet_writer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -249,6 +249,7 @@ impl ParquetOptions {
skip_arrow_metadata: _,
max_predicate_cache_size: _,
max_in_list_size: _,
enable_rle_to_dictionary: _,
} = self;

let mut builder = WriterProperties::builder()
Expand Down Expand Up @@ -509,6 +510,7 @@ mod tests {
coerce_int96_tz: None,
max_predicate_cache_size: defaults.max_predicate_cache_size,
content_defined_chunking: defaults.content_defined_chunking.clone(),
enable_rle_to_dictionary: defaults.enable_rle_to_dictionary,
}
}

Expand Down Expand Up @@ -631,6 +633,8 @@ mod tests {
coerce_int96: None,
coerce_int96_tz: None,
content_defined_chunking: props.content_defined_chunking().into(),
enable_rle_to_dictionary: global_options_defaults
.enable_rle_to_dictionary,
},
column_specific_options,
key_value_metadata,
Expand Down
18 changes: 13 additions & 5 deletions datafusion/datasource-parquet/src/file_format.rs
Original file line number Diff line number Diff line change
Expand Up @@ -360,14 +360,16 @@ impl FileFormat for ParquetFormat {
&object.location,
)
.await?;
let result = DFParquetMetadata::new(store.as_ref(), object)
let meta = 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?;
.with_enable_rle_to_dictionary(
self.options.global.enable_rle_to_dictionary,
);
let result = meta.fetch_schema_with_location().await?;
Ok::<_, DataFusionError>(result)
})
.boxed() // Workaround https://github.com/rust-lang/rust/issues/64552
Expand Down Expand Up @@ -402,7 +404,12 @@ impl FileFormat for ParquetFormat {
}
drop(seen);

let schemas = schemas.into_iter().map(|(_, schema)| schema);
// Normalize dict-promoted schemas before merging so mixed dict/plain files merge cleanly.
let mut schemas: Vec<Schema> =
schemas.into_iter().map(|(_, schema)| schema).collect();
Comment thread
Rich-T-kid marked this conversation as resolved.
if self.options.global.enable_rle_to_dictionary {
schemas = crate::schema_coercion::uniform_dict_schemas(schemas);
}

let schema = if self.skip_metadata() {
Schema::try_merge(clear_metadata(schemas))
Expand Down Expand Up @@ -524,7 +531,7 @@ impl FileFormat for ParquetFormat {
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 = source.with_metadata_size_hint(metadata_size_hint);
}

source = self.set_source_encryption_factory(source, state)?;
Expand Down Expand Up @@ -738,6 +745,7 @@ impl From<&ParquetFormatFactory> for protobuf::TableParquetOptions {
compression_opt: global_options.global.compression.map(|compression| {
parquet_options::CompressionOpt::Compression(compression)
}),
enable_rle_to_dictionary: global_options.global.enable_rle_to_dictionary,
dictionary_enabled_opt: global_options.global.dictionary_enabled.map(|enabled| {
parquet_options::DictionaryEnabledOpt::DictionaryEnabled(enabled)
}),
Expand Down
75 changes: 73 additions & 2 deletions datafusion/datasource-parquet/src/metadata.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,11 +23,11 @@ use crate::{Int96Coercer, apply_file_schema_type_coercions};
use arrow::array::{Array, ArrayRef, BooleanArray};
use arrow::compute::kernels::cmp::eq;
use arrow::compute::{and, sum};
use arrow::datatypes::{DataType, Schema, SchemaRef, TimeUnit};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
use datafusion_common::encryption::FileDecryptionProperties;
use datafusion_common::stats::Precision;
use datafusion_common::{
ColumnStatistics, DataFusionError, HashMap, Result, ScalarValue, Statistics,
ColumnStatistics, DataFusionError, HashMap, HashSet, Result, ScalarValue, Statistics,
internal_datafusion_err,
};
use datafusion_execution::cache::cache_manager::{
Expand Down Expand Up @@ -146,6 +146,8 @@ pub struct DFParquetMetadata<'a> {
pub coerce_int96: Option<TimeUnit>,
/// Optional timezone applied to INT96-coerced timestamps.
pub coerce_int96_tz: Option<Arc<str>>,
/// If true, promote string/binary columns with dictionary pages to `Dictionary(Int32, ...)`.
enable_rle_to_dictionary: bool,
}

impl<'a> DFParquetMetadata<'a> {
Expand All @@ -163,9 +165,16 @@ impl<'a> DFParquetMetadata<'a> {
page_index_policy: None,
coerce_int96: None,
coerce_int96_tz: None,
enable_rle_to_dictionary: false,
}
}

/// Promote string/binary columns with dictionary pages to `Dictionary(Int32, ...)`.
pub fn with_enable_rle_to_dictionary(mut self, enable: bool) -> Self {
self.enable_rle_to_dictionary = enable;
self
}

/// Set a hint for the number of trailing bytes to prefetch from the end
/// of the file, equivalent to
/// [`ParquetMetaDataReader::with_prefetch_hint`].
Expand Down Expand Up @@ -439,6 +448,68 @@ impl<'a> DFParquetMetadata<'a> {
.coerce()
})
.unwrap_or(schema);

let schema = if self.enable_rle_to_dictionary {
let schema_descr = file_metadata.schema_descr();
// Top-level columns that have a dictionary page in at least one row group.
let dict_cols: HashSet<String> = metadata
.row_groups()
.iter()
.flat_map(|rg| {
rg.columns()
.iter()
.enumerate()
.filter_map(|(col_idx, col)| {
col.dictionary_page_offset()?;
let col_desc = schema_descr.column(col_idx);
let parts = col_desc.path().parts();
// Skip nested columns: their leaf name doesn't match the
// Arrow top-level field name.
(parts.len() == 1).then(|| parts[0].clone())
})
})
.collect();
if dict_cols.is_empty() {
schema
} else {
let promoted: Vec<_> = schema
.fields()
.iter()
.map(|field| {
if !dict_cols.contains(field.name()) {
return Arc::clone(field);
}
let dict_value_type = match field.data_type() {
DataType::Utf8 => Some(DataType::Utf8),
DataType::LargeUtf8 => Some(DataType::LargeUtf8),
DataType::Binary => Some(DataType::Binary),
DataType::LargeBinary => Some(DataType::LargeBinary),
_ => None,
};
dict_value_type.map_or_else(
|| Arc::clone(field),
|value_type| {
Arc::new(
Field::new(
field.name(),
DataType::Dictionary(
Box::new(DataType::Int32),
Box::new(value_type),
),
field.is_nullable(),
)
.with_metadata(field.metadata().clone()),
)
},
)
})
.collect();
Schema::new_with_metadata(promoted, schema.metadata().clone())
}
} else {
schema
};

Ok(schema)
}

Expand Down
72 changes: 70 additions & 2 deletions datafusion/datasource-parquet/src/opener/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ use crate::row_group_filter::{RowGroupAccessPlanFilter, row_group_in_range};
use crate::{
BloomFilterStatistics, Int96Coercer, ParquetAccessPlan, ParquetFileMetrics,
ParquetFileReaderFactory, ParquetRowSelection, ParquetVirtualColumn,
apply_file_schema_type_coercions,
schema_coercion::apply_file_schema_type_coercions_with_rle,
};
use arrow::array::RecordBatch;
use arrow::datatypes::DataType;
Expand Down Expand Up @@ -294,6 +294,8 @@ pub(super) struct ParquetMorselizer {
/// lists skip container-level pruning. Sourced from
/// `datafusion.execution.parquet.max_in_list_size`.
pub max_in_list_size: usize,
/// Whether to ask arrow-rs to read promoted dictionary columns directly.
pub enable_rle_to_dictionary: bool,
/// Whether to read row groups in reverse order
pub reverse_row_groups: bool,
/// Optional sort order used to reorder row groups by their min/max statistics.
Expand Down Expand Up @@ -467,6 +469,7 @@ struct PreparedParquetOpen {
predicate_creation_errors: Count,
max_predicate_cache_size: Option<usize>,
max_in_list_size: usize,
enable_rle_to_dictionary: bool,
reverse_row_groups: bool,
sort_order_for_reorder: Option<LexOrdering>,
preserve_order: bool,
Expand Down Expand Up @@ -874,6 +877,7 @@ impl ParquetMorselizer {
predicate_creation_errors,
max_predicate_cache_size: self.max_predicate_cache_size,
max_in_list_size: self.max_in_list_size,
enable_rle_to_dictionary: self.enable_rle_to_dictionary,
reverse_row_groups: self.reverse_row_groups,
sort_order_for_reorder: self.sort_order_for_reorder.clone(),
preserve_order: self.preserve_order,
Expand Down Expand Up @@ -984,9 +988,10 @@ impl MetadataLoadedParquetOpen {
// desired schema (for example if we want to instruct the parquet
// reader to read strings using Utf8View instead). Update if necessary
let mut metadata_dirty = false;
if let Some(merged) = apply_file_schema_type_coercions(
if let Some(merged) = apply_file_schema_type_coercions_with_rle(
&prepared.logical_file_schema,
&physical_file_schema,
prepared.enable_rle_to_dictionary,
) {
physical_file_schema = Arc::new(merged);
options = options.with_schema(Arc::clone(&physical_file_schema));
Expand Down Expand Up @@ -1937,6 +1942,7 @@ mod test {
coerce_int96: Option<TimeUnit>,
max_predicate_cache_size: Option<usize>,
max_in_list_size: usize,
enable_rle_to_dictionary: bool,
reverse_row_groups: bool,
preserve_order: bool,
}
Expand Down Expand Up @@ -2148,6 +2154,7 @@ mod test {
coerce_int96: None,
max_predicate_cache_size: None,
max_in_list_size: MAX_IN_LIST_SIZE,
enable_rle_to_dictionary: false,
reverse_row_groups: false,
preserve_order: false,
}
Expand Down Expand Up @@ -2223,6 +2230,11 @@ mod test {
self
}

fn with_enable_rle_to_dictionary(mut self, enable: bool) -> Self {
self.enable_rle_to_dictionary = enable;
self
}

fn with_metrics(mut self, metrics: ExecutionPlanMetricsSet) -> Self {
self.metrics = metrics;
self
Expand Down Expand Up @@ -2326,6 +2338,7 @@ mod test {
encryption_factory: None,
max_predicate_cache_size: self.max_predicate_cache_size,
max_in_list_size: self.max_in_list_size,
enable_rle_to_dictionary: self.enable_rle_to_dictionary,
reverse_row_groups: self.reverse_row_groups,
sort_order_for_reorder: None,
virtual_state,
Expand Down Expand Up @@ -4333,4 +4346,59 @@ mod test {
assert_eq!(rows, 5);
}
}

async fn collect_batches(
morselizer: &ParquetMorselizer,
file: PartitionedFile,
) -> Vec<RecordBatch> {
let mut stream = open_file(morselizer, file).await.unwrap();
let mut batches = Vec::new();
while let Some(batch) = stream.next().await {
batches.push(batch.unwrap());
}
batches
}

// Proves the opener passes a promoted binary Dictionary schema to arrow-rs.
#[tokio::test]
async fn test_rle_binary_column_promotion() {
let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
let bin_schema = Arc::new(Schema::new(vec![Field::new(
"payload",
DataType::Binary,
true,
)]));
let values =
Arc::new(arrow::array::BinaryArray::from_vec(vec![b"a", b"b", b"a"]));
let batch = RecordBatch::try_new(Arc::clone(&bin_schema), vec![values]).unwrap();
let props = WriterProperties::builder()
.set_dictionary_enabled(true)
.build();
let bin_size = write_parquet_batches(
Arc::clone(&store),
"binary.parquet",
vec![batch],
Some(props),
)
.await;
let dict_bin_schema = Arc::new(Schema::new(vec![Field::new(
"payload",
DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary)),
true,
)]));
let morselizer = ParquetMorselizerBuilder::new()
.with_store(Arc::clone(&store))
.with_schema(Arc::clone(&dict_bin_schema))
.with_enable_rle_to_dictionary(true)
.build();
let batches = collect_batches(
&morselizer,
PartitionedFile::new("binary.parquet".to_string(), bin_size as u64),
)
.await;
assert_eq!(
batches[0].schema().field(0).data_type(),
&DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary))
);
}
}
Loading