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 f1168c4a57..7c517a0a63 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, @@ -254,15 +279,6 @@ 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), - } - } - /// 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 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 4ae73565c6..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") + 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 0a33a50450..fbe7ce15cb 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} @@ -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,