feat: support gzip compression in native Parquet writes

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"] }

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;

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<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
@@ -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
@@ -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>,
@@ -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
@@ -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());
@@ -576,6 +592,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 +860,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

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

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 {

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")
+  private val supportedCompressionCodes =
+    Set("none", "uncompressed", "snappy", "lz4", "zstd", "gzip")
 
   override def enabledConfig: Option[ConfigEntry[Boolean]] =
     Some(CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED)
@@ -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
@@ -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))

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}
@@ -32,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 {
@@ -138,6 +143,71 @@ class CometParquetWriterSuite extends CometTestBase {
     }
   }
 
+  test("parquet write with each supported compression 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")
+
+        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 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") {
+    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")
+
+      withSQLConf(
+        CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true",
+        CometConf.getOperatorAllowIncompatConfigKey(classOf[DataWritingCommandExec]) -> "true",
+        CometConf.COMET_EXEC_ENABLED.key -> "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)
+    }
+  }
+
   test("parquet write with array type") {
     withTempPath { dir =>
       val outputPath = new File(dir, "output.parquet").getAbsolutePath
@@ -437,6 +507,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,
@@ -470,6 +556,40 @@ class CometParquetWriterSuite extends CometTestBase {
     compareRows(schema, sparkRows, cometRows)
   }
 
+  private def expectedCodecName(sparkCodec: String): CompressionCodecName = sparkCodec match {
+    case "none" | "uncompressed" => 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,