Skip to content

Commit af47ffe

Browse files
authored
Merge pull request #132 from Embucket/df55-copy-compression-auto
Support AUTO and additional compression codecs for CSV/JSON scans
2 parents d02410a + 6295c85 commit af47ffe

20 files changed

Lines changed: 1041 additions & 44 deletions

File tree

‎Cargo.lock‎

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎datafusion/catalog-listing/src/config.rs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,24 @@ pub enum SchemaSource {
4040
Specified,
4141
}
4242

43+
#[cfg(test)]
44+
mod compression_extension_tests {
45+
use super::ListingTableConfig;
46+
47+
#[test]
48+
fn new_codec_suffixes_preserve_csv_extension() {
49+
for suffix in ["br", "zlib", "raw_deflate"] {
50+
assert_eq!(
51+
ListingTableConfig::infer_file_extension_and_compression_type(&format!(
52+
"data.csv.{suffix}"
53+
))
54+
.unwrap(),
55+
("csv".to_owned(), Some(suffix.to_owned()))
56+
);
57+
}
58+
}
59+
}
60+
4361
/// Configuration for creating a [`crate::ListingTable`]
4462
///
4563
/// # Schema Evolution Support

‎datafusion/catalog-listing/src/helpers.rs‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -389,6 +389,28 @@ pub async fn pruned_partition_list<'a>(
389389
filters: &'a [Expr],
390390
file_extension: &'a str,
391391
partition_cols: &'a [(String, DataType)],
392+
) -> Result<BoxStream<'a, Result<PartitionedFile>>> {
393+
pruned_partition_list_with_suffixes(
394+
ctx,
395+
store,
396+
table_path,
397+
filters,
398+
file_extension,
399+
None,
400+
partition_cols,
401+
)
402+
.await
403+
}
404+
405+
/// As [`pruned_partition_list`], but accepts several filename suffixes.
406+
pub async fn pruned_partition_list_with_suffixes<'a>(
407+
ctx: &'a dyn Session,
408+
store: &'a dyn ObjectStore,
409+
table_path: &'a ListingTableUrl,
410+
filters: &'a [Expr],
411+
file_extension: &'a str,
412+
file_suffixes: Option<&'a [String]>,
413+
partition_cols: &'a [(String, DataType)],
392414
) -> Result<BoxStream<'a, Result<PartitionedFile>>> {
393415
let prefix = if !partition_cols.is_empty() {
394416
evaluate_partition_prefix(partition_cols, filters)
@@ -397,7 +419,13 @@ pub async fn pruned_partition_list<'a>(
397419
};
398420

399421
let objects = table_path
400-
.list_prefixed_files(ctx, store, prefix, file_extension)
422+
.list_prefixed_files_with_suffixes(
423+
ctx,
424+
store,
425+
prefix,
426+
file_extension,
427+
file_suffixes,
428+
)
401429
.await?
402430
.try_filter(|object_meta| futures::future::ready(object_meta.size > 0));
403431

‎datafusion/catalog-listing/src/options.rs‎

Lines changed: 52 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,9 @@
1717

