szehon-ho commented on code in PR #58296:
URL: https://github.com/apache/spark/pull/58296#discussion_r3867696966
##########
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:
Done -- `InMemoryTableWithV2Filter` and `InMemoryCatalystRuntimeFilterTable`
now take the full table-creation metadata (constraints, distribution, ordering,
partition counts, advisory size, strictness) like `InMemoryTable`, and
`InMemoryCatalystRuntimeFilterCatalog.createTable` forwards all of it.
##########
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:
Done -- partition columns now keep their `fieldNames` as an unflattened
`Seq[String]`, and `partitionAttrFor` compares path parts component-wise with
the resolver, so a quoted top-level `` `a.b` `` no longer collides with nested
`a`.`b`. Added a test (`dotted top-level and nested partition columns -> bound
to the correct partition slot`) covering the collision.
##########
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:
Done -- `runtimeFilterAttrs` and `fullyPushedRuntimeFilterAttrs` now reject
any multi-part reference, including one over a struct, instead of widening. The
assertion now expects the rejection for both the int-parent and struct-parent
cases, and the `SupportsRuntimeCatalystFiltering` doc is updated to match.
##########
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:
Done. I removed the temporal word so the comment states the current
invariant rather than a before/after transition.
##########
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:
Done.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala:
##########
@@ -398,3 +579,73 @@ private class BothRuntimeFilteringInterfacesScan
override def filter(expressions: Array[Expression]): Unit = {}
}
+
+/** A scan declaring a filter attribute the read schema does not carry. */
+private class MissingFilterAttributeScan extends Scan with
SupportsRuntimeCatalystFiltering {
+
+ override def readSchema(): StructType = new StructType().add("part",
IntegerType)
+
+ override def filterAttributes(): Array[NamedReference] =
Array(FieldReference("missing"))
+
+ override def filter(expressions: Array[Expression]): Unit = {}
+}
+
+/**
+ * A scan breaking the rule that a filter attribute must be a top level read
schema column: it
+ * reports `part.nested` over the int column `part`, so resolving it fails on
the extract base.
+ */
+private class NestedFilterAttributeScan extends Scan with
SupportsRuntimeCatalystFiltering {
+
+ override def readSchema(): StructType = new StructType().add("part",
IntegerType)
+
+ override def filterAttributes(): Array[NamedReference] =
+ Array(FieldReference(Seq("part", "nested")))
+
+ override def filter(expressions: Array[Expression]): Unit = {}
+}
+
+/**
+ * A scan breaking the same rule over a struct column: it reports `s.tz` where
`s` is a struct, so
+ * resolving it succeeds and widens to `s` rather than failing.
+ */
+private class StructNestedFilterAttributeScan extends Scan with
SupportsRuntimeCatalystFiltering {
+
+ override def readSchema(): StructType =
+ new StructType().add("s", new StructType().add("tz", StringType))
+
+ override def filterAttributes(): Array[NamedReference] =
Array(FieldReference(Seq("s", "tz")))
+
+ override def filter(expressions: Array[Expression]): Unit = {}
+}
+
+private case class KeyedInputPartition(key: Int) extends InputPartition with
HasPartitionKey {
+ override def partitionKey(): InternalRow = InternalRow(key)
+}
+
+/**
+ * A scan reporting one set of partitions before filtering and another after,
so it can break the
+ * requirement to preserve the partitioning it originally reported.
+ */
+private class PartitioningBreakingScan(
+ initialPartitions: Seq[InputPartition],
+ afterFilter: Seq[InputPartition])
+ extends Scan with Batch with SupportsRuntimeCatalystFiltering {
+
+ private var filtered = false
+
+ override def readSchema(): StructType = new StructType().add("part",
IntegerType)
+
+ override def toBatch: Batch = this
+
+ override def planInputPartitions(): Array[InputPartition] =
+ if (filtered) afterFilter.toArray else initialPartitions.toArray
+
+ override def createReaderFactory(): PartitionReaderFactory =
+ throw new UnsupportedOperationException()
+
+ override def filterAttributes(): Array[NamedReference] =
Array(FieldReference("part"))
+
+ override def filter(expressions: Array[Expression]): Unit = {
+ filtered = true
+ }
+}
Review Comment:
Done -- added the trailing newline.
--
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]
