cloud-fan commented on code in PR #58296:
URL: https://github.com/apache/spark/pull/58296#discussion_r3859456715
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala:
##########
@@ -755,15 +756,51 @@ abstract class InMemoryBaseTable(
/** Predicates recorded by [[filter]], for test assertions only. */
def pushedCatalystPredicates: Seq[CatalystExpression] =
catalystPredicates.toSeq
- /** AttributeReferences matching the partition-key InternalRow field
order. */
+ def filterCallCount: Int = filterCalls
+
+ /**
+ * The `AttributeReference`s standing for the partition key InternalRow
fields, in its field
+ * order, each named by the dotted path of its partition column. Example:
+ * - `PARTITIONED BY (part, s.nested)` -> `AttributeReference(part)`,
then
+ * `AttributeReference(s.nested)`
+ */
private def partitionAttributes: Seq[AttributeReference] = {
partitioning.flatMap(_.references()).flatMap { ref =>
- val name = ref.fieldNames.mkString(".")
- readSchema.find(_.name == name).orElse(tableSchema.find(_.name ==
name)).map { f =>
- AttributeReference(f.name, f.dataType, f.nullable)()
+ val path = ref.fieldNames.toImmutableArraySeq
+
readSchema.findNestedField(path).orElse(tableSchema.findNestedField(path)).map {
+ case (_, f) =>
+ AttributeReference(ref.fieldNames.mkString("."), f.dataType,
f.nullable)()
Review Comment:
**Non-blocking:**
Please preserve `fieldNames` as parts here. Flattening them makes a quoted
top-level `a.b` indistinguishable from nested `a`.`b`; if both are partition
columns, `partitionAttrFor` can bind a predicate to the first, wrong
partition-key slot. Compare path parts component-wise with the resolver and add
that collision case.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala:
##########
@@ -247,26 +291,126 @@ class DataSourceV2CatalystRuntimeFilterSuite extends
SharedSparkSession {
}
}
- test("filter on column outside filterAttributes -> not pushed") {
+ test("filter on column outside filterAttributes -> not pushed, even if
declared fully pushed") {
val tbl = s"$catalogName.tbl4"
val dim = s"$catalogName.dim4"
withTable(tbl, dim) {
+ // p2 is a partition column but is not declared filterable, so no
runtime filter is derived
+ // for it. Declaring it fully pushed as well, which the interface
forbids for an attribute
+ // that is not filterable, must not cost it the post-scan filter:
nothing was pushed, so the
+ // scan prunes nothing and the nonmatching rows would come back.
sql(s"CREATE TABLE $tbl (id INT, p1 INT, p2 INT) USING $v2Source " +
"PARTITIONED BY (p1, p2) " +
- "TBLPROPERTIES('filter-attributes' = 'p1')")
+ "TBLPROPERTIES('filter-attributes' = 'p1',
'fully-pushed-filter-attributes' = 'p2')")
for (i <- 0 until 5) {
- sql(s"INSERT INTO $tbl VALUES ($i, $i, 10)")
+ sql(s"INSERT INTO $tbl VALUES ($i, $i, $i)")
}
sql(s"CREATE TABLE $dim (val INT) USING $v2Source")
- sql(s"INSERT INTO $dim VALUES (10)")
+ sql(s"INSERT INTO $dim VALUES (3)")
- // p2 is a partition column but is not declared filterable, so no
runtime filter is derived.
val df = sql(s"SELECT * FROM $tbl WHERE p2 = (SELECT max(val) FROM
$dim)")
- checkAnswer(df, (0 until 5).map(i => Row(i, i, 10)))
+ checkAnswer(df, Row(3, 3, 3))
assert(collectBatchScan(df).runtimeFilters.isEmpty,
"Expected no runtime filters for a column outside filterAttributes")
assertPushedCatalystPredicates(df, 0)
+ assertScalarSubqueryEvaluatedAfterScan(df, expected = true)
+ }
+ }
+
+ test("two predicates on filter attributes -> pushed together in a single
filter() call") {
+ val tbl = s"$catalogName.tbl_two_predicates"
+ val dim1 = s"$catalogName.dim_two_predicates1"
+ val dim2 = s"$catalogName.dim_two_predicates2"
+ withTable(tbl, dim1, dim2) {
+ sql(s"CREATE TABLE $tbl (id INT, p1 INT, p2 INT) USING $v2Source
PARTITIONED BY (p1, p2)")
+ for (i <- 0 until 5) {
+ sql(s"INSERT INTO $tbl VALUES ($i, $i, ${i * 10})")
+ }
+ sql(s"CREATE TABLE $dim1 (val INT) USING $v2Source")
+ sql(s"INSERT INTO $dim1 VALUES (3)")
+ sql(s"CREATE TABLE $dim2 (val INT) USING $v2Source")
+ sql(s"INSERT INTO $dim2 VALUES (30)")
+
+ val df = sql(s"SELECT * FROM $tbl WHERE p1 = (SELECT max(val) FROM
$dim1) " +
+ s"AND p2 = (SELECT max(val) FROM $dim2)")
+ checkAnswer(df, Row(3, 3, 30))
+
+ assertScalarSubqueryRuntimeFilters(df, expectedCount = 2)
+ val p1 = AttributeReference("p1", IntegerType, nullable = false)()
+ val p2 = AttributeReference("p2", IntegerType, nullable = false)()
+ assertPushedCatalystPredicatesEqual(
+ df, EqualTo(p1, Literal(3)), EqualTo(p2, Literal(30)))
+ assert(getCatalystScan(df).filterCallCount === 1,
+ "expected both predicates pushed in a single filter() call")
+ }
+ }
+
+ test("nested field of a filter attribute -> pushed with the nested access
intact") {
+ val tbl = s"$catalogName.tbl_nested"
+ val dim = s"$catalogName.dim_nested"
+ withTable(tbl, dim) {
+ sql(s"CREATE TABLE $tbl (id INT, s STRUCT<tz: STRING>) USING $v2Source "
+
+ "PARTITIONED BY (s.tz)")
+ for (i <- 0 until 3) {
+ sql(s"INSERT INTO $tbl VALUES ($i, named_struct('tz', 'tz$i'))")
+ }
+ sql(s"CREATE TABLE $dim (val STRING) USING $v2Source")
+ sql(s"INSERT INTO $dim VALUES ('tz1')")
+
+ // The scan declares the top-level struct column `s` as its filter
attribute, so the
+ // predicate qualifies for pushdown even though it reaches into `s.tz`.
Matching the nested
+ // access against the partition layout is left to the scan, which the
fixture does not do.
+ val df = sql(s"SELECT * FROM $tbl WHERE s.tz = (SELECT max(val) FROM
$dim)")
+ checkAnswer(df, Row(1, Row("tz1")))
+
+ assertScalarSubqueryRuntimeFilters(df)
+ val pushed = getPushedCatalystPredicates(df)
+ assert(pushed.size === 1, s"expected a single pushed predicate, got
$pushed")
+ val nestedAccesses = pushed.head.collect { case g: GetStructField => g }
+ assert(nestedAccesses.size === 1,
+ s"expected the pushed predicate to keep the nested access, got
${pushed.head}")
+ assert(nestedAccesses.head.childSchema.fieldNames.contains("tz"))
+ }
+ }
+
+ test("filterAttributes that is not a top-level scan attribute") {
+ val tbl = s"$catalogName.tbl_unresolvable_attr"
+ withTable(tbl) {
+ sql(s"CREATE TABLE $tbl (id INT, part INT, s STRUCT<tz: STRING>) USING
$v2Source " +
+ "PARTITIONED BY (part)")
+ sql(s"INSERT INTO $tbl VALUES (1, 1, named_struct('tz', 'a'))")
+
+ val scanRelation = sql(s"SELECT * FROM
$tbl").queryExecution.optimizedPlan.collectFirst {
+ case r: DataSourceV2ScanRelation => r
+ }.getOrElse(fail("Expected a DataSourceV2ScanRelation"))
+
+ // An attribute the read schema does not carry, such as one pruned out
of the projection.
+ val missing = intercept[AnalysisException] {
+ scanRelation.copy(scan = new
MissingFilterAttributeScan).runtimeFilterAttrs
+ }
+ checkError(
+ exception = missing,
+ condition = "_LEGACY_ERROR_TEMP_1137",
+ parameters = Map("name" -> "missing", "outputStr" -> "id,part,s"))
+
+ // A nested reference is resolved as a field extraction on the top-level
attribute, which
+ // fails as soon as that attribute is not a struct.
+ val nested = intercept[AnalysisException] {
+ scanRelation.copy(scan = new
NestedFilterAttributeScan).runtimeFilterAttrs
+ }
+ checkError(
+ exception = nested,
+ condition = "INVALID_EXTRACT_BASE_FIELD_TYPE",
+ parameters = Map("base" -> "\"part\"", "other" -> "\"INT\""))
+
+ // Over a struct it does not fail at all: the reference widens to the
struct column, so
+ // declaring `s.tz` makes filters over every field of `s` eligible, not
just `s.tz`. This
+ // is shared with the two predicate-based interfaces, which resolve the
same way.
+ val widened = scanRelation.copy(scan = new
StructNestedFilterAttributeScan)
+ .runtimeFilterAttrs
+ assert(widened.map(_.name).toSeq === Seq("s"),
Review Comment:
**Non-blocking:**
`filterAttributes()` requires top-level read-schema attributes. Please make
`runtimeFilterAttrs` reject nested references even when the parent is a struct
and change this assertion to expect that error; accepting `s.tz` here widens it
to `s` and makes filters over every field eligible.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalystRuntimeFilterCatalog.scala:
##########
@@ -48,3 +50,40 @@ class InMemoryTableCatalystRuntimeFilterCatalog extends
InMemoryTableCatalog {
createTable(ident, tableInfo.columns(), tableInfo.partitions(),
tableInfo.properties)
}
}
+
+/**
+ * The [[InMemoryCatalog]] counterpart of
[[InMemoryTableCatalystRuntimeFilterCatalog]]: it hands
+ * out tables whose scans take runtime filters as Catalyst expressions, and
honors
+ * `numRowsPerSplit` so that a partition key can have several splits.
+ */
+class InMemoryCatalystRuntimeFilterCatalog extends InMemoryCatalog {
+ import CatalogV2Implicits._
+
+ // scalastyle:off argcount
+ override def createTable(
+ ident: Identifier,
+ columns: Array[Column],
+ partitions: Array[Transform],
+ properties: util.Map[String, String],
+ distribution: Distribution,
+ ordering: Array[SortOrder],
+ requiredNumPartitions: Option[Int],
+ advisoryPartitionSize: Option[Long],
+ constraints: Array[Constraint],
+ distributionStrictlyRequired: Boolean,
+ numRowsPerSplit: Int): Table = {
+ // scalastyle:on argcount
+ if (tables.containsKey(ident)) {
+ throw new TableAlreadyExistsException(ident.asMultipartIdentifier)
+ }
+
+ InMemoryTableCatalog.maybeSimulateFailedTableCreation(properties)
+
+ val tableName = s"$name.${ident.quoted}"
+ val table = new InMemoryCatalystRuntimeFilterTable(
+ tableName, columns, partitions, properties, numRowsPerSplit)
Review Comment:
**Non-blocking:**
This override accepts distribution, ordering, partition counts, advisory
size, constraints, and strictness, but this constructor call drops all of them.
Please thread the full metadata through the Catalyst table constructors, as
`InMemoryTableCatalog` does, so shared suites do not silently exercise defaults.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -47,8 +47,285 @@ import org.apache.spark.sql.internal.SQLConf._
import org.apache.spark.sql.types._
import org.apache.spark.tags.ExtendedSQLTest
+abstract class KeyGroupedPartitioningSuiteBase extends
DistributionAndOrderingSuiteBase {
+
+ protected val emptyProps: java.util.Map[String, String] = {
+ Collections.emptyMap[String, String]
+ }
+
+ protected val items: String = "items"
+ protected val itemsColumns: Array[Column] = Array(
+ Column.create("id", LongType),
+ Column.create("name", StringType),
+ Column.create("price", FloatType),
+ Column.create("arrive_time", TimestampType))
+
+ protected val purchases: String = "purchases"
+ protected val purchasesColumns: Array[Column] = Array(
+ Column.create("item_id", LongType),
+ Column.create("price", FloatType),
+ Column.create("time", TimestampType))
+
+ protected def createTable(
+ table: String,
+ columns: Array[Column],
+ partitions: Array[Transform],
+ ordering: Array[SortOrder] = Array.empty,
+ catalog: InMemoryTableCatalog = catalog): Unit = {
+ catalog.createTable(Identifier.of(Array("ns"), table),
+ columns, partitions, emptyProps, Distributions.unspecified(), ordering,
None, None,
+ numRowsPerSplit = 1)
+ }
+
+ protected def collectShuffles(plan: SparkPlan): Seq[ShuffleExchangeLike] = {
+ // here we skip collecting shuffle operators that are not associated with
SMJ
+ collect(plan) {
+ case s: SortMergeJoinExec => s
+ }.flatMap(smj =>
+ collect(smj) {
+ case s: ShuffleExchangeExec => s
+ })
+ }.toSet.toSeq
+
+ protected def collectGroupPartitions(plan: SparkPlan):
Seq[GroupPartitionsExec] = {
+ // here we skip collecting shuffle operators that are not associated with
SMJ
Review Comment:
**Nit:**
This helper returns `GroupPartitionsExec` nodes, not shuffle operators.
```suggestion
// here we skip collecting group-partition operators that are not
associated with SMJ
```
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala:
##########
@@ -247,26 +291,126 @@ class DataSourceV2CatalystRuntimeFilterSuite extends
SharedSparkSession {
}
}
- test("filter on column outside filterAttributes -> not pushed") {
+ test("filter on column outside filterAttributes -> not pushed, even if
declared fully pushed") {
val tbl = s"$catalogName.tbl4"
val dim = s"$catalogName.dim4"
withTable(tbl, dim) {
+ // p2 is a partition column but is not declared filterable, so no
runtime filter is derived
+ // for it. Declaring it fully pushed as well, which the interface
forbids for an attribute
+ // that is not filterable, must not cost it the post-scan filter:
nothing was pushed, so the
+ // scan prunes nothing and the nonmatching rows would come back.
sql(s"CREATE TABLE $tbl (id INT, p1 INT, p2 INT) USING $v2Source " +
"PARTITIONED BY (p1, p2) " +
- "TBLPROPERTIES('filter-attributes' = 'p1')")
+ "TBLPROPERTIES('filter-attributes' = 'p1',
'fully-pushed-filter-attributes' = 'p2')")
for (i <- 0 until 5) {
- sql(s"INSERT INTO $tbl VALUES ($i, $i, 10)")
+ sql(s"INSERT INTO $tbl VALUES ($i, $i, $i)")
}
sql(s"CREATE TABLE $dim (val INT) USING $v2Source")
- sql(s"INSERT INTO $dim VALUES (10)")
+ sql(s"INSERT INTO $dim VALUES (3)")
- // p2 is a partition column but is not declared filterable, so no
runtime filter is derived.
val df = sql(s"SELECT * FROM $tbl WHERE p2 = (SELECT max(val) FROM
$dim)")
- checkAnswer(df, (0 until 5).map(i => Row(i, i, 10)))
+ checkAnswer(df, Row(3, 3, 3))
assert(collectBatchScan(df).runtimeFilters.isEmpty,
"Expected no runtime filters for a column outside filterAttributes")
assertPushedCatalystPredicates(df, 0)
+ assertScalarSubqueryEvaluatedAfterScan(df, expected = true)
+ }
+ }
+
+ test("two predicates on filter attributes -> pushed together in a single
filter() call") {
+ val tbl = s"$catalogName.tbl_two_predicates"
+ val dim1 = s"$catalogName.dim_two_predicates1"
+ val dim2 = s"$catalogName.dim_two_predicates2"
+ withTable(tbl, dim1, dim2) {
+ sql(s"CREATE TABLE $tbl (id INT, p1 INT, p2 INT) USING $v2Source
PARTITIONED BY (p1, p2)")
+ for (i <- 0 until 5) {
+ sql(s"INSERT INTO $tbl VALUES ($i, $i, ${i * 10})")
+ }
+ sql(s"CREATE TABLE $dim1 (val INT) USING $v2Source")
+ sql(s"INSERT INTO $dim1 VALUES (3)")
+ sql(s"CREATE TABLE $dim2 (val INT) USING $v2Source")
+ sql(s"INSERT INTO $dim2 VALUES (30)")
+
+ val df = sql(s"SELECT * FROM $tbl WHERE p1 = (SELECT max(val) FROM
$dim1) " +
+ s"AND p2 = (SELECT max(val) FROM $dim2)")
+ checkAnswer(df, Row(3, 3, 30))
+
+ assertScalarSubqueryRuntimeFilters(df, expectedCount = 2)
+ val p1 = AttributeReference("p1", IntegerType, nullable = false)()
+ val p2 = AttributeReference("p2", IntegerType, nullable = false)()
+ assertPushedCatalystPredicatesEqual(
+ df, EqualTo(p1, Literal(3)), EqualTo(p2, Literal(30)))
+ assert(getCatalystScan(df).filterCallCount === 1,
+ "expected both predicates pushed in a single filter() call")
+ }
+ }
+
+ test("nested field of a filter attribute -> pushed with the nested access
intact") {
+ val tbl = s"$catalogName.tbl_nested"
+ val dim = s"$catalogName.dim_nested"
+ withTable(tbl, dim) {
+ sql(s"CREATE TABLE $tbl (id INT, s STRUCT<tz: STRING>) USING $v2Source "
+
+ "PARTITIONED BY (s.tz)")
+ for (i <- 0 until 3) {
+ sql(s"INSERT INTO $tbl VALUES ($i, named_struct('tz', 'tz$i'))")
+ }
+ sql(s"CREATE TABLE $dim (val STRING) USING $v2Source")
+ sql(s"INSERT INTO $dim VALUES ('tz1')")
+
+ // The scan declares the top-level struct column `s` as its filter
attribute, so the
+ // predicate qualifies for pushdown even though it reaches into `s.tz`.
Matching the nested
+ // access against the partition layout is left to the scan, which the
fixture does not do.
Review Comment:
**Nit:**
The current sentence describes the pre-change fixture behavior.
```suggestion
// access against the partition layout is left to the scan, which this
fixture now does.
```
--
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]
