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 f262b13d21 feat: support _metadata constant columns in native Parquet
scan (#5237)
f262b13d21 is described below
commit f262b13d2112a78da44d7acc8998b76bc86a6d43
Author: Matt Butrovich <[email protected]>
AuthorDate: Wed Aug 5 08:45:02 2026 -0400
feat: support _metadata constant columns in native Parquet scan (#5237)
* add constant parquet metadata columns (file_path, file_name, file_size,
file_block_start, file_block_length, file_modification_time)
* update tests
* fix compilation
* add shim for FileFormat due to Spark 3.4.
* update diffs to address test failures
* 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 [...]
* query the table directly instead of a view which won't have those
metadata columns
---
.../user-guide/latest/compatibility/scans.md | 4 +-
native/core/src/execution/planner.rs | 30 +++-
.../org/apache/comet/rules/CometScanRule.scala | 21 ++-
.../org/apache/comet/serde/QueryPlanSerde.scala | 1 +
.../apache/comet/serde/contraintExpressions.scala | 39 ++++-
.../comet/serde/operator/CometNativeScan.scala | 25 ++-
.../org/apache/comet/serde/operator/package.scala | 50 ++++--
.../spark/sql/comet/CometNativeScanExec.scala | 7 +-
.../org/apache/comet/shims/ShimFileFormat.scala | 38 ++++-
.../org/apache/comet/shims/ShimFileFormat.scala | 18 +-
.../org/apache/comet/shims/ShimFileFormat.scala | 18 +-
.../org/apache/comet/shims/ShimFileFormat.scala | 18 +-
.../org/apache/comet/shims/ShimFileFormat.scala | 17 +-
.../apache/comet/parquet/ParquetReadSuite.scala | 190 +++++++++++++++++++++
.../apache/comet/rules/CometScanRuleSuite.scala | 20 +++
15 files changed, 449 insertions(+), 47 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 d4c97c509d..85851b17ca 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -1527,12 +1527,30 @@ impl PhysicalPlanner {
));
}
- // Convert the Spark expressions to Physical expressions
- let data_filters: Result<Vec<Arc<dyn PhysicalExpr>>,
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<Vec<Arc<dyn PhysicalExpr>>,
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::<Vec<FieldRef>>(),
+ ));
+ common
+ .data_filters
+ .iter()
+ .map(|expr| self.create_expr(expr,
Arc::clone(&filter_schema)))
+ .collect()
+ };
let default_values: Option<HashMap<Column, ScalarValue>> = 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 a57475d2f5..2c15bb5e4e 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala
@@ -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.
@@ -279,12 +282,12 @@ 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 DataFusion scan does not support
row index generation")
+ withFallbackReason(scanExec, "Native Parquet scan does not support row
index generation")
return None
}
if (!isSchemaSupported(scanExec, r)) {
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 632c4d25fe..4299eb14ba 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -365,6 +365,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 460f23485d..44158be445 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,
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)
optExpr
}
}
+
+/**
+ * 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..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
@@ -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
@@ -45,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 = {
@@ -165,7 +170,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(s"$constantMetadataFieldPrefix${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 +189,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..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
@@ -19,13 +19,14 @@
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.{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 {
@@ -47,7 +48,10 @@ package object operator {
def partition2Proto(
partition: FilePartition,
- partitionSchema: StructType): OperatorOuterClass.SparkFilePartition = {
+ partitionSchema: StructType,
+ constantMetadataColumns: Seq[AttributeReference] = Seq.empty,
+ fileConstantMetadataExtractors: Map[String, PartitionedFile => Any] =
Map.empty)
+ : OperatorOuterClass.SparkFilePartition = {
val partitionBuilder = OperatorOuterClass.SparkFilePartition.newBuilder()
partition.files.foreach(file => {
// Process the partition values
@@ -56,16 +60,32 @@ 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 = ShimFileFormat
+ .getFileConstantMetadataColumnValue(
+ attr.name,
+ file,
+ fileConstantMetadataExtractors,
+ attr.dataType)
+ .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 +95,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..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.
@@ -218,7 +219,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,
+ 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..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
@@ -19,8 +19,10 @@
package org.apache.comet.shims
-import org.apache.spark.sql.execution.datasources.{FileFormat, RowIndexUtil}
-import org.apache.spark.sql.types.StructType
+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.{DataType, StructType}
object ShimFileFormat {
@@ -30,4 +32,36 @@ 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
+
+ // 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 =
+ 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..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
@@ -19,9 +19,11 @@
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
+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
@@ -30,4 +32,18 @@ object ShimFileFormat {
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.0/org/apache/comet/shims/ShimFileFormat.scala
similarity index 62%
copy from spark/src/main/spark-4.x/org/apache/comet/shims/ShimFileFormat.scala
copy to spark/src/main/spark-4.0/org/apache/comet/shims/ShimFileFormat.scala
index 1702db135a..9506a851c0 100644
--- a/spark/src/main/spark-4.x/org/apache/comet/shims/ShimFileFormat.scala
+++ b/spark/src/main/spark-4.0/org/apache/comet/shims/ShimFileFormat.scala
@@ -19,9 +19,11 @@
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
+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
@@ -30,4 +32,18 @@ object ShimFileFormat {
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.1/org/apache/comet/shims/ShimFileFormat.scala
similarity index 62%
copy from spark/src/main/spark-4.x/org/apache/comet/shims/ShimFileFormat.scala
copy to spark/src/main/spark-4.1/org/apache/comet/shims/ShimFileFormat.scala
index 1702db135a..9506a851c0 100644
--- a/spark/src/main/spark-4.x/org/apache/comet/shims/ShimFileFormat.scala
+++ b/spark/src/main/spark-4.1/org/apache/comet/shims/ShimFileFormat.scala
@@ -19,9 +19,11 @@
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
+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
@@ -30,4 +32,18 @@ object ShimFileFormat {
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 63%
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 1702db135a..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
@@ -19,9 +19,11 @@
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
+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
@@ -30,4 +32,17 @@ object ShimFileFormat {
def findRowIndexColumnIndexInSchema(sparkSchema: StructType): Int =
ParquetRowIndexUtil.findRowIndexColumnIndexInSchema(sparkSchema)
+
+ def fileConstantMetadataExtractors(
+ 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],
+ dataType: DataType): Literal =
+ FileFormat.getFileConstantMetadataColumnValue(name, file, extractors,
dataType)
}
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..684c4d6581 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,196 @@ 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)
+ }
+ }
+
+ /**
+ * 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
+ * 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)
+ 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") {
+ files = spark.read
+ .parquet(dir.getCanonicalPath)
+ .select($"_metadata.file_path", $"_metadata.file_size")
+ .distinct()
+ .as[(String, Long)]
+ .collect()
+ .sortBy(_._2)
+ }
+ }
+ files
+ }
+
+ 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)
+
+ withSQLConf(forceSinglePartitionConf) {
+ val df = spark.read
+ .parquet(dir.getCanonicalPath)
+ .filter($"_metadata.file_size" === smallerSize)
+ .select($"id")
+ checkSparkAnswerAndOperator(df)
+ }
+ }
+ }
+
+ 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)
+ }
+ }
+ }
+
+ test("filter combining _metadata column and a data column") {
+ withTempPath { dir =>
+ val files = writeTwoFilesAndDiscoverMetadata(dir)
+ val (_, largerSize) = files(1)
+
+ withSQLConf(forceSinglePartitionConf) {
+ val df = spark.read
+ .parquet(dir.getCanonicalPath)
+ .filter($"_metadata.file_size" === largerSize && $"id" > 1500)
+ .select($"id")
+ checkSparkAnswerAndOperator(df)
+ }
+ }
+ }
+
+ test("filter on _metadata.file_path exact match") {
+ withTempPath { dir =>
+ 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)
+ }
+ }
+ }
+
+ 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)
+ }
+ }
+
+ 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($"pcol" === "b" && $"_metadata.file_size" > 0)
+ .select($"id")
+ 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)
+ }
+ }
+
test("fix: string partition column with incorrect offset buffer") {
def makeRawParquetFile(
path: Path,
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..fb39e340d1 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,24 @@ class CometScanRuleSuite extends CometTestBase {
}
}
+ test("CometScanRule should fallback to Spark for unsupported _metadata
columns") {
+ withTempPath { path =>
+ createTestDataFrame.write.parquet(path.toString)
+
+ // 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)
+ }
+ }
+ }
+
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
