diff --git a/Cargo.lock b/Cargo.lock index f12ba856b9fbe..35d74540b6cdc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1382,6 +1382,7 @@ version = "0.4.38" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ce2548391e9c1929c21bf6aa2680af86fe4c1b33e6cea9ac1cfeec0bd11218cf" dependencies = [ + "brotli", "bzip2", "compression-core", "flate2", @@ -1919,6 +1920,7 @@ dependencies = [ "arrow", "async-compression", "async-trait", + "brotli", "bytes", "bzip2", "chrono", diff --git a/datafusion/catalog-listing/src/config.rs b/datafusion/catalog-listing/src/config.rs index 2b83c8ec92b2c..a3791c23da041 100644 --- a/datafusion/catalog-listing/src/config.rs +++ b/datafusion/catalog-listing/src/config.rs @@ -40,6 +40,24 @@ pub enum SchemaSource { Specified, } +#[cfg(test)] +mod compression_extension_tests { + use super::ListingTableConfig; + + #[test] + fn new_codec_suffixes_preserve_csv_extension() { + for suffix in ["br", "zlib", "raw_deflate"] { + assert_eq!( + ListingTableConfig::infer_file_extension_and_compression_type(&format!( + "data.csv.{suffix}" + )) + .unwrap(), + ("csv".to_owned(), Some(suffix.to_owned())) + ); + } + } +} + /// Configuration for creating a [`crate::ListingTable`] /// /// # Schema Evolution Support diff --git a/datafusion/catalog-listing/src/helpers.rs b/datafusion/catalog-listing/src/helpers.rs index 098f3d51ef911..ca6f84c91eb78 100644 --- a/datafusion/catalog-listing/src/helpers.rs +++ b/datafusion/catalog-listing/src/helpers.rs @@ -389,6 +389,28 @@ pub async fn pruned_partition_list<'a>( filters: &'a [Expr], file_extension: &'a str, partition_cols: &'a [(String, DataType)], +) -> Result>> { + pruned_partition_list_with_suffixes( + ctx, + store, + table_path, + filters, + file_extension, + None, + partition_cols, + ) + .await +} + +/// As [`pruned_partition_list`], but accepts several filename suffixes. +pub async fn pruned_partition_list_with_suffixes<'a>( + ctx: &'a dyn Session, + store: &'a dyn ObjectStore, + table_path: &'a ListingTableUrl, + filters: &'a [Expr], + file_extension: &'a str, + file_suffixes: Option<&'a [String]>, + partition_cols: &'a [(String, DataType)], ) -> Result>> { let prefix = if !partition_cols.is_empty() { evaluate_partition_prefix(partition_cols, filters) @@ -397,7 +419,13 @@ pub async fn pruned_partition_list<'a>( }; let objects = table_path - .list_prefixed_files(ctx, store, prefix, file_extension) + .list_prefixed_files_with_suffixes( + ctx, + store, + prefix, + file_extension, + file_suffixes, + ) .await? .try_filter(|object_meta| futures::future::ready(object_meta.size > 0)); diff --git a/datafusion/catalog-listing/src/options.rs b/datafusion/catalog-listing/src/options.rs index 44337e52a1e05..e8ce423156d16 100644 --- a/datafusion/catalog-listing/src/options.rs +++ b/datafusion/catalog-listing/src/options.rs @@ -17,8 +17,9 @@ use arrow::datatypes::{DataType, SchemaRef}; use datafusion_catalog::Session; -use datafusion_common::plan_err; +use datafusion_common::{GetExt, plan_err}; use datafusion_datasource::ListingTableUrl; +use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_datasource::file_format::FileFormat; use datafusion_expr::{Partitioning, SortExpr}; use futures::StreamExt; @@ -33,6 +34,9 @@ pub struct ListingOptions { /// A suffix on which files should be filtered (leave empty to /// keep all files on the path) pub file_extension: String, + /// Accepted filename suffixes for header-detected compression. An explicit + /// file-extension override disables this filter. + pub auto_file_suffixes: Option>, /// The file format pub format: Arc, /// The expected partition column names in the folder structure. @@ -96,14 +100,38 @@ pub struct ListingOptions { pub output_partitioning: Option, } +fn auto_compression_suffixes(base_extension: &str) -> Vec { + let base = format!(".{}", base_extension.trim_start_matches('.')); + [ + FileCompressionType::UNCOMPRESSED, + FileCompressionType::GZIP, + FileCompressionType::BZIP2, + FileCompressionType::XZ, + FileCompressionType::ZSTD, + FileCompressionType::DEFLATE, + ] + .into_iter() + .map(|compression| format!("{base}{}", compression.get_ext())) + .collect() +} + impl ListingOptions { /// Creates an options instance with the given format /// Default values: /// - use default file extension filter /// - no input partition to discover pub fn new(format: Arc) -> Self { + let auto = format.compression_type() == Some(FileCompressionType::AUTO); + let file_extension = if auto { + String::new() + } else { + format.get_ext() + }; + let auto_file_suffixes = + auto.then(|| auto_compression_suffixes(&format.get_ext())); Self { - file_extension: format.get_ext(), + file_extension, + auto_file_suffixes, format, table_partition_cols: vec![], file_sort_order: vec![], @@ -126,6 +154,15 @@ impl ListingOptions { /// ``` pub fn with_file_extension(mut self, file_extension: impl Into) -> Self { self.file_extension = file_extension.into(); + self.auto_file_suffixes = None; + self + } + + /// Detect compression while restricting a listing to files of this format. + pub fn with_auto_compression_file_extension(mut self, base_extension: &str) -> Self { + self.file_extension.clear(); + self.auto_file_suffixes = (!base_extension.is_empty()) + .then(|| auto_compression_suffixes(base_extension)); self } @@ -150,7 +187,7 @@ impl ListingOptions { S: Into, { if let Some(file_extension) = file_extension { - self.file_extension = file_extension.into(); + self = self.with_file_extension(file_extension); } self } @@ -280,7 +317,12 @@ impl ListingOptions { let store = state.runtime_env().object_store(table_path)?; let all_files: Vec<_> = table_path - .list_all_files(state, store.as_ref(), &self.file_extension) + .list_all_files_with_suffixes( + state, + store.as_ref(), + &self.file_extension, + self.auto_file_suffixes.as_deref(), + ) .await? .try_collect() .await?; @@ -375,7 +417,12 @@ impl ListingOptions { // This can fail to detect inconsistent partition keys // A DFS traversal approach of the store can help here let files: Vec<_> = table_path - .list_all_files(state, store.as_ref(), &self.file_extension) + .list_all_files_with_suffixes( + state, + store.as_ref(), + &self.file_extension, + self.auto_file_suffixes.as_deref(), + ) .await? .take(10) .try_collect() diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index b3328cc06303d..e0c349e04ea67 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -17,7 +17,8 @@ use crate::config::SchemaSource; use crate::helpers::{ - expr_applicable_for_cols, filter_partitioned_file, pruned_partition_list, + expr_applicable_for_cols, filter_partitioned_file, + pruned_partition_list_with_suffixes as pruned_partition_list, }; use crate::{ListingOptions, ListingTableConfig}; use arrow::datatypes::{Field, Schema, SchemaBuilder, SchemaRef}; @@ -725,6 +726,7 @@ impl TableProvider for ListingTable { table_path, &[], &self.options.file_extension, + self.options.auto_file_suffixes.as_deref(), &self.options.table_partition_cols, ) .await?; @@ -812,6 +814,7 @@ impl ListingTable { table_path, listing_time_filters, &self.options.file_extension, + self.options.auto_file_suffixes.as_deref(), &self.options.table_partition_cols, ) })) diff --git a/datafusion/common/src/parsers.rs b/datafusion/common/src/parsers.rs index 6b930d110f47b..da16ce049c01c 100644 --- a/datafusion/common/src/parsers.rs +++ b/datafusion/common/src/parsers.rs @@ -35,6 +35,15 @@ pub enum CompressionTypeVariant { ZSTD, /// Uncompressed file UNCOMPRESSED, + /// Detect file compression from its header when reading + AUTO, + /// Brotli-compressed file + BROTLI, + /// Deflate-compressed file with a zlib header + DEFLATE, + /// Deflate-compressed file without a zlib header + #[expect(non_camel_case_types, reason = "match the external codec name")] + RAW_DEFLATE, } impl FromStr for CompressionTypeVariant { @@ -48,6 +57,10 @@ impl FromStr for CompressionTypeVariant { "XZ" => Ok(Self::XZ), "ZST" | "ZSTD" => Ok(Self::ZSTD), "" | "UNCOMPRESSED" => Ok(Self::UNCOMPRESSED), + "AUTO" => Ok(Self::AUTO), + "BROTLI" | "BR" => Ok(Self::BROTLI), + "DEFLATE" | "ZLIB" => Ok(Self::DEFLATE), + "RAW_DEFLATE" => Ok(Self::RAW_DEFLATE), _ => Err(DataFusionError::NotImplemented(format!( "Unsupported file compression type {s}" ))), @@ -63,6 +76,10 @@ impl Display for CompressionTypeVariant { Self::XZ => "XZ", Self::ZSTD => "ZSTD", Self::UNCOMPRESSED => "", + Self::AUTO => "AUTO", + Self::BROTLI => "BROTLI", + Self::DEFLATE => "DEFLATE", + Self::RAW_DEFLATE => "RAW_DEFLATE", }; write!(f, "{str}") } diff --git a/datafusion/core/src/datasource/file_format/options.rs b/datafusion/core/src/datasource/file_format/options.rs index 8ef780ea1e972..e6c157b5e69b0 100644 --- a/datafusion/core/src/datasource/file_format/options.rs +++ b/datafusion/core/src/datasource/file_format/options.rs @@ -643,8 +643,13 @@ impl ReadOptions<'_> for CsvReadOptions<'_> { .with_null_regex(self.null_regex.clone()) .with_truncated_rows(self.truncated_rows); - ListingOptions::new(Arc::new(file_format)) - .with_file_extension(self.file_extension) + let listing = ListingOptions::new(Arc::new(file_format)); + let listing = if self.file_compression_type == FileCompressionType::AUTO { + listing.with_auto_compression_file_extension(self.file_extension) + } else { + listing.with_file_extension(self.file_extension) + }; + listing .with_table_partition_cols(self.table_partition_cols.clone()) .with_file_sort_order(self.file_sort_order.clone()) } @@ -724,8 +729,13 @@ impl ReadOptions<'_> for JsonReadOptions<'_> { .with_file_compression_type(self.file_compression_type.to_owned()) .with_newline_delimited(self.newline_delimited); - ListingOptions::new(Arc::new(file_format)) - .with_file_extension(self.file_extension) + let listing = ListingOptions::new(Arc::new(file_format)); + let listing = if self.file_compression_type == FileCompressionType::AUTO { + listing.with_auto_compression_file_extension(self.file_extension) + } else { + listing.with_file_extension(self.file_extension) + }; + listing .with_table_partition_cols(self.table_partition_cols.clone()) .with_file_sort_order(self.file_sort_order.clone()) } diff --git a/datafusion/core/src/datasource/listing_table_factory.rs b/datafusion/core/src/datasource/listing_table_factory.rs index f012add5d1a2c..a9074195ae59d 100644 --- a/datafusion/core/src/datasource/listing_table_factory.rs +++ b/datafusion/core/src/datasource/listing_table_factory.rs @@ -32,6 +32,7 @@ use datafusion_common::{Result, config_datafusion_err}; use datafusion_common::{ ToDFSchema, arrow_datafusion_err, internal_datafusion_err, plan_err, }; +use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_expr::CreateExternalTable; use async_trait::async_trait; @@ -130,8 +131,14 @@ impl TableProviderFactory for ListingTableFactory { } else { String::new() }; + let auto_extension = (file_format.compression_type() + == Some(FileCompressionType::AUTO)) + .then(|| file_format.get_ext()); let mut options = ListingOptions::new(file_format).with_file_extension(file_extension); + if let Some(extension) = auto_extension { + options = options.with_auto_compression_file_extension(&extension); + } // Partition columns are derived from the first location; all locations // are expected to share the same partitioning. @@ -215,6 +222,9 @@ impl TableProviderFactory for ListingTableFactory { // derive the pattern based on compression type. // So for gzipped CSV the pattern is `*.csv.gz` let glob = match options.format.compression_type() { + Some(FileCompressionType::AUTO) => { + format!("*.{}*", options.format.get_ext()) + } Some(compression) => { match options .format @@ -492,6 +502,380 @@ mod tests { ); } + #[cfg(feature = "compression")] + #[tokio::test] + async fn test_create_using_folder_with_auto_compression() { + use crate::prelude::CsvReadOptions; + use arrow::array::Int64Array; + use flate2::{Compression, write::GzEncoder}; + use std::io::Write; + + let dir = tempfile::tempdir().unwrap(); + fs::write(dir.path().join("plain.csv"), "id\n1\n").unwrap(); + let mut gzip = GzEncoder::new(Vec::new(), Compression::default()); + gzip.write_all(b"id\n2\n").unwrap(); + fs::write(dir.path().join("compressed.csv.gz"), gzip.finish().unwrap()).unwrap(); + fs::write(dir.path().join("ignored.csv.bak"), "not,a,csv\n").unwrap(); + fs::write(dir.path().join("ignored.json"), "{\"id\":999}\n").unwrap(); + + let factory = ListingTableFactory::new(); + let context = SessionContext::new(); + let state = context.state(); + let cmd = CreateExternalTable::builder( + TableReference::bare("auto_csv"), + dir.path().to_str().unwrap().to_owned(), + "csv", + Arc::new(DFSchema::empty()), + ) + .with_options(HashMap::from([ + ("format.has_header".into(), "true".into()), + ("format.compression".into(), "auto".into()), + ])) + .build(); + let provider = factory.create(&state, &cmd).await.unwrap(); + let table = provider.downcast_ref::().unwrap(); + + assert_eq!(table.options().file_extension, ""); + assert_eq!( + table.options().format.compression_type(), + Some(FileCompressionType::AUTO) + ); + assert_eq!( + table.table_paths()[0].get_glob().as_ref().unwrap(), + &Pattern::new("*.csv*").unwrap() + ); + context.register_table("auto_csv", provider).unwrap(); + let batches = context + .sql("SELECT id FROM auto_csv ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + let values = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.values(), &[1, 2]); + + let options = + CsvReadOptions::new().file_compression_type(FileCompressionType::AUTO); + let direct = context + .read_csv( + dir.path().join("compressed.csv.gz").to_str().unwrap(), + options.clone(), + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + direct[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 2 + ); + let folder = context + .read_csv(dir.path().to_str().unwrap(), options) + .await + .unwrap() + .collect() + .await + .unwrap(); + let mut ids = folder + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + ids.sort_unstable(); + assert_eq!(ids, [1, 2]); + } + + #[cfg(feature = "compression")] + #[tokio::test] + async fn test_read_json_with_auto_compression_filters_other_files() { + use crate::prelude::JsonReadOptions; + use arrow::array::Int64Array; + use flate2::{Compression, write::GzEncoder}; + use std::io::Write; + + let dir = tempfile::tempdir().unwrap(); + fs::write(dir.path().join("plain.json"), "{\"id\":1}\n").unwrap(); + let mut gzip = GzEncoder::new(Vec::new(), Compression::default()); + gzip.write_all(b"{\"id\":2}\n").unwrap(); + fs::write( + dir.path().join("compressed.json.gz"), + gzip.finish().unwrap(), + ) + .unwrap(); + fs::write(dir.path().join("ignored.json.bak"), "bad json\n").unwrap(); + fs::write(dir.path().join("ignored.csv"), "id\n999\n").unwrap(); + + let context = SessionContext::new(); + let options = + JsonReadOptions::default().file_compression_type(FileCompressionType::AUTO); + let direct = context + .read_json( + dir.path().join("compressed.json.gz").to_str().unwrap(), + options.clone(), + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + direct[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 2 + ); + let folder = context + .read_json(dir.path().to_str().unwrap(), options) + .await + .unwrap() + .collect() + .await + .unwrap(); + let mut ids = folder + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + ids.sort_unstable(); + assert_eq!(ids, [1, 2]); + } + + #[cfg(feature = "compression")] + #[tokio::test] + async fn test_auto_compression_filters_multiple_folders_but_keeps_explicit_files() { + use arrow::array::Int64Array; + use flate2::{Compression, write::GzEncoder}; + use std::io::Write; + + let first = tempfile::tempdir().unwrap(); + let second = tempfile::tempdir().unwrap(); + let explicit = tempfile::tempdir().unwrap(); + fs::write(first.path().join("plain.csv"), "id\n1\n").unwrap(); + let mut gzip = GzEncoder::new(Vec::new(), Compression::default()); + gzip.write_all(b"id\n2\n").unwrap(); + fs::write( + second.path().join("compressed.csv.gz"), + gzip.finish().unwrap(), + ) + .unwrap(); + fs::write(second.path().join("ignored.csv.bak"), "bad,data\n").unwrap(); + let explicit_path = explicit.path().join("custom.data"); + fs::write(&explicit_path, "id\n3\n").unwrap(); + + let factory = ListingTableFactory::new(); + let context = SessionContext::new(); + let state = context.state(); + let locations = vec![ + first.path().to_str().unwrap().to_owned(), + second.path().to_str().unwrap().to_owned(), + explicit_path.to_str().unwrap().to_owned(), + ]; + let cmd = CreateExternalTable::builder( + TableReference::bare("auto_multi"), + locations[0].clone(), + "csv", + Arc::new(DFSchema::empty()), + ) + .with_locations(locations) + .with_options(HashMap::from([ + ("format.has_header".into(), "true".into()), + ("format.compression".into(), "auto".into()), + ])) + .build(); + let provider = factory.create(&state, &cmd).await.unwrap(); + context.register_table("auto_multi", provider).unwrap(); + let batches = context + .sql("SELECT id FROM auto_multi ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + let values = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.values(), &[1, 2, 3]); + } + + #[cfg(feature = "compression")] + #[tokio::test] + async fn test_large_plain_csv_auto_preserves_parallel_range_scan() { + use crate::prelude::CsvReadOptions; + use datafusion_datasource::source::DataSourceExec; + use datafusion_datasource_csv::source::CsvSource; + use datafusion_physical_plan::ExecutionPlan; + + fn find_scan(plan: &dyn ExecutionPlan) -> Option<&DataSourceExec> { + plan.downcast_ref::().or_else(|| { + plan.children() + .into_iter() + .find_map(|child| find_scan(child.as_ref())) + }) + } + + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("large.csv"); + fs::write(&path, "1\n".repeat(9_000_000)).unwrap(); + let schema = Schema::new(vec![Field::new("id", DataType::Int64, true)]); + let context = SessionContext::new_with_config( + SessionConfig::new() + .with_target_partitions(4) + .with_repartition_file_scans(true), + ); + let plan = context + .read_csv( + path.to_str().unwrap(), + CsvReadOptions::new() + .has_header(false) + .schema(&schema) + .file_compression_type(FileCompressionType::AUTO), + ) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let scan = find_scan(plan.as_ref()).unwrap(); + let (config, _) = scan.downcast_to_file_source::().unwrap(); + assert_eq!( + config.file_compression_type, + FileCompressionType::UNCOMPRESSED + ); + assert!(config.file_groups.len() > 1); + let batches = datafusion_physical_plan::collect(plan, context.task_ctx()) + .await + .unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 9_000_000 + ); + } + + #[cfg(feature = "compression")] + #[tokio::test] + async fn test_multiple_large_plain_csv_auto_preserves_parallel_range_scan() { + use crate::prelude::CsvReadOptions; + use datafusion_datasource::source::DataSourceExec; + use datafusion_datasource_csv::source::CsvSource; + use datafusion_physical_plan::ExecutionPlan; + + fn find_scan(plan: &dyn ExecutionPlan) -> Option<&DataSourceExec> { + plan.downcast_ref::().or_else(|| { + plan.children() + .into_iter() + .find_map(|child| find_scan(child.as_ref())) + }) + } + + let dir = tempfile::tempdir().unwrap(); + let contents = "10000000000000000000000000000000\n".repeat(600_000); + fs::write(dir.path().join("first.csv"), &contents).unwrap(); + fs::write(dir.path().join("second.csv"), contents).unwrap(); + let schema = Schema::new(vec![Field::new("value", DataType::Utf8, true)]); + let context = SessionContext::new_with_config( + SessionConfig::new() + .with_target_partitions(4) + .with_repartition_file_scans(true), + ); + let plan = context + .read_csv( + dir.path().to_str().unwrap(), + CsvReadOptions::new() + .has_header(false) + .schema(&schema) + .file_compression_type(FileCompressionType::AUTO), + ) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let scan = find_scan(plan.as_ref()).unwrap(); + let (config, _) = scan.downcast_to_file_source::().unwrap(); + assert_eq!( + config.file_compression_type, + FileCompressionType::UNCOMPRESSED + ); + assert!(config.file_groups.len() > 2); + let batches = datafusion_physical_plan::collect(plan, context.task_ctx()) + .await + .unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 1_200_000 + ); + } + + #[cfg(not(feature = "compression"))] + #[tokio::test] + async fn test_large_plain_csv_auto_requires_compression_feature() { + use crate::prelude::CsvReadOptions; + use datafusion_datasource::source::DataSourceExec; + use datafusion_datasource_csv::source::CsvSource; + + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("large.csv"); + fs::write(&path, "1\n".repeat(9_000_000)).unwrap(); + let schema = Schema::new(vec![Field::new("id", DataType::Int64, true)]); + let context = SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(4), + ); + let frame = context + .read_csv( + path.to_str().unwrap(), + CsvReadOptions::new() + .has_header(false) + .schema(&schema) + .file_compression_type(FileCompressionType::AUTO), + ) + .await + .unwrap(); + let plan = frame.create_physical_plan().await.unwrap(); + let scan = plan.downcast_ref::().unwrap(); + let (config, _) = scan.downcast_to_file_source::().unwrap(); + assert_eq!(config.file_compression_type, FileCompressionType::AUTO); + assert!( + frame + .collect() + .await + .unwrap_err() + .to_string() + .contains("Compression feature is not enabled") + ); + } + /// Validates that CreateExternalTable without compression /// searches for normal files in a directory location #[tokio::test] diff --git a/datafusion/core/src/test/mod.rs b/datafusion/core/src/test/mod.rs index f46a5a0749065..00904853bee04 100644 --- a/datafusion/core/src/test/mod.rs +++ b/datafusion/core/src/test/mod.rs @@ -156,6 +156,12 @@ pub fn partitioned_file_groups( FileCompressionType::BZIP2 => { Box::new(BzEncoder::new(file, BzCompression::default())) } + FileCompressionType::AUTO + | FileCompressionType::BROTLI + | FileCompressionType::DEFLATE + | FileCompressionType::RAW_DEFLATE => { + panic!("This test helper cannot write the requested compression type") + } #[cfg(not(feature = "compression"))] FileCompressionType::GZIP | FileCompressionType::BZIP2 diff --git a/datafusion/datasource-csv/src/file_format.rs b/datafusion/datasource-csv/src/file_format.rs index 1edce45102ec0..8943897dc8eca 100644 --- a/datafusion/datasource-csv/src/file_format.rs +++ b/datafusion/datasource-csv/src/file_format.rs @@ -450,8 +450,60 @@ impl FileFormat for CsvFormat { .expect("file_source should be a CsvSource"); let source = Arc::new(csv_source.clone().with_csv_options(csv_options)); + let mut compression: FileCompressionType = self.options.compression.into(); + if compression == FileCompressionType::AUTO + && FileCompressionType::compression_enabled() + && source.supports_repartitioning() + && state.config_options().optimizer.repartition_file_scans + && state.config_options().execution.target_partitions > 1 + { + let files = conf + .file_groups + .iter() + .flat_map(|group| group.iter()) + .take(33) + .collect::>(); + let min_size = u64::try_from( + state + .config_options() + .optimizer + .repartition_file_min_size + .max(16 * 1024 * 1024), + ) + .unwrap_or(u64::MAX); + if !files.is_empty() + && files.len() <= 32 + && files.iter().all(|file| file.range.is_none()) + && files.iter().fold(0_u64, |size, file| { + size.saturating_add(file.object_meta.size) + }) >= min_size + { + let store = state.runtime_env().object_store(&conf.object_store_url)?; + let plain = futures::future::join_all(files.into_iter().map(|file| { + let store = Arc::clone(&store); + async move { + if file.object_meta.size == 0 { + return true; + } + let len = file.object_meta.size.min(6); + store + .get_range(&file.object_meta.location, 0..len) + .await + .is_ok_and(|header| { + FileCompressionType::detect_from_header(&header) + == FileCompressionType::UNCOMPRESSED + }) + } + })) + .await; + if plain.into_iter().all(|is_plain| is_plain) { + compression = FileCompressionType::UNCOMPRESSED; + } + } + } + let config = FileScanConfigBuilder::from(conf) - .with_file_compression_type(self.options.compression.into()) + .with_file_compression_type(compression) .with_source(source) .build(); diff --git a/datafusion/datasource/Cargo.toml b/datafusion/datasource/Cargo.toml index 98bf9ac13f1ff..9216feb02115b 100644 --- a/datafusion/datasource/Cargo.toml +++ b/datafusion/datasource/Cargo.toml @@ -32,7 +32,7 @@ all-features = true [features] backtrace = ["datafusion-common/backtrace"] -compression = ["async-compression", "liblzma", "bzip2", "flate2", "zstd", "tokio-util"] +compression = ["async-compression", "brotli", "liblzma", "bzip2", "flate2", "zstd", "tokio-util"] default = ["compression"] # Enables protobuf conversions for datasource types, source serialization hooks, # and the shared `FileScanConfig` <-> proto conversion. Off by default so @@ -47,14 +47,18 @@ proto = [ arrow = { workspace = true } async-compression = { version = "0.4.40", features = [ "bzip2", + "brotli", + "deflate", "gzip", "xz", + "zlib", "zstd", "tokio", ], optional = true } async-trait = { workspace = true } bytes = { workspace = true } bzip2 = { workspace = true, optional = true } +brotli = { version = "8.0.2", optional = true } chrono = { workspace = true } datafusion-common = { workspace = true, features = ["object_store"] } datafusion-common-runtime = { workspace = true } diff --git a/datafusion/datasource/src/file_compression_type.rs b/datafusion/datasource/src/file_compression_type.rs index 89efb580652b1..8ae2d6a4c6bc1 100644 --- a/datafusion/datasource/src/file_compression_type.rs +++ b/datafusion/datasource/src/file_compression_type.rs @@ -17,6 +17,7 @@ //! File Compression type abstraction +use std::io::Read; use std::str::FromStr; use datafusion_common::error::{DataFusionError, Result}; @@ -26,19 +27,27 @@ use datafusion_common::parsers::CompressionTypeVariant::{self, *}; #[cfg(feature = "compression")] use async_compression::tokio::bufread::{ + BrotliDecoder as AsyncBrotliDecoder, BrotliEncoder as AsyncBrotliEncoder, BzDecoder as AsyncBzDecoder, BzEncoder as AsyncBzEncoder, + DeflateDecoder as AsyncDeflateDecoder, DeflateEncoder as AsyncDeflateEncoder, GzipDecoder as AsyncGzDecoder, GzipEncoder as AsyncGzEncoder, XzDecoder as AsyncXzDecoder, XzEncoder as AsyncXzEncoder, + ZlibDecoder as AsyncZlibDecoder, ZlibEncoder as AsyncZlibEncoder, ZstdDecoder as AsyncZstdDecoer, ZstdEncoder as AsyncZstdEncoder, }; #[cfg(feature = "compression")] -use async_compression::tokio::write::{BzEncoder, GzipEncoder, XzEncoder, ZstdEncoder}; +use async_compression::tokio::write::{ + BrotliEncoder, BzEncoder, DeflateEncoder, GzipEncoder, XzEncoder, ZlibEncoder, + ZstdEncoder, +}; +#[cfg(feature = "compression")] +use brotli::Decompressor as BrotliReader; use bytes::Bytes; #[cfg(feature = "compression")] use bzip2::read::MultiBzDecoder; #[cfg(feature = "compression")] -use flate2::read::MultiGzDecoder; +use flate2::read::{DeflateDecoder, MultiGzDecoder, ZlibDecoder}; use futures::StreamExt; #[cfg(feature = "compression")] use futures::TryStreamExt; @@ -66,6 +75,10 @@ impl GetExt for FileCompressionType { XZ => ".xz".to_owned(), ZSTD => ".zst".to_owned(), UNCOMPRESSED => "".to_owned(), + AUTO => "".to_owned(), + BROTLI => ".br".to_owned(), + DEFLATE => ".zlib".to_owned(), + RAW_DEFLATE => ".raw_deflate".to_owned(), } } } @@ -95,6 +108,11 @@ impl FromStr for FileCompressionType { /// `FileCompressionType` implementation impl FileCompressionType { + /// Whether this build can decode compressed input or use AUTO detection. + pub const fn compression_enabled() -> bool { + cfg!(feature = "compression") + } + /// Gzip-ed file pub const GZIP: Self = Self { variant: GZIP }; @@ -112,6 +130,20 @@ impl FileCompressionType { variant: UNCOMPRESSED, }; + /// Detect the compression type from the file header while reading + pub const AUTO: Self = Self { variant: AUTO }; + + /// Brotli-compressed file + pub const BROTLI: Self = Self { variant: BROTLI }; + + /// Deflate-compressed file with a zlib header + pub const DEFLATE: Self = Self { variant: DEFLATE }; + + /// Deflate-compressed file without a zlib header + pub const RAW_DEFLATE: Self = Self { + variant: RAW_DEFLATE, + }; + /// Read only access to self.variant pub fn get_variant(&self) -> &CompressionTypeVariant { &self.variant @@ -144,13 +176,32 @@ impl FileCompressionType { ZSTD => ReaderStream::new(AsyncZstdEncoder::new(StreamReader::new(s))) .map_err(DataFusionError::from) .boxed(), + #[cfg(feature = "compression")] + BROTLI => ReaderStream::new(AsyncBrotliEncoder::new(StreamReader::new(s))) + .map_err(DataFusionError::from) + .boxed(), + #[cfg(feature = "compression")] + DEFLATE => ReaderStream::new(AsyncZlibEncoder::new(StreamReader::new(s))) + .map_err(DataFusionError::from) + .boxed(), + #[cfg(feature = "compression")] + RAW_DEFLATE => { + ReaderStream::new(AsyncDeflateEncoder::new(StreamReader::new(s))) + .map_err(DataFusionError::from) + .boxed() + } #[cfg(not(feature = "compression"))] - GZIP | BZIP2 | XZ | ZSTD => { + GZIP | BZIP2 | XZ | ZSTD | BROTLI | DEFLATE | RAW_DEFLATE => { return Err(DataFusionError::NotImplemented( "Compression feature is not enabled".to_owned(), )); } UNCOMPRESSED => s.boxed(), + AUTO => { + return Err(DataFusionError::NotImplemented( + "AUTO compression is only supported for reading".to_owned(), + )); + } }) } @@ -205,8 +256,30 @@ impl FileCompressionType { } None => Box::new(ZstdEncoder::new(w)), }, + #[cfg(feature = "compression")] + BROTLI => match compression_level { + Some(level) => { + Box::new(BrotliEncoder::with_quality(w, Level::Precise(level as i32))) + } + None => Box::new(BrotliEncoder::new(w)), + }, + #[cfg(feature = "compression")] + DEFLATE => match compression_level { + Some(level) => { + Box::new(ZlibEncoder::with_quality(w, Level::Precise(level as i32))) + } + None => Box::new(ZlibEncoder::new(w)), + }, + #[cfg(feature = "compression")] + RAW_DEFLATE => match compression_level { + Some(level) => Box::new(DeflateEncoder::with_quality( + w, + Level::Precise(level as i32), + )), + None => Box::new(DeflateEncoder::new(w)), + }, #[cfg(not(feature = "compression"))] - GZIP | BZIP2 | XZ | ZSTD => { + GZIP | BZIP2 | XZ | ZSTD | BROTLI | DEFLATE | RAW_DEFLATE => { // compression_level is not used when compression feature is disabled let _ = compression_level; return Err(DataFusionError::NotImplemented( @@ -214,6 +287,11 @@ impl FileCompressionType { )); } UNCOMPRESSED => Box::new(w), + AUTO => { + return Err(DataFusionError::NotImplemented( + "AUTO compression is only supported for reading".to_owned(), + )); + } }) } @@ -244,21 +322,66 @@ impl FileCompressionType { ZSTD => ReaderStream::new(AsyncZstdDecoer::new(StreamReader::new(s))) .map_err(DataFusionError::from) .boxed(), + #[cfg(feature = "compression")] + BROTLI => ReaderStream::new(AsyncBrotliDecoder::new(StreamReader::new(s))) + .map_err(DataFusionError::from) + .boxed(), + #[cfg(feature = "compression")] + DEFLATE => ReaderStream::new(AsyncZlibDecoder::new(StreamReader::new(s))) + .map_err(DataFusionError::from) + .boxed(), + #[cfg(feature = "compression")] + RAW_DEFLATE => { + ReaderStream::new(AsyncDeflateDecoder::new(StreamReader::new(s))) + .map_err(DataFusionError::from) + .boxed() + } #[cfg(not(feature = "compression"))] - GZIP | BZIP2 | XZ | ZSTD => { + GZIP | BZIP2 | XZ | ZSTD | BROTLI | DEFLATE | RAW_DEFLATE => { return Err(DataFusionError::NotImplemented( "Compression feature is not enabled".to_owned(), )); } UNCOMPRESSED => s.boxed(), + #[cfg(feature = "compression")] + AUTO => futures::stream::once(async move { + let mut source = s; + let mut header = [0_u8; 6]; + let mut header_len = 0; + let mut initial_chunks = Vec::new(); + while header_len < header.len() { + match source.next().await { + Some(Ok(bytes)) if bytes.is_empty() => continue, + Some(Ok(bytes)) => { + let copied = (header.len() - header_len).min(bytes.len()); + header[header_len..header_len + copied] + .copy_from_slice(&bytes[..copied]); + header_len += copied; + initial_chunks.push(Ok(bytes)); + } + Some(Err(error)) => return Err(error), + None => break, + } + } + let replay = futures::stream::iter(initial_chunks).chain(source).boxed(); + Self::detect_from_header(&header[..header_len]).convert_stream(replay) + }) + .try_flatten() + .boxed(), + #[cfg(not(feature = "compression"))] + AUTO => { + return Err(DataFusionError::NotImplemented( + "Compression feature is not enabled".to_owned(), + )); + } }) } /// Given a `Read`, create a `Read` which data are decompressed with `FileCompressionType`. - pub fn convert_read( + pub fn convert_read( &self, r: T, - ) -> Result> { + ) -> Result> { Ok(match self.variant { #[cfg(feature = "compression")] GZIP => Box::new(MultiGzDecoder::new(r)), @@ -271,17 +394,70 @@ impl FileCompressionType { Ok(decoder) => Box::new(decoder), Err(e) => return Err(DataFusionError::External(Box::new(e))), }, + #[cfg(feature = "compression")] + BROTLI => Box::new(BrotliReader::new(r, 4096)), + #[cfg(feature = "compression")] + DEFLATE => Box::new(ZlibDecoder::new(r)), + #[cfg(feature = "compression")] + RAW_DEFLATE => Box::new(DeflateDecoder::new(r)), #[cfg(not(feature = "compression"))] - GZIP | BZIP2 | XZ | ZSTD => { + GZIP | BZIP2 | XZ | ZSTD | BROTLI | DEFLATE | RAW_DEFLATE => { return Err(DataFusionError::NotImplemented( "Compression feature is not enabled".to_owned(), )); } UNCOMPRESSED => Box::new(r), + #[cfg(feature = "compression")] + AUTO => { + let mut source = r; + let mut header = [0_u8; 6]; + let mut header_len = 0; + while header_len < header.len() { + let read = source.read(&mut header[header_len..])?; + if read == 0 { + break; + } + header_len += read; + } + let replay: Box = Box::new( + std::io::Cursor::new(header[..header_len].to_vec()).chain(source), + ); + Self::detect_from_header(&header[..header_len]).convert_read(replay)? + } + #[cfg(not(feature = "compression"))] + AUTO => { + return Err(DataFusionError::NotImplemented( + "Compression feature is not enabled".to_owned(), + )); + } }) } } +impl FileCompressionType { + /// Identify codecs supported by AUTO from up to the first six file bytes. + /// Headerless codecs such as Brotli and raw Deflate require explicit selection. + pub fn detect_from_header(header: &[u8]) -> Self { + if header.starts_with(&[0x1f, 0x8b]) { + Self::GZIP + } else if header.starts_with(b"BZh") { + Self::BZIP2 + } else if header.starts_with(&[0x28, 0xb5, 0x2f, 0xfd]) { + Self::ZSTD + } else if header.starts_with(&[0xfd, b'7', b'z', b'X', b'Z', 0x00]) { + Self::XZ + } else if header.len() >= 2 + && header[0] & 0x0f == 8 + && header[0] >> 4 <= 7 + && u16::from_be_bytes([header[0], header[1]]).is_multiple_of(31) + { + Self::DEFLATE + } else { + Self::UNCOMPRESSED + } + } +} + /// Trait for extending the functionality of the `FileType` enum. pub trait FileTypeExt { /// Given a `FileCompressionType`, return the `FileType`'s extension with compression suffix @@ -315,6 +491,12 @@ mod tests { ("ZST", FileCompressionType::ZSTD), ("zstd", FileCompressionType::ZSTD), ("ZSTD", FileCompressionType::ZSTD), + ("AUTO", FileCompressionType::AUTO), + ("BROTLI", FileCompressionType::BROTLI), + ("BR", FileCompressionType::BROTLI), + ("DEFLATE", FileCompressionType::DEFLATE), + ("ZLIB", FileCompressionType::DEFLATE), + ("RAW_DEFLATE", FileCompressionType::RAW_DEFLATE), ("", FileCompressionType::UNCOMPRESSED), ] { assert_eq!( @@ -329,6 +511,93 @@ mod tests { )); } + #[cfg(feature = "compression")] + #[tokio::test] + async fn auto_detects_compression_across_stream_chunks() -> Result<(), DataFusionError> + { + use futures::TryStreamExt; + use std::io::Read; + + let plain = b"1,alice\n2,bob\n"; + for codec in [ + FileCompressionType::UNCOMPRESSED, + FileCompressionType::GZIP, + FileCompressionType::BZIP2, + FileCompressionType::ZSTD, + FileCompressionType::DEFLATE, + FileCompressionType::XZ, + ] { + let input = futures::stream::once(async { + Ok::(Bytes::from_static(plain)) + }); + let encoded = codec + .convert_to_compress_stream(input.boxed())? + .try_collect::>() + .await? + .concat(); + let chunks = std::iter::repeat_with(|| Ok(Bytes::new())) + .take(100) + .chain( + encoded + .chunks(1) + .map(|chunk| Ok(Bytes::copy_from_slice(chunk))), + ) + .collect::>>(); + let decoded = FileCompressionType::AUTO + .convert_stream(futures::stream::iter(chunks).boxed())? + .try_collect::>() + .await? + .concat(); + assert_eq!(decoded, plain, "stream decoder for {codec:?}"); + + let mut reader = + FileCompressionType::AUTO.convert_read(std::io::Cursor::new(encoded))?; + let mut decoded = Vec::new(); + reader.read_to_end(&mut decoded)?; + assert_eq!(decoded, plain, "sync decoder for {codec:?}"); + } + Ok(()) + } + + #[cfg(feature = "compression")] + #[tokio::test] + async fn explicit_brotli_and_deflate_roundtrip() -> Result<(), DataFusionError> { + use futures::TryStreamExt; + use std::io::Read; + + let plain = b"1,alice\n2,bob\n"; + for codec in [ + FileCompressionType::BROTLI, + FileCompressionType::DEFLATE, + FileCompressionType::RAW_DEFLATE, + ] { + let input = futures::stream::once(async { + Ok::(Bytes::from_static(plain)) + }); + let encoded = codec + .convert_to_compress_stream(input.boxed())? + .try_collect::>() + .await? + .concat(); + let chunks = encoded + .chunks(3) + .map(|chunk| Ok(Bytes::copy_from_slice(chunk))) + .collect::>>(); + let decoded = codec + .convert_stream(futures::stream::iter(chunks).boxed())? + .try_collect::>() + .await? + .concat(); + assert_eq!(decoded, plain, "stream decoder for {codec:?}"); + + let mut reader = codec.convert_read(std::io::Cursor::new(encoded))?; + let mut decoded = Vec::new(); + reader.read_to_end(&mut decoded)?; + assert_eq!(decoded, plain, "sync decoder for {codec:?}"); + } + Ok(()) + } + #[tokio::test] async fn test_bgzip_stream_decoding() -> Result<(), DataFusionError> { // As described in https://samtools.github.io/hts-specs/SAMv1.pdf ("The BGZF compression format") diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs index e865e7ee6d519..b3a6c2831b3e7 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -279,6 +279,20 @@ impl ListingTableUrl { store: &'a dyn ObjectStore, prefix: Option, file_extension: &'a str, + ) -> Result>> { + self.list_prefixed_files_with_suffixes(ctx, store, prefix, file_extension, None) + .await + } + + /// List files with optional accepted suffixes for collections. Explicit + /// paths retain their existing ability to name files with arbitrary suffixes. + pub async fn list_prefixed_files_with_suffixes<'a>( + &'a self, + ctx: &'a dyn Session, + store: &'a dyn ObjectStore, + prefix: Option, + file_extension: &'a str, + file_suffixes: Option<&'a [String]>, ) -> Result>> { let exec_options = &ctx.config_options().execution; let ignore_subdirectory = exec_options.listing_table_ignore_subdirectory; @@ -292,23 +306,9 @@ impl ListingTableUrl { self.prefix.clone() }; - let list: BoxStream<'a, Result> = if self.is_collection() { - list_with_cache( - ctx, - store, - self.table_ref.as_ref(), - &self.prefix, - prefix.as_ref(), - ) - .await? - } else { - match store.head(&full_prefix).await { - Ok(meta) => futures::stream::once(async { Ok(meta) }) - .map_err(|e| DataFusionError::ObjectStore(Box::new(e))) - .boxed(), - // If the head command fails, it is likely that object doesn't exist. - // Retry as though it were a prefix (aka a collection) - Err(object_store::Error::NotFound { .. }) => { + let (list, is_collection): (BoxStream<'a, Result>, bool) = + if self.is_collection() { + ( list_with_cache( ctx, store, @@ -316,17 +316,47 @@ impl ListingTableUrl { &self.prefix, prefix.as_ref(), ) - .await? + .await?, + true, + ) + } else { + match store.head(&full_prefix).await { + Ok(meta) => ( + futures::stream::once(async { Ok(meta) }) + .map_err(|e| DataFusionError::ObjectStore(Box::new(e))) + .boxed(), + false, + ), + // If the head command fails, it is likely that object doesn't exist. + // Retry as though it were a prefix (aka a collection) + Err(object_store::Error::NotFound { .. }) => ( + list_with_cache( + ctx, + store, + self.table_ref.as_ref(), + &self.prefix, + prefix.as_ref(), + ) + .await?, + true, + ), + Err(e) => return Err(e.into()), } - Err(e) => return Err(e.into()), - } - }; + }; if let Some(regex) = &self.full_path_regex { Ok(list .try_filter(move |meta| { let path = &meta.location; - let extension_match = path.as_ref().ends_with(file_extension); + let extension_match = file_suffixes.map_or_else( + || path.as_ref().ends_with(file_extension), + |suffixes| { + !is_collection + || suffixes + .iter() + .any(|suffix| path.as_ref().ends_with(suffix)) + }, + ); let regex_match = self.strip_prefix(path).is_some() && regex.0.is_match(path.as_ref()); futures::future::ready(extension_match && regex_match) @@ -336,7 +366,15 @@ impl ListingTableUrl { Ok(list .try_filter(move |meta| { let path = &meta.location; - let extension_match = path.as_ref().ends_with(file_extension); + let extension_match = file_suffixes.map_or_else( + || path.as_ref().ends_with(file_extension), + |suffixes| { + !is_collection + || suffixes + .iter() + .any(|suffix| path.as_ref().ends_with(suffix)) + }, + ); let glob_match = self.contains_without_regex(path, ignore_subdirectory); futures::future::ready(extension_match && glob_match) @@ -356,6 +394,24 @@ impl ListingTableUrl { .await } + /// List files matching any of the supplied suffixes. + pub async fn list_all_files_with_suffixes<'a>( + &'a self, + ctx: &'a dyn Session, + store: &'a dyn ObjectStore, + file_extension: &'a str, + file_suffixes: Option<&'a [String]>, + ) -> Result>> { + self.list_prefixed_files_with_suffixes( + ctx, + store, + None, + file_extension, + file_suffixes, + ) + .await + } + /// Returns this [`ListingTableUrl`] as a string pub fn as_str(&self) -> &str { self.as_ref() @@ -846,6 +902,22 @@ mod tests { vec!["t/c.csv", "t/d.csv"], ); + create_file(&store, "/t/e.csv.gz").await; + create_file(&store, "/t/ignored.csv.bak").await; + let suffixes = vec![".csv".to_owned(), ".csv.gz".to_owned()]; + let session = MockSession::new(); + for url in ["/t", "/t/"] { + let files = ListingTableUrl::parse(url)? + .list_all_files_with_suffixes(&session, &store, "", Some(&suffixes)) + .await? + .try_collect::>() + .await? + .into_iter() + .map(|meta| meta.location.to_string()) + .collect::>(); + assert_eq!(files, ["t/c.csv", "t/d.csv", "t/e.csv.gz"]); + } + let session = MockSession::new(); let regex_url = ListingTableUrl::parse("memory:///t/")?.with_full_path_regex( Regex::new(r"\At/(?:c|nested/e)[.]csv\z") diff --git a/datafusion/proto-common/proto/datafusion_common.proto b/datafusion/proto-common/proto/datafusion_common.proto index 2ceeb6eb47f96..47e60f63dfde6 100644 --- a/datafusion/proto-common/proto/datafusion_common.proto +++ b/datafusion/proto-common/proto/datafusion_common.proto @@ -427,6 +427,10 @@ enum CompressionTypeVariant { XZ = 2; ZSTD = 3; UNCOMPRESSED = 4; + AUTO = 5; + BROTLI = 6; + DEFLATE = 7; + RAW_DEFLATE = 8; } message JsonWriterOptions { diff --git a/datafusion/proto-common/src/from_proto/mod.rs b/datafusion/proto-common/src/from_proto/mod.rs index 506dd0b7545d3..1006f562e23fc 100644 --- a/datafusion/proto-common/src/from_proto/mod.rs +++ b/datafusion/proto-common/src/from_proto/mod.rs @@ -948,6 +948,10 @@ impl From for CompressionTypeVariant { protobuf::CompressionTypeVariant::Xz => Self::XZ, protobuf::CompressionTypeVariant::Zstd => Self::ZSTD, protobuf::CompressionTypeVariant::Uncompressed => Self::UNCOMPRESSED, + protobuf::CompressionTypeVariant::Auto => Self::AUTO, + protobuf::CompressionTypeVariant::Brotli => Self::BROTLI, + protobuf::CompressionTypeVariant::Deflate => Self::DEFLATE, + protobuf::CompressionTypeVariant::RawDeflate => Self::RAW_DEFLATE, } } } @@ -960,6 +964,35 @@ impl From for protobuf::CompressionTypeVariant { CompressionTypeVariant::XZ => Self::Xz, CompressionTypeVariant::ZSTD => Self::Zstd, CompressionTypeVariant::UNCOMPRESSED => Self::Uncompressed, + CompressionTypeVariant::AUTO => Self::Auto, + CompressionTypeVariant::BROTLI => Self::Brotli, + CompressionTypeVariant::DEFLATE => Self::Deflate, + CompressionTypeVariant::RAW_DEFLATE => Self::RawDeflate, + } + } +} + +#[cfg(test)] +mod compression_type_tests { + use super::*; + + #[test] + fn compression_type_wire_values_roundtrip() { + for (variant, tag) in [ + (CompressionTypeVariant::GZIP, 0), + (CompressionTypeVariant::BZIP2, 1), + (CompressionTypeVariant::XZ, 2), + (CompressionTypeVariant::ZSTD, 3), + (CompressionTypeVariant::UNCOMPRESSED, 4), + (CompressionTypeVariant::AUTO, 5), + (CompressionTypeVariant::BROTLI, 6), + (CompressionTypeVariant::DEFLATE, 7), + (CompressionTypeVariant::RAW_DEFLATE, 8), + ] { + let protobuf = protobuf::CompressionTypeVariant::from(variant); + assert_eq!(protobuf as i32, tag); + assert_eq!(CompressionTypeVariant::from(protobuf), variant); + assert_eq!(protobuf::CompressionTypeVariant::from(&variant), protobuf); } } } diff --git a/datafusion/proto-common/src/generated/pbjson.rs b/datafusion/proto-common/src/generated/pbjson.rs index b4a622d2e648a..1fc692a4e74d0 100644 --- a/datafusion/proto-common/src/generated/pbjson.rs +++ b/datafusion/proto-common/src/generated/pbjson.rs @@ -1304,6 +1304,10 @@ impl serde::Serialize for CompressionTypeVariant { Self::Xz => "XZ", Self::Zstd => "ZSTD", Self::Uncompressed => "UNCOMPRESSED", + Self::Auto => "AUTO", + Self::Brotli => "BROTLI", + Self::Deflate => "DEFLATE", + Self::RawDeflate => "RAW_DEFLATE", }; serializer.serialize_str(variant) } @@ -1320,6 +1324,10 @@ impl<'de> serde::Deserialize<'de> for CompressionTypeVariant { "XZ", "ZSTD", "UNCOMPRESSED", + "AUTO", + "BROTLI", + "DEFLATE", + "RAW_DEFLATE", ]; struct GeneratedVisitor; @@ -1365,6 +1373,10 @@ impl<'de> serde::Deserialize<'de> for CompressionTypeVariant { "XZ" => Ok(CompressionTypeVariant::Xz), "ZSTD" => Ok(CompressionTypeVariant::Zstd), "UNCOMPRESSED" => Ok(CompressionTypeVariant::Uncompressed), + "AUTO" => Ok(CompressionTypeVariant::Auto), + "BROTLI" => Ok(CompressionTypeVariant::Brotli), + "DEFLATE" => Ok(CompressionTypeVariant::Deflate), + "RAW_DEFLATE" => Ok(CompressionTypeVariant::RawDeflate), _ => Err(serde::de::Error::unknown_variant(value, FIELDS)), } } diff --git a/datafusion/proto-common/src/generated/prost.rs b/datafusion/proto-common/src/generated/prost.rs index d9a0ce71469bf..84548d1283241 100644 --- a/datafusion/proto-common/src/generated/prost.rs +++ b/datafusion/proto-common/src/generated/prost.rs @@ -1239,6 +1239,10 @@ pub enum CompressionTypeVariant { Xz = 2, Zstd = 3, Uncompressed = 4, + Auto = 5, + Brotli = 6, + Deflate = 7, + RawDeflate = 8, } impl CompressionTypeVariant { /// String value of the enum field names used in the ProtoBuf definition. @@ -1252,6 +1256,10 @@ impl CompressionTypeVariant { Self::Xz => "XZ", Self::Zstd => "ZSTD", Self::Uncompressed => "UNCOMPRESSED", + Self::Auto => "AUTO", + Self::Brotli => "BROTLI", + Self::Deflate => "DEFLATE", + Self::RawDeflate => "RAW_DEFLATE", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -1262,6 +1270,10 @@ impl CompressionTypeVariant { "XZ" => Some(Self::Xz), "ZSTD" => Some(Self::Zstd), "UNCOMPRESSED" => Some(Self::Uncompressed), + "AUTO" => Some(Self::Auto), + "BROTLI" => Some(Self::Brotli), + "DEFLATE" => Some(Self::Deflate), + "RAW_DEFLATE" => Some(Self::RawDeflate), _ => None, } } diff --git a/datafusion/proto-common/src/to_proto/mod.rs b/datafusion/proto-common/src/to_proto/mod.rs index fe79657c85a60..2f06127324e0a 100644 --- a/datafusion/proto-common/src/to_proto/mod.rs +++ b/datafusion/proto-common/src/to_proto/mod.rs @@ -841,6 +841,10 @@ impl From<&CompressionTypeVariant> for protobuf::CompressionTypeVariant { CompressionTypeVariant::XZ => Self::Xz, CompressionTypeVariant::ZSTD => Self::Zstd, CompressionTypeVariant::UNCOMPRESSED => Self::Uncompressed, + CompressionTypeVariant::AUTO => Self::Auto, + CompressionTypeVariant::BROTLI => Self::Brotli, + CompressionTypeVariant::DEFLATE => Self::Deflate, + CompressionTypeVariant::RAW_DEFLATE => Self::RawDeflate, } } } diff --git a/datafusion/proto-models/src/from_proto.rs b/datafusion/proto-models/src/from_proto.rs index 7003e45455d27..34bbc55582816 100644 --- a/datafusion/proto-models/src/from_proto.rs +++ b/datafusion/proto-models/src/from_proto.rs @@ -218,6 +218,10 @@ impl From<&CsvOptionsProto> for CsvOptions { 1 => CompressionTypeVariant::BZIP2, 2 => CompressionTypeVariant::XZ, 3 => CompressionTypeVariant::ZSTD, + 5 => CompressionTypeVariant::AUTO, + 6 => CompressionTypeVariant::BROTLI, + 7 => CompressionTypeVariant::DEFLATE, + 8 => CompressionTypeVariant::RAW_DEFLATE, _ => CompressionTypeVariant::UNCOMPRESSED, }, schema_infer_max_rec: proto.schema_infer_max_rec.map(|v| v as usize), @@ -302,6 +306,10 @@ impl From<&JsonOptionsProto> for JsonOptions { 1 => CompressionTypeVariant::BZIP2, 2 => CompressionTypeVariant::XZ, 3 => CompressionTypeVariant::ZSTD, + 5 => CompressionTypeVariant::AUTO, + 6 => CompressionTypeVariant::BROTLI, + 7 => CompressionTypeVariant::DEFLATE, + 8 => CompressionTypeVariant::RAW_DEFLATE, _ => CompressionTypeVariant::UNCOMPRESSED, }, schema_infer_max_rec: proto.schema_infer_max_rec.map(|v| v as usize), diff --git a/datafusion/proto-models/src/generated/datafusion_proto_common.rs b/datafusion/proto-models/src/generated/datafusion_proto_common.rs index d9a0ce71469bf..84548d1283241 100644 --- a/datafusion/proto-models/src/generated/datafusion_proto_common.rs +++ b/datafusion/proto-models/src/generated/datafusion_proto_common.rs @@ -1239,6 +1239,10 @@ pub enum CompressionTypeVariant { Xz = 2, Zstd = 3, Uncompressed = 4, + Auto = 5, + Brotli = 6, + Deflate = 7, + RawDeflate = 8, } impl CompressionTypeVariant { /// String value of the enum field names used in the ProtoBuf definition. @@ -1252,6 +1256,10 @@ impl CompressionTypeVariant { Self::Xz => "XZ", Self::Zstd => "ZSTD", Self::Uncompressed => "UNCOMPRESSED", + Self::Auto => "AUTO", + Self::Brotli => "BROTLI", + Self::Deflate => "DEFLATE", + Self::RawDeflate => "RAW_DEFLATE", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -1262,6 +1270,10 @@ impl CompressionTypeVariant { "XZ" => Some(Self::Xz), "ZSTD" => Some(Self::Zstd), "UNCOMPRESSED" => Some(Self::Uncompressed), + "AUTO" => Some(Self::Auto), + "BROTLI" => Some(Self::Brotli), + "DEFLATE" => Some(Self::Deflate), + "RAW_DEFLATE" => Some(Self::RawDeflate), _ => None, } }