1818
use arrow::datatypes::{DataType, SchemaRef};
1919
use datafusion_catalog::Session;
20-
use datafusion_common::plan_err;
20+
use datafusion_common::{GetExt, plan_err};
2121
use datafusion_datasource::ListingTableUrl;
22+
use datafusion_datasource::file_compression_type::FileCompressionType;
2223
use datafusion_datasource::file_format::FileFormat;
2324
use datafusion_expr::{Partitioning, SortExpr};
2425
use futures::StreamExt;
@@ -33,6 +34,9 @@ pub struct ListingOptions {
3334
/// A suffix on which files should be filtered (leave empty to
3435
/// keep all files on the path)
3536
pub file_extension: String,
37+
/// Accepted filename suffixes for header-detected compression. An explicit
38+
/// file-extension override disables this filter.
39+
pub auto_file_suffixes: Option<Vec<String>>,
3640
/// The file format
3741
pub format: Arc<dyn FileFormat>,
3842
/// The expected partition column names in the folder structure.
@@ -96,14 +100,38 @@ pub struct ListingOptions {
96100
pub output_partitioning: Option<Partitioning>,
97101
}
98102

103+
fn auto_compression_suffixes(base_extension: &str) -> Vec<String> {
104+
let base = format!(".{}", base_extension.trim_start_matches('.'));
105+
[
106+
FileCompressionType::UNCOMPRESSED,
107+
FileCompressionType::GZIP,
108+
FileCompressionType::BZIP2,
109+
FileCompressionType::XZ,
110+
FileCompressionType::ZSTD,
111+
FileCompressionType::DEFLATE,
112+
]
113+
.into_iter()
114+
.map(|compression| format!("{base}{}", compression.get_ext()))
115+
.collect()
116+
}
117+
99118
impl ListingOptions {
100119
/// Creates an options instance with the given format
101120
/// Default values:
102121
/// - use default file extension filter
103122
/// - no input partition to discover
104123
pub fn new(format: Arc<dyn FileFormat>) -> Self {
124+
let auto = format.compression_type() == Some(FileCompressionType::AUTO);
125+
let file_extension = if auto {
126+
String::new()
127+
} else {
128+
format.get_ext()
129+
};
130+
let auto_file_suffixes =
131+
auto.then(|| auto_compression_suffixes(&format.get_ext()));
105132
Self {
106-
file_extension: format.get_ext(),
133+
file_extension,
134+
auto_file_suffixes,
107135
format,
108136
table_partition_cols: vec![],
109137
file_sort_order: vec![],
@@ -126,6 +154,15 @@ impl ListingOptions {
126154
/// ```
127155
pub fn with_file_extension(mut self, file_extension: impl Into<String>) -> Self {
128156
self.file_extension = file_extension.into();
157+
self.auto_file_suffixes = None;
158+
self
159+
}
160+
161+
/// Detect compression while restricting a listing to files of this format.
162+
pub fn with_auto_compression_file_extension(mut self, base_extension: &str) -> Self {
163+
self.file_extension.clear();
164+
self.auto_file_suffixes = (!base_extension.is_empty())
165+
.then(|| auto_compression_suffixes(base_extension));
129166
self
130167
}
131168

@@ -150,7 +187,7 @@ impl ListingOptions {
150187
S: Into<String>,
151188
{
152189
if let Some(file_extension) = file_extension {
153-
self.file_extension = file_extension.into();
190+
self = self.with_file_extension(file_extension);
154191
}
155192
self
156193
}
@@ -280,7 +317,12 @@ impl ListingOptions {
280317
let store = state.runtime_env().object_store(table_path)?;
281318

282319
let all_files: Vec<_> = table_path
283-
.list_all_files(state, store.as_ref(), &self.file_extension)
320+
.list_all_files_with_suffixes(
321+
state,
322+
store.as_ref(),
323+
&self.file_extension,
324+
self.auto_file_suffixes.as_deref(),
325+
)
284326
.await?
285327
.try_collect()
286328
.await?;
@@ -375,7 +417,12 @@ impl ListingOptions {
375417
// This can fail to detect inconsistent partition keys
376418
// A DFS traversal approach of the store can help here
377419
let files: Vec<_> = table_path
378-
.list_all_files(state, store.as_ref(), &self.file_extension)
420+
.list_all_files_with_suffixes(
421+
state,
422+
store.as_ref(),
423+
&self.file_extension,
424+
self.auto_file_suffixes.as_deref(),
425+
)
379426
.await?
380427
.take(10)
381428
.try_collect()

‎datafusion/catalog-listing/src/table.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,8 @@
1717

1818
use crate::config::SchemaSource;
1919
use crate::helpers::{
20-
expr_applicable_for_cols, filter_partitioned_file, pruned_partition_list,
20+
expr_applicable_for_cols, filter_partitioned_file,
21+
pruned_partition_list_with_suffixes as pruned_partition_list,
2122
};
2223
use crate::{ListingOptions, ListingTableConfig};
2324
use arrow::datatypes::{Field, Schema, SchemaBuilder, SchemaRef};
@@ -725,6 +726,7 @@ impl TableProvider for ListingTable {
725726
table_path,
726727
&[],
727728
&self.options.file_extension,
729+
self.options.auto_file_suffixes.as_deref(),
728730
&self.options.table_partition_cols,
729731
)
730732
.await?;
@@ -812,6 +814,7 @@ impl ListingTable {
812814
table_path,
813815
listing_time_filters,
814816
&self.options.file_extension,
817+
self.options.auto_file_suffixes.as_deref(),
815818
&self.options.table_partition_cols,
816819
)
817820
}))

‎datafusion/common/src/parsers.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,15 @@ pub enum CompressionTypeVariant {
3535
ZSTD,
3636
/// Uncompressed file
3737
UNCOMPRESSED,
38+
/// Detect file compression from its header when reading
39+
AUTO,
40+
/// Brotli-compressed file
41+
BROTLI,
42+
/// Deflate-compressed file with a zlib header
43+
DEFLATE,
44+
/// Deflate-compressed file without a zlib header
45+
#[expect(non_camel_case_types, reason = "match the external codec name")]
46+
RAW_DEFLATE,
3847
}
3948

4049
impl FromStr for CompressionTypeVariant {
@@ -48,6 +57,10 @@ impl FromStr for CompressionTypeVariant {
4857
"XZ" => Ok(Self::XZ),
4958
"ZST" | "ZSTD" => Ok(Self::ZSTD),
5059
"" | "UNCOMPRESSED" => Ok(Self::UNCOMPRESSED),
60+
"AUTO" => Ok(Self::AUTO),
61+
"BROTLI" | "BR" => Ok(Self::BROTLI),
62+
"DEFLATE" | "ZLIB" => Ok(Self::DEFLATE),
63+
"RAW_DEFLATE" => Ok(Self::RAW_DEFLATE),
5164
_ => Err(DataFusionError::NotImplemented(format!(
5265
"Unsupported file compression type {s}"
5366
))),
@@ -63,6 +76,10 @@ impl Display for CompressionTypeVariant {
6376
Self::XZ => "XZ",
6477
Self::ZSTD => "ZSTD",
6578
Self::UNCOMPRESSED => "",
79+
Self::AUTO => "AUTO",
80+
Self::BROTLI => "BROTLI",
81+
Self::DEFLATE => "DEFLATE",
82+
Self::RAW_DEFLATE => "RAW_DEFLATE",
6683
};
6784
write!(f, "{str}")
6885
}

‎datafusion/core/src/datasource/file_format/options.rs‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -643,8 +643,13 @@ impl ReadOptions<'_> for CsvReadOptions<'_> {
643643
.with_null_regex(self.null_regex.clone())
644644
.with_truncated_rows(self.truncated_rows);
645645

646-
ListingOptions::new(Arc::new(file_format))
647-
.with_file_extension(self.file_extension)
646+
let listing = ListingOptions::new(Arc::new(file_format));
647+
let listing = if self.file_compression_type == FileCompressionType::AUTO {
648+
listing.with_auto_compression_file_extension(self.file_extension)
649+
} else {
650+
listing.with_file_extension(self.file_extension)
651+
};
652+
listing
648653
.with_table_partition_cols(self.table_partition_cols.clone())
649654
.with_file_sort_order(self.file_sort_order.clone())
650655
}
@@ -724,8 +729,13 @@ impl ReadOptions<'_> for JsonReadOptions<'_> {
724729
.with_file_compression_type(self.file_compression_type.to_owned())
725730
.with_newline_delimited(self.newline_delimited);
726731

727-
ListingOptions::new(Arc::new(file_format))
728-
.with_file_extension(self.file_extension)
732+
let listing = ListingOptions::new(Arc::new(file_format));
733+
let listing = if self.file_compression_type == FileCompressionType::AUTO {
734+
listing.with_auto_compression_file_extension(self.file_extension)
735+
} else {
736+
listing.with_file_extension(self.file_extension)
737+
};
738+
listing
729739
.with_table_partition_cols(self.table_partition_cols.clone())
730740
.with_file_sort_order(self.file_sort_order.clone())
731741
}

0 commit comments

Comments
 (0)