andygrove commented on code in PR #5293:
URL: https://github.com/apache/datafusion-comet/pull/5293#discussion_r3789664250
##########
native/proto/src/proto/operator.proto:
##########
@@ -458,16 +458,17 @@ message ShuffleWriter {
}
message ParquetWriter {
+ // Fully-qualified path of the Parquet file that this task must write, set
per task by
+ // CometWriteFilesExec from FileCommitProtocol.newTaskTempFile. Naming and
staging are owned by
+ // Spark's commit protocol so that task-attempt isolation, speculative
execution, and committers
+ // that track individual files (S3A magic, streaming manifest) all behave as
they do for Spark's
+ // own writer. The native writer uses this path verbatim.
string output_path = 1;
CompressionCodec compression = 2;
repeated string column_names = 4;
- // Working directory for temporary files (used by FileCommitProtocol)
- // If not set, files are written directly to output_path
- optional string work_dir = 5;
- // Job ID for tracking this write operation
- optional string job_id = 6;
- // Task attempt ID for this specific task
- optional int32 task_attempt_id = 7;
+ // Was work_dir / job_id / task_attempt_id, used when the native writer
derived its own file
+ // names from the task context. File naming now comes from the commit
protocol instead.
+ reserved 5, 6, 7;
Review Comment:
Done in 551c604d8 - dropped the `reserved` block and renumbered so the tags
are contiguous:
```protobuf
message ParquetWriter {
string output_path = 1;
CompressionCodec compression = 2;
repeated string column_names = 3;
map<string, string> object_store_options = 4;
}
```
Agreed on the reasoning: plans are serialized on the driver and deserialized
on executors within a single jar version and are never persisted, so there is
no wire-compatibility surface the reserved tags were protecting. Generated
accessors key off field names rather than tags, so no Rust or Scala call site
changed. `cargo check` and a full build pass, and `CometParquetWriterSuite` is
37/37 on Spark 4.1.
One consistency note for a possible follow-up: `HashAggregate` still carries
`reserved 3, 8` from #4507 for the same removed-field situation, so the file
now does this two ways.
##########
spark/src/main/scala/org/apache/comet/serde/operator/CometWriteFiles.scala:
##########
@@ -0,0 +1,178 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.serde.operator
+
+import java.net.URI
+import java.util.Locale
+
+import org.apache.parquet.hadoop.ParquetOutputFormat
+import org.apache.spark.sql.comet.{CometNativeExec, CometWriteFilesExec}
+import org.apache.spark.sql.execution.datasources.WriteFilesExec
+import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
+import org.apache.spark.sql.internal.SQLConf
+
+import org.apache.comet.{CometConf, ConfigEntry}
+import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus,
withFallbackReason}
+import org.apache.comet.objectstore.NativeConfig
+import org.apache.comet.rules.CometExecRule
+import org.apache.comet.serde.{CometOperatorSerde, Incompatible,
OperatorOuterClass, SupportLevel, Unsupported}
+import org.apache.comet.serde.OperatorOuterClass.Operator
+import org.apache.comet.serde.QueryPlanSerde.serializeDataType
+
+/**
+ * Serde for Spark's `WriteFilesExec`, replacing the per-task Parquet write
with Comet's native
+ * writer while leaving the surrounding write framework (commit protocol,
stats trackers, SaveMode
+ * handling, `_SUCCESS`) to Spark. See [[CometWriteFilesExec]] for how the two
fit together.
+ */
+object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] {
+
+ private val supportedCompressionCodecs =
+ Set("none", "uncompressed", "snappy", "lz4", "zstd", "gzip")
+
+ override def enabledConfig: Option[ConfigEntry[Boolean]] =
+ Some(CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED)
+
+ // Native writes require Arrow-formatted input data. If the query falls back
to Spark
+ // (e.g., due to unsupported complex types), the write must also fall back.
+ override def requiresNativeChildren: Boolean = true
+
+ override def getSupportLevel(op: WriteFilesExec): SupportLevel = {
+ // `V1WritesUtils.getWriteFilesOpt` matches the `WriteFilesExecBase` trait
on Spark 4.0+, which
+ // is what makes Spark route the write through CometWriteFilesExec. Spark
3.x matches the
+ // concrete `WriteFilesExec` case class instead, so a Comet node would be
silently ignored and
+ // the write would fall into FileFormatWriter's non-planned, row-based
branch.
+ if (!isSpark40Plus) {
+ return Unsupported(Some("Native Parquet writes require Spark 4.0 or
later"))
Review Comment:
Yes. On 3.x `getSupportLevel` returns `Unsupported("Native Parquet writes
require Spark 4.0 or later")` and the write runs through Spark's own writer
unchanged. It is reported as a fallback reason so users on 3.x can see why,
rather than silently getting a different path.
##########
spark/src/main/scala/org/apache/comet/serde/operator/CometWriteFiles.scala:
##########
@@ -0,0 +1,178 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.serde.operator
+
+import java.net.URI
+import java.util.Locale
+
+import org.apache.parquet.hadoop.ParquetOutputFormat
+import org.apache.spark.sql.comet.{CometNativeExec, CometWriteFilesExec}
+import org.apache.spark.sql.execution.datasources.WriteFilesExec
+import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
+import org.apache.spark.sql.internal.SQLConf
+
+import org.apache.comet.{CometConf, ConfigEntry}
+import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus,
withFallbackReason}
+import org.apache.comet.objectstore.NativeConfig
+import org.apache.comet.rules.CometExecRule
+import org.apache.comet.serde.{CometOperatorSerde, Incompatible,
OperatorOuterClass, SupportLevel, Unsupported}
+import org.apache.comet.serde.OperatorOuterClass.Operator
+import org.apache.comet.serde.QueryPlanSerde.serializeDataType
+
+/**
+ * Serde for Spark's `WriteFilesExec`, replacing the per-task Parquet write
with Comet's native
+ * writer while leaving the surrounding write framework (commit protocol,
stats trackers, SaveMode
+ * handling, `_SUCCESS`) to Spark. See [[CometWriteFilesExec]] for how the two
fit together.
+ */
+object CometWriteFiles extends CometOperatorSerde[WriteFilesExec] {
+
+ private val supportedCompressionCodecs =
+ Set("none", "uncompressed", "snappy", "lz4", "zstd", "gzip")
+
+ override def enabledConfig: Option[ConfigEntry[Boolean]] =
+ Some(CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED)
+
+ // Native writes require Arrow-formatted input data. If the query falls back
to Spark
+ // (e.g., due to unsupported complex types), the write must also fall back.
+ override def requiresNativeChildren: Boolean = true
+
+ override def getSupportLevel(op: WriteFilesExec): SupportLevel = {
+ // `V1WritesUtils.getWriteFilesOpt` matches the `WriteFilesExecBase` trait
on Spark 4.0+, which
+ // is what makes Spark route the write through CometWriteFilesExec. Spark
3.x matches the
+ // concrete `WriteFilesExec` case class instead, so a Comet node would be
silently ignored and
+ // the write would fall into FileFormatWriter's non-planned, row-based
branch.
+ if (!isSpark40Plus) {
+ return Unsupported(Some("Native Parquet writes require Spark 4.0 or
later"))
+ }
+
+ if (!op.fileFormat.isInstanceOf[ParquetFileFormat]) {
+ return Unsupported(Some("Only Parquet writes are supported"))
+ }
+
+ // The write node does not carry the output path, so CometExecRule tags it
from the enclosing
+ // InsertIntoHadoopFsRelationCommand. An absent tag means this write
belongs to some other V1
+ // write command (a Hive insert, for example) whose semantics Comet has
not been verified
+ // against, so decline it.
+ val outputPath = outputPathOf(op) match {
+ case Some(path) => path
+ case None =>
+ return Unsupported(Some("Only InsertIntoHadoopFsRelationCommand writes
are supported"))
+ }
+
+ if (!outputPath.startsWith("file:") && !outputPath.startsWith("hdfs:")) {
Review Comment:
Covered in 43bce31a6. `InsertIntoHadoopFsRelationCommand` only sets
`dynamicPartitionOverwrite` when `staticPartitions.size <
partitionColumns.length`, which implies partition columns, so a dynamic
overwrite always hits the existing `partitionColumns.nonEmpty ||
staticPartitions.nonEmpty` guard and is declined. I made that explicit in a
comment rather than adding a redundant condition.
##########
spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala:
##########
@@ -24,25 +24,41 @@ import java.io.File
import scala.jdk.CollectionConverters._
import scala.util.{Random, Using}
+import org.scalactic.source.Position
+import org.scalatest.Tag
+
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.comet.{CometBatchScanExec, CometNativeScanExec,
CometScanExec, CometWriteFilesExec}
import org.apache.spark.sql.execution.{FileSourceScanExec, QueryExecution,
SparkPlan}
import org.apache.spark.sql.execution.command.DataWritingCommandExec
+import org.apache.spark.sql.execution.datasources.WriteFilesExec
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.{CometConf, CometExplainInfo}
+import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus,
isSpark40Plus}
import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator,
SchemaGenOptions}
class CometParquetWriterSuite extends CometTestBase {
import testImplicits._
+ /**
Review Comment:
Added in 43bce31a6: a dynamic-overwrite test that asserts the untouched
partitions survive, alongside tests for the abort-and-retry path (injected
failing commit protocol), the `maxRecordsPerFile` fallback (both the write
option and the conf, verifying Spark's writer rolls 10 files), and the
schema-only empty-input write.
On the fallback itself: `dynamicPartitionOverwrite` is only set when
`staticPartitions.size < partitionColumns.length`, which implies partition
columns, so it always lands on the existing `partitionColumns.nonEmpty ||
staticPartitions.nonEmpty` guard and declines.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]