Skip to content

feat: support gzip compression in native Parquet writes #4930

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
andygrove wants to merge 9 commits into
base: main
Choose a base branch
from
Open

feat: support gzip compression in native Parquet writes #4930

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
29e1f72
feat: support gzip compression in the native Parquet writer
andygrove Jul 14, 2026
4f40cb3
feat: write Parquet natively when the compression codec is gzip
andygrove Jul 14, 2026
7b449bd
fix: request parquet-rs codec features explicitly for native Parquet …
andygrove Jul 14, 2026
45c08fb
fix: honor parquet.compression precedence and uncompressed codec alias
andygrove Jul 14, 2026
8a1fa5e
refactor: inline single-call-site compression_to_parquet delegate
andygrove Jul 14, 2026
be81af0
test: replace brittle brotli fallback test with lz4_raw round trip
andygrove Jul 14, 2026
3e0572f
test: skip lz4_raw fallback test on Spark 3.4
andygrove Jul 15, 2026
fb444ca
Merge remote-tracking branch 'apache/main' into feat-parquet-write-gzip
andygrove Jul 16, 2026
15b01f3
test: cover full three-tier codec precedence for Parquet writes
andygrove Jul 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion native/core/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ publish = false

[dependencies]
arrow = { workspace = true }
parquet = { workspace = true, default-features = false, features = ["experimental", "arrow"] }
parquet = { workspace = true, default-features = false, features = ["experimental", "arrow", "snap", "lz4", "zstd", "flate2-zlib-rs"] }
futures = { workspace = true }
mimalloc = { version = "*", default-features = false, optional = true }
tikv-jemallocator = { version = "0.6.1", optional = true, features = ["disable_initial_exec_tls"] }
Expand Down
2 changes: 1 addition & 1 deletion native/core/src/execution/operators/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ mod expand;
pub use expand::ExpandExec;
mod iceberg_scan;
mod parquet_writer;
pub use parquet_writer::ParquetWriterExec;
pub use parquet_writer::{ParquetCompression, ParquetWriterExec};
mod csv_scan;
pub mod projection;
mod scan;
Expand Down
71 changes: 56 additions & 15 deletions native/core/src/execution/operators/parquet_writer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@ use opendal::Operator;
#[cfg(feature = "hdfs-opendal")]
use std::io::Cursor;

