From 3466779af8fa1eb86c3cda969873eec3acb260c1 Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Mon, 3 Aug 2026 16:13:11 -0400 Subject: [PATCH 1/7] add constant parquet metadata columns (file_path, file_name, file_size, file_block_start, file_block_length, file_modification_time) --- .../user-guide/latest/compatibility/scans.md | 4 +- native/core/src/execution/planner.rs | 30 +++++-- .../apache/comet/rules/CometScanRule.scala | 23 +++--- .../apache/comet/serde/QueryPlanSerde.scala | 1 + .../comet/serde/contraintExpressions.scala | 39 +++++++-- .../serde/operator/CometNativeScan.scala | 21 ++++- .../apache/comet/serde/operator/package.scala | 45 ++++++++--- .../spark/sql/comet/CometNativeScanExec.scala | 6 +- .../comet/parquet/ParquetReadSuite.scala | 79 +++++++++++++++++++ 9 files changed, 204 insertions(+), 44 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/scans.md b/docs/source/user-guide/latest/compatibility/scans.md index 00b3284e10..071b83ba5a 100644 --- a/docs/source/user-guide/latest/compatibility/scans.md +++ b/docs/source/user-guide/latest/compatibility/scans.md @@ -38,8 +38,8 @@ The following features are not supported and cause Comet to fall back to Spark: - Default values that are nested types (e.g., maps, arrays, structs). Literal default values are supported. - Spark's Datasource V2 API. When `spark.sql.sources.useV1SourceList` does not include `parquet`, Spark uses the V2 API for Parquet scans. Comet's Parquet scan only supports the V1 API. -- Spark metadata columns (e.g., `_metadata.file_path`) -- No support for row indexes +- `_metadata.row_index`. Other `_metadata` columns (`file_path`, `file_name`, `file_size`, `file_block_start`, + `file_block_length`, `file_modification_time`) are supported. - No support for `input_file_name()`, `input_file_block_start()`, or `input_file_block_length()` SQL functions. Comet's Parquet scan does not use Spark's `FileScanRDD`, so these functions cannot populate their values. - No support for `ignoreMissingFiles` or `ignoreCorruptFiles` being set to `true` diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index f20dadf7f3..cd57737dd8 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1526,12 +1526,30 @@ impl PhysicalPlanner { )); } - // Convert the Spark expressions to Physical expressions - let data_filters: Result>, ExecutionError> = common - .data_filters - .iter() - .map(|expr| self.create_expr(expr, Arc::clone(&required_schema))) - .collect(); + // data_filters may reference partition columns and constant metadata columns + // (e.g. `_metadata.file_size`), which the Parquet reader appends after + // required_schema's columns once partition_values are projected into the + // batch. Bind against the combined schema so `Bound` indices resolve + // correctly -- Scala's `exprToProto(filter, scan.output)` + // (CometNativeScan.scala) numbers columns against that same ordering. + let data_filters: Result>, ExecutionError> = + if common.data_filters.is_empty() { + Ok(vec![]) + } else { + let filter_schema: SchemaRef = Arc::new(Schema::new( + required_schema + .fields() + .iter() + .chain(partition_schema.fields().iter()) + .cloned() + .collect::>(), + )); + common + .data_filters + .iter() + .map(|expr| self.create_expr(expr, Arc::clone(&filter_schema))) + .collect() + }; let default_values: Option> = if !common .default_values diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index ca917335da..39c347bfa3 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -51,7 +51,7 @@ import org.apache.comet.iceberg.{CometIcebergNativeScanMetadata, IcebergReflecti import org.apache.comet.objectstore.NativeConfig import org.apache.comet.parquet.CometParquetUtils.{encryptionEnabled, isEncryptionConfigSupported} import org.apache.comet.serde.operator.{CometIcebergNativeScan, CometNativeScan} -import org.apache.comet.shims.{CometTypeShim, ShimCometStreaming, ShimFileFormat, ShimSubqueryBroadcast} +import org.apache.comet.shims.{CometTypeShim, ShimCometStreaming, ShimSubqueryBroadcast} /** * Spark physical optimizer rule for replacing Spark scans with Comet scans. @@ -143,11 +143,18 @@ case class CometScanRule(session: SparkSession) } private def transformV1Scan(plan: SparkPlan, scanExec: FileSourceScanExec): SparkPlan = { - val metadataColNames = metadataCols(scanExec) - if (metadataColNames.nonEmpty) { + // fileConstantMetadataColumns (file_path, file_name, file_size, file_block_start, + // file_block_length, file_modification_time) are known before opening the file and + // supported below via the same projection mechanism as partition columns. Any other + // metadata column (currently only `_metadata.row_index`, generated per row by the reader) + // is not. + val constantMetadataColNames = scanExec.fileConstantMetadataColumns.map(_.name).toSet + val unsupportedMetadataColNames = + metadataCols(scanExec).filterNot(constantMetadataColNames.contains) + if (unsupportedMetadataColNames.nonEmpty) { return withFallbackReason( scanExec, - s"Metadata column(s) ${metadataColNames.mkString(", ")} is not supported") + s"Metadata column(s) ${unsupportedMetadataColNames.mkString(", ")} is not supported") } // On Spark 3.4, injectQueryStageOptimizerRule is unavailable, so @@ -265,10 +272,6 @@ case class CometScanRule(session: SparkSession) withFallbackReason(scanExec, "Native Parquet scan does not support encryption") return None } - if (scanExec.fileConstantMetadataColumns.nonEmpty) { - withFallbackReason(scanExec, "Native DataFusion scan does not support metadata columns") - return None - } // input_file_name, input_file_block_start, and input_file_block_length read from // InputFileBlockHolder, a thread-local set by Spark's FileScanRDD. The native DataFusion // scan does not use FileScanRDD, so these expressions would return empty/default values. @@ -283,10 +286,6 @@ case class CometScanRule(session: SparkSession) "input_file_block_start, or input_file_block_length") return None } - if (ShimFileFormat.findRowIndexColumnIndexInSchema(scanExec.requiredSchema) >= 0) { - withFallbackReason(scanExec, "Native DataFusion scan does not support row index generation") - return None - } if (!isSchemaSupported(scanExec, r)) { return None } diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 4ae8144648..77bc7153f6 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -361,6 +361,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { classOf[CheckOverflow] -> CometCheckOverflow, classOf[Coalesce] -> CometCoalesce, classOf[KnownFloatingPointNormalized] -> CometKnownFloatingPointNormalized, + classOf[KnownNotNull] -> CometKnownNotNull, classOf[KnownNullable] -> CometKnownNullable, classOf[Literal] -> CometLiteral, classOf[MakeDecimal] -> CometMakeDecimal, diff --git a/spark/src/main/scala/org/apache/comet/serde/contraintExpressions.scala b/spark/src/main/scala/org/apache/comet/serde/contraintExpressions.scala index 9b30357f9f..abc6babdd5 100644 --- a/spark/src/main/scala/org/apache/comet/serde/contraintExpressions.scala +++ b/spark/src/main/scala/org/apache/comet/serde/contraintExpressions.scala @@ -19,7 +19,7 @@ package org.apache.comet.serde -import org.apache.spark.sql.catalyst.expressions.{Attribute, KnownFloatingPointNormalized, KnownNullable} +import org.apache.spark.sql.catalyst.expressions.{Attribute, KnownFloatingPointNormalized, KnownNotNull, KnownNullable, TaggingExpression} import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithFallbackReason, serializeDataType} @@ -78,17 +78,40 @@ object CometKnownFloatingPointNormalized } /** - * `KnownNullable` is a tagging expression that only marks its child as nullable; it is a runtime - * no-op (`eval` returns the child's value unchanged). Spark's time-window resolution wraps window - * bounds in `KnownNullable`, so supporting it lets those grouping queries run natively. We simply - * serialize the child and drop the tag. + * `KnownNullable` and `KnownNotNull` (below) are Spark `TaggingExpression`s that only annotate a + * child's nullability; both are runtime no-ops whose `eval` returns the child's value unchanged. + * We serialize the child directly and drop the tag. */ -object CometKnownNullable extends CometExpressionSerde[KnownNullable] { - override def convert( - expr: KnownNullable, +private object CometTaggingExpression { + def convert( + expr: TaggingExpression, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { val optExpr = exprToProtoInternal(expr.child, inputs, binding) optExprWithFallbackReason(optExpr, expr, expr.child) } } + +/** + * Spark's time-window resolution wraps window bounds in `KnownNullable`, so supporting it lets + * those grouping queries run natively. + */ +object CometKnownNullable extends CometExpressionSerde[KnownNullable] { + override def convert( + expr: KnownNullable, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = + CometTaggingExpression.convert(expr, inputs, binding) +} + +/** + * Spark's `FileSourceStrategy` wraps the `_metadata` struct in `KnownNotNull` so the schema + * advertises it as non-nullable without relying on `CreateStruct`'s own nullability inference. + */ +object CometKnownNotNull extends CometExpressionSerde[KnownNotNull] { + override def convert( + expr: KnownNotNull, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = + CometTaggingExpression.convert(expr, inputs, binding) +} diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala index f017acfcd2..3dca22502a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala @@ -29,6 +29,7 @@ import org.apache.spark.sql.comet.{CometNativeExec, CometNativeScanExec, CometSc import org.apache.spark.sql.execution.{FileSourceScanExec, InSubqueryExec, SubqueryAdaptiveBroadcastExec} import org.apache.spark.sql.execution.datasources.parquet.ParquetUtils import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StructField import org.apache.comet.{CometConf, ConfigEntry} import org.apache.comet.CometConf.COMET_EXEC_ENABLED @@ -165,7 +166,18 @@ object CometNativeScan extends CometOperatorSerde[CometScanExec] with Logging { .headOption .map(_.getPath.toUri) - val partitionSchema = schema2Proto(scan.relation.partitionSchema) + // Constant metadata columns (file_path, file_name, file_size, file_block_start, + // file_block_length, file_modification_time) are known before opening the file and + // constant for every row read from it, exactly like partition columns. Spark places + // them immediately after partition columns in `scan.output` + // (FileSourceStrategy.scala: readDataColumns ++ generatedMetadataColumns ++ + // partitionColumns ++ constantMetadataColumns), so appending them after the real + // partition schema here keeps the two in lockstep. + val constantMetadataFields = scan.wrapped.fileConstantMetadataColumns.map(attr => + StructField(attr.name, attr.dataType, attr.nullable)) + val partitionSchemaFields = scan.relation.partitionSchema.fields.toSeq ++ + constantMetadataFields + val partitionSchema = schema2Proto(partitionSchemaFields) val requiredSchema = schema2Proto(scan.requiredSchema) val dataSchema = schema2Proto(scan.relation.dataSchema) @@ -173,15 +185,16 @@ object CometNativeScan extends CometOperatorSerde[CometScanExec] with Logging { scan.relation.dataSchema.fieldIndex(field.name) }) val partitionSchemaIndexes = scan.relation.dataSchema.fields.length until - (scan.relation.dataSchema.length + scan.relation.partitionSchema.fields.length) + (scan.relation.dataSchema.length + partitionSchemaFields.length) val projectionVector = (dataSchemaIndexes ++ partitionSchemaIndexes).map(idx => idx.toLong.asInstanceOf[java.lang.Long]) commonBuilder.addAllProjectionVector(projectionVector.asJava) - // In `CometScanRule`, we ensure partitionSchema is supported. - assert(partitionSchema.length == scan.relation.partitionSchema.fields.length) + // In `CometScanRule`, we ensure partitionSchema (including constant metadata columns) + // is supported. + assert(partitionSchema.length == partitionSchemaFields.length) commonBuilder.addAllDataSchema(dataSchema.asJava) commonBuilder.addAllRequiredSchema(requiredSchema.asJava) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/package.scala b/spark/src/main/scala/org/apache/comet/serde/operator/package.scala index 8bf6832f26..2a386d0aa9 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/package.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/package.scala @@ -19,8 +19,8 @@ package org.apache.comet.serde -import org.apache.spark.sql.catalyst.expressions.Literal -import org.apache.spark.sql.execution.datasources.FilePartition +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Literal} +import org.apache.spark.sql.execution.datasources.{FileFormat, FilePartition, PartitionedFile} import org.apache.spark.sql.execution.datasources.parquet.ParquetUtils import org.apache.spark.sql.types.{StructField, StructType} @@ -47,7 +47,10 @@ package object operator { def partition2Proto( partition: FilePartition, - partitionSchema: StructType): OperatorOuterClass.SparkFilePartition = { + partitionSchema: StructType, + constantMetadataColumns: Seq[AttributeReference] = Seq.empty, + fileConstantMetadataExtractors: Map[String, PartitionedFile => Any] = + FileFormat.BASE_METADATA_EXTRACTORS): OperatorOuterClass.SparkFilePartition = { val partitionBuilder = OperatorOuterClass.SparkFilePartition.newBuilder() partition.files.foreach(file => { // Process the partition values @@ -56,16 +59,28 @@ package object operator { val partitionVals = partitionValues.toSeq(partitionSchema).zipWithIndex.map { case (value, i) => val attr = partitionSchema(i) - val valueProto = exprToProto(Literal(value, attr.dataType), Seq.empty) - // In `CometScanRule`, we have already checked that all partition values are - // supported. So, we can safely use `get` here. - assert( - valueProto.isDefined, - s"Unsupported partition value: $value, type: ${attr.dataType}") - valueProto.get + literalToProto( + Literal(value, attr.dataType), + s"partition value: $value, type: ${attr.dataType}") } + // Constant metadata columns (file_path, file_name, file_size, file_block_start, + // file_block_length, file_modification_time) are, like partition columns, known before + // opening the file and constant for every row read from it. Reuse the same + // partition-value wire format and projection mechanism, and Spark's own extractor + // dispatch (which also covers custom per-format overrides), rather than a bespoke one. + // getFileConstantMetadataColumnValue's Literal has an inferred, not declared, dataType + // (e.g. file_modification_time's raw micros value infers as LongType, not + // TimestampType) -- take only its value and retype against the attribute's actual + // dataType, exactly as Spark's own FileFormat.updateMetadataInternalRow does + // (`row.update(i, literal.value)`). + val metadataVals = constantMetadataColumns.map { attr => + val value = FileFormat + .getFileConstantMetadataColumnValue(attr.name, file, fileConstantMetadataExtractors) + .value + literalToProto(Literal(value, attr.dataType), s"metadata column value for ${attr.name}") + } val fileBuilder = OperatorOuterClass.SparkPartitionedFile.newBuilder() - partitionVals.foreach(fileBuilder.addPartitionValues) + (partitionVals ++ metadataVals).foreach(fileBuilder.addPartitionValues) fileBuilder .setFilePath(file.filePath.toString) .setStart(file.start) @@ -75,4 +90,12 @@ package object operator { }) partitionBuilder.build() } + + // In `CometScanRule`, we have already checked that all partition and metadata column values + // are supported. So, we can safely use `get` here. + private def literalToProto(literal: Literal, description: String): ExprOuterClass.Expr = { + val valueProto = exprToProto(literal, Seq.empty) + assert(valueProto.isDefined, s"Unsupported $description") + valueProto.get + } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala index 52990cfa3d..50ee7ec21c 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala @@ -218,7 +218,11 @@ case class CometNativeScanExec( // Serialize each partition's files import org.apache.comet.serde.operator.partition2Proto val perPartitionBytes = filePartitions.map { filePartition => - val partitionProto = partition2Proto(filePartition, relation.partitionSchema) + val partitionProto = partition2Proto( + filePartition, + relation.partitionSchema, + originalPlan.fileConstantMetadataColumns, + relation.fileFormat.fileConstantMetadataExtractors) val partitionNativeScan = org.apache.comet.serde.OperatorOuterClass.NativeScan .newBuilder() .setFilePartition(partitionProto) diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala index a5fc92a5eb..9aa6a174e0 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala @@ -521,6 +521,85 @@ abstract class ParquetReadSuite extends CometTestBase { } } + test("read _metadata constant columns via native scan") { + // TODO(https://github.com/apache/datafusion-comet/issues/3432): `_metadata.row_index` is + // generated per row by the reader, not constant per file, so it needs DataFusion's + // virtual-column mechanism rather than the partition-value path used here. Not covered. + withTempPath { dir => + (1 to 100).toDF("id").repartition(1).write.parquet(dir.getCanonicalPath) + val df = spark.read + .parquet(dir.getCanonicalPath) + .select( + $"id", + $"_metadata.file_path", + $"_metadata.file_name", + $"_metadata.file_size", + $"_metadata.file_block_start", + $"_metadata.file_block_length", + $"_metadata.file_modification_time") + checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + } + } + + /** + * Writes two Parquet files with distinct row counts (and thus distinct file sizes) into `dir`, + * and returns their (file_path, file_size) pairs sorted by size. Discovers them with Comet + * disabled so the tests below filter on ground truth rather than hardcoded assumptions. + */ + private def writeTwoFilesAndDiscoverMetadata(dir: File): Array[(String, Long)] = { + (1 to 5).toDF("id").repartition(1).write.mode("overwrite").parquet(dir.getCanonicalPath) + (1000 to 1999).toDF("id").repartition(1).write.mode("append").parquet(dir.getCanonicalPath) + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.read + .parquet(dir.getCanonicalPath) + .select($"_metadata.file_path", $"_metadata.file_size") + .distinct() + .as[(String, Long)] + .collect() + .sortBy(_._2) + } + } + + test("filter on _metadata.file_size selects rows from the matching file only") { + withTempPath { dir => + val files = writeTwoFilesAndDiscoverMetadata(dir) + assert(files.length == 2, s"expected two distinct files, got ${files.toSeq}") + val (_, smallerSize) = files(0) + + val df = spark.read + .parquet(dir.getCanonicalPath) + .filter($"_metadata.file_size" === smallerSize) + .select($"id") + checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + } + } + + test("filter combining _metadata column and a data column") { + withTempPath { dir => + val files = writeTwoFilesAndDiscoverMetadata(dir) + val (_, largerSize) = files(1) + + val df = spark.read + .parquet(dir.getCanonicalPath) + .filter($"_metadata.file_size" === largerSize && $"id" > 1500) + .select($"id") + checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + } + } + + test("filter on _metadata.file_path exact match") { + withTempPath { dir => + val files = writeTwoFilesAndDiscoverMetadata(dir) + val (targetPath, _) = files(0) + + val df = spark.read + .parquet(dir.getCanonicalPath) + .filter($"_metadata.file_path" === targetPath) + .select($"id") + checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + } + } + test("fix: string partition column with incorrect offset buffer") { def makeRawParquetFile( path: Path, From 940641e3ac335729d86db70224ed9d54e3782367 Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Mon, 3 Aug 2026 16:58:53 -0400 Subject: [PATCH 2/7] update tests --- .../comet/parquet/ParquetReadSuite.scala | 107 ++++++++++++++---- 1 file changed, 88 insertions(+), 19 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala index 9aa6a174e0..4941095a63 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala @@ -541,6 +541,17 @@ abstract class ParquetReadSuite extends CometTestBase { } } + /** + * Two tiny files easily fit in one Spark partition by data volume, but + * spark.sql.files.minPartitionNum defaults to the session's target parallelism, which pushes + * maxSplitBytes down to spark.sql.files.openCostInBytes; a single file's own virtual open-cost + * surcharge then already consumes that budget, so Spark schedules one file per task instead of + * packing both together. Forcing minPartitionNum to 1 is what actually gets both files into one + * partition, so partition2Proto's multi-file loop is exercised instead of one file per task. + */ + private val forceSinglePartitionConf: (String, String) = + SQLConf.FILES_MIN_PARTITION_NUM.key -> "1" + /** * Writes two Parquet files with distinct row counts (and thus distinct file sizes) into `dir`, * and returns their (file_path, file_size) pairs sorted by size. Discovers them with Comet @@ -549,14 +560,20 @@ abstract class ParquetReadSuite extends CometTestBase { private def writeTwoFilesAndDiscoverMetadata(dir: File): Array[(String, Long)] = { (1 to 5).toDF("id").repartition(1).write.mode("overwrite").parquet(dir.getCanonicalPath) (1000 to 1999).toDF("id").repartition(1).write.mode("append").parquet(dir.getCanonicalPath) - withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - spark.read - .parquet(dir.getCanonicalPath) - .select($"_metadata.file_path", $"_metadata.file_size") - .distinct() - .as[(String, Long)] - .collect() - .sortBy(_._2) + withSQLConf(forceSinglePartitionConf) { + assert( + spark.read.parquet(dir.getCanonicalPath).rdd.getNumPartitions == 1, + "expected both files to be packed into the same Spark partition, to exercise " + + "partition2Proto's multi-file loop rather than one file per task") + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.read + .parquet(dir.getCanonicalPath) + .select($"_metadata.file_path", $"_metadata.file_size") + .distinct() + .as[(String, Long)] + .collect() + .sortBy(_._2) + } } } @@ -566,11 +583,28 @@ abstract class ParquetReadSuite extends CometTestBase { assert(files.length == 2, s"expected two distinct files, got ${files.toSeq}") val (_, smallerSize) = files(0) - val df = spark.read - .parquet(dir.getCanonicalPath) - .filter($"_metadata.file_size" === smallerSize) - .select($"id") - checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + withSQLConf(forceSinglePartitionConf) { + val df = spark.read + .parquet(dir.getCanonicalPath) + .filter($"_metadata.file_size" === smallerSize) + .select($"id") + checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + } + } + } + + test("filter on _metadata.file_size using a range predicate") { + withTempPath { dir => + val files = writeTwoFilesAndDiscoverMetadata(dir) + val (_, smallerSize) = files(0) + + withSQLConf(forceSinglePartitionConf) { + val df = spark.read + .parquet(dir.getCanonicalPath) + .filter($"_metadata.file_size" > smallerSize) + .select($"id") + checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + } } } @@ -579,11 +613,13 @@ abstract class ParquetReadSuite extends CometTestBase { val files = writeTwoFilesAndDiscoverMetadata(dir) val (_, largerSize) = files(1) - val df = spark.read - .parquet(dir.getCanonicalPath) - .filter($"_metadata.file_size" === largerSize && $"id" > 1500) - .select($"id") - checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + withSQLConf(forceSinglePartitionConf) { + val df = spark.read + .parquet(dir.getCanonicalPath) + .filter($"_metadata.file_size" === largerSize && $"id" > 1500) + .select($"id") + checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + } } } @@ -592,9 +628,42 @@ abstract class ParquetReadSuite extends CometTestBase { val files = writeTwoFilesAndDiscoverMetadata(dir) val (targetPath, _) = files(0) + withSQLConf(forceSinglePartitionConf) { + val df = spark.read + .parquet(dir.getCanonicalPath) + .filter($"_metadata.file_path" === targetPath) + .select($"id") + checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + } + } + } + + test("read _metadata constant columns together with a real Hive partition column") { + withTempPath { dir => + Seq((1, "a"), (2, "a"), (3, "b"), (4, "b")) + .toDF("id", "pcol") + .repartition(1) + .write + .partitionBy("pcol") + .parquet(dir.getCanonicalPath) + val df = spark.read + .parquet(dir.getCanonicalPath) + .select($"id", $"pcol", $"_metadata.file_path", $"_metadata.file_size") + checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + } + } + + test("filter combining a real Hive partition column and a metadata column") { + withTempPath { dir => + Seq((1, "a"), (2, "a"), (3, "b"), (4, "b")) + .toDF("id", "pcol") + .repartition(1) + .write + .partitionBy("pcol") + .parquet(dir.getCanonicalPath) val df = spark.read .parquet(dir.getCanonicalPath) - .filter($"_metadata.file_path" === targetPath) + .filter($"pcol" === "b" && $"_metadata.file_size" > 0) .select($"id") checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) } From a005b404ffe71f8d9da360329eca889568bae11b Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Mon, 3 Aug 2026 17:21:29 -0400 Subject: [PATCH 3/7] fix compilation --- .../scala/org/apache/comet/parquet/ParquetReadSuite.scala | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala index 4941095a63..703ef73f50 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala @@ -560,13 +560,14 @@ abstract class ParquetReadSuite extends CometTestBase { private def writeTwoFilesAndDiscoverMetadata(dir: File): Array[(String, Long)] = { (1 to 5).toDF("id").repartition(1).write.mode("overwrite").parquet(dir.getCanonicalPath) (1000 to 1999).toDF("id").repartition(1).write.mode("append").parquet(dir.getCanonicalPath) + var files: Array[(String, Long)] = null withSQLConf(forceSinglePartitionConf) { assert( spark.read.parquet(dir.getCanonicalPath).rdd.getNumPartitions == 1, "expected both files to be packed into the same Spark partition, to exercise " + "partition2Proto's multi-file loop rather than one file per task") withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - spark.read + files = spark.read .parquet(dir.getCanonicalPath) .select($"_metadata.file_path", $"_metadata.file_size") .distinct() @@ -575,6 +576,7 @@ abstract class ParquetReadSuite extends CometTestBase { .sortBy(_._2) } } + files } test("filter on _metadata.file_size selects rows from the matching file only") { From 3a131f676753e850e9cc875b86ebf6276a677c44 Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Mon, 3 Aug 2026 17:58:36 -0400 Subject: [PATCH 4/7] add shim for FileFormat due to Spark 3.4. --- .../apache/comet/serde/operator/package.scala | 9 +++--- .../spark/sql/comet/CometNativeScanExec.scala | 3 +- .../apache/comet/shims/ShimFileFormat.scala | 32 ++++++++++++++++++- .../apache/comet/shims/ShimFileFormat.scala | 12 +++++++ .../apache/comet/shims/ShimFileFormat.scala | 12 +++++++ 5 files changed, 62 insertions(+), 6 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/package.scala b/spark/src/main/scala/org/apache/comet/serde/operator/package.scala index 2a386d0aa9..d2d674e485 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/package.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/package.scala @@ -20,12 +20,13 @@ package org.apache.comet.serde import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Literal} -import org.apache.spark.sql.execution.datasources.{FileFormat, FilePartition, PartitionedFile} +import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionedFile} import org.apache.spark.sql.execution.datasources.parquet.ParquetUtils import org.apache.spark.sql.types.{StructField, StructType} import org.apache.comet.parquet.CometParquetUtils import org.apache.comet.serde.QueryPlanSerde.{exprToProto, serializeDataType} +import org.apache.comet.shims.ShimFileFormat package object operator { @@ -49,8 +50,8 @@ package object operator { partition: FilePartition, partitionSchema: StructType, constantMetadataColumns: Seq[AttributeReference] = Seq.empty, - fileConstantMetadataExtractors: Map[String, PartitionedFile => Any] = - FileFormat.BASE_METADATA_EXTRACTORS): OperatorOuterClass.SparkFilePartition = { + fileConstantMetadataExtractors: Map[String, PartitionedFile => Any] = Map.empty) + : OperatorOuterClass.SparkFilePartition = { val partitionBuilder = OperatorOuterClass.SparkFilePartition.newBuilder() partition.files.foreach(file => { // Process the partition values @@ -74,7 +75,7 @@ package object operator { // dataType, exactly as Spark's own FileFormat.updateMetadataInternalRow does // (`row.update(i, literal.value)`). val metadataVals = constantMetadataColumns.map { attr => - val value = FileFormat + val value = ShimFileFormat .getFileConstantMetadataColumnValue(attr.name, file, fileConstantMetadataExtractors) .value literalToProto(Literal(value, attr.dataType), s"metadata column value for ${attr.name}") diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala index 50ee7ec21c..fbebc42004 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala @@ -40,6 +40,7 @@ import com.google.common.base.Objects import org.apache.comet.parquet.CometParquetUtils import org.apache.comet.serde.OperatorOuterClass.Operator import org.apache.comet.serde.QueryPlanSerde.exprToProto +import org.apache.comet.shims.ShimFileFormat /** * Native scan operator for DataSource V1 Parquet files using DataFusion's ParquetExec. @@ -222,7 +223,7 @@ case class CometNativeScanExec( filePartition, relation.partitionSchema, originalPlan.fileConstantMetadataColumns, - relation.fileFormat.fileConstantMetadataExtractors) + ShimFileFormat.fileConstantMetadataExtractors(relation.fileFormat)) val partitionNativeScan = org.apache.comet.serde.OperatorOuterClass.NativeScan .newBuilder() .setFilePartition(partitionProto) diff --git a/spark/src/main/spark-3.4/org/apache/comet/shims/ShimFileFormat.scala b/spark/src/main/spark-3.4/org/apache/comet/shims/ShimFileFormat.scala index 7b4911e81f..3656c9c348 100644 --- a/spark/src/main/spark-3.4/org/apache/comet/shims/ShimFileFormat.scala +++ b/spark/src/main/spark-3.4/org/apache/comet/shims/ShimFileFormat.scala @@ -19,7 +19,9 @@ package org.apache.comet.shims -import org.apache.spark.sql.execution.datasources.{FileFormat, RowIndexUtil} +import org.apache.hadoop.fs.Path +import org.apache.spark.sql.catalyst.expressions.Literal +import org.apache.spark.sql.execution.datasources.{FileFormat, PartitionedFile, RowIndexUtil} import org.apache.spark.sql.types.StructType object ShimFileFormat { @@ -30,4 +32,32 @@ object ShimFileFormat { def findRowIndexColumnIndexInSchema(sparkSchema: StructType): Int = RowIndexUtil.findRowIndexColumnIndexInSchema(sparkSchema) + + // Spark 3.4 has no per-format metadata extractor concept (added in Spark 3.5, SPARK-43868); + // it derives these values inline in FileFormat.updateMetadataInternalRow. Replicate that + // logic here so callers can use the same extractor-map shape across all Spark versions. + private val baseMetadataExtractors: Map[String, PartitionedFile => Any] = Map( + FileFormat.FILE_PATH -> { pf: PartitionedFile => + // Use `new Path(Path.toString)` as a form of canonicalization + new Path(pf.filePath.toPath.toString).toUri.toString + }, + FileFormat.FILE_NAME -> { pf: PartitionedFile => + pf.filePath.toUri.getRawPath.split("/").lastOption.getOrElse("") + }, + FileFormat.FILE_SIZE -> { pf: PartitionedFile => pf.fileSize }, + FileFormat.FILE_BLOCK_START -> { pf: PartitionedFile => pf.start }, + FileFormat.FILE_BLOCK_LENGTH -> { pf: PartitionedFile => pf.length }, + // The modificationTime from the file has millisecond granularity, but the TimestampType for + // `file_modification_time` has microsecond granularity. + FileFormat.FILE_MODIFICATION_TIME -> { pf: PartitionedFile => pf.modificationTime * 1000 }) + + def fileConstantMetadataExtractors( + fileFormat: FileFormat): Map[String, PartitionedFile => Any] = + baseMetadataExtractors + + def getFileConstantMetadataColumnValue( + name: String, + file: PartitionedFile, + extractors: Map[String, PartitionedFile => Any]): Literal = + Literal(extractors(name)(file)) } diff --git a/spark/src/main/spark-3.5/org/apache/comet/shims/ShimFileFormat.scala b/spark/src/main/spark-3.5/org/apache/comet/shims/ShimFileFormat.scala index 1702db135a..6e827b1d7f 100644 --- a/spark/src/main/spark-3.5/org/apache/comet/shims/ShimFileFormat.scala +++ b/spark/src/main/spark-3.5/org/apache/comet/shims/ShimFileFormat.scala @@ -19,6 +19,8 @@ package org.apache.comet.shims +import org.apache.spark.sql.catalyst.expressions.Literal +import org.apache.spark.sql.execution.datasources.{FileFormat, PartitionedFile} import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat import org.apache.spark.sql.execution.datasources.parquet.ParquetRowIndexUtil import org.apache.spark.sql.types.StructType @@ -30,4 +32,14 @@ object ShimFileFormat { def findRowIndexColumnIndexInSchema(sparkSchema: StructType): Int = ParquetRowIndexUtil.findRowIndexColumnIndexInSchema(sparkSchema) + + def fileConstantMetadataExtractors( + fileFormat: FileFormat): Map[String, PartitionedFile => Any] = + fileFormat.fileConstantMetadataExtractors + + def getFileConstantMetadataColumnValue( + name: String, + file: PartitionedFile, + extractors: Map[String, PartitionedFile => Any]): Literal = + FileFormat.getFileConstantMetadataColumnValue(name, file, extractors) } diff --git a/spark/src/main/spark-4.x/org/apache/comet/shims/ShimFileFormat.scala b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimFileFormat.scala index 1702db135a..6e827b1d7f 100644 --- a/spark/src/main/spark-4.x/org/apache/comet/shims/ShimFileFormat.scala +++ b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimFileFormat.scala @@ -19,6 +19,8 @@ package org.apache.comet.shims +import org.apache.spark.sql.catalyst.expressions.Literal +import org.apache.spark.sql.execution.datasources.{FileFormat, PartitionedFile} import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat import org.apache.spark.sql.execution.datasources.parquet.ParquetRowIndexUtil import org.apache.spark.sql.types.StructType @@ -30,4 +32,14 @@ object ShimFileFormat { def findRowIndexColumnIndexInSchema(sparkSchema: StructType): Int = ParquetRowIndexUtil.findRowIndexColumnIndexInSchema(sparkSchema) + + def fileConstantMetadataExtractors( + fileFormat: FileFormat): Map[String, PartitionedFile => Any] = + fileFormat.fileConstantMetadataExtractors + + def getFileConstantMetadataColumnValue( + name: String, + file: PartitionedFile, + extractors: Map[String, PartitionedFile => Any]): Literal = + FileFormat.getFileConstantMetadataColumnValue(name, file, extractors) } From 3bae5a685f741e1b84364ccafa3c8745e667a331 Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Tue, 4 Aug 2026 07:38:00 -0400 Subject: [PATCH 5/7] update diffs to address test failures --- .../apache/comet/rules/CometScanRule.scala | 8 ++- .../apache/comet/serde/operator/package.scala | 6 ++- .../apache/comet/shims/ShimFileFormat.scala | 8 ++- .../apache/comet/shims/ShimFileFormat.scala | 8 ++- .../apache/comet/shims/ShimFileFormat.scala | 49 +++++++++++++++++++ .../apache/comet/shims/ShimFileFormat.scala | 49 +++++++++++++++++++ .../apache/comet/shims/ShimFileFormat.scala | 9 ++-- 7 files changed, 127 insertions(+), 10 deletions(-) create mode 100644 spark/src/main/spark-4.0/org/apache/comet/shims/ShimFileFormat.scala create mode 100644 spark/src/main/spark-4.1/org/apache/comet/shims/ShimFileFormat.scala rename spark/src/main/{spark-4.x => spark-4.2}/org/apache/comet/shims/ShimFileFormat.scala (84%) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index 39c347bfa3..0ceaf00431 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -51,7 +51,7 @@ import org.apache.comet.iceberg.{CometIcebergNativeScanMetadata, IcebergReflecti import org.apache.comet.objectstore.NativeConfig import org.apache.comet.parquet.CometParquetUtils.{encryptionEnabled, isEncryptionConfigSupported} import org.apache.comet.serde.operator.{CometIcebergNativeScan, CometNativeScan} -import org.apache.comet.shims.{CometTypeShim, ShimCometStreaming, ShimSubqueryBroadcast} +import org.apache.comet.shims.{CometTypeShim, ShimCometStreaming, ShimFileFormat, ShimSubqueryBroadcast} /** * Spark physical optimizer rule for replacing Spark scans with Comet scans. @@ -282,10 +282,14 @@ case class CometScanRule(session: SparkSession) }))) { withFallbackReason( scanExec, - "Native DataFusion scan is not compatible with input_file_name, " + + "Native Parquet scan is not compatible with input_file_name, " + "input_file_block_start, or input_file_block_length") return None } + if (ShimFileFormat.findRowIndexColumnIndexInSchema(scanExec.requiredSchema) >= 0) { + withFallbackReason(scanExec, "Native Parquet scan does not support row index generation") + return None + } if (!isSchemaSupported(scanExec, r)) { return None } diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/package.scala b/spark/src/main/scala/org/apache/comet/serde/operator/package.scala index d2d674e485..cb7702083b 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/package.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/package.scala @@ -76,7 +76,11 @@ package object operator { // (`row.update(i, literal.value)`). val metadataVals = constantMetadataColumns.map { attr => val value = ShimFileFormat - .getFileConstantMetadataColumnValue(attr.name, file, fileConstantMetadataExtractors) + .getFileConstantMetadataColumnValue( + attr.name, + file, + fileConstantMetadataExtractors, + attr.dataType) .value literalToProto(Literal(value, attr.dataType), s"metadata column value for ${attr.name}") } diff --git a/spark/src/main/spark-3.4/org/apache/comet/shims/ShimFileFormat.scala b/spark/src/main/spark-3.4/org/apache/comet/shims/ShimFileFormat.scala index 3656c9c348..6906a39fbf 100644 --- a/spark/src/main/spark-3.4/org/apache/comet/shims/ShimFileFormat.scala +++ b/spark/src/main/spark-3.4/org/apache/comet/shims/ShimFileFormat.scala @@ -22,7 +22,7 @@ package org.apache.comet.shims import org.apache.hadoop.fs.Path import org.apache.spark.sql.catalyst.expressions.Literal import org.apache.spark.sql.execution.datasources.{FileFormat, PartitionedFile, RowIndexUtil} -import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.types.{DataType, StructType} object ShimFileFormat { @@ -55,9 +55,13 @@ object ShimFileFormat { fileFormat: FileFormat): Map[String, PartitionedFile => Any] = baseMetadataExtractors + // dataType is unused on this Spark version; getFileConstantMetadataColumnValue only gained + // a dataType parameter in Spark 4.2 (SPARK-56931). Accepted here so callers can pass it + // uniformly across Spark versions. def getFileConstantMetadataColumnValue( name: String, file: PartitionedFile, - extractors: Map[String, PartitionedFile => Any]): Literal = + extractors: Map[String, PartitionedFile => Any], + dataType: DataType): Literal = Literal(extractors(name)(file)) } diff --git a/spark/src/main/spark-3.5/org/apache/comet/shims/ShimFileFormat.scala b/spark/src/main/spark-3.5/org/apache/comet/shims/ShimFileFormat.scala index 6e827b1d7f..9506a851c0 100644 --- a/spark/src/main/spark-3.5/org/apache/comet/shims/ShimFileFormat.scala +++ b/spark/src/main/spark-3.5/org/apache/comet/shims/ShimFileFormat.scala @@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.expressions.Literal import org.apache.spark.sql.execution.datasources.{FileFormat, PartitionedFile} import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat import org.apache.spark.sql.execution.datasources.parquet.ParquetRowIndexUtil -import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.types.{DataType, StructType} object ShimFileFormat { // A name for a temporary column that holds row indexes computed by the file format reader @@ -37,9 +37,13 @@ object ShimFileFormat { fileFormat: FileFormat): Map[String, PartitionedFile => Any] = fileFormat.fileConstantMetadataExtractors + // dataType is unused on this Spark version; getFileConstantMetadataColumnValue only gained + // a dataType parameter in Spark 4.2 (SPARK-56931). Accepted here so callers can pass it + // uniformly across Spark versions. def getFileConstantMetadataColumnValue( name: String, file: PartitionedFile, - extractors: Map[String, PartitionedFile => Any]): Literal = + extractors: Map[String, PartitionedFile => Any], + dataType: DataType): Literal = FileFormat.getFileConstantMetadataColumnValue(name, file, extractors) } diff --git a/spark/src/main/spark-4.0/org/apache/comet/shims/ShimFileFormat.scala b/spark/src/main/spark-4.0/org/apache/comet/shims/ShimFileFormat.scala new file mode 100644 index 0000000000..9506a851c0 --- /dev/null +++ b/spark/src/main/spark-4.0/org/apache/comet/shims/ShimFileFormat.scala @@ -0,0 +1,49 @@ +/* + * 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.shims + +import org.apache.spark.sql.catalyst.expressions.Literal +import org.apache.spark.sql.execution.datasources.{FileFormat, PartitionedFile} +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.execution.datasources.parquet.ParquetRowIndexUtil +import org.apache.spark.sql.types.{DataType, StructType} + +object ShimFileFormat { + // A name for a temporary column that holds row indexes computed by the file format reader + // until they can be placed in the _metadata struct. + val ROW_INDEX_TEMPORARY_COLUMN_NAME = ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME + + def findRowIndexColumnIndexInSchema(sparkSchema: StructType): Int = + ParquetRowIndexUtil.findRowIndexColumnIndexInSchema(sparkSchema) + + def fileConstantMetadataExtractors( + fileFormat: FileFormat): Map[String, PartitionedFile => Any] = + fileFormat.fileConstantMetadataExtractors + + // dataType is unused on this Spark version; getFileConstantMetadataColumnValue only gained + // a dataType parameter in Spark 4.2 (SPARK-56931). Accepted here so callers can pass it + // uniformly across Spark versions. + def getFileConstantMetadataColumnValue( + name: String, + file: PartitionedFile, + extractors: Map[String, PartitionedFile => Any], + dataType: DataType): Literal = + FileFormat.getFileConstantMetadataColumnValue(name, file, extractors) +} diff --git a/spark/src/main/spark-4.1/org/apache/comet/shims/ShimFileFormat.scala b/spark/src/main/spark-4.1/org/apache/comet/shims/ShimFileFormat.scala new file mode 100644 index 0000000000..9506a851c0 --- /dev/null +++ b/spark/src/main/spark-4.1/org/apache/comet/shims/ShimFileFormat.scala @@ -0,0 +1,49 @@ +/* + * 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.shims + +import org.apache.spark.sql.catalyst.expressions.Literal +import org.apache.spark.sql.execution.datasources.{FileFormat, PartitionedFile} +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.execution.datasources.parquet.ParquetRowIndexUtil +import org.apache.spark.sql.types.{DataType, StructType} + +object ShimFileFormat { + // A name for a temporary column that holds row indexes computed by the file format reader + // until they can be placed in the _metadata struct. + val ROW_INDEX_TEMPORARY_COLUMN_NAME = ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME + + def findRowIndexColumnIndexInSchema(sparkSchema: StructType): Int = + ParquetRowIndexUtil.findRowIndexColumnIndexInSchema(sparkSchema) + + def fileConstantMetadataExtractors( + fileFormat: FileFormat): Map[String, PartitionedFile => Any] = + fileFormat.fileConstantMetadataExtractors + + // dataType is unused on this Spark version; getFileConstantMetadataColumnValue only gained + // a dataType parameter in Spark 4.2 (SPARK-56931). Accepted here so callers can pass it + // uniformly across Spark versions. + def getFileConstantMetadataColumnValue( + name: String, + file: PartitionedFile, + extractors: Map[String, PartitionedFile => Any], + dataType: DataType): Literal = + FileFormat.getFileConstantMetadataColumnValue(name, file, extractors) +} diff --git a/spark/src/main/spark-4.x/org/apache/comet/shims/ShimFileFormat.scala b/spark/src/main/spark-4.2/org/apache/comet/shims/ShimFileFormat.scala similarity index 84% rename from spark/src/main/spark-4.x/org/apache/comet/shims/ShimFileFormat.scala rename to spark/src/main/spark-4.2/org/apache/comet/shims/ShimFileFormat.scala index 6e827b1d7f..4aee5671de 100644 --- a/spark/src/main/spark-4.x/org/apache/comet/shims/ShimFileFormat.scala +++ b/spark/src/main/spark-4.2/org/apache/comet/shims/ShimFileFormat.scala @@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.expressions.Literal import org.apache.spark.sql.execution.datasources.{FileFormat, PartitionedFile} import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat import org.apache.spark.sql.execution.datasources.parquet.ParquetRowIndexUtil -import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.types.{DataType, StructType} object ShimFileFormat { // A name for a temporary column that holds row indexes computed by the file format reader @@ -37,9 +37,12 @@ object ShimFileFormat { fileFormat: FileFormat): Map[String, PartitionedFile => Any] = fileFormat.fileConstantMetadataExtractors + // Spark 4.2 (SPARK-56931) added a required dataType parameter so complex constant metadata + // values go through Literal.create instead of Literal.apply's type inference. def getFileConstantMetadataColumnValue( name: String, file: PartitionedFile, - extractors: Map[String, PartitionedFile => Any]): Literal = - FileFormat.getFileConstantMetadataColumnValue(name, file, extractors) + extractors: Map[String, PartitionedFile => Any], + dataType: DataType): Literal = + FileFormat.getFileConstantMetadataColumnValue(name, file, extractors, dataType) } From 6739cde503ef558510dab64eadd9a82f493f1f6c Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Tue, 4 Aug 2026 14:14:13 -0400 Subject: [PATCH 6/7] fix: avoid _metadata column name collisions with real columns in native scan DataFusion's `table_partition_cols` literal substitution matches columns by name, not position, so a real data or Hive partition column sharing a name with one of the six constant metadata columns (e.g. `file_size`) could be silently overwritten with the metadata value. Spark itself has no such hazard: metadata attributes get a fresh ExprId regardless of name and are always resolved by ExprId/ordinal, never by name. Prefix the wire field names with `_comet_metadata_` so the collision is structurally impossible, matching Spark's guarantee instead of falling back. Also adds test coverage requested in review: fallback still applies to `_metadata` (whole struct) and `_metadata.row_index`, file_path/file_name are correctly url-encoded on Spark 3.4 for a directory containing a space, and drops redundant `includeClasses` assertions from `checkSparkAnswerAndOperator` calls now that `CometNativeScanExec` is the only native scan implementation. --- .../serde/operator/CometNativeScan.scala | 6 ++- .../comet/parquet/ParquetReadSuite.scala | 54 ++++++++++++++++--- .../comet/rules/CometScanRuleSuite.scala | 17 ++++++ 3 files changed, 69 insertions(+), 8 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala index 3dca22502a..0240bd4a07 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala @@ -46,6 +46,10 @@ import org.apache.comet.serde.QueryPlanSerde.{exprToProto, serializeDataType} */ object CometNativeScan extends CometOperatorSerde[CometScanExec] with Logging { + // DataFusion's table_partition_cols literal substitution matches by name, so a bare name + // like "file_size" could collide with a real column of the same name. Prefix to avoid it. + private val constantMetadataFieldPrefix = "_comet_metadata_" + /** Determine whether the scan is supported and tag the Spark plan with any fallback reasons */ def isSupported(scanExec: FileSourceScanExec): Boolean = { @@ -174,7 +178,7 @@ object CometNativeScan extends CometOperatorSerde[CometScanExec] with Logging { // partitionColumns ++ constantMetadataColumns), so appending them after the real // partition schema here keeps the two in lockstep. val constantMetadataFields = scan.wrapped.fileConstantMetadataColumns.map(attr => - StructField(attr.name, attr.dataType, attr.nullable)) + StructField(s"$constantMetadataFieldPrefix${attr.name}", attr.dataType, attr.nullable)) val partitionSchemaFields = scan.relation.partitionSchema.fields.toSeq ++ constantMetadataFields val partitionSchema = schema2Proto(partitionSchemaFields) diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala index 703ef73f50..684c4d6581 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala @@ -537,7 +537,7 @@ abstract class ParquetReadSuite extends CometTestBase { $"_metadata.file_block_start", $"_metadata.file_block_length", $"_metadata.file_modification_time") - checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + checkSparkAnswerAndOperator(df) } } @@ -590,7 +590,7 @@ abstract class ParquetReadSuite extends CometTestBase { .parquet(dir.getCanonicalPath) .filter($"_metadata.file_size" === smallerSize) .select($"id") - checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + checkSparkAnswerAndOperator(df) } } } @@ -605,7 +605,7 @@ abstract class ParquetReadSuite extends CometTestBase { .parquet(dir.getCanonicalPath) .filter($"_metadata.file_size" > smallerSize) .select($"id") - checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + checkSparkAnswerAndOperator(df) } } } @@ -620,7 +620,7 @@ abstract class ParquetReadSuite extends CometTestBase { .parquet(dir.getCanonicalPath) .filter($"_metadata.file_size" === largerSize && $"id" > 1500) .select($"id") - checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + checkSparkAnswerAndOperator(df) } } } @@ -635,7 +635,7 @@ abstract class ParquetReadSuite extends CometTestBase { .parquet(dir.getCanonicalPath) .filter($"_metadata.file_path" === targetPath) .select($"id") - checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + checkSparkAnswerAndOperator(df) } } } @@ -651,7 +651,7 @@ abstract class ParquetReadSuite extends CometTestBase { val df = spark.read .parquet(dir.getCanonicalPath) .select($"id", $"pcol", $"_metadata.file_path", $"_metadata.file_size") - checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + checkSparkAnswerAndOperator(df) } } @@ -667,7 +667,47 @@ abstract class ParquetReadSuite extends CometTestBase { .parquet(dir.getCanonicalPath) .filter($"pcol" === "b" && $"_metadata.file_size" > 0) .select($"id") - checkSparkAnswerAndOperator(df, Seq(classOf[CometNativeScanExec])) + checkSparkAnswerAndOperator(df) + } + } + + test("_metadata.file_path and file_name are url-encoded for a directory with a space") { + withTempDir { parent => + val dir = new File(parent, "dir with space") + (1 to 10).toDF("id").repartition(1).write.parquet(dir.getCanonicalPath) + val df = spark.read + .parquet(dir.getCanonicalPath) + .select($"id", $"_metadata.file_path", $"_metadata.file_name") + checkSparkAnswerAndOperator(df) + } + } + + test("_metadata column does not collide with a data column of the same name") { + withTempPath { dir => + Seq((1L, 10), (2L, 20)) + .toDF("file_size", "x") + .repartition(1) + .write + .parquet(dir.getCanonicalPath) + val df = spark.read + .parquet(dir.getCanonicalPath) + .select($"file_size", $"x", $"_metadata.file_size".as("meta_size")) + checkSparkAnswerAndOperator(df) + } + } + + test("_metadata column does not collide with a partition column of the same name") { + withTempPath { dir => + Seq((1, 100L), (2, 200L)) + .toDF("id", "file_size") + .repartition(1) + .write + .partitionBy("file_size") + .parquet(dir.getCanonicalPath) + val df = spark.read + .parquet(dir.getCanonicalPath) + .select($"id", $"file_size", $"_metadata.file_size") + checkSparkAnswerAndOperator(df) } } diff --git a/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala index f444fe62c9..41fdc8009b 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala @@ -134,4 +134,21 @@ class CometScanRuleSuite extends CometTestBase { } } + test("CometScanRule should fallback to Spark for unsupported _metadata columns") { + withTempPath { path => + createTestDataFrame.write.parquet(path.toString) + withTempView("test_data") { + spark.read.parquet(path.toString).createOrReplaceTempView("test_data") + + for (query <- Seq( + "SELECT id, _metadata FROM test_data", + "SELECT id, _metadata.row_index FROM test_data")) { + val transformedPlan = applyCometScanRule(createSparkPlan(spark, query)) + assert(countOperators(transformedPlan, classOf[FileSourceScanExec]) == 1) + assert(countOperators(transformedPlan, classOf[CometScanExec]) == 0) + } + } + } + } + } From 705c7bed522a21bec9e074c2264ff4f05bb85e34 Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Tue, 4 Aug 2026 15:12:51 -0400 Subject: [PATCH 7/7] query the table directly instead of a view which won't have those metadata columns --- .../comet/rules/CometScanRuleSuite.scala | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala index 41fdc8009b..fb39e340d1 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala @@ -137,16 +137,19 @@ class CometScanRuleSuite extends CometTestBase { test("CometScanRule should fallback to Spark for unsupported _metadata columns") { withTempPath { path => createTestDataFrame.write.parquet(path.toString) - withTempView("test_data") { - spark.read.parquet(path.toString).createOrReplaceTempView("test_data") - for (query <- Seq( - "SELECT id, _metadata FROM test_data", - "SELECT id, _metadata.row_index FROM test_data")) { - val transformedPlan = applyCometScanRule(createSparkPlan(spark, query)) - assert(countOperators(transformedPlan, classOf[FileSourceScanExec]) == 1) - assert(countOperators(transformedPlan, classOf[CometScanExec]) == 0) + // A temp view's output schema is fixed at creation time (here [id, name]), so + // `_metadata` cannot resolve through one; query the relation directly. + for (df <- Seq( + spark.read.parquet(path.toString).select("id", "_metadata"), + spark.read.parquet(path.toString).selectExpr("id", "_metadata.row_index"))) { + var sparkPlan: SparkPlan = null + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkPlan = df.queryExecution.executedPlan } + val transformedPlan = applyCometScanRule(stripAQEPlan(sparkPlan)) + assert(countOperators(transformedPlan, classOf[FileSourceScanExec]) == 1) + assert(countOperators(transformedPlan, classOf[CometScanExec]) == 0) } } }