diff --git a/Cargo.toml b/Cargo.toml index 7bbfe851fdab6..f3cc9aed9734f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -371,7 +371,7 @@ unused_qualifications = "deny" # `CARGO_PROFILE_DEV_DEBUG=2 cargo build` when you need to step through code in a # debugger. [profile.dev] -debug = "line-tables-only" +#debug = "line-tables-only" [profile.release] codegen-units = 1 diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index b28709dd096a2..2e8426d9cf57d 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -2947,6 +2947,7 @@ pub enum ConfigFileType { #[cfg(feature = "parquet")] PARQUET, JSON, + ARROW, } /// Represents the configuration options available for handling different table formats within a data processing application. @@ -2965,6 +2966,10 @@ pub struct TableOptions { /// Configuration options for JSON file handling. pub json: JsonOptions, + /// Configuration options for Arrow IPC file handling. This includes settings for the + /// compression codec used for record batches written to disk. + pub arrow: ArrowOptions, + /// The current file format that the table operations should assume. This option allows /// for dynamic switching between the supported file types (e.g., CSV, Parquet, JSON). pub current_format: Option, @@ -2988,11 +2993,13 @@ impl ConfigField for TableOptions { ConfigFileType::PARQUET => self.parquet.visit(v, "format", ""), ConfigFileType::CSV => self.csv.visit(v, "format", ""), ConfigFileType::JSON => self.json.visit(v, "format", ""), + ConfigFileType::ARROW => self.arrow.visit(v, "format", ""), } } else { self.csv.visit(v, "csv", ""); self.parquet.visit(v, "parquet", ""); self.json.visit(v, "json", ""); + self.arrow.visit(v, "arrow", ""); } } @@ -3024,6 +3031,7 @@ impl ConfigField for TableOptions { ConfigFileType::PARQUET => self.parquet.set(rem, value), ConfigFileType::CSV => self.csv.set(rem, value), ConfigFileType::JSON => self.json.set(rem, value), + ConfigFileType::ARROW => self.arrow.set(rem, value), } } _ => _config_err!("Config value \"{key}\" not found on TableOptions"), @@ -3885,6 +3893,80 @@ impl EncryptionFactoryOptions { } } +/// Compression codec used for the record batch bodies of Arrow IPC files. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ArrowIpcCompression { + /// No compression + Uncompressed, + /// LZ4 (frame format) + Lz4Frame, + /// Zstd + Zstd, +} + +impl FromStr for ArrowIpcCompression { + type Err = DataFusionError; + + fn from_str(s: &str) -> Result { + match s.replace('\'', "").to_ascii_lowercase().as_str() { + "uncompressed" | "" => Ok(Self::Uncompressed), + "lz4_frame" => Ok(Self::Lz4Frame), + "zstd" => Ok(Self::Zstd), + other => Err(DataFusionError::Configuration(format!( + "Invalid Arrow IPC compression type: {other}. \ + Expected one of: uncompressed, lz4_frame, zstd" + ))), + } + } +} + +impl Display for ArrowIpcCompression { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let str = match self { + Self::Uncompressed => "uncompressed", + Self::Lz4Frame => "lz4_frame", + Self::Zstd => "zstd", + }; + write!(f, "{str}") + } +} + +impl ConfigField for ArrowIpcCompression { + fn visit(&self, v: &mut V, key: &str, description: &'static str) { + v.some(key, self, description) + } + + fn set(&mut self, _: &str, value: &str) -> Result<()> { + *self = ArrowIpcCompression::from_str(value)?; + Ok(()) + } +} + +impl From for Option { + fn from(c: ArrowIpcCompression) -> Self { + match c { + ArrowIpcCompression::Uncompressed => None, + ArrowIpcCompression::Lz4Frame => Some(CompressionType::LZ4_FRAME), + ArrowIpcCompression::Zstd => Some(CompressionType::ZSTD), + } + } +} + +config_namespace! { + /// Options controlling Arrow IPC format + pub struct ArrowOptions { + /// Compression codec used for the record batch bodies of Arrow IPC files. + /// One of: "uncompressed", "lz4_frame", "zstd" + pub compression: ArrowIpcCompression, default = ArrowIpcCompression::Lz4Frame + /// Compression level to use with the "zstd" codec. Not supported by "lz4_frame" + /// or "uncompressed". If not specified, the default zstd compression level is used. + pub compression_level: Option, default = None + /// Byte alignment to pad buffers to when writing Arrow IPC files. + /// Must be one of 8, 16, 32, or 64. + pub alignment: usize, default = 64 + } +} + config_namespace! { /// Options controlling CSV format pub struct CsvOptions { @@ -4140,7 +4222,7 @@ pub enum OutputFormat { #[cfg(feature = "parquet")] PARQUET(TableParquetOptions), AVRO, - ARROW, + ARROW(ArrowOptions), } impl Display for OutputFormat { @@ -4151,7 +4233,7 @@ impl Display for OutputFormat { #[cfg(feature = "parquet")] OutputFormat::PARQUET(_) => "parquet", OutputFormat::AVRO => "avro", - OutputFormat::ARROW => "arrow", + OutputFormat::ARROW(_) => "arrow", }; write!(f, "{out}") } diff --git a/datafusion/core/src/datasource/file_format/options.rs b/datafusion/core/src/datasource/file_format/options.rs index 8250ef45be2f7..a3743b07a01ef 100644 --- a/datafusion/core/src/datasource/file_format/options.rs +++ b/datafusion/core/src/datasource/file_format/options.rs @@ -821,7 +821,8 @@ impl ReadOptions<'_> for ArrowReadOptions<'_> { _config: &SessionConfig, _table_options: TableOptions, ) -> ListingOptions { - let file_format = ArrowFormat; + // There are no read options for Arrow, so we can ignore the `TableOptions` + let file_format = ArrowFormat::default(); ListingOptions::new(Arc::new(file_format)) .with_file_extension(self.file_extension) diff --git a/datafusion/datasource-arrow/src/file_format.rs b/datafusion/datasource-arrow/src/file_format.rs index b874466e2f7d7..e9453224aab2a 100644 --- a/datafusion/datasource-arrow/src/file_format.rs +++ b/datafusion/datasource-arrow/src/file_format.rs @@ -28,8 +28,9 @@ use arrow::datatypes::{Schema, SchemaRef}; use arrow::error::ArrowError; use arrow::ipc::convert::try_fb_to_schema; use arrow::ipc::reader::{FileReader, StreamReader}; +use arrow::ipc::root_as_message; use arrow::ipc::writer::IpcWriteOptions; -use arrow::ipc::{CompressionType, root_as_message}; +use datafusion_common::config::{ArrowOptions, ConfigField, ConfigFileType}; use datafusion_common::error::Result; use datafusion_common::parsers::CompressionTypeVariant; use datafusion_common::{ @@ -76,26 +77,52 @@ const BUFFER_FLUSH_BYTES: usize = 1024000; /// Factory struct used to create [`ArrowFormat`] #[derive(Default, Debug)] -pub struct ArrowFormatFactory; +pub struct ArrowFormatFactory { + /// the write options for arrow + pub options: Option, +} impl ArrowFormatFactory { /// Creates an instance of [ArrowFormatFactory] pub fn new() -> Self { - Self {} + Self { options: None } + } + + /// Creates an instance of [`ArrowFormatFactory`] with customized default options + pub fn new_with_options(options: ArrowOptions) -> Self { + Self { + options: Some(options), + } } } impl FileFormatFactory for ArrowFormatFactory { fn create( &self, - _state: &dyn Session, - _format_options: &HashMap, + state: &dyn Session, + format_options: &HashMap, ) -> Result> { - Ok(Arc::new(ArrowFormat)) + let arrow_options = match &self.options { + None => { + let mut table_options = state.default_table_options(); + table_options.set_config_format(ConfigFileType::ARROW); + table_options.alter_with_string_hash_map(format_options)?; + table_options.arrow + } + Some(arrow_options) => { + let mut arrow_options = arrow_options.clone(); + for (k, v) in format_options { + arrow_options.set(k, v)?; + } + arrow_options + } + }; + + Ok(Arc::new(ArrowFormat::default().with_options(arrow_options))) } fn default(&self) -> Arc { - Arc::new(ArrowFormat) + Arc::new(ArrowFormat::default()) } } @@ -108,7 +135,22 @@ impl GetExt for ArrowFormatFactory { /// Arrow [`FileFormat`] implementation. #[derive(Default, Debug)] -pub struct ArrowFormat; +pub struct ArrowFormat { + options: ArrowOptions, +} + +impl ArrowFormat { + /// Set the arrow options + pub fn with_options(mut self, options: ArrowOptions) -> Self { + self.options = options; + self + } + + /// Retrieve the arrow options + pub fn options(&self) -> &ArrowOptions { + &self.options + } +} #[async_trait] impl FileFormat for ArrowFormat { @@ -234,7 +276,7 @@ impl FileFormat for ArrowFormat { return not_impl_err!("Overwrites are not implemented yet for Arrow format"); } - let sink = Arc::new(ArrowFileSink::new(conf)); + let sink = Arc::new(ArrowFileSink::new(conf, self.options.clone())); Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) } @@ -247,11 +289,12 @@ impl FileFormat for ArrowFormat { /// Implements [`FileSink`] for Arrow IPC files struct ArrowFileSink { config: FileSinkConfig, + options: ArrowOptions, } impl ArrowFileSink { - fn new(config: FileSinkConfig) -> Self { - Self { config } + fn new(config: FileSinkConfig, options: ArrowOptions) -> Self { + Self { config, options } } } @@ -271,9 +314,13 @@ impl FileSink for ArrowFileSink { let mut file_write_tasks: JoinSet> = JoinSet::new(); - let ipc_options = - IpcWriteOptions::try_new(64, false, arrow_ipc::MetadataVersion::V5)? - .try_with_compression(Some(CompressionType::LZ4_FRAME))?; + let ipc_options = IpcWriteOptions::try_new( + self.options.alignment, + false, + arrow_ipc::MetadataVersion::V5, + )? + .try_with_compression(self.options.compression.into())? + .try_with_compression_level(self.options.compression_level)?; while let Some((file_metadata, mut rx)) = file_stream_rx.recv().await { let shared_buffer = SharedBuffer::new(INITIAL_BUFFER_BYTES); let mut arrow_writer = arrow_ipc::writer::FileWriter::try_new_with_options( @@ -660,7 +707,7 @@ mod tests { version: None, }; - let arrow_format = ArrowFormat {}; + let arrow_format = ArrowFormat::default(); let expected = vec!["f0: Int64", "f1: Utf8", "f2: Boolean"]; // Test chunk sizes where too small so we keep having to read more bytes @@ -704,7 +751,7 @@ mod tests { version: None, }; - let arrow_format = ArrowFormat {}; + let arrow_format = ArrowFormat::default(); let store = Arc::new(ChunkedStore::new(in_memory_store.clone(), 7)); let err = arrow_format diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index c3a7519cf584e..fcf546f43803e 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -329,6 +329,7 @@ fn table_options_to_rhash(mut options: TableOptions) -> SVec<(SString, SString)> #[cfg(feature = "parquet")] ConfigFileType::PARQUET => "parquet", ConfigFileType::CSV => "csv", + ConfigFileType::ARROW => "arrow", } .into(), ); @@ -633,6 +634,7 @@ fn table_options_from_rhashmap(options: SVec<(SString, SString)>) -> TableOption ConfigFileType::JSON, #[cfg(feature = "parquet")] ConfigFileType::PARQUET, + ConfigFileType::ARROW, ]; for format in formats { // It is imperative that if new enum variants are added below that they be @@ -642,6 +644,7 @@ fn table_options_from_rhashmap(options: SVec<(SString, SString)>) -> TableOption #[cfg(feature = "parquet")] ConfigFileType::PARQUET => "parquet", ConfigFileType::JSON => "json", + ConfigFileType::ARROW => "arrow", }; let format_options: HashMap = options .iter() @@ -665,7 +668,7 @@ fn table_options_from_rhashmap(options: SVec<(SString, SString)>) -> TableOption .iter() .filter_map(|(k, v)| { let (prefix, _) = k.split_once('.')?; - if !["json", "parquet", "csv"].contains(&prefix) { + if !["json", "parquet", "csv", "arrow"].contains(&prefix) { Some((k.to_owned(), v.to_owned())) } else { None @@ -684,6 +687,7 @@ fn table_options_from_rhashmap(options: SVec<(SString, SString)>) -> TableOption #[cfg(feature = "parquet")] "parquet" => Some(ConfigFileType::PARQUET), "json" => Some(ConfigFileType::JSON), + "arrow" => Some(ConfigFileType::ARROW), _ => None, }); table_options diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 2e1811f76b3a7..03065a33e93fe 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -649,7 +649,7 @@ impl AsLogicalPlan for LogicalPlanNode { } } FileFormatType::Arrow(..) => { - Arc::new(ArrowFormat) + Arc::new(ArrowFormat::default()) } }; diff --git a/datafusion/sqllogictest/test_files/copy.slt b/datafusion/sqllogictest/test_files/copy.slt index c569dcb64c0de..6f62827412738 100644 --- a/datafusion/sqllogictest/test_files/copy.slt +++ b/datafusion/sqllogictest/test_files/copy.slt @@ -661,6 +661,86 @@ select * from validate_arrow; 1 Foo 2 Bar +# Copy to arrow file with zstd compression +query I +COPY source_table to 'test_files/scratch/copy/table_zstd.arrow' STORED AS ARROW OPTIONS ('format.compression' 'zstd'); +---- +2 + +statement ok +CREATE EXTERNAL TABLE validate_arrow_zstd +STORED AS arrow +LOCATION 'test_files/scratch/copy/table_zstd.arrow'; + +query IT +select * from validate_arrow_zstd; +---- +1 Foo +2 Bar + +# Copy to arrow file with no compression +query I +COPY source_table to 'test_files/scratch/copy/table_uncompressed.arrow' STORED AS ARROW OPTIONS ('format.compression' 'uncompressed'); +---- +2 + +statement ok +CREATE EXTERNAL TABLE validate_arrow_uncompressed +STORED AS arrow +LOCATION 'test_files/scratch/copy/table_uncompressed.arrow'; + +query IT +select * from validate_arrow_uncompressed; +---- +1 Foo +2 Bar + +# Copy to arrow file with an invalid compression codec +query error DataFusion error: Invalid or Unsupported Configuration: Invalid Arrow IPC compression type: snappy\. Expected one of: uncompressed, lz4_frame, zstd +COPY source_table to 'test_files/scratch/copy/table_invalid_compression.arrow' STORED AS ARROW OPTIONS ('format.compression' 'snappy'); + +# Copy to arrow file with zstd compression and an explicit compression level +query I +COPY source_table to 'test_files/scratch/copy/table_zstd_level.arrow' STORED AS ARROW OPTIONS ('format.compression' 'zstd', 'format.compression_level' '15'); +---- +2 + +statement ok +CREATE EXTERNAL TABLE validate_arrow_zstd_level +STORED AS arrow +LOCATION 'test_files/scratch/copy/table_zstd_level.arrow'; + +query IT +select * from validate_arrow_zstd_level; +---- +1 Foo +2 Bar + +# compression_level is not supported outside of zstd +query error DataFusion error: Arrow error: Invalid argument error: LZ4 Frame compression does not support configurable compression levels +COPY source_table to 'test_files/scratch/copy/table_lz4_level.arrow' STORED AS ARROW OPTIONS ('format.compression' 'lz4_frame', 'format.compression_level' '5'); + +# Copy to arrow file with a non-default byte alignment +query I +COPY source_table to 'test_files/scratch/copy/table_alignment.arrow' STORED AS ARROW OPTIONS ('format.alignment' '32'); +---- +2 + +statement ok +CREATE EXTERNAL TABLE validate_arrow_alignment +STORED AS arrow +LOCATION 'test_files/scratch/copy/table_alignment.arrow'; + +query IT +select * from validate_arrow_alignment; +---- +1 Foo +2 Bar + +# alignment must be one of 8, 16, 32, or 64 +query error DataFusion error: Arrow error: Invalid argument error: Alignment should be 8, 16, 32, or 64\. +COPY source_table to 'test_files/scratch/copy/table_bad_alignment.arrow' STORED AS ARROW OPTIONS ('format.alignment' '10'); + # Format Options Support without the 'format.' prefix # Copy with format options for Parquet without the 'format.' prefix diff --git a/docs/source/user-guide/sql/format_options.md b/docs/source/user-guide/sql/format_options.md index 719fd5bd6b1ee..738ed8b5de638 100644 --- a/docs/source/user-guide/sql/format_options.md +++ b/docs/source/user-guide/sql/format_options.md @@ -78,6 +78,25 @@ In this example, we write the entire `source_table` out to a folder of Parquet f # Available Options +## Arrow Format Options + +The following options are available when writing [Arrow IPC](https://arrow.apache.org/docs/format/Columnar.html#ipc-file-format) files. Note: If any unsupported option is specified, an error will be raised and the query will fail. + +| Option | Description | Default Value | +| ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------- | +| COMPRESSION | Sets the compression codec used for the record batch bodies within the Arrow IPC file. Supported values are UNCOMPRESSED, LZ4_FRAME, and ZSTD. | LZ4_FRAME | +| COMPRESSION_LEVEL | Sets the compression level to use with the `ZSTD` codec. Not supported by `LZ4_FRAME` or `UNCOMPRESSED`. If not specified, the default zstd compression level is used. | None | +| ALIGNMENT | Sets the byte alignment used to pad buffers when writing the Arrow IPC file. Must be one of 8, 16, 32, or 64. | 64 | + +**Example:** + +```sql +COPY source_table + TO '/tmp/foo.arrow' + STORED AS ARROW + OPTIONS('COMPRESSION' 'zstd', 'COMPRESSION_LEVEL' '15'); +``` + ## JSON Format Options The following options are available when reading or writing JSON files. Note: If any unsupported option is specified, an error will be raised and the query will fail.