use crate::execution::shuffle::CompressionCodec;
use crate::parquet::parquet_support::is_hdfs_scheme;
#[cfg(feature = "hdfs-opendal")]
use crate::parquet::parquet_support::{create_hdfs_operator, prepare_object_store_with_configs};
Expand All @@ -52,11 +51,37 @@ use datafusion::{
use futures::TryStreamExt;
use parquet::{
arrow::ArrowWriter,
basic::{Compression, ZstdLevel},
basic::{Compression, GzipLevel, ZstdLevel},
file::properties::WriterProperties,
};
use url::Url;

/// Compression codecs supported by the native Parquet writer.
///
/// This is deliberately separate from the shuffle `CompressionCodec`: gzip is a valid Parquet
/// codec but is not a codec we compress shuffle blocks with.
#[derive(Debug, Clone, PartialEq)]
pub enum ParquetCompression {
None,
Snappy,
Lz4,
Zstd(i32),
Gzip,
}

impl ParquetCompression {

@comphead comphead Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we need this enum, cant we reuse Compression ?

@andygrove andygrove Jul 15, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

From the PR description:

"Give the native Parquet writer its own ParquetCompression enum instead of borrowing the shuffle crate's CompressionCodec. gzip is a Parquet codec, not a shuffle codec, and the shuffle path has no business carrying it."

@comphead comphead Jul 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in this case we prob can just add GZIP variant into CompressionCodec and reuse, however if we follow Apache Spark there are 2 enums indeed, for parquet and internal compression like shuffle.

@andygrove andygrove Jul 16, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it would be better to maintain separate enums for shuffle vs parquet writes, since they support different codecs.

pub fn to_parquet(&self) -> Result<Compression> {
match self {
ParquetCompression::None => Ok(Compression::UNCOMPRESSED),
ParquetCompression::Snappy => Ok(Compression::SNAPPY),
ParquetCompression::Lz4 => Ok(Compression::LZ4),
ParquetCompression::Zstd(level) => Ok(Compression::ZSTD(ZstdLevel::try_new(*level)?)),
// Level 6, matching the default of parquet-mr's zlib codec
ParquetCompression::Gzip => Ok(Compression::GZIP(GzipLevel::default())),
}
}
}

/// Enum representing different types of Arrow writers based on storage backend
enum ParquetWriter {
/// Writer for local file system
Expand Down Expand Up @@ -202,7 +227,7 @@ pub struct ParquetWriterExec {
/// Task attempt ID for this specific task
task_attempt_id: Option<i32>,
/// Compression codec
compression: CompressionCodec,
compression: ParquetCompression,
/// Partition ID (from Spark TaskContext)
partition_id: i32,
/// Column names to use in the output Parquet file
Expand All @@ -224,7 +249,7 @@ impl ParquetWriterExec {
work_dir: String,
job_id: Option<String>,
task_attempt_id: Option<i32>,
compression: CompressionCodec,
compression: ParquetCompression,
partition_id: i32,
column_names: Vec<String>,
object_store_options: HashMap<String, String>,
Expand Down Expand Up @@ -254,15 +279,6 @@ impl ParquetWriterExec {
})
}

fn compression_to_parquet(&self) -> Result<Compression> {
match self.compression {
CompressionCodec::None => Ok(Compression::UNCOMPRESSED),
CompressionCodec::Zstd(level) => Ok(Compression::ZSTD(ZstdLevel::try_new(level)?)),
CompressionCodec::Lz4Frame => Ok(Compression::LZ4),
CompressionCodec::Snappy => Ok(Compression::SNAPPY),
}
}

/// Create an Arrow writer based on the storage scheme
///
/// # Arguments
Expand Down Expand Up @@ -462,7 +478,7 @@ impl ExecutionPlan for ParquetWriterExec {
let input_schema = self.input.schema();
let work_dir = self.work_dir.clone();
let task_attempt_id = self.task_attempt_id;
let compression = self.compression_to_parquet()?;
let compression = self.compression.to_parquet()?;
let column_names = self.column_names.clone();

assert_eq!(input_schema.fields().len(), column_names.len());
Expand Down Expand Up @@ -577,6 +593,31 @@ mod tests {
use arrow::datatypes::{DataType, Field, Schema};
use std::sync::Arc;

#[test]
fn test_parquet_compression_to_parquet() {
assert_eq!(
ParquetCompression::None.to_parquet().unwrap(),
Compression::UNCOMPRESSED
);
assert_eq!(
ParquetCompression::Snappy.to_parquet().unwrap(),
Compression::SNAPPY
);
assert_eq!(
ParquetCompression::Lz4.to_parquet().unwrap(),
Compression::LZ4
);
assert_eq!(
ParquetCompression::Zstd(3).to_parquet().unwrap(),
Compression::ZSTD(ZstdLevel::try_new(3).unwrap())
);
// gzip level 6 matches the default of parquet-mr's zlib codec
assert_eq!(
ParquetCompression::Gzip.to_parquet().unwrap(),
Compression::GZIP(GzipLevel::default())
);
}

/// Helper function to create a test RecordBatch with 1000 rows of (int, string) data
/// Example batch_id 1 -> 0..1000, 2 -> 1001..2000
#[allow(dead_code)]
Expand Down Expand Up @@ -820,7 +861,7 @@ mod tests {
work_dir,
None, // job_id
Some(123), // task_attempt_id
CompressionCodec::None,
ParquetCompression::None,
0, // partition_id
column_names,
HashMap::new(), // object_store_options
Expand Down
14 changes: 9 additions & 5 deletions native/core/src/execution/planner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,10 @@ use crate::execution::operators::IcebergScanExec;
use crate::execution::{
expressions::list_positions::ListPositionsExpr,
expressions::subquery::Subquery,
operators::{ExecutionError, ExpandExec, ParquetWriterExec, ScanExec, ShuffleScanExec},
operators::{
ExecutionError, ExpandExec, ParquetCompression, ParquetWriterExec, ScanExec,
ShuffleScanExec,
},
planner::expression_registry::ExpressionRegistry,
planner::operator_registry::OperatorRegistry,
serde::to_arrow_datatype,
Expand Down Expand Up @@ -1656,10 +1659,11 @@ impl PhysicalPlanner {
self.create_plan(&children[0], inputs, partition_count)?;

let codec = match writer.compression.try_into() {
Ok(SparkCompressionCodec::None) => Ok(CompressionCodec::None),
Ok(SparkCompressionCodec::Snappy) => Ok(CompressionCodec::Snappy),
Ok(SparkCompressionCodec::Zstd) => Ok(CompressionCodec::Zstd(3)),
Ok(SparkCompressionCodec::Lz4) => Ok(CompressionCodec::Lz4Frame),
Ok(SparkCompressionCodec::None) => Ok(ParquetCompression::None),
Ok(SparkCompressionCodec::Snappy) => Ok(ParquetCompression::Snappy),
Ok(SparkCompressionCodec::Zstd) => Ok(ParquetCompression::Zstd(3)),
Ok(SparkCompressionCodec::Lz4) => Ok(ParquetCompression::Lz4),
Ok(SparkCompressionCodec::Gzip) => Ok(ParquetCompression::Gzip),
_ => Err(GeneralError(format!(
"Unsupported parquet compression codec: {:?}",
writer.compression
Expand Down
2 changes: 2 additions & 0 deletions native/proto/src/proto/operator.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -320,6 +320,8 @@ enum CompressionCodec {
Zstd = 1;
Lz4 = 2;
Snappy = 3;
// Parquet write only. The shuffle writer rejects this codec.
Gzip = 4;
}

message ShuffleWriter {
Expand Down
13 changes: 10 additions & 3 deletions spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import java.util.Locale

import scala.jdk.CollectionConverters._

import org.apache.parquet.hadoop.ParquetOutputFormat
import org.apache.spark.SparkException
import org.apache.spark.sql.comet.{CometNativeExec, CometNativeWriteExec}
import org.apache.spark.sql.execution.command.DataWritingCommandExec
Expand All @@ -44,7 +45,8 @@ import org.apache.comet.serde.QueryPlanSerde.serializeDataType
*/
object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec] {

private val supportedCompressionCodes = Set("none", "snappy", "lz4", "zstd")
private val supportedCompressionCodes =
Set("none", "uncompressed", "snappy", "lz4", "zstd", "gzip")

override def enabledConfig: Option[ConfigEntry[Boolean]] =
Some(CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED)
Expand Down Expand Up @@ -121,7 +123,8 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec
case "snappy" => OperatorOuterClass.CompressionCodec.Snappy
case "lz4" => OperatorOuterClass.CompressionCodec.Lz4
case "zstd" => OperatorOuterClass.CompressionCodec.Zstd
case "none" => OperatorOuterClass.CompressionCodec.None
case "gzip" => OperatorOuterClass.CompressionCodec.Gzip
case "none" | "uncompressed" => OperatorOuterClass.CompressionCodec.None
case other =>
withFallbackReason(op, s"Unsupported compression codec: $other")
return None
Expand Down Expand Up @@ -204,9 +207,13 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec
}

private def parseCompressionCodec(cmd: InsertIntoHadoopFsRelationCommand) = {
// `compression`, `parquet.compression` (i.e., ParquetOutputFormat.COMPRESSION), and
// `spark.sql.parquet.compression.codec` are in order of precedence from highest to
// lowest, matching Spark's own ParquetOptions.compressionCodecClassName.
cmd.options
.get("compression")
.orElse(cmd.options.get(ParquetOutputFormat.COMPRESSION))
.getOrElse(
"compression",
SQLConf.get.getConfString(
SQLConf.PARQUET_COMPRESSION.key,
SQLConf.PARQUET_COMPRESSION.defaultValueString))
Expand Down
Loading
Loading