This is an automated email from the ASF dual-hosted git repository.

andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git


The following commit(s) were added to refs/heads/main by this push:
     new c65a5ee803 feat: support gzip compression in native Parquet writes 
(#4930)
c65a5ee803 is described below

commit c65a5ee8032a17ea1597eae0986945e81b787ad3
Author: Andy Grove <[email protected]>
AuthorDate: Fri Jul 17 07:42:28 2026 -0600

    feat: support gzip compression in native Parquet writes (#4930)
    
    * 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.
    
    * feat: write Parquet natively when the compression codec is gzip
    
    * 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.
    
    * 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.
    
    * 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.
    
    * 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.
    
    * 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.
    
    * 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.
---
 native/core/Cargo.toml                             |   2 +-
 native/core/src/execution/operators/mod.rs         |   2 +-
 .../core/src/execution/operators/parquet_writer.rs |  71 +++++++---
 native/core/src/execution/planner.rs               |  14 +-
 native/proto/src/proto/operator.proto              |   2 +
 .../serde/operator/CometDataWritingCommand.scala   |  13 +-
 .../comet/parquet/CometParquetWriterSuite.scala    | 149 +++++++++++++++++++++
 7 files changed, 228 insertions(+), 25 deletions(-)

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"] }
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 e0bd7fc887..dbbee713ae 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<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());
@@ -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)]
@@ -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
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 {
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 b07a458dcb..8157f28682 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")
+  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))
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 da19f31d7c..a1ae1af1d1 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.{AnalysisException, CometTestBase, DataFrame, Row, 
SaveMode}
 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,100 @@ 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 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 =>
+      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
@@ -671,6 +770,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,
@@ -704,6 +819,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,


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to