sam-1112 commented on code in PR #7196:
URL: https://github.com/apache/datafusion-comet/pull/7196#discussion_r4369360873
##########
spark/src/test/scala/org/apache/comet/exec/CometNativeReaderSuite.scala:
##########
@@ -1269,6 +1270,373 @@ class CometNativeReaderSuite extends CometTestBase with
AdaptiveSparkPlanHelper
nullable = true,
new MetadataBuilder().putLong("parquet.field.id", id.toLong).build())
+ private def withFieldClaimFile(nested: Boolean)(f: (String, StructType) =>
Unit): Unit = {
+ val id10 = new MetadataBuilder().putLong("parquet.field.id", 10L).build()
+ def wrapSchema(fields: StructType): StructType =
+ if (nested) StructType(Seq(StructField("s", fields, nullable = true,
id10)))
+ else fields
+ val writeSchema = wrapSchema(StructType(Seq(withFieldId("A", 1))))
+ val readSchema = wrapSchema(
+ StructType(Seq(withFieldId("x", 1), StructField("a", LongType, nullable
= true))))
+ withSQLConf(
+ SQLConf.CASE_SENSITIVE.key -> "false",
+ SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key -> "true",
+ SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true") {
+ withTempPath { path =>
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ val rows = (1L to 3L).map(value => if (nested) Row(Row(value)) else
Row(value))
+ spark
+ .createDataFrame(spark.sparkContext.parallelize(rows, 1),
writeSchema)
+ .write
+ .parquet(path.getCanonicalPath)
+ }
+ f(path.getCanonicalPath, readSchema)
+ }
+ }
+ }
+
+ private def fieldClaimReference(query: => DataFrame): Seq[Row] = {
+ var rows = Seq.empty[Row]
+ withSQLConf(
+ CometConf.COMET_ENABLED.key -> "false",
+ SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "false",
+ SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "false") {
+ val df = query
+ val scans = collect(df.queryExecution.executedPlan) { case scan:
FileSourceScanExec =>
+ scan
+ }
+ assert(scans.nonEmpty && scans.forall(!_.supportsColumnar))
+ rows = df.collect().toSeq
+ }
+ rows
+ }
+
+ private def checkNativeFieldClaimAnswer(df: DataFrame, expected: Seq[Row]):
Unit = {
+ assert(
+ collect(df.queryExecution.executedPlan) { case scan: CometNativeScanExec
=> scan }.nonEmpty,
+ s"expected native field claim
resolution:\n${df.queryExecution.executedPlan}")
+ checkAnswer(df, expected)
+ }
+
+ for {
+ nested <- Seq(false, true)
+ pruning <- (if (nested) Seq(false, true) else Seq(false))
+ selected <- Seq("both", "id", "name")
+ } {
+ test(
+ s"Parquet field claims projection - nested=$nested pruning=$pruning
selected=$selected") {
+ withFieldClaimFile(nested) { (path, schema) =>
+ withSQLConf(SQLConf.NESTED_SCHEMA_PRUNING_ENABLED.key ->
pruning.toString) {
+ def query: DataFrame = {
+ val df = spark.read.schema(schema).parquet(path)
+ val prefix = if (nested) "s." else ""
+ val names = selected match {
+ case "both" => Seq("x", "a")
+ case "id" => Seq("x")
+ case "name" => Seq("a")
+ }
+ df.select(names.map(name => col(prefix + name)): _*)
+ }
+ val expected = fieldClaimReference(query)
+ checkNativeFieldClaimAnswer(query, expected)
+ }
+ }
+ }
+ }
+
+ for {
+ nested <- Seq(false, true)
+ pushdown <- Seq(false, true)
+ predicate <- Seq("missing is null", "missing is not null", "id greater
than one")
+ } {
+ test(
+ s"Parquet field claims filters - nested=$nested pushdown=$pushdown
predicate=$predicate") {
+ withFieldClaimFile(nested) { (path, schema) =>
+ withSQLConf(
+ SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.key ->
pushdown.toString) {
+ def query: DataFrame = {
+ val df = spark.read.schema(schema).parquet(path)
+ val prefix = if (nested) "s." else ""
+ val x = col(prefix + "x")
+ val a = col(prefix + "a")
+ val condition = predicate match {
+ case "missing is null" => a.isNull
+ case "missing is not null" => a.isNotNull
+ case "id greater than one" => x > 1L
+ }
+ df.where(condition).select(x, a)
+ }
+ val expected = fieldClaimReference(query)
+ val values = predicate match {
+ case "missing is null" => Seq(1L, 2L, 3L)
+ case "missing is not null" => Seq.empty[Long]
+ case "id greater than one" => Seq(2L, 3L)
+ }
+ assert(expected.sortBy(_.toString) == values.map(value => Row(value,
null)))
+ val df = query
+ checkNativeFieldClaimAnswer(df, expected)
+ if (!nested && pushdown && predicate == "id greater than one") {
+ val scans = collect(df.queryExecution.executedPlan) {
+ case scan: CometNativeScanExec =>
+ scan
+ }
+ assert(
+ scans.exists(scan =>
+ scan.metrics("pushdown_rows_pruned").value +
+ scan.metrics("pushdown_rows_matched").value > 0))
+ }
+ }
+ }
+ }
+ }
+
+ for {
+ nested <- Seq(false, true)
+ reverse <- Seq(false, true)
+ kind <- Seq("struct", "array", "map")
+ } {
+ test(s"Parquet field claims complex source - nested=$nested
reverse=$reverse kind=$kind") {
+ val (dataType, values): (DataType, Seq[Any]) = kind match {
+ case "struct" =>
+ (new StructType().add("v", LongType), Seq(Row(11L), Row(null), null))
+ case "array" =>
+ (
+ ArrayType(LongType, containsNull = true),
+ Seq(Seq(11L, null, 33L), Seq.empty[Long], null))
+ case "map" =>
+ (
+ MapType(StringType, LongType, valueContainsNull = true),
+ Seq(Map("k" -> 11L, "nil" -> null), Map.empty[String, Long], null))
+ }
+ val id1 = new MetadataBuilder().putLong("parquet.field.id", 1L).build()
+ val id10 = new MetadataBuilder().putLong("parquet.field.id", 10L).build()
+ def wrapSchema(fields: StructType): StructType =
+ if (nested) StructType(Seq(StructField("s", fields, nullable = true,
id10))) else fields
+ def wrapRow(fields: Row): Row = if (nested) Row(fields) else fields
+ val writeSchema =
+ wrapSchema(StructType(Seq(StructField("A", dataType, nullable = true,
id1))))
+ val fields = Seq(
+ StructField("x", dataType, nullable = true, id1),
+ StructField("a", dataType, nullable = true))
+ val readSchema = wrapSchema(StructType(if (reverse) fields.reverse else
fields))
+ val expected =
+ values.map(value => wrapRow(if (reverse) Row(null, value) else
Row(value, null)))
+ withSQLConf(
+ SQLConf.CASE_SENSITIVE.key -> "false",
+ SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key -> "true",
+ SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true") {
+ withTempPath { path =>
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark
+ .createDataFrame(
+ spark.sparkContext.parallelize(values.map(value =>
wrapRow(Row(value))), 1),
+ writeSchema)
+ .write
+ .parquet(path.getCanonicalPath)
+ }
+ def query: DataFrame =
spark.read.schema(readSchema).parquet(path.getCanonicalPath)
+ withSQLConf(
+ CometConf.COMET_ENABLED.key -> "false",
+ SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "false") {
+ val df = query
+ val scans = collect(df.queryExecution.executedPlan) { case scan:
FileSourceScanExec =>
+ scan
+ }
+ assert(scans.nonEmpty && scans.forall(!_.supportsColumnar))
+ checkAnswer(df, expected)
+ }
+ checkNativeFieldClaimAnswer(query, expected)
+ }
+ }
+ }
+ }
+
+ for {
+ nested <- Seq(false, true)
+ reverse <- Seq(false, true)
+ count <- Seq(2, 3, 4)
+ } {
+ test(
+ s"Parquet field claims - repeated requested Parquet ids match parquet-mr
- nested=$nested reverse=$reverse count=$count") {
+ val id10 = new MetadataBuilder().putLong("parquet.field.id", 10L).build()
+ def wrapSchema(fields: StructType): StructType =
+ if (nested) StructType(Seq(StructField("s", fields, nullable = true,
id10)))
+ else fields
+ def wrapRow(fields: Row): Row = if (nested) Row(fields) else fields
+
+ val fields = Seq("x", "y", "z", "w").take(count).map(name =>
withFieldId(name, 1))
+ val readSchema = wrapSchema(StructType(if (reverse) fields.reverse else
fields))
+ val writeSchema = wrapSchema(StructType(Seq(withFieldId("A", 1))))
+ val expected = Seq(11L, 22L, 33L).map(value =>
+ wrapRow(Row.fromSeq(Seq.fill[Any](count - 1)(null) :+ value)))
+ withSQLConf(
+ SQLConf.CASE_SENSITIVE.key -> "false",
+ SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key -> "true",
+ SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true",
+ SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "false") {
+ withTempPath { path =>
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ val rows = Seq(11L, 22L, 33L).map(value => wrapRow(Row(value)))
+ spark
+ .createDataFrame(spark.sparkContext.parallelize(rows, 1),
writeSchema)
+ .write
+ .parquet(path.getCanonicalPath)
+ }
+ // ParquetRowConverter builds its requested-id-to-index map with
toMap, so the last
+ // field with the repeated id owns the value. Check the Spark result
in both orders.
+ val sparkRows = fieldClaimReference {
+ spark.read.schema(readSchema).parquet(path.getCanonicalPath)
+ }
+ assert(sparkRows.sortBy(_.toString) == expected.sortBy(_.toString))
+ withSQLConf(
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_SCAN_ALLOW_DISABLED_PARQUET_VECTORIZED_READER.key
-> "true") {
+ val df =
spark.read.schema(readSchema).parquet(path.getCanonicalPath)
+ // Keep the existing JVM guard: native matching is tested
separately by the Rust
+ // Parquet roundtrip, which bypasses this guard and exercises the
actual converter.
+ val plan = df.queryExecution.executedPlan
+ assert(collect(plan) { case scan: CometNativeScanExec =>
+ scan
+ }.isEmpty)
+ val reasons = new ExtendedExplainInfo().getFallbackReasons(plan)
+ assert(reasons.exists(_.contains("duplicate Parquet field ids")),
s"$reasons\n$plan")
+ assert(df.collect().toSeq.sortBy(_.toString) ==
sparkRows.sortBy(_.toString))
+ }
+ }
+ }
+ }
+ }
+
+ for {
+ nested <- Seq(false, true)
+ offheap <- Seq(false, true)
+ } {
+ test(
+ s"Parquet field claims - issue 6994 original reproduction -
nested=$nested offheap=$offheap") {
+ val id1 = new MetadataBuilder().putLong("parquet.field.id", 1L).build()
+ val id10 = new MetadataBuilder().putLong("parquet.field.id", 10L).build()
+ def wrapSchema(fields: StructType): StructType =
+ if (nested) StructType(Seq(StructField("s", fields, nullable = true,
id10)))
+ else fields
+ def wrapRow(fields: Row): Row = if (nested) Row(fields) else fields
+
+ val writeSchema =
+ wrapSchema(StructType(Seq(StructField("A", IntegerType, nullable =
true, id1))))
+ val readSchema = wrapSchema(
+ StructType(
+ Seq(
+ StructField("x", IntegerType, nullable = true, id1),
+ StructField("a", IntegerType, nullable = true))))
+ val expected = (1 to 3).map(value => wrapRow(Row(value, null)))
+
+ withSQLConf(
+ SQLConf.CASE_SENSITIVE.key -> "false",
+ SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key -> "true",
+ SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true",
+ "spark.sql.parquet.enableNestedColumnVectorizedReader" -> "true",
+ "spark.sql.columnVector.offheap.enabled" -> offheap.toString) {
+ withTempPath { path =>
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ val rows = (1 to 3).map(value => wrapRow(Row(value)))
+ spark
+ .createDataFrame(spark.sparkContext.parallelize(rows, 1),
writeSchema)
+ .write
+ .parquet(path.getCanonicalPath)
+ }
+ for ((reader, comet, vectorized) <- Seq(
+ ("Spark vectorized", false, true),
+ ("Spark parquet-mr", false, false),
+ ("Comet native", true, true))) {
+ withSQLConf(
+ CometConf.COMET_ENABLED.key -> comet.toString,
+ SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key ->
vectorized.toString) {
+ val df =
spark.read.schema(readSchema).parquet(path.getCanonicalPath)
+ val plan = df.queryExecution.executedPlan
+ val native = collect(plan) { case scan: CometNativeScanExec =>
scan }.nonEmpty
+ assert(native == comet, s"unexpected scan for $reader:\n$plan")
+ if (!comet) {
+ val scans = collect(plan) { case scan: FileSourceScanExec =>
scan }
+ assert(scans.nonEmpty && scans.forall(_.supportsColumnar ==
vectorized))
+ }
+ val rows = df.collect().toSeq
+ info(
+ s"Spark ${spark.version}, nested=$nested, offheap=$offheap,
$reader: ${rows.mkString(", ")}")
Review Comment:
The affected Spark vectorized branch now checks only the plan. Collection,
row formatting, and value assertions are inside the `if (comet || !vectorized)`
guard, so only Comet and parquet-mr results are collected and logged.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
