Skip to content

[KYUUBI #6912][LINEAGE] Properly handle empty attribute set on mergeRelationColumnLineage #6911

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
xglv1985 wants to merge 7 commits into from
Closed

[KYUUBI #6912][LINEAGE] Properly handle empty attribute set on mergeRelationColumnLineage #6911

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
fbb4df8
fix a runtime exception when generate column lineage tuple
Feb 7, 2025
9018790
fix a runtime exception when generate column lineage tuple--unit test
Feb 8, 2025
fea6bbc
fix a runtime exception when generate column lineage tuple--remove ta…
Feb 10, 2025
52bc028
fix a runtime exception when generate column lineage tuple--spotless …
Feb 11, 2025
59b350b
fix a runtime exception when generate column lineage tuple--more read…
Feb 11, 2025
4e89b95
Update extensions/spark/kyuubi-spark-lineage/src/test/scala/org/apach…
pan3793 Feb 11, 2025
13a7107
Update extensions/spark/kyuubi-spark-lineage/src/test/scala/org/apach…
pan3793 Feb 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions ...e/src/main/scala/org/apache/kyuubi/plugin/lineage/helper/SparkSQLLineageParseHelper.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,8 @@ trait LineageParser {
relationColumnLineage: AttributeMap[AttributeSet]): AttributeMap[AttributeSet] = {
val mergedRelationColumnLineage = {
relationOutput.foldLeft((ListMap[Attribute, AttributeSet](), relationColumnLineage)) {
case ((acc, x), attr) if x.isEmpty =>
(acc + (attr -> AttributeSet.empty), x.empty)
case ((acc, x), attr) =>
(acc + (attr -> x.head._2), x.tail)
}._1
Expand Down
36 changes: 35 additions & 1 deletion ...test/scala/org/apache/kyuubi/plugin/lineage/helper/SparkSQLLineageParserHelperSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,10 +59,16 @@ class SparkSQLLineageParserHelperSuite extends KyuubiFunSuite
" partitioned by(pid)")
spark.sql("create table if not exists test_db0.test_table1" +
" (key int, value string) using parquet")
spark.sql("create table test_db.test_table_from_dir" +
" (`a0` string, `b0` string) using parquet")
}

override def afterAll(): Unit = {
Seq("test_db0.test_table0", "test_db0.test_table1", "test_db0.test_table_part0").foreach { t =>
Seq(
"test_db0.test_table0",
"test_db0.test_table1",
"test_db0.test_table_part0",
"test_db.test_table_from_dir").foreach { t =>
spark.sql(s"drop table if exists $t")
}
spark.sql("drop database if exists test_db")
Expand Down Expand Up @@ -1442,6 +1448,34 @@ class SparkSQLLineageParserHelperSuite extends KyuubiFunSuite
}
}

test("test directory to table") {
val inputFile = getClass.getResource("/").getPath + "input_file"
val sourceFile = File(inputFile).createFile()
spark.sql(
s"""
|CREATE OR REPLACE TEMPORARY VIEW temp_view (
| `a` STRING COMMENT '',
| `b` STRING COMMENT ''
|) USING csv OPTIONS(
| sep='\t',
| path='${sourceFile.path}'
Comment thread
pan3793 marked this conversation as resolved.
Outdated
|);
|""".stripMargin).collect()
Comment on lines +1455 to +1463

@pan3793 pan3793 Feb 11, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
s"""
| CREATE OR REPLACE TEMPORARY VIEW temp_view
|(
| `a` STRING COMMENT '',
| `b` STRING COMMENT ''
|)
|USING csv OPTIONS(
| sep='\t',
| path='${sourceFile.path}'
|);
|""".stripMargin).collect()
s"""
|CREATE OR REPLACE TEMPORARY VIEW temp_view (
| `a` STRING COMMENT '',
| `b` STRING COMMENT ''
|) USING csv OPTIONS(
| sep='\t',
| path='${sourceFile.path}'
|)
|""".stripMargin).collect()


val ret0 = extractLineageWithoutExecuting(
s"""
|INSERT OVERWRITE TABLE test_db.test_table_from_dir
|SELECT `a`, `b` FROM temp_view
|""".stripMargin)
Comment on lines +1468 to +1469

@pan3793 pan3793 Feb 11, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: use UPPER CASE for SQL keywords, and format the SQL in a short readable way

Suggested change
s"""
|insert overwrite table test_db.test_table_from_dir
|SELECT
| `a`,
| `b`
|FROM temp_view
|""".stripMargin)
s"""
|INSERT OVERWRITE TABLE test_db.test_table_from_dir
|SELECT `a`, `b` FROM temp_view
|""".stripMargin)


assert(ret0 == Lineage(
List(),
List(s"spark_catalog.test_db.test_table_from_dir"),
List(
(s"spark_catalog.test_db.test_table_from_dir.a0", Set()),
(s"spark_catalog.test_db.test_table_from_dir.b0", Set()))))
}

private def extractLineageWithoutExecuting(sql: String): Lineage = {
val parsed = spark.sessionState.sqlParser.parsePlan(sql)
val analyzed = spark.sessionState.analyzer.execute(parsed)
Expand Down