feat: add complex type support to native Parquet writer (#3214)

Author: andygrove

benchmarks/pyspark/run_all_benchmarks.sh

@@ -25,7 +25,7 @@ set -e
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 DATA_PATH="${1:-/tmp/shuffle-benchmark-data}"
-COMET_JAR="${COMET_JAR:-$SCRIPT_DIR/../spark/target/comet-spark-spark3.5_2.12-0.13.0-SNAPSHOT.jar}"
+COMET_JAR="${COMET_JAR:-$SCRIPT_DIR/../../spark/target/comet-spark-spark3.5_2.12-0.13.0-SNAPSHOT.jar}"
 SPARK_MASTER="${SPARK_MASTER:-local[*]}"
 EXECUTOR_MEMORY="${EXECUTOR_MEMORY:-16g}"
 EVENT_LOG_DIR="${EVENT_LOG_DIR:-/tmp/spark-events}"
@@ -71,9 +71,10 @@ $SPARK_HOME/bin/spark-submit \
   --conf spark.memory.offHeap.enabled=true \
   --conf spark.memory.offHeap.size=16g \
   --conf spark.comet.enabled=true \
-  --conf spark.comet.exec.enabled=true \
-  --conf spark.comet.exec.all.enabled=true \
-  --conf spark.comet.exec.shuffle.enabled=true \
+  --conf spark.comet.operator.DataWritingCommandExec.allowIncompatible=true \
+  --conf spark.comet.parquet.write.enabled=true \
+  --conf spark.comet.logFallbackReasons.enabled=true \
+  --conf spark.comet.explainFallback.enabled=true \
   --conf spark.comet.shuffle.mode=jvm \
   --conf spark.comet.exec.shuffle.mode=jvm \
   --conf spark.comet.exec.replaceSortMergeJoin=true \
@@ -98,9 +99,10 @@ $SPARK_HOME/bin/spark-submit \
   --conf spark.memory.offHeap.enabled=true \
   --conf spark.memory.offHeap.size=16g \
   --conf spark.comet.enabled=true \
-  --conf spark.comet.exec.enabled=true \
-  --conf spark.comet.exec.all.enabled=true \
-  --conf spark.comet.exec.shuffle.enabled=true \
+  --conf spark.comet.operator.DataWritingCommandExec.allowIncompatible=true \
+  --conf spark.comet.parquet.write.enabled=true \
+  --conf spark.comet.logFallbackReasons.enabled=true \
+  --conf spark.comet.explainFallback.enabled=true \
   --conf spark.comet.exec.shuffle.mode=native \
   --conf spark.comet.exec.replaceSortMergeJoin=true \
   --conf spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager \

native/core/src/execution/operators/parquet_writer.rs

@@ -535,8 +535,12 @@ impl ExecutionPlan for ParquetWriterExec {
                 DataFusionError::Execution(format!("Failed to close writer: {}", e))
             })?;
 
-            // Get file size
-            let file_size = std::fs::metadata(&part_file)
+            // Get file size - strip file:// prefix if present for local filesystem access
+            let local_path = part_file
+                .strip_prefix("file://")
+                .or_else(|| part_file.strip_prefix("file:"))
+                .unwrap_or(&part_file);
+            let file_size = std::fs::metadata(local_path)
                 .map(|m| m.len() as i64)
                 .unwrap_or(0);
 

spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala

@@ -31,7 +31,7 @@ import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCom
 import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
 import org.apache.spark.sql.internal.SQLConf
 
-import org.apache.comet.{CometConf, ConfigEntry, DataTypeSupport}
+import org.apache.comet.{CometConf, ConfigEntry}
 import org.apache.comet.CometSparkSessionExtensions.withInfo
 import org.apache.comet.objectstore.NativeConfig
 import org.apache.comet.serde.{CometOperatorSerde, Incompatible, OperatorOuterClass, SupportLevel, Unsupported}
@@ -67,10 +67,6 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec
               return Unsupported(Some("Partitioned writes are not supported"))
             }
 
-            if (cmd.query.output.exists(attr => DataTypeSupport.isComplexType(attr.dataType))) {
-              return Unsupported(Some("Complex types are not supported"))
-            }
-
             val codec = parseCompressionCodec(cmd)
             if (!supportedCompressionCodes.contains(codec)) {
               return Unsupported(Some(s"Unsupported compression codec: $codec"))