pan3793 commented on code in PR #7792:
URL: https://github.com/apache/kyuubi/pull/7792#discussion_r4396101072
##########
extensions/spark/kyuubi-spark-lineage/src/main/scala/org/apache/kyuubi/plugin/lineage/helper/SparkSQLLineageParseHelper.scala:
##########
@@ -455,6 +456,17 @@ trait LineageParser {
nextColumnsLineage,
inputTablesByPlan)).reduce(mergeColumnsLineage)
+ case p: Filter =>
Review Comment:
This only collects subquery input tables from `Filter` conditions.
Subqueries elsewhere still go through `case p =>`, which only visits
`p.children`, so their tables are still missing with legacy=false, e.g.
1. `INSERT OVERWRITE tb3 SELECT t1.col1, t1.col2 IN (SELECT col2 FROM tb2)
FROM tb1 t1` (`getExpressionSubqueryPlans` only matches `ScalarSubquery`)
2. `... JOIN t2 ON t1.col1 = t2.col1 AND EXISTS (SELECT 1 FROM tb4 ...)`
3. `INSERT INTO tb3 REPLACE WHERE col1 IN (SELECT col1 FROM tb2) SELECT ...`
So the same subquery reports `tb2` in `WHERE` but not in the select list.
Could we collect subquery tables generically instead, e.g. visit
`plan.subqueries` (or `plan.expressions.flatMap(_.collect { case s:
SubqueryExpression => s.plan })`) once in `extractColumnsLineage` for every
node? Then `Filter` can fall back to the default path, and
`getExpressionSubqueryPlans` stays only for the column lineage of select-list
scalar subqueries.
##########
extensions/spark/kyuubi-spark-lineage/src/test/scala/org/apache/kyuubi/plugin/lineage/helper/SparkSQLLineageParserHelperSuite.scala:
##########
@@ -1494,6 +1494,163 @@ abstract class SparkSQLLineageParserHelperSuite extends
KyuubiFunSuite
}
}
+ test("columns lineage extract - collect input tables from a correlated
EXISTS filter") {
Review Comment:
This test repeats the parameterized `EXISTS` legacy=false case below. Could
we drop it?
##########
extensions/spark/kyuubi-spark-lineage/src/main/scala/org/apache/kyuubi/plugin/lineage/helper/SparkSQLLineageParseHelper.scala:
##########
@@ -455,6 +456,17 @@ trait LineageParser {
nextColumnsLineage,
inputTablesByPlan)).reduce(mergeColumnsLineage)
+ case p: Filter =>
+ p.condition.foreach {
+ case expression: SubqueryExpression =>
+ extractColumnsLineage(
Review Comment:
Two things here:
1. This runs even when `LEGACY_COLLECT_INPUT_TABLES_ENABLED=true`, where
`inputTablesByPlan` is ignored. Any exception from code that is newly reached
here (e.g. `identifier.get` in `getV2TableName`, an unmatched `NamedExpression`
in `getSelectColumnLineage`) makes `transformToLineage` return `None`, so a
statement that used to get lineage now gets none, in legacy mode too. Please
skip this when legacy is enabled.
2. The return value is ignored on purpose (we only want the input tables;
predicate columns must not feed column lineage). Please add a short comment
saying so.
##########
extensions/spark/kyuubi-spark-lineage/src/main/scala/org/apache/kyuubi/plugin/lineage/helper/SparkSQLLineageParseHelper.scala:
##########
@@ -144,6 +144,7 @@ trait LineageParser {
if (!containsCountAll(exp.child)) references
else references +
exp.toAttribute.withName(AGGREGATE_COUNT_COLUMN_IDENTIFIER))
case a: Attribute => a -> AttributeSet(a)
+ case ref: OuterReference => ref.toAttribute ->
AttributeSet(ref.toAttribute)
Review Comment:
This also changes select-list scalar subqueries such as `INSERT OVERWRITE
tb3 SELECT (SELECT t1.col1) AS col1, ... FROM tb1 t1`, which used to throw
`MatchError` here. Please add a test asserting the column lineage for that case
too.
##########
extensions/spark/kyuubi-spark-lineage/src/test/scala/org/apache/kyuubi/plugin/lineage/helper/SparkSQLLineageParserHelperSuite.scala:
##########
@@ -1494,6 +1494,163 @@ abstract class SparkSQLLineageParserHelperSuite extends
KyuubiFunSuite
}
}
+ test("columns lineage extract - collect input tables from a correlated
EXISTS filter") {
+ val tables = Seq("v2_catalog.db.tb1", "v2_catalog.db.tb2",
"v2_catalog.db.tb3")
+ withTable(tables: _*) { _ =>
+ tables.foreach { table =>
+ spark.sql(s"CREATE TABLE $table (col1 STRING, col2 STRING, col3
STRING)").collect()
+ }
+
assert(!SparkContextHelper.getConf(LineageConf.LEGACY_COLLECT_INPUT_TABLES_ENABLED))
+ val sql =
+ """
+ |INSERT OVERWRITE v2_catalog.db.tb3
+ |SELECT t1.* FROM v2_catalog.db.tb1 t1
+ |WHERE EXISTS (
+ | SELECT 1 FROM v2_catalog.db.tb2 t2 WHERE t2.col1 = t1.col1
+ |)
+ |""".stripMargin
+ val parsed = spark.sessionState.sqlParser.parsePlan(sql)
+ val analyzed = spark.sessionState.analyzer.execute(parsed)
+ spark.sessionState.analyzer.checkAnalysis(analyzed)
+ val helper = SparkSQLLineageParseHelper(spark)
+ val lineage = helper.parse(analyzed)
+ info(s"Correlated EXISTS SQL:\n$sql")
+ info(s"Correlated EXISTS analyzed plan:\n${analyzed.treeString}")
+ info(s"Correlated EXISTS actual lineage: $lineage")
+ assert(helper.transformToLineage(0, analyzed).contains(lineage))
+ assert(lineage.outputTables == List("v2_catalog.db.tb3"))
+ assert(lineage.columnLineage.map(c => c.column -> c.originalColumns) ==
List(
+ ("v2_catalog.db.tb3.col1", Set("v2_catalog.db.tb1.col1")),
+ ("v2_catalog.db.tb3.col2", Set("v2_catalog.db.tb1.col2")),
+ ("v2_catalog.db.tb3.col3", Set("v2_catalog.db.tb1.col3"))))
+ assert(lineage.inputTables.toSet == Set("v2_catalog.db.tb1",
"v2_catalog.db.tb2"))
+ }
+ }
+
+ Seq(
+ (
+ "EXISTS",
+ "WHERE EXISTS (SELECT 1 FROM v2_catalog.db.tb2 t2 WHERE t2.col1 =
t1.col1)",
+ Set("v2_catalog.db.tb2")),
+ (
+ "NOT EXISTS",
+ "WHERE NOT EXISTS (SELECT 1 FROM v2_catalog.db.tb2 t2 WHERE t2.col1 =
t1.col1)",
+ Set("v2_catalog.db.tb2")),
+ (
+ "IN",
+ "WHERE t1.col1 IN (SELECT col1 FROM v2_catalog.db.tb2)",
+ Set("v2_catalog.db.tb2")),
+ (
+ "NOT IN",
+ "WHERE t1.col1 NOT IN (SELECT col1 FROM v2_catalog.db.tb2)",
+ Set("v2_catalog.db.tb2")),
+ (
+ "scalar subquery",
+ "WHERE t1.col1 > (SELECT max(col1) FROM v2_catalog.db.tb2)",
+ Set("v2_catalog.db.tb2")),
+ (
+ "HAVING",
+ "GROUP BY t1.col1, t1.col2, t1.col3 " +
+ "HAVING count(*) > (SELECT count(*) FROM v2_catalog.db.tb2)",
+ Set("v2_catalog.db.tb2")),
+ (
+ "nested IN",
+ "WHERE t1.col1 IN (SELECT t2.col1 FROM v2_catalog.db.tb2 t2 " +
+ "WHERE t2.col2 IN (SELECT col2 FROM v2_catalog.db.tb4))",
+ Set("v2_catalog.db.tb2", "v2_catalog.db.tb4")),
+ (
+ "multiple subqueries",
+ "WHERE EXISTS (SELECT 1 FROM v2_catalog.db.tb2 t2 WHERE t2.col1 =
t1.col1) " +
+ "AND (t1.col2 IN (SELECT col2 FROM v2_catalog.db.tb4) OR t1.col3 =
'keep')",
+ Set("v2_catalog.db.tb2", "v2_catalog.db.tb4")),
+ (
+ "no subquery",
+ "WHERE t1.col1 = 'keep'",
+ Set.empty[String])).foreach { case (conditionName, condition,
filterTables) =>
+ Seq(false, true).foreach { legacy =>
+ test(s"columns lineage extract - filter $conditionName with
legacy=$legacy") {
+ val legacyEntry = LineageConf.LEGACY_COLLECT_INPUT_TABLES_ENABLED
+ val previousLegacy = SparkContextHelper.getConf(legacyEntry)
+ try {
+ SparkContextHelper.setConf(legacyEntry, legacy)
+ val tables = Seq("tb1", "tb2", "tb3", "tb4").map(t =>
s"v2_catalog.db.$t")
+ withTable(tables: _*) { _ =>
+ tables.foreach { table =>
+ spark.sql(s"CREATE TABLE $table (col1 STRING, col2 STRING, col3
STRING)").collect()
+ }
+ val sql = "INSERT OVERWRITE v2_catalog.db.tb3 " +
+ s"SELECT t1.* FROM v2_catalog.db.tb1 t1 $condition"
+ val lineage = extractLineage(sql)
+ val expectedInputs =
+ if (legacy) Set("v2_catalog.db.tb1")
+ else Set("v2_catalog.db.tb1") ++ filterTables
+ info(s"Filter $conditionName with legacy=$legacy SQL: $sql")
+ info(s"Filter $conditionName with legacy=$legacy actual lineage:
$lineage")
+ assert(lineage == Lineage(
+ expectedInputs.toList,
+ List("v2_catalog.db.tb3"),
+ List(
+ ("v2_catalog.db.tb3.col1", Set("v2_catalog.db.tb1.col1")),
+ ("v2_catalog.db.tb3.col2", Set("v2_catalog.db.tb1.col2")),
+ ("v2_catalog.db.tb3.col3", Set("v2_catalog.db.tb1.col3")))))
+ }
+ } finally {
+ SparkContextHelper.setConf(legacyEntry, previousLegacy)
+ }
+ assert(SparkContextHelper.getConf(legacyEntry) == previousLegacy)
+ }
+ }
+ }
+
+ Seq(false, true).foreach { legacy =>
+ Seq("SELECT", "INSERT").foreach { statement =>
+ test(s"columns lineage extract - bare outer reference $statement with
legacy=$legacy") {
+ val legacyEntry = LineageConf.LEGACY_COLLECT_INPUT_TABLES_ENABLED
+ val previousLegacy = SparkContextHelper.getConf(legacyEntry)
+ try {
+ SparkContextHelper.setConf(legacyEntry, legacy)
+ withTable("v2_catalog.db.tb1", "v2_catalog.db.tb3") { tables =>
+ tables.foreach { table =>
+ spark.sql(s"CREATE TABLE $table (col1 INT, col2 STRING, col3
INT)").collect()
+ }
+ spark.sql("INSERT INTO v2_catalog.db.tb1 VALUES (1, 'one', 10),
(2, 'two', 20)")
+ .collect()
+ val select = "SELECT * FROM v2_catalog.db.tb1 t1 " +
+ "WHERE t1.col1 = (SELECT t1.col1)"
+ val sql = if (statement == "SELECT") select
+ else s"INSERT OVERWRITE v2_catalog.db.tb3 $select"
+ val parsed = spark.sessionState.sqlParser.parsePlan(sql)
+ val analyzed = spark.sessionState.analyzer.execute(parsed)
+ spark.sessionState.analyzer.checkAnalysis(analyzed)
+ info(s"Bare outer reference $statement with legacy=$legacy SQL:
$sql")
+ info(s"Bare outer reference analyzed
plan:\n${analyzed.treeString}")
+ val result = spark.sql(sql).collect()
Review Comment:
Inserting data, running the query and checking result rows tests Spark
rather than lineage. The lineage assertion on `analyzed` is enough. Can we drop
the insert/execute/row checks?
##########
extensions/spark/kyuubi-spark-lineage/src/test/scala/org/apache/kyuubi/plugin/lineage/helper/SparkSQLLineageParserHelperSuite.scala:
##########
@@ -1494,6 +1494,163 @@ abstract class SparkSQLLineageParserHelperSuite extends
KyuubiFunSuite
}
}
+ test("columns lineage extract - collect input tables from a correlated
EXISTS filter") {
+ val tables = Seq("v2_catalog.db.tb1", "v2_catalog.db.tb2",
"v2_catalog.db.tb3")
+ withTable(tables: _*) { _ =>
+ tables.foreach { table =>
+ spark.sql(s"CREATE TABLE $table (col1 STRING, col2 STRING, col3
STRING)").collect()
+ }
+
assert(!SparkContextHelper.getConf(LineageConf.LEGACY_COLLECT_INPUT_TABLES_ENABLED))
+ val sql =
+ """
+ |INSERT OVERWRITE v2_catalog.db.tb3
+ |SELECT t1.* FROM v2_catalog.db.tb1 t1
+ |WHERE EXISTS (
+ | SELECT 1 FROM v2_catalog.db.tb2 t2 WHERE t2.col1 = t1.col1
+ |)
+ |""".stripMargin
+ val parsed = spark.sessionState.sqlParser.parsePlan(sql)
+ val analyzed = spark.sessionState.analyzer.execute(parsed)
+ spark.sessionState.analyzer.checkAnalysis(analyzed)
+ val helper = SparkSQLLineageParseHelper(spark)
+ val lineage = helper.parse(analyzed)
+ info(s"Correlated EXISTS SQL:\n$sql")
+ info(s"Correlated EXISTS analyzed plan:\n${analyzed.treeString}")
+ info(s"Correlated EXISTS actual lineage: $lineage")
+ assert(helper.transformToLineage(0, analyzed).contains(lineage))
+ assert(lineage.outputTables == List("v2_catalog.db.tb3"))
+ assert(lineage.columnLineage.map(c => c.column -> c.originalColumns) ==
List(
+ ("v2_catalog.db.tb3.col1", Set("v2_catalog.db.tb1.col1")),
+ ("v2_catalog.db.tb3.col2", Set("v2_catalog.db.tb1.col2")),
+ ("v2_catalog.db.tb3.col3", Set("v2_catalog.db.tb1.col3"))))
+ assert(lineage.inputTables.toSet == Set("v2_catalog.db.tb1",
"v2_catalog.db.tb2"))
+ }
+ }
+
+ Seq(
+ (
+ "EXISTS",
+ "WHERE EXISTS (SELECT 1 FROM v2_catalog.db.tb2 t2 WHERE t2.col1 =
t1.col1)",
+ Set("v2_catalog.db.tb2")),
+ (
+ "NOT EXISTS",
+ "WHERE NOT EXISTS (SELECT 1 FROM v2_catalog.db.tb2 t2 WHERE t2.col1 =
t1.col1)",
+ Set("v2_catalog.db.tb2")),
+ (
+ "IN",
+ "WHERE t1.col1 IN (SELECT col1 FROM v2_catalog.db.tb2)",
+ Set("v2_catalog.db.tb2")),
+ (
+ "NOT IN",
+ "WHERE t1.col1 NOT IN (SELECT col1 FROM v2_catalog.db.tb2)",
+ Set("v2_catalog.db.tb2")),
+ (
+ "scalar subquery",
+ "WHERE t1.col1 > (SELECT max(col1) FROM v2_catalog.db.tb2)",
+ Set("v2_catalog.db.tb2")),
+ (
+ "HAVING",
+ "GROUP BY t1.col1, t1.col2, t1.col3 " +
+ "HAVING count(*) > (SELECT count(*) FROM v2_catalog.db.tb2)",
+ Set("v2_catalog.db.tb2")),
+ (
+ "nested IN",
+ "WHERE t1.col1 IN (SELECT t2.col1 FROM v2_catalog.db.tb2 t2 " +
+ "WHERE t2.col2 IN (SELECT col2 FROM v2_catalog.db.tb4))",
+ Set("v2_catalog.db.tb2", "v2_catalog.db.tb4")),
+ (
+ "multiple subqueries",
+ "WHERE EXISTS (SELECT 1 FROM v2_catalog.db.tb2 t2 WHERE t2.col1 =
t1.col1) " +
+ "AND (t1.col2 IN (SELECT col2 FROM v2_catalog.db.tb4) OR t1.col3 =
'keep')",
+ Set("v2_catalog.db.tb2", "v2_catalog.db.tb4")),
+ (
+ "no subquery",
+ "WHERE t1.col1 = 'keep'",
+ Set.empty[String])).foreach { case (conditionName, condition,
filterTables) =>
+ Seq(false, true).foreach { legacy =>
+ test(s"columns lineage extract - filter $conditionName with
legacy=$legacy") {
+ val legacyEntry = LineageConf.LEGACY_COLLECT_INPUT_TABLES_ENABLED
+ val previousLegacy = SparkContextHelper.getConf(legacyEntry)
+ try {
+ SparkContextHelper.setConf(legacyEntry, legacy)
+ val tables = Seq("tb1", "tb2", "tb3", "tb4").map(t =>
s"v2_catalog.db.$t")
+ withTable(tables: _*) { _ =>
Review Comment:
Each of the 18 generated cases creates and drops the same 4 tables, but the
check only analyzes the plan. Could we create the tables once with a single
`withTable` and loop over the conditions inside it? Also:
1. Please remove the `info(...)` dumps here and in the other new tests
(`:1517-1519`, `:1586-1587`, `:1625-1626`, `:1631`, `:1641`).
2. The `assert(SparkContextHelper.getConf(legacyEntry) == previousLegacy)`
after `finally` (`:1600`, `:1649`) only re-checks the line above it, so it can
go.
--
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]
