diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/SqlGraphRegistrationContext.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/SqlGraphRegistrationContext.scala index 625844de74cb0..83c168bdf42b9 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/SqlGraphRegistrationContext.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/SqlGraphRegistrationContext.scala @@ -20,12 +20,16 @@ import scala.collection.mutable import org.apache.spark.{SparkException, SparkRuntimeException} import org.apache.spark.sql.{AnalysisException, SparkSession} -import org.apache.spark.sql.catalyst.QueryPlanningTracker -import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation -import org.apache.spark.sql.catalyst.plans.logical.{CreateFlowCommand, CreateMaterializedViewAsSelect, CreateStreamingTable, CreateStreamingTableAsSelect, CreateView, InsertIntoStatement, LogicalPlan} +import org.apache.spark.sql.Column +import org.apache.spark.sql.catalyst.{QueryPlanningTracker, TableIdentifier} +import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedRelation} +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.plans.logical.{AutoCdcIntoCommand, CreateFlowCommand, CreateMaterializedViewAsSelect, CreateStreamingTable, CreateStreamingTableAsSelect, CreateStreamingTableAutoCdc, CreateView, InsertIntoStatement, LogicalPlan} import org.apache.spark.sql.catalyst.util.StringUtils +import org.apache.spark.sql.classic.ClassicConversions._ import org.apache.spark.sql.execution.command.{CreateViewCommand, SetCatalogCommand, SetCommand, SetNamespaceCommand} import org.apache.spark.sql.pipelines.Language +import org.apache.spark.sql.pipelines.autocdc.{ChangeArgs, ColumnSelection, ScdType, UnqualifiedColumnName} import org.apache.spark.sql.types.StructType /** @@ -162,6 +166,10 @@ class SqlGraphRegistrationContext( case createStreamingTableCommand: CreateStreamingTable => // CREATE STREAMING TABLE [ streaming_table_name ] [ options ] CreateStreamingTableHandler.handle(createStreamingTableCommand, queryOrigin) + case createStreamingTableAutoCdcCommand: CreateStreamingTableAutoCdc => + // CREATE STREAMING TABLE [ streaming_table_name ] [ options ] + // FLOW AUTO CDC FROM [ source ] KEYS ( ... ) SEQUENCE BY [ expr ] ... + CreateStreamingTableAutoCdcHandler.handle(createStreamingTableAutoCdcCommand, queryOrigin) case createFlowCommand: CreateFlowCommand => // CREATE FLOW [ flow_name ] AS INSERT INTO [ destination_name ] BY NAME CreateFlowHandler.handle(createFlowCommand, queryOrigin) @@ -256,6 +264,108 @@ class SqlGraphRegistrationContext( } } + /** + * Converts the parse-time AUTO CDC parameters (catalyst expressions and unresolved attributes) + * into the [[ChangeArgs]] consumed by an [[AutoCdcFlow]]. Shared by the two SQL AUTO CDC entry + * points: `CREATE STREAMING TABLE ... FLOW AUTO CDC ...` and `CREATE FLOW ... AS AUTO CDC INTO`. + * + * SQL AUTO CDC syntax only supports SCD Type 1, so [[ChangeArgs.storedAsScdType]] is always + * [[ScdType.Type1]]. [[includeColumns]] and [[excludeColumns]] are mutually exclusive at the + * grammar level; the guard here is defensive. + */ + private def buildChangeArgs( + keys: Seq[UnresolvedAttribute], + sequenceByExpr: Expression, + deleteCondition: Option[Expression], + includeColumns: Option[Seq[UnresolvedAttribute]], + excludeColumns: Option[Seq[UnresolvedAttribute]], + queryOrigin: QueryOrigin): ChangeArgs = { + val columnSelection: Option[ColumnSelection] = (includeColumns, excludeColumns) match { + case (Some(_), Some(_)) => + throw SqlGraphElementRegistrationException( + msg = "AUTO CDC cannot specify both COLUMNS and COLUMNS * EXCEPT.", + queryOrigin = queryOrigin + ) + case (Some(included), None) => + Option(ColumnSelection.IncludeColumns(included.map(toUnqualifiedColumnName))) + case (None, Some(excluded)) => + Option(ColumnSelection.ExcludeColumns(excluded.map(toUnqualifiedColumnName))) + case (None, None) => + None + } + + ChangeArgs( + keys = keys.map(toUnqualifiedColumnName), + sequencing = Column(sequenceByExpr), + storedAsScdType = ScdType.Type1, + deleteCondition = deleteCondition.map(Column(_)), + columnSelection = columnSelection + ) + } + + private def toUnqualifiedColumnName(attr: UnresolvedAttribute): UnqualifiedColumnName = + UnqualifiedColumnName(attr.nameParts) + + private object CreateStreamingTableAutoCdcHandler { + def handle(cst: CreateStreamingTableAutoCdc, queryOrigin: QueryOrigin): Unit = { + val stIdentifier = GraphIdentifierManager + .parseAndQualifyTableIdentifier( + rawTableIdentifier = IdentifierHelper.toTableIdentifier(cst.name), + currentCatalog = context.getCurrentCatalogOpt, + currentDatabase = context.getCurrentDatabaseOpt + ) + .identifier + + // Register the streaming table as a table. The streaming table is itself the target of the + // CDC operation. + graphRegistrationContext.registerTable( + Table( + identifier = stIdentifier, + comment = cst.tableSpec.comment, + specifiedSchema = + Option.when(cst.columns.nonEmpty)(StructType(cst.columns.map(_.toV1Column))), + partitionCols = Option(PartitionHelper.applyPartitioning(cst.partitioning, queryOrigin)), + clusterCols = None, + properties = cst.tableSpec.properties, + origin = queryOrigin.copy( + objectName = Option(stIdentifier.unquotedString), + objectType = Option(QueryOriginType.Table.toString) + ), + format = cst.tableSpec.provider, + normalizedPath = None, + isStreamingTable = true + ) + ) + + // Register the AutoCDC flow that backs this streaming table. Both the flow and its + // destination are the streaming table itself. + graphRegistrationContext.registerFlow( + AutoCdcFlow( + identifier = stIdentifier, + destinationIdentifier = stIdentifier, + func = FlowAnalysis.createFlowFunctionFromLogicalPlan(cst.source), + sqlConf = context.getSqlConf, + queryContext = QueryContext( + currentCatalog = context.getCurrentCatalogOpt, + currentDatabase = context.getCurrentDatabaseOpt + ), + origin = queryOrigin.copy( + objectName = Option(stIdentifier.unquotedString), + objectType = Option(QueryOriginType.Flow.toString) + ), + changeArgs = buildChangeArgs( + keys = cst.keys, + sequenceByExpr = cst.sequenceByExpr, + deleteCondition = cst.deleteCondition, + includeColumns = cst.includeColumns, + excludeColumns = cst.excludeColumns, + queryOrigin = queryOrigin + ) + ) + ) + } + } + private object CreateMaterializedViewAsSelectHandler { def handle(cmv: CreateMaterializedViewAsSelect, queryOrigin: QueryOrigin): Unit = { val mvIdentifier = GraphIdentifierManager @@ -415,7 +525,7 @@ class SqlGraphRegistrationContext( ) .identifier - val (flowTargetDatasetIdentifier, flowQueryLogicalPlan) = cf.flowOperation match { + cf.flowOperation match { case i: InsertIntoStatement => validateInsertIntoFlow(i, queryOrigin) val flowTargetDatasetName = i.table match { @@ -427,22 +537,55 @@ class SqlGraphRegistrationContext( queryOrigin = queryOrigin ) } - val qualifiedFlowTargetDatasetName = GraphIdentifierManager - .parseAndQualifyTableIdentifier( - rawTableIdentifier = flowTargetDatasetName, - currentCatalog = context.getCurrentCatalogOpt, - currentDatabase = context.getCurrentDatabaseOpt + graphRegistrationContext.registerFlow( + UntypedFlow( + identifier = flowIdentifier, + destinationIdentifier = qualifyDestinationIdentifier(flowTargetDatasetName), + func = FlowAnalysis.createFlowFunctionFromLogicalPlan(i.query), + sqlConf = context.getSqlConf, + once = false, + queryContext = QueryContext( + currentCatalog = context.getCurrentCatalogOpt, + currentDatabase = context.getCurrentDatabaseOpt + ), + origin = queryOrigin + ) + ) + case a: AutoCdcIntoCommand => + val flowTargetDatasetName = IdentifierHelper.toTableIdentifier(a.targetTable) + graphRegistrationContext.registerFlow( + AutoCdcFlow( + identifier = flowIdentifier, + destinationIdentifier = qualifyDestinationIdentifier(flowTargetDatasetName), + func = FlowAnalysis.createFlowFunctionFromLogicalPlan(a.source), + sqlConf = context.getSqlConf, + queryContext = QueryContext( + currentCatalog = context.getCurrentCatalogOpt, + currentDatabase = context.getCurrentDatabaseOpt + ), + origin = queryOrigin, + changeArgs = buildChangeArgs( + keys = a.keys, + sequenceByExpr = a.sequenceByExpr, + deleteCondition = a.deleteCondition, + includeColumns = a.includeColumns, + excludeColumns = a.excludeColumns, + queryOrigin = queryOrigin + ) ) - .identifier - (qualifiedFlowTargetDatasetName, i.query) + ) case _ => throw SqlGraphElementRegistrationException( - msg = "Unable flow type. Only INSERT INTO flows are supported.", + msg = "Unable flow type. Only INSERT INTO and AUTO CDC INTO flows are supported.", queryOrigin = queryOrigin ) } + } - val qualifiedDestinationIdentifier = GraphIdentifierManager + /** Qualifies a raw flow target dataset identifier against the current catalog/database. */ + private def qualifyDestinationIdentifier( + flowTargetDatasetIdentifier: TableIdentifier): TableIdentifier = + GraphIdentifierManager .parseAndQualifyFlowIdentifier( rawFlowIdentifier = flowTargetDatasetIdentifier, currentCatalog = context.getCurrentCatalogOpt, @@ -450,22 +593,6 @@ class SqlGraphRegistrationContext( ) .identifier - graphRegistrationContext.registerFlow( - UntypedFlow( - identifier = flowIdentifier, - destinationIdentifier = qualifiedDestinationIdentifier, - func = FlowAnalysis.createFlowFunctionFromLogicalPlan(flowQueryLogicalPlan), - sqlConf = context.getSqlConf, - once = false, - queryContext = QueryContext( - currentCatalog = context.getCurrentCatalogOpt, - currentDatabase = context.getCurrentDatabaseOpt - ), - origin = queryOrigin - ) - ) - } - private def validateInsertIntoFlow( insertIntoStatement: InsertIntoStatement, queryOrigin: QueryOrigin diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/SqlPipelineSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/SqlPipelineSuite.scala index 90132c7ea7bbf..855426983c634 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/SqlPipelineSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/SqlPipelineSuite.scala @@ -18,7 +18,9 @@ package org.apache.spark.sql.pipelines.graph import org.apache.spark.sql.{AnalysisException, Row} import org.apache.spark.sql.catalyst.parser.ParseException -import org.apache.spark.sql.connector.catalog.{Identifier, TableCatalog} +import org.apache.spark.sql.connector.catalog.{Identifier, SharedTablesInMemoryRowLevelOperationTableCatalog, TableCatalog} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.pipelines.autocdc.{ColumnSelection, ScdType} import org.apache.spark.sql.pipelines.utils.{PipelineTest, TestGraphRegistrationContext} import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{LongType, StructType} @@ -1080,4 +1082,229 @@ class SqlPipelineSuite extends PipelineTest with SharedSparkSession { } } } + + // =========================================================================== + // AUTO CDC syntax registration and execution. + // + // Two SQL forms register an [[AutoCdcFlow]] into the dataflow graph: + // 1. CREATE STREAMING TABLE FLOW AUTO CDC FROM KEYS (...) SEQUENCE BY + // 2. CREATE FLOW AS AUTO CDC INTO FROM KEYS (...) SEQUENCE BY + // SQL AUTO CDC only supports SCD Type 1. + // =========================================================================== + + /** Returns the single unresolved [[AutoCdcFlow]] registered for the given flow identifier. */ + private def autoCdcFlowFor(graph: DataflowGraph, name: String): AutoCdcFlow = { + val ident = fullyQualifiedIdentifier(name) + graph.flows.collect { + case f: AutoCdcFlow if f.identifier == ident => f + }.headOption.getOrElse( + fail(s"No AutoCdcFlow registered for identifier ${ident.unquotedString}") + ) + } + + test("CREATE STREAMING TABLE FLOW AUTO CDC registers a streaming table and an AutoCDC flow") { + val graph = unresolvedDataflowGraphFromSql( + sqlText = s""" + |CREATE STREAMING TABLE st + |FLOW AUTO CDC + |FROM STREAM $externalTable1Ident + |KEYS (id) + |SEQUENCE BY id + |""".stripMargin + ) + + // The streaming table is registered as a table. + assert(graph.tables.exists(_.identifier == fullyQualifiedIdentifier("st"))) + + // The backing flow is an AutoCdcFlow whose flow and destination are the streaming table. + val flow = autoCdcFlowFor(graph, "st") + assert(flow.destinationIdentifier == fullyQualifiedIdentifier("st")) + assert(flow.changeArgs.keys.map(_.name) == Seq("id")) + assert(flow.changeArgs.storedAsScdType == ScdType.Type1) + assert(flow.changeArgs.deleteCondition.isEmpty) + assert(flow.changeArgs.columnSelection.isEmpty) + } + + test("CREATE FLOW AS AUTO CDC INTO registers an AutoCDC flow targeting a streaming table") { + val graph = unresolvedDataflowGraphFromSql( + sqlText = s""" + |CREATE STREAMING TABLE target; + |CREATE FLOW f AS AUTO CDC INTO target + |FROM STREAM $externalTable1Ident + |KEYS (id) + |SEQUENCE BY id + |""".stripMargin + ) + + assert(graph.tables.exists(_.identifier == fullyQualifiedIdentifier("target"))) + + // The flow identifier is the named flow, and its destination is the target streaming table. + val flow = autoCdcFlowFor(graph, "f") + assert(flow.destinationIdentifier == fullyQualifiedIdentifier("target")) + assert(flow.changeArgs.keys.map(_.name) == Seq("id")) + assert(flow.changeArgs.storedAsScdType == ScdType.Type1) + } + + test("AUTO CDC optional clauses map onto ChangeArgs") { + val graph = unresolvedDataflowGraphFromSql( + sqlText = s""" + |CREATE STREAMING TABLE target; + |CREATE FLOW f AS AUTO CDC INTO target + |FROM STREAM $externalTable1Ident + |KEYS (id) + |APPLY AS DELETE WHEN id = 0 + |SEQUENCE BY id + |COLUMNS * EXCEPT (id) + |""".stripMargin + ) + + val flow = autoCdcFlowFor(graph, "f") + assert(flow.changeArgs.deleteCondition.isDefined) + flow.changeArgs.columnSelection match { + case Some(ColumnSelection.ExcludeColumns(cols)) => assert(cols.map(_.name) == Seq("id")) + case other => fail(s"Expected ExcludeColumns(id), got $other") + } + } + + test("AUTO CDC COLUMNS include list maps onto ChangeArgs") { + val graph = unresolvedDataflowGraphFromSql( + sqlText = s""" + |CREATE STREAMING TABLE st + |FLOW AUTO CDC + |FROM STREAM $externalTable1Ident + |KEYS (id) + |SEQUENCE BY id + |COLUMNS (id) + |""".stripMargin + ) + + autoCdcFlowFor(graph, "st").changeArgs.columnSelection match { + case Some(ColumnSelection.IncludeColumns(cols)) => assert(cols.map(_.name) == Seq("id")) + case other => fail(s"Expected IncludeColumns(id), got $other") + } + } + + test("Multipart AUTO CDC flow name is not supported") { + Seq("a.b", "a.b.c").foreach { flowIdentifier => + val ex = intercept[AnalysisException] { + unresolvedDataflowGraphFromSql( + sqlText = s""" + |CREATE STREAMING TABLE target; + |CREATE FLOW $flowIdentifier AS AUTO CDC INTO target + |FROM STREAM $externalTable1Ident + |KEYS (id) + |SEQUENCE BY id + |""".stripMargin + ) + } + checkError( + exception = ex, + condition = "MULTIPART_FLOW_NAME_NOT_SUPPORTED", + parameters = Map("flowName" -> flowIdentifier) + ) + } + } + + // --------------------------------------------------------------------------- + // End-to-end execution. AutoCDC applies its microbatch to the target via the DataFrame + // MERGE API, which requires a row-level-operation-capable v2 catalog, so these tests scaffold + // one inline (mirroring the AutoCDC E2E suites) rather than using the default `spark_catalog`. + // --------------------------------------------------------------------------- + + private val rowLevelCatalog: String = "cat" + private val rowLevelNamespace: String = "ns1" + + private def withRowLevelAutoCdcCatalog(testBody: => Unit): Unit = { + spark.conf.set( + s"spark.sql.catalog.$rowLevelCatalog", + classOf[SharedTablesInMemoryRowLevelOperationTableCatalog].getName + ) + // Surface flow failures on the first attempt instead of retrying. + spark.conf.set(SQLConf.PIPELINES_MAX_FLOW_RETRY_ATTEMPTS.key, "0") + spark.sql(s"CREATE NAMESPACE IF NOT EXISTS $rowLevelCatalog.$rowLevelNamespace") + try { + testBody + } finally { + SharedTablesInMemoryRowLevelOperationTableCatalog.reset() + spark.sessionState.catalogManager.reset() + spark.sessionState.conf.unsetConf(s"spark.sql.catalog.$rowLevelCatalog") + spark.sessionState.conf.unsetConf(SQLConf.PIPELINES_MAX_FLOW_RETRY_ATTEMPTS.key) + } + } + + /** Build a target row's `_cdc_metadata` struct value (deleteSequence, upsertSequence). */ + private def cdcMeta(deleteSeq: Option[Long], upsertSeq: Option[Long]): Row = + Row(deleteSeq.orNull, upsertSeq.orNull) + + test("CREATE STREAMING TABLE FLOW AUTO CDC upserts rows into the target end-to-end") { + withRowLevelAutoCdcCatalog { + // Source and target both live in the row-level catalog. Streaming over a static table + // replays all rows in one microbatch, exercising the SCD1 upsert path. Two versions of + // key 1 test latest-wins. + val source = s"$rowLevelCatalog.$rowLevelNamespace.cdc_source" + spark.sql(s"CREATE TABLE $source (id INT, name STRING, version BIGINT)") + spark.sql( + s"INSERT INTO $source VALUES (1, 'alice', 1), (1, 'alice2', 2), (2, 'bob', 1)") + + val target = s"$rowLevelCatalog.$rowLevelNamespace.target" + val graph = unresolvedDataflowGraphFromSql( + sqlText = s""" + |CREATE STREAMING TABLE $target + |FLOW AUTO CDC + |FROM STREAM $source + |KEYS (id) + |SEQUENCE BY version + |""".stripMargin + ) + + startPipelineAndWaitForCompletion(graph) + + // Key 1 converges to the highest-sequenced upsert (alice2 @ v2); key 2 lands as-is. + checkAnswer( + spark.table(target), + Seq( + Row(1, "alice2", 2L, cdcMeta(None, Some(2L))), + Row(2, "bob", 1L, cdcMeta(None, Some(1L))) + ) + ) + } + } + + test("CREATE FLOW AS AUTO CDC INTO applies deletes and column exclusion end-to-end") { + withRowLevelAutoCdcCatalog { + // Source carries a control column `op` that drives the delete condition and is excluded + // from the target projection. + val source = s"$rowLevelCatalog.$rowLevelNamespace.cdc_source" + spark.sql(s"CREATE TABLE $source (id INT, name STRING, version BIGINT, op STRING)") + spark.sql( + s"""INSERT INTO $source VALUES + | (1, 'alice', 1, 'UPSERT'), + | (2, 'bob', 1, 'UPSERT'), + | (2, 'bob', 2, 'DELETE') + |""".stripMargin) + + val target = s"$rowLevelCatalog.$rowLevelNamespace.target" + val graph = unresolvedDataflowGraphFromSql( + sqlText = s""" + |CREATE STREAMING TABLE $target; + |CREATE FLOW f AS AUTO CDC INTO $target + |FROM STREAM $source + |KEYS (id) + |APPLY AS DELETE WHEN op = 'DELETE' + |SEQUENCE BY version + |COLUMNS * EXCEPT (op) + |""".stripMargin + ) + + startPipelineAndWaitForCompletion(graph) + + // Key 2's delete @ v2 supersedes its upsert @ v1; only key 1 remains. The `op` column is + // excluded from the target schema. + checkAnswer( + spark.table(target), + Seq(Row(1, "alice", 1L, cdcMeta(None, Some(1L)))) + ) + } + } + }