feat: support gzip compression in native Parquet writes#4930
Conversation
Add Gzip to the CompressionCodec proto enum and introduce a Parquet-specific ParquetCompression enum in parquet_writer.rs so the native writer can honor gzip without affecting the shuffle codec. The planner maps the new proto variant to ParquetCompression::Gzip, which writes with parquet's default GzipLevel (6), matching parquet-mr's zlib default. The shuffle writer path is unchanged and still rejects Gzip via its catch-all error arm.
…writer The native writer's gzip, snappy, lz4, and zstd support only worked because Cargo feature unification happened to pull in parquet-rs's flate2/snap/lz4/zstd features via other workspace dependencies. Declare these features directly on the parquet dependency so the native writer does not silently lose codec support if that unification changes.
parseCompressionCodec only checked the compression option and the SQLConf default, skipping the parquet.compression option that Spark's own ParquetOptions treats as the middle rung of precedence. Fix the lookup to match Spark's compression, parquet.compression, SQLConf order, and accept "uncompressed" as an alias for "none" since Spark does the same. Add coverage for the parquet.compression option taking precedence over the SQLConf default, for uncompressed as a none alias, and for an unsupported codec (brotli) causing the write to fall back to Spark's own writer instead of CometNativeWriteExec.
compression_to_parquet was a one-line wrapper around ParquetCompression::to_parquet with a single caller; call it directly instead.
The unsupported-codec fallback test relied on Spark's write failing with ClassNotFoundException for BrotliCodec, so it only passed because the environment lacks a Brotli codec class rather than because Comet routed the write through the fallback path. Use lz4_raw instead, which Spark can write directly via parquet-mr and Comet does not support, so the test can assert a real successful round trip and the correct footer codec instead of intercepting an unrelated failure. Revert the allowFailure plumbing added to captureWritePlan for that test since it is no longer needed.
parthchandra
left a comment
There was a problem hiding this comment.
lgtm. approved pending ci.
| } | ||
| } | ||
|
|
||
| test("parquet write with each supported compression codec") { |
There was a problem hiding this comment.
Not a requirement but do you think it might be useful to add a test to cover codec precedence? I.e. specify zstd in the conf, but then override it to gzip
|
Spark 3.4 failure: |
Spark 3.4's PARQUET_COMPRESSION config validates against a fixed set of codecs (brotli, uncompressed, lz4, gzip, lzo, snappy, none, zstd) that does not include lz4_raw, so withSQLConf threw before the fallback path ran. Guard the test with assume(isSpark35Plus, ...) so it runs on versions where lz4_raw is a valid codec value.
| Gzip, | ||
| } | ||
|
|
||
| impl ParquetCompression { |
There was a problem hiding this comment.
do we need this enum, cant we reuse Compression ?
There was a problem hiding this comment.
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."
There was a problem hiding this comment.
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.
Which issue does this PR close?
Closes #4929. Part of #1625.
Rationale for this change
Comet's native Parquet writer supported
none,snappy,lz4, andzstd, so a write configured withspark.sql.parquet.compression.codec=gzipfell back to Spark's writer. gzip is a valid Spark Parquet codec and is widely used, so this was an avoidable loss of native write coverage.Reviewing that code path surfaced two related gaps, both fixed here:
parseCompressionCodecdid not follow Spark's codec precedence. Spark'sParquetOptionsresolvescompression>parquet.compression(ParquetOutputFormat.COMPRESSION) >spark.sql.parquet.compression.codec, but Comet checked only the first and third.df.write.option("parquet.compression", "gzip")with the default SQLConf was therefore written natively as snappy, silently substituting the user's requested codec.uncompressedas a spelling ofnone, but the supported-codec set listed onlynone, souncompressedfell back to Spark's writer unnecessarily.native/corealso requested theparquetcrate withdefault-features = falsewhile naming onlyexperimentalandarrow. The codec features (snap,lz4,zstd, and gzip'sflate2) were enabled only incidentally, through Cargo feature unification with other crates that pullparquetwith default features. A dependency bump that changed that would have left the writer compiling fine but failing at execution time with "Disabled feature at compile time", so the features are now requested explicitly.What changes are included in this PR?
Gzip = 4to the protobufCompressionCodecenum.ParquetCompressionenum instead of borrowing the shuffle crate'sCompressionCodec. gzip is a Parquet codec, not a shuffle codec, and the shuffle path has no business carrying it. The shuffle branch of the planner is unchanged and still rejects any codec it does not map. gzip maps toCompression::GZIP(GzipLevel::default()), i.e. level 6, matching parquet-mr's zlib default.parquetcodec features explicitly innative/core/Cargo.toml.CometDataWritingCommand, honor Spark'sparquet.compressionprecedence, and acceptuncompressedas an alias ofnone.Compression levels remain non-configurable (gzip and zstd use fixed defaults), so parquet-mr's
parquet.compression.codec.{gzip,zstd}.levelconfigs are still ignored. That is out of scope here and noted in the issue.How are these changes tested?
CometParquetWriterSuitehad no compression coverage at all, so this adds it:none,uncompressed,snappy,lz4,zstd, andgzip. For each, it assertsCometNativeWriteExecis in the executed plan (so a silent fallback to Spark's writer cannot make the test pass vacuously), reads the data back and compares against the source, and opens the Parquet footer to assert every column chunk reports the expected codec (so a write that ignored the requested codec cannot pass either)..option("parquet.compression", "gzip")beats asnappySQLConf default and produces a natively written GZIP file.lz4_raw, a codec Spark writes but Comet does not support: the plan must contain noCometNativeWriteExec, and the data must still round-trip. This catches a codec being added to the supported set without a matching serde case.A Rust unit test covers the
ParquetCompressiontoparquet::basic::Compressionmapping for all five variants.This work was scaffolded with the brainstorming, writing-plans, and subagent-driven-development skills.