From 29e1f725d4ac95e3450de4cb200815fb4f00d4bc Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 14 Jul 2026 13:48:57 -0600 Subject: [PATCH 1/8] feat: support gzip compression in the native Parquet writer 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. --- native/core/src/execution/operators/mod.rs | 2 +- .../src/execution/operators/parquet_writer.rs | 67 ++++++++++++++++--- native/core/src/execution/planner.rs | 14 ++-- native/proto/src/proto/operator.proto | 2 + 4 files changed, 68 insertions(+), 17 deletions(-) diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index d68252bd9b..6d27a17042 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -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; diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index f1168c4a57..c13676d58d 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -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}; @@ -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 { + pub fn to_parquet(&self) -> Result { + 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 @@ -202,7 +227,7 @@ pub struct ParquetWriterExec { /// Task attempt ID for this specific task task_attempt_id: Option, /// Compression codec - compression: CompressionCodec, + compression: ParquetCompression, /// Partition ID (from Spark TaskContext) partition_id: i32, /// Column names to use in the output Parquet file @@ -224,7 +249,7 @@ impl ParquetWriterExec { work_dir: String, job_id: Option, task_attempt_id: Option, - compression: CompressionCodec, + compression: ParquetCompression, partition_id: i32, column_names: Vec, object_store_options: HashMap, @@ -255,12 +280,7 @@ impl ParquetWriterExec { } fn compression_to_parquet(&self) -> Result { - 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), - } + self.compression.to_parquet() } /// Create an Arrow writer based on the storage scheme @@ -576,6 +596,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)] @@ -819,7 +864,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 diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 7ed2b3331c..a23a7b1f67 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -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, @@ -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 diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 2fcfe7f25b..00158644d1 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -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 { From 4f40cb3fd233b2838ee8df9b9a59e58b24e6f6bf Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 14 Jul 2026 13:57:22 -0600 Subject: [PATCH 2/8] feat: write Parquet natively when the compression codec is gzip --- .../operator/CometDataWritingCommand.scala | 3 +- .../parquet/CometParquetWriterSuite.scala | 60 +++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala index 4ae73565c6..0cf22a86c6 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala @@ -44,7 +44,7 @@ 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", "snappy", "lz4", "zstd", "gzip") override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED) @@ -121,6 +121,7 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec case "snappy" => OperatorOuterClass.CompressionCodec.Snappy case "lz4" => OperatorOuterClass.CompressionCodec.Lz4 case "zstd" => OperatorOuterClass.CompressionCodec.Zstd + case "gzip" => OperatorOuterClass.CompressionCodec.Gzip case "none" => OperatorOuterClass.CompressionCodec.None case other => withFallbackReason(op, s"Unsupported compression codec: $other") diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index 0a33a50450..cae92c2f6a 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -21,9 +21,13 @@ package org.apache.comet.parquet import java.io.File +import scala.jdk.CollectionConverters._ import scala.util.{Random, Using} import org.apache.hadoop.fs.{FileSystem, Path} +import org.apache.parquet.hadoop.ParquetFileReader +import org.apache.parquet.hadoop.metadata.CompressionCodecName +import org.apache.parquet.hadoop.util.HadoopInputFile import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeScanExec, CometNativeWriteExec, CometScanExec} import org.apache.spark.sql.execution.{FileSourceScanExec, QueryExecution, SparkPlan} @@ -138,6 +142,28 @@ class CometParquetWriterSuite extends CometTestBase { } } + test("parquet write with each supported compression codec") { + Seq("none", "snappy", "lz4", "zstd", "gzip").foreach { codec => + withTempPath { dir => + val outputPath = new File(dir, s"output_$codec.parquet").getAbsolutePath + val df = spark.range(0, 100).selectExpr("id", "cast(id as string) as name") + + withSQLConf( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + SQLConf.PARQUET_COMPRESSION.key -> codec) { + + val plan = captureWritePlan(path => df.write.parquet(path), outputPath) + assertHasCometNativeWriteExec(plan) + } + + checkAnswer(spark.read.parquet(outputPath), df.collect()) + assertParquetCodec(outputPath, expectedCodecName(codec)) + } + } + } + test("parquet write with array type") { withTempPath { dir => val outputPath = new File(dir, "output.parquet").getAbsolutePath @@ -470,6 +496,40 @@ class CometParquetWriterSuite extends CometTestBase { compareRows(schema, sparkRows, cometRows) } + private def expectedCodecName(sparkCodec: String): CompressionCodecName = sparkCodec match { + case "none" => CompressionCodecName.UNCOMPRESSED + case "snappy" => CompressionCodecName.SNAPPY + case "lz4" => CompressionCodecName.LZ4 + case "zstd" => CompressionCodecName.ZSTD + case "gzip" => CompressionCodecName.GZIP + case other => fail(s"unexpected codec: $other") + } + + /** + * Asserts that every column chunk in every part file under `outputPath` reports `expected` as + * its compression codec. Reading the data back is not enough on its own: a Parquet reader + * honors whatever the footer says, so a write that silently ignored the requested codec would + * still round-trip. + */ + private def assertParquetCodec(outputPath: String, expected: CompressionCodecName): Unit = { + val conf = spark.sparkContext.hadoopConfiguration + val partFiles = new File(outputPath).listFiles().filter(_.getName.startsWith("part-")) + assert(partFiles.nonEmpty, s"No part files found under $outputPath") + + partFiles.foreach { partFile => + val inputFile = HadoopInputFile.fromPath(new Path(partFile.getAbsolutePath), conf) + Using.resource(ParquetFileReader.open(inputFile)) { reader => + val codecs = reader.getFooter.getBlocks.asScala + .flatMap(_.getColumns.asScala) + .map(_.getCodec) + .toSet + assert( + codecs == Set(expected), + s"Expected all column chunks in ${partFile.getName} to use $expected, found $codecs") + } + } + } + private def writeComplexTypeData( inputDf: DataFrame, outputPath: String, From 7b449bd4bd9d14778bff7fcdef78b15991cc3128 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 14 Jul 2026 14:23:21 -0600 Subject: [PATCH 3/8] fix: request parquet-rs codec features explicitly for native Parquet 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. --- native/core/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index d78418a7ae..7be91683b8 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -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"] } From 45c08fb703aa08d7410149ae0a3f88db18a87f4f Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 14 Jul 2026 14:23:28 -0600 Subject: [PATCH 4/8] fix: honor parquet.compression precedence and uncompressed codec alias 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. --- .../operator/CometDataWritingCommand.scala | 12 ++- .../parquet/CometParquetWriterSuite.scala | 86 +++++++++++++++++-- 2 files changed, 87 insertions(+), 11 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala index 0cf22a86c6..99bae56017 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala @@ -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 @@ -44,7 +45,8 @@ import org.apache.comet.serde.QueryPlanSerde.serializeDataType */ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec] { - private val supportedCompressionCodes = Set("none", "snappy", "lz4", "zstd", "gzip") + private val supportedCompressionCodes = + Set("none", "uncompressed", "snappy", "lz4", "zstd", "gzip") override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED) @@ -122,7 +124,7 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec case "lz4" => OperatorOuterClass.CompressionCodec.Lz4 case "zstd" => OperatorOuterClass.CompressionCodec.Zstd case "gzip" => OperatorOuterClass.CompressionCodec.Gzip - case "none" => OperatorOuterClass.CompressionCodec.None + case "none" | "uncompressed" => OperatorOuterClass.CompressionCodec.None case other => withFallbackReason(op, s"Unsupported compression codec: $other") return None @@ -205,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)) diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index cae92c2f6a..42c09a59ec 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -143,7 +143,7 @@ class CometParquetWriterSuite extends CometTestBase { } test("parquet write with each supported compression codec") { - Seq("none", "snappy", "lz4", "zstd", "gzip").foreach { codec => + Seq("none", "uncompressed", "snappy", "lz4", "zstd", "gzip").foreach { codec => withTempPath { dir => val outputPath = new File(dir, s"output_$codec.parquet").getAbsolutePath val df = spark.range(0, 100).selectExpr("id", "cast(id as string) as name") @@ -164,6 +164,52 @@ class CometParquetWriterSuite extends CometTestBase { } } + test("parquet write honors parquet.compression option over SQLConf default") { + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + val df = spark.range(0, 100).selectExpr("id", "cast(id as string) as name") + + withSQLConf( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + SQLConf.PARQUET_COMPRESSION.key -> "snappy") { + + val plan = captureWritePlan( + path => df.write.option("parquet.compression", "gzip").parquet(path), + outputPath) + assertHasCometNativeWriteExec(plan) + } + + checkAnswer(spark.read.parquet(outputPath), df.collect()) + assertParquetCodec(outputPath, CompressionCodecName.GZIP) + } + } + + test("parquet write with unsupported compression codec falls back to Spark") { + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + val df = spark.range(0, 100).selectExpr("id", "cast(id as string) as name") + + withSQLConf( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + SQLConf.PARQUET_COMPRESSION.key -> "brotli") { + + // brotli is a real Parquet codec that Comet does not support, so the write must + // fall back to Spark's own writer instead of using CometNativeWriteExec. This + // project's Hadoop dependency does not bundle a BrotliCodec implementation, so + // Spark's writer itself fails here for reasons unrelated to Comet; that failure + // (rather than a successful round trip) is what proves the write actually reached + // Spark's normal code path instead of Comet's native one. + val plan = + captureWritePlan(path => df.write.parquet(path), outputPath, allowFailure = true) + assertNoCometNativeWriteExec(plan) + } + } + } + test("parquet write with array type") { withTempPath { dir => val outputPath = new File(dir, "output.parquet").getAbsolutePath @@ -402,7 +448,10 @@ class CometParquetWriterSuite extends CometTestBase { * @return * The captured execution plan */ - private def captureWritePlan(writeOp: String => Unit, outputPath: String): SparkPlan = { + private def captureWritePlan( + writeOp: String => Unit, + outputPath: String, + allowFailure: Boolean = false): SparkPlan = { var capturedPlan: Option[QueryExecution] = None val listener = new org.apache.spark.sql.util.QueryExecutionListener { @@ -412,16 +461,21 @@ class CometParquetWriterSuite extends CometTestBase { } } - override def onFailure( - funcName: String, - qe: QueryExecution, - exception: Exception): Unit = {} + override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = { + if (allowFailure && (funcName == "save" || funcName.contains("command"))) { + capturedPlan = Some(qe) + } + } } spark.listenerManager.register(listener) try { - writeOp(outputPath) + if (allowFailure) { + intercept[Exception](writeOp(outputPath)) + } else { + writeOp(outputPath) + } // Wait for listener to be called with timeout val maxWaitTimeMs = 15000 @@ -463,6 +517,22 @@ class CometParquetWriterSuite extends CometTestBase { s"Expected exactly one CometNativeWriteExec in the plan, but found $nativeWriteCount:\n${plan.treeString}") } + private def assertNoCometNativeWriteExec(plan: SparkPlan): Unit = { + val hasNativeWrite = plan.exists { + case _: CometNativeWriteExec => true + case d: DataWritingCommandExec => + d.child.exists { + case _: CometNativeWriteExec => true + case _ => false + } + case _ => false + } + + assert( + !hasNativeWrite, + s"Expected no CometNativeWriteExec in the plan, but found one:\n${plan.treeString}") + } + private def writeWithCometNativeWriteExec( inputPath: String, outputPath: String, @@ -497,7 +567,7 @@ class CometParquetWriterSuite extends CometTestBase { } private def expectedCodecName(sparkCodec: String): CompressionCodecName = sparkCodec match { - case "none" => CompressionCodecName.UNCOMPRESSED + case "none" | "uncompressed" => CompressionCodecName.UNCOMPRESSED case "snappy" => CompressionCodecName.SNAPPY case "lz4" => CompressionCodecName.LZ4 case "zstd" => CompressionCodecName.ZSTD From 8a1fa5ef3f36776b345a68cc9aa22438a4ceb1b0 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 14 Jul 2026 14:23:32 -0600 Subject: [PATCH 5/8] refactor: inline single-call-site compression_to_parquet delegate compression_to_parquet was a one-line wrapper around ParquetCompression::to_parquet with a single caller; call it directly instead. --- native/core/src/execution/operators/parquet_writer.rs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index c13676d58d..7c517a0a63 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -279,10 +279,6 @@ impl ParquetWriterExec { }) } - fn compression_to_parquet(&self) -> Result { - self.compression.to_parquet() - } - /// Create an Arrow writer based on the storage scheme /// /// # Arguments @@ -482,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()); From be81af07c6cd53d166d733f81948cbf2725d455d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 14 Jul 2026 14:29:11 -0600 Subject: [PATCH 6/8] test: replace brittle brotli fallback test with lz4_raw round trip 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. --- .../parquet/CometParquetWriterSuite.scala | 36 +++++++------------ 1 file changed, 12 insertions(+), 24 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index 42c09a59ec..d88aa06ff5 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -195,18 +195,14 @@ class CometParquetWriterSuite extends CometTestBase { CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true", - SQLConf.PARQUET_COMPRESSION.key -> "brotli") { - - // brotli is a real Parquet codec that Comet does not support, so the write must - // fall back to Spark's own writer instead of using CometNativeWriteExec. This - // project's Hadoop dependency does not bundle a BrotliCodec implementation, so - // Spark's writer itself fails here for reasons unrelated to Comet; that failure - // (rather than a successful round trip) is what proves the write actually reached - // Spark's normal code path instead of Comet's native one. - val plan = - captureWritePlan(path => df.write.parquet(path), outputPath, allowFailure = true) + SQLConf.PARQUET_COMPRESSION.key -> "lz4_raw") { + + val plan = captureWritePlan(path => df.write.parquet(path), outputPath) assertNoCometNativeWriteExec(plan) } + + checkAnswer(spark.read.parquet(outputPath), df.collect()) + assertParquetCodec(outputPath, CompressionCodecName.LZ4_RAW) } } @@ -448,10 +444,7 @@ class CometParquetWriterSuite extends CometTestBase { * @return * The captured execution plan */ - private def captureWritePlan( - writeOp: String => Unit, - outputPath: String, - allowFailure: Boolean = false): SparkPlan = { + private def captureWritePlan(writeOp: String => Unit, outputPath: String): SparkPlan = { var capturedPlan: Option[QueryExecution] = None val listener = new org.apache.spark.sql.util.QueryExecutionListener { @@ -461,21 +454,16 @@ class CometParquetWriterSuite extends CometTestBase { } } - override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = { - if (allowFailure && (funcName == "save" || funcName.contains("command"))) { - capturedPlan = Some(qe) - } - } + override def onFailure( + funcName: String, + qe: QueryExecution, + exception: Exception): Unit = {} } spark.listenerManager.register(listener) try { - if (allowFailure) { - intercept[Exception](writeOp(outputPath)) - } else { - writeOp(outputPath) - } + writeOp(outputPath) // Wait for listener to be called with timeout val maxWaitTimeMs = 15000 From 3e0572f0fc655954528b993f949e3f987a58bd0e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 15 Jul 2026 13:59:31 -0600 Subject: [PATCH 7/8] test: skip lz4_raw fallback test on Spark 3.4 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. --- .../org/apache/comet/parquet/CometParquetWriterSuite.scala | 2 ++ 1 file changed, 2 insertions(+) diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index d88aa06ff5..fbe7ce15cb 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -36,6 +36,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, SchemaGenOptions} class CometParquetWriterSuite extends CometTestBase { @@ -187,6 +188,7 @@ class CometParquetWriterSuite extends CometTestBase { } test("parquet write with unsupported compression codec falls back to Spark") { + assume(isSpark35Plus, "lz4_raw was added in Spark 3.5") withTempPath { dir => val outputPath = new File(dir, "output.parquet").getAbsolutePath val df = spark.range(0, 100).selectExpr("id", "cast(id as string) as name") From 15b01f382f05912d83a204cb655bae5402dd7eb5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 16 Jul 2026 15:45:26 -0600 Subject: [PATCH 8/8] test: cover full three-tier codec precedence for Parquet writes Adds a test that pins the `compression` > `parquet.compression` > `spark.sql.parquet.compression.codec` precedence: SQLConf is `zstd`, `parquet.compression` is `snappy`, and `compression` is `gzip`; the written file must report GZIP. A leak from either lower layer would surface as a codec mismatch. The existing test still covers `parquet.compression` beating the SQLConf default when `compression` is absent. --- .../parquet/CometParquetWriterSuite.scala | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index 8377b95c22..a1ae1af1d1 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -187,6 +187,35 @@ class CometParquetWriterSuite extends CometTestBase { } } + test("parquet write honors compression option over parquet.compression and SQLConf") { + // Precedence, highest to lowest, matches Spark's ParquetOptions: + // `compression` write option > `parquet.compression` write option > spark.sql.parquet.compression.codec + // Use a distinct wrong codec at each lower layer so any leak surfaces as a codec mismatch. + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + val df = spark.range(0, 100).selectExpr("id", "cast(id as string) as name") + + withSQLConf( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + SQLConf.PARQUET_COMPRESSION.key -> "zstd") { + + val plan = captureWritePlan( + path => + df.write + .option("compression", "gzip") + .option("parquet.compression", "snappy") + .parquet(path), + outputPath) + assertHasCometNativeWriteExec(plan) + } + + checkAnswer(spark.read.parquet(outputPath), df.collect()) + assertParquetCodec(outputPath, CompressionCodecName.GZIP) + } + } + test("parquet write with unsupported compression codec falls back to Spark") { assume(isSpark35Plus, "lz4_raw was added in Spark 3.5") withTempPath { dir =>