hudi-agent commented on code in PR #19406:
URL: https://github.com/apache/hudi/pull/19406#discussion_r3685725145
##########
hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/common/TestNestedSchemaPruningOptimization.scala:
##########
@@ -167,6 +143,134 @@ class TestNestedSchemaPruningOptimization extends
HoodieSparkSqlTestBase with Sp
}
}
+ test("Test NestedSchemaPruning prunes nested struct when array and map
columns are present") {
+ withTempDir { tmp =>
+ val tableName = generateTableName
+ val tablePath = s"${tmp.getCanonicalPath}/$tableName"
+
+ createTableWithComplexNestedSchema(tableName, tablePath)
+
+ // Only a single nested sub-field is projected, so "item" is pruned down
to just "name" and the
+ // unreferenced "tags" (array<struct>) and "props" (map<string,struct>)
columns are dropped.
+ // Traversing the full data schema exercises the array and map branches
of countLeaves.
+ val selectDF = spark.sql(s"SELECT id, item.name FROM $tableName")
+
+ val expectedSchema = StructType(Seq(
+ StructField("id", IntegerType, nullable = true),
+ StructField("item", StructType(Seq(StructField("name", StringType,
nullable = false))), nullable = true)
+ ))
+ val expectedReadSchemaClause = "ReadSchema:
struct<id:int,item:struct<name:string>>"
+
+ assertPrunedReadSchema(selectDF, tableName, expectedSchema,
expectedReadSchemaClause)
+
+ // Execute the query to make sure it's working as expected (smoke test)
+ selectDF.count
+ }
+ }
+
+ test("Test NestedSchemaPruning is a no-op when all nested sub-fields are
selected") {
+ withTempDir { tmp =>
+ val tableName = generateTableName
+ val tablePath = s"${tmp.getCanonicalPath}/$tableName"
+
+ createTableWithNestedStructSchema("mor", tableName, tablePath)
+
+ // Every leaf is projected, so the pruned schema has the same leaf count
as the data schema and
+ // the rule leaves "item" untouched (the countLeaves comparison is an
equality, not a >).
+ val selectDF = spark.sql(s"SELECT id, item.name, item.price, ts FROM
$tableName")
+
+ val expectedItemStruct = StructType(Seq(
+ StructField("name", StringType, nullable = false),
+ StructField("price", IntegerType, nullable = false)
+ ))
+ assertEquals(expectedItemStruct, nestedFieldOf(selectDF, "item"))
+
+ selectDF.count
+ }
+ }
+
+ test("Test NestedSchemaPruning is a no-op when nested schema pruning is
disabled") {
+ withTempDir { tmp =>
+ val tableName = generateTableName
+ val tablePath = s"${tmp.getCanonicalPath}/$tableName"
+
+ createTableWithNestedStructSchema("mor", tableName, tablePath)
+
+ // With the optimizer flag off the rule short-circuits, so "item" keeps
"price" even though only
+ // "item.name" is projected.
+ val expectedItemStruct = StructType(Seq(
+ StructField("name", StringType, nullable = false),
+ StructField("price", IntegerType, nullable = false)
+ ))
+
+ withSQLConf(SQLConf.NESTED_SCHEMA_PRUNING_ENABLED.key -> "false") {
+ val selectDF = spark.sql(s"SELECT id, item.name FROM $tableName")
+ assertEquals(expectedItemStruct, nestedFieldOf(selectDF, "item"))
+ selectDF.count
+ }
+ }
+ }
+
+ private def assertPrunedReadSchema(selectDF: DataFrame,
+ tableName: String,
+ expectedSchema: StructType,
+ expectedReadSchemaClause: String,
+ hint: String = ""): Unit = {
+ // NOTE: We're disabling WCE to simplify resulting plan
+ spark.sessionState.conf.setConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED, false)
+
+ // NOTE: Unfortunately, we can't use pattern-matching to extract required
fields, due to a need to maintain
+ // compatibility w/ Spark 2.4
+ selectDF.queryExecution.executedPlan match {
+ // COW
+ case ProjectExec(_, fileScan: FileSourceScanExec) =>
+ assertEquals(tableName, fileScan.tableIdentifier.get.table)
+ assertEquals(expectedSchema, fileScan.requiredSchema, hint)
+
+ // MOR
+ case ProjectExec(_, dataScan: RowDataSourceScanExec) =>
+ // NOTE: This is temporary solution to assert for Spark 2.4, until
it's deprecated
+ val explainedPlan = explain(selectDF.queryExecution.logical)
+ assertTrue(explainedPlan.contains(expectedReadSchemaClause))
+
+ assertEquals(tableName, dataScan.tableIdentifier.get.table)
Review Comment:
🤖 nit: `nestedFieldOf` reads as "give me a field belonging to X", but it
actually inspects the physical plan and returns a `StructType`. Something like
`prunedStructTypeOf` or `requiredStructTypeOf` would make the return type and
the plan-inspection nature immediately obvious.
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag
quality.</i></sub>
##########
hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/common/TestNestedSchemaPruningOptimization.scala:
##########
@@ -167,6 +143,134 @@ class TestNestedSchemaPruningOptimization extends
HoodieSparkSqlTestBase with Sp
}
}
+ test("Test NestedSchemaPruning prunes nested struct when array and map
columns are present") {
+ withTempDir { tmp =>
+ val tableName = generateTableName
+ val tablePath = s"${tmp.getCanonicalPath}/$tableName"
+
+ createTableWithComplexNestedSchema(tableName, tablePath)
+
+ // Only a single nested sub-field is projected, so "item" is pruned down
to just "name" and the
+ // unreferenced "tags" (array<struct>) and "props" (map<string,struct>)
columns are dropped.
+ // Traversing the full data schema exercises the array and map branches
of countLeaves.
+ val selectDF = spark.sql(s"SELECT id, item.name FROM $tableName")
+
+ val expectedSchema = StructType(Seq(
+ StructField("id", IntegerType, nullable = true),
+ StructField("item", StructType(Seq(StructField("name", StringType,
nullable = false))), nullable = true)
+ ))
+ val expectedReadSchemaClause = "ReadSchema:
struct<id:int,item:struct<name:string>>"
+
+ assertPrunedReadSchema(selectDF, tableName, expectedSchema,
expectedReadSchemaClause)
+
+ // Execute the query to make sure it's working as expected (smoke test)
+ selectDF.count
+ }
+ }
+
+ test("Test NestedSchemaPruning is a no-op when all nested sub-fields are
selected") {
+ withTempDir { tmp =>
+ val tableName = generateTableName
+ val tablePath = s"${tmp.getCanonicalPath}/$tableName"
+
+ createTableWithNestedStructSchema("mor", tableName, tablePath)
+
+ // Every leaf is projected, so the pruned schema has the same leaf count
as the data schema and
+ // the rule leaves "item" untouched (the countLeaves comparison is an
equality, not a >).
+ val selectDF = spark.sql(s"SELECT id, item.name, item.price, ts FROM
$tableName")
+
+ val expectedItemStruct = StructType(Seq(
+ StructField("name", StringType, nullable = false),
+ StructField("price", IntegerType, nullable = false)
+ ))
+ assertEquals(expectedItemStruct, nestedFieldOf(selectDF, "item"))
+
+ selectDF.count
+ }
+ }
+
+ test("Test NestedSchemaPruning is a no-op when nested schema pruning is
disabled") {
+ withTempDir { tmp =>
+ val tableName = generateTableName
+ val tablePath = s"${tmp.getCanonicalPath}/$tableName"
+
+ createTableWithNestedStructSchema("mor", tableName, tablePath)
+
+ // With the optimizer flag off the rule short-circuits, so "item" keeps
"price" even though only
+ // "item.name" is projected.
+ val expectedItemStruct = StructType(Seq(
+ StructField("name", StringType, nullable = false),
+ StructField("price", IntegerType, nullable = false)
+ ))
+
+ withSQLConf(SQLConf.NESTED_SCHEMA_PRUNING_ENABLED.key -> "false") {
+ val selectDF = spark.sql(s"SELECT id, item.name FROM $tableName")
+ assertEquals(expectedItemStruct, nestedFieldOf(selectDF, "item"))
+ selectDF.count
+ }
+ }
+ }
+
+ private def assertPrunedReadSchema(selectDF: DataFrame,
+ tableName: String,
+ expectedSchema: StructType,
+ expectedReadSchemaClause: String,
+ hint: String = ""): Unit = {
+ // NOTE: We're disabling WCE to simplify resulting plan
+ spark.sessionState.conf.setConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED, false)
+
+ // NOTE: Unfortunately, we can't use pattern-matching to extract required
fields, due to a need to maintain
+ // compatibility w/ Spark 2.4
+ selectDF.queryExecution.executedPlan match {
+ // COW
+ case ProjectExec(_, fileScan: FileSourceScanExec) =>
+ assertEquals(tableName, fileScan.tableIdentifier.get.table)
+ assertEquals(expectedSchema, fileScan.requiredSchema, hint)
+
+ // MOR
+ case ProjectExec(_, dataScan: RowDataSourceScanExec) =>
+ // NOTE: This is temporary solution to assert for Spark 2.4, until
it's deprecated
+ val explainedPlan = explain(selectDF.queryExecution.logical)
+ assertTrue(explainedPlan.contains(expectedReadSchemaClause))
+
+ assertEquals(tableName, dataScan.tableIdentifier.get.table)
+ assertEquals(expectedSchema, dataScan.requiredSchema, hint)
Review Comment:
🤖 nit: disabling `WHOLESTAGE_CODEGEN` inside what looks like a pure
schema-accessor is a surprising side effect -- could you move this `setConf`
call to the test setup or wrap it in `withSQLConf` at the call site, so the
method body is a straightforward plan inspection?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag
quality.</i></sub>
--
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]
