refactor: centralize default Parquet writer properties (delta-io#4335)

ethan-tyler · web-flow · commit c8f9184eecaf · 2026-04-07T23:00:43.000+02:00
diff --git a/crates/core/src/delta_datafusion/mod.rs b/crates/core/src/delta_datafusion/mod.rs
@@ -633,6 +633,7 @@ mod tests {
     use delta_kernel::schema::ArrayType;
     use futures::{StreamExt, TryStreamExt};
     use serde_json::json;
+    use std::ops::Range;
     use url::Url;
 
     use super::*;
@@ -1451,28 +1452,65 @@ mod tests {
         let small = batch.column_by_name("small").unwrap().as_string::<i32>();
         assert_eq!("a", small.iter().next().unwrap().unwrap());
 
+        let files = table.get_files_by_partitions(&[]).await.unwrap();
+        assert_eq!(1, files.len());
+        let object_store = table.object_store();
+        let file_meta = object_store.head(&files[0]).await.unwrap();
+        let file_reader = parquet::arrow::async_reader::ParquetObjectReader::new(
+            object_store,
+            file_meta.location.clone(),
+        )
+        .with_file_size(file_meta.size);
+        let parquet_metadata =
+            parquet::arrow::async_reader::ParquetRecordBatchStreamBuilder::new(file_reader)
+                .await
+                .unwrap()
+                .metadata()
+                .as_ref()
+                .clone();
+        let (small_start, small_len) = parquet_metadata.row_group(0).column(0).byte_range();
+        let small_range = small_start..small_start + small_len;
+        let (large_start, large_len) = parquet_metadata.row_group(0).column(1).byte_range();
+        let large_range = large_start..large_start + large_len;
+
         let actual = drain_recorded_ops(&mut operations).await;
+
+        let data_ranges = actual
+            .iter()
+            .flat_map(|operation| match operation {
+                ObjectStoreOperation::GetRange(PathKind::Data, range) => vec![range.clone()],
+                ObjectStoreOperation::GetRanges(PathKind::Data, ranges) => ranges.clone(),
+                _ => Vec::new(),
+            })
+            .collect::<Vec<_>>();
+
+        let overlaps = |left: &Range<u64>, right: &Range<u64>| {
+            left.start < right.end && right.start < left.end
+        };
+
         assert!(
-            actual.iter().any(|op| matches!(
-                op,
-                ObjectStoreOperation::GetRange(PathKind::Data, range)
-                    if range == &(957_u64..965_u64)
-            )),
-            "expected footer range read for the selected column, saw {actual:?}",
+            !data_ranges.is_empty(),
+            "expected ranged parquet data reads, saw {actual:?}"
         );
         assert!(
-            actual.iter().any(|op| matches!(
-                op,
-                ObjectStoreOperation::GetRange(PathKind::Data, range)
-                    if range == &(326_u64..957_u64)
-            )),
-            "expected column chunk range read for the selected column, saw {actual:?}",
+            data_ranges
+                .iter()
+                .any(|range| overlaps(range, &small_range)),
+            "expected selected column chunk {small_range:?} to be read, saw {actual:?}"
         );
         assert!(
-            !actual
+            data_ranges
                 .iter()
-                .any(|op| matches!(op, ObjectStoreOperation::Get(PathKind::Data))),
-            "expected column pruning to avoid full parquet object reads, saw {actual:?}",
+                .all(|range| !overlaps(range, &large_range)),
+            "expected unselected column chunk {large_range:?} to be pruned, saw {actual:?}"
+        );
+        assert!(
+            !actual.iter().any(|operation| matches!(
+                operation,
+                ObjectStoreOperation::Get(PathKind::Data)
+                    | ObjectStoreOperation::GetOpts(PathKind::Data)
+            )),
+            "expected no full data file reads, saw {actual:?}"
         );
     }
 
diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs
@@ -83,6 +83,7 @@ pub mod errors;
 pub mod kernel;
 pub mod logstore;
 pub mod operations;
+pub(crate) mod parquet_utils;
 pub mod protocol;
 pub use kernel::schema;
 pub mod table;
diff --git a/crates/core/src/operations/optimize.rs b/crates/core/src/operations/optimize.rs
@@ -58,11 +58,12 @@ use crate::kernel::transaction::{CommitBuilder, CommitProperties, DEFAULT_RETRIE
 use crate::kernel::{Action, Add, PartitionsExt, Remove, Version, scalars::ScalarExt};
 use crate::kernel::{EagerSnapshot, resolve_snapshot};
 use crate::logstore::{LogStore, LogStoreRef, ObjectStoreRef};
+use crate::parquet_utils::default_writer_properties;
 use crate::protocol::DeltaOperation;
 use crate::table::config::TablePropertiesExt as _;
 use crate::table::state::DeltaTableState;
 use crate::writer::utils::arrow_schema_without_partitions;
-use crate::{DeltaTable, ObjectMeta, PartitionFilter, crate_version, to_kernel_predicate};
+use crate::{DeltaTable, ObjectMeta, PartitionFilter, to_kernel_predicate};
 
 /// Planner used by optimize.
 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
@@ -421,10 +422,7 @@ impl<'a> std::future::IntoFuture for OptimizeBuilder<'a> {
             this.pre_execute(operation_id).await?;
 
             let writer_properties = this.writer_properties.unwrap_or_else(|| {
-                WriterProperties::builder()
-                    .set_compression(Compression::ZSTD(ZstdLevel::try_new(4).unwrap()))
-                    .set_created_by(format!("delta-rs version {}", crate_version()))
-                    .build()
+                default_writer_properties(Compression::ZSTD(ZstdLevel::try_new(4).unwrap()))
             });
             let (session, _) = resolve_session_state(
                 this.session.as_deref(),
diff --git a/crates/core/src/operations/write/writer.rs b/crates/core/src/operations/write/writer.rs
@@ -19,11 +19,10 @@ use parquet::file::properties::WriterProperties;
 use tokio::task::JoinSet;
 use tracing::*;
 
-use crate::crate_version;
-
 use crate::errors::{DeltaResult, DeltaTableError};
 use crate::kernel::{Add, PartitionsExt};
 use crate::logstore::ObjectStoreRef;
+use crate::parquet_utils::default_writer_properties;
 use crate::writer::record_batch::{PartitionResult, divide_by_partition_values};
 use crate::writer::stats::create_add;
 use crate::writer::utils::{
@@ -88,17 +87,6 @@ fn sort_completed_writes_by_path<T>(results: &mut [(Path, usize, T)]) {
     results.sort_unstable_by(|a, b| a.0.cmp(&b.0));
 }
 
-fn default_writer_properties(include_created_by: bool) -> WriterProperties {
-    let builder = WriterProperties::builder().set_compression(Compression::SNAPPY);
-    if include_created_by {
-        builder
-            .set_created_by(format!("delta-rs version {}", crate_version()))
-            .build()
-    } else {
-        builder.build()
-    }
-}
-
 #[derive(thiserror::Error, Debug)]
 enum WriteError {
     #[error("Unexpected Arrow schema: got: {schema}, expected: {expected_schema}")]
@@ -169,7 +157,7 @@ impl WriterConfig {
         stats_columns: Option<Vec<String>>,
     ) -> Self {
         let writer_properties =
-            writer_properties.unwrap_or_else(|| default_writer_properties(false));
+            writer_properties.unwrap_or_else(|| default_writer_properties(Compression::SNAPPY));
         let write_batch_size = write_batch_size.unwrap_or(DEFAULT_WRITE_BATCH_SIZE);
 
         Self {
@@ -335,7 +323,7 @@ impl PartitionWriterConfig {
         let part_path = partition_values.hive_partition_path();
         let prefix = Path::parse(part_path)?;
         let writer_properties =
-            writer_properties.unwrap_or_else(|| default_writer_properties(true));
+            writer_properties.unwrap_or_else(|| default_writer_properties(Compression::SNAPPY));
         let write_batch_size = write_batch_size.unwrap_or(DEFAULT_WRITE_BATCH_SIZE);
 
         Ok(Self {
@@ -550,11 +538,13 @@ impl PartitionWriter {
 mod tests {
     use super::*;
     use crate::DeltaTableBuilder;
+    use crate::crate_version;
     use crate::logstore::tests::flatten_list_stream as list;
     use crate::table::config::DEFAULT_NUM_INDEX_COLS;
     use crate::writer::test_utils::*;
     use arrow::array::{Int32Array, StringArray};
     use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
+    use parquet::schema::types::ColumnPath;
     use std::sync::Arc;
 
     fn get_delta_writer(
@@ -601,6 +591,59 @@ mod tests {
         .unwrap()
     }
 
+    fn assert_default_created_by(writer_properties: &WriterProperties) {
+        assert_eq!(
+            writer_properties.created_by(),
+            format!("delta-rs version {}", crate_version())
+        );
+    }
+
+    #[test]
+    fn test_writer_config_defaults_include_delta_rs_created_by() {
+        let schema = Arc::new(ArrowSchema::new(vec![Field::new(
+            "id",
+            DataType::Int32,
+            true,
+        )]));
+        let config = WriterConfig::new(
+            schema,
+            vec![],
+            None,
+            None,
+            None,
+            DataSkippingNumIndexedCols::NumColumns(DEFAULT_NUM_INDEX_COLS),
+            None,
+        );
+
+        assert_default_created_by(&config.writer_properties);
+        assert_eq!(
+            config
+                .writer_properties
+                .compression(&ColumnPath::from("id")),
+            Compression::SNAPPY
+        );
+    }
+
+    #[test]
+    fn test_partition_writer_config_defaults_include_delta_rs_created_by() {
+        let schema = Arc::new(ArrowSchema::new(vec![Field::new(
+            "id",
+            DataType::Int32,
+            true,
+        )]));
+        let config =
+            PartitionWriterConfig::try_new(schema, IndexMap::new(), None, None, None, None)
+                .unwrap();
+
+        assert_default_created_by(&config.writer_properties);
+        assert_eq!(
+            config
+                .writer_properties
+                .compression(&ColumnPath::from("id")),
+            Compression::SNAPPY
+        );
+    }
+
     #[tokio::test]
     async fn test_write_partition() {
         let log_store = DeltaTableBuilder::from_url(url::Url::parse("memory:///").unwrap())
diff --git a/crates/core/src/parquet_utils.rs b/crates/core/src/parquet_utils.rs
@@ -0,0 +1,29 @@
+use parquet::basic::Compression;
+use parquet::file::properties::WriterProperties;
+
+pub(crate) fn default_writer_properties(compression: Compression) -> WriterProperties {
+    WriterProperties::builder()
+        .set_created_by(format!("delta-rs version {}", crate::crate_version()))
+        .set_compression(compression)
+        .build()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use parquet::schema::types::ColumnPath;
+
+    #[test]
+    fn default_writer_properties_sets_created_by_and_compression() {
+        let writer_properties = default_writer_properties(Compression::SNAPPY);
+
+        assert_eq!(
+            writer_properties.created_by(),
+            format!("delta-rs version {}", crate::crate_version())
+        );
+        assert_eq!(
+            writer_properties.compression(&ColumnPath::from("id")),
+            Compression::SNAPPY
+        );
+    }
+}
diff --git a/crates/core/src/protocol/checkpoints.rs b/crates/core/src/protocol/checkpoints.rs
@@ -2,7 +2,6 @@
 
 use std::sync::LazyLock;
 
-use parquet::file::properties::WriterProperties;
 use url::Url;
 
 use chrono::{TimeZone, Utc};
@@ -20,6 +19,7 @@ use uuid::Uuid;
 
 use crate::kernel::{Version, spawn_blocking_with_span};
 use crate::logstore::{DELTA_LOG_REGEX, LogStore};
+use crate::parquet_utils::default_writer_properties;
 use crate::protocol::to_rb;
 use crate::table::config::TablePropertiesExt as _;
 use crate::{DeltaResult, DeltaTableError};
@@ -67,11 +67,9 @@ pub(crate) async fn create_checkpoint_for(
     let mut writer = AsyncArrowWriter::try_new(
         object_store_writer,
         first_batch.schema(),
-        Some(
-            WriterProperties::builder()
-                .set_compression(parquet::basic::Compression::SNAPPY)
-                .build(),
-        ),
+        Some(default_writer_properties(
+            parquet::basic::Compression::SNAPPY,
+        )),
     )?;
     writer.write(&first_batch).await?;
 
diff --git a/crates/core/src/writer/json.rs b/crates/core/src/writer/json.rs
@@ -10,10 +10,7 @@ use delta_kernel::expressions::Scalar;
 use indexmap::IndexMap;
 use itertools::Itertools;
 use object_store::path::Path;
-use parquet::{
-    arrow::ArrowWriter, basic::Compression, errors::ParquetError,
-    file::properties::WriterProperties,
-};
+use parquet::{arrow::ArrowWriter, errors::ParquetError, file::properties::WriterProperties};
 use serde_json::Value;
 use tracing::*;
 use url::Url;
@@ -29,6 +26,7 @@ use crate::DeltaTable;
 use crate::errors::DeltaTableError;
 use crate::kernel::{Add, PartitionsExt, scalars::ScalarExt};
 use crate::logstore::ObjectStoreRetryExt;
+use crate::parquet_utils::default_writer_properties;
 use crate::table::builder::DeltaTableBuilder;
 use crate::table::config::TablePropertiesExt as _;
 use crate::writer::utils::ShareableBuffer;
@@ -198,10 +196,7 @@ impl JsonWriter {
             .load()
             .await?;
         // Initialize writer properties for the underlying arrow writer
-        let writer_properties = WriterProperties::builder()
-            // NOTE: Consider extracting config for writer properties and setting more than just compression
-            .set_compression(Compression::SNAPPY)
-            .build();
+        let writer_properties = default_writer_properties(parquet::basic::Compression::SNAPPY);
 
         Ok(Self {
             table,
@@ -219,10 +214,7 @@ impl JsonWriter {
         let partition_columns = metadata.partition_columns().clone();
 
         // Initialize writer properties for the underlying arrow writer
-        let writer_properties = WriterProperties::builder()
-            // NOTE: Consider extracting config for writer properties and setting more than just compression
-            .set_compression(Compression::SNAPPY)
-            .build();
+        let writer_properties = default_writer_properties(parquet::basic::Compression::SNAPPY);
 
         Ok(Self {
             table: table.clone(),
@@ -547,6 +539,40 @@ mod tests {
         assert_eq!(columns, vec!["id".to_string(), "value".to_string()]);
     }
 
+    #[tokio::test]
+    async fn test_json_writer_for_table_defaults_include_delta_rs_created_by() {
+        let table_dir = tempfile::tempdir().unwrap();
+        let table = get_test_table(&table_dir).await;
+
+        let writer = JsonWriter::for_table(&table).unwrap();
+
+        assert_eq!(
+            writer.writer_properties.created_by(),
+            format!("delta-rs version {}", crate::crate_version())
+        );
+    }
+
+    #[tokio::test]
+    async fn test_json_writer_try_new_defaults_include_delta_rs_created_by() {
+        let table_dir = tempfile::tempdir().unwrap();
+        let table = get_test_table(&table_dir).await;
+        let arrow_schema = table.snapshot().unwrap().snapshot().arrow_schema();
+
+        let writer = JsonWriter::try_new(
+            table.table_url().clone(),
+            arrow_schema,
+            Some(vec!["modified".to_string()]),
+            None,
+        )
+        .await
+        .unwrap();
+
+        assert_eq!(
+            writer.writer_properties.created_by(),
+            format!("delta-rs version {}", crate::crate_version())
+        );
+    }
+
     #[test]
     fn test_extract_partition_values() {
         let record_batch = RecordBatch::try_new(
diff --git a/crates/core/src/writer/record_batch.rs b/crates/core/src/writer/record_batch.rs
diff --git a/crates/core/tests/checkpoint_writer.rs b/crates/core/tests/checkpoint_writer.rs
diff --git a/crates/core/tests/command_optimize.rs b/crates/core/tests/command_optimize.rs
