[SPARK-60270][CONNECT] Preserve streaming source names for table reads

python/pyspark/sql/connect/plan.py

@@ -421,15 +421,19 @@ def __init__(
         table_name: str,
         options: Optional[Dict[str, str]] = None,
         is_streaming: Optional[bool] = None,
+        source_name: Optional[str] = None,
     ) -> None:
         super().__init__(None)
         self.table_name = table_name
         self.options = options or {}
         self._is_streaming = is_streaming
+        self._source_name = source_name
 
     def plan(self, session: "SparkConnectClient") -> proto.Relation:
         plan = self._create_proto_relation()
         plan.read.named_table.unparsed_identifier = self.table_name
+        if self._source_name is not None:
+            plan.read.named_table.source_name = self._source_name
         if self._is_streaming is not None:
             plan.read.is_streaming = self._is_streaming
         for k, v in self.options.items():

python/pyspark/sql/connect/proto/relations_pb2.pyi

@@ -1144,25 +1144,44 @@ class Read(google.protobuf.message.Message):
 
         UNPARSED_IDENTIFIER_FIELD_NUMBER: builtins.int
         OPTIONS_FIELD_NUMBER: builtins.int
+        SOURCE_NAME_FIELD_NUMBER: builtins.int
         unparsed_identifier: builtins.str
         """(Required) Unparsed identifier for the table."""
         @property
         def options(
             self,
         ) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]:
             """Options for the named table. The map key is case insensitive."""
+        source_name: builtins.str
         def __init__(
             self,
             *,
             unparsed_identifier: builtins.str = ...,
             options: collections.abc.Mapping[builtins.str, builtins.str] | None = ...,
+            source_name: builtins.str | None = ...,
         ) -> None: ...
+        def HasField(
+            self,
+            field_name: typing_extensions.Literal[
+                "_source_name", b"_source_name", "source_name", b"source_name"
+            ],
+        ) -> builtins.bool: ...
         def ClearField(
             self,
             field_name: typing_extensions.Literal[
-                "options", b"options", "unparsed_identifier", b"unparsed_identifier"
+                "_source_name",
+                b"_source_name",
+                "options",
+                b"options",
+                "source_name",
+                b"source_name",
+                "unparsed_identifier",
+                b"unparsed_identifier",
             ],
         ) -> None: ...
+        def WhichOneof(
+            self, oneof_group: typing_extensions.Literal["_source_name", b"_source_name"]
+        ) -> typing_extensions.Literal["source_name"] | None: ...
 
     class DataSource(google.protobuf.message.Message):
         DESCRIPTOR: google.protobuf.descriptor.Descriptor

python/pyspark/sql/connect/streaming/readwriter.py

@@ -454,7 +454,9 @@ def xml(
     xml.__doc__ = PySparkDataStreamReader.xml.__doc__
 
     def table(self, tableName: str) -> "DataFrame":
-        return self._df(Read(tableName, self._options, is_streaming=True))
+        return self._df(
+            Read(tableName, self._options, is_streaming=True, source_name=self._source_name)
+        )
 
     table.__doc__ = PySparkDataStreamReader.table.__doc__
 

python/pyspark/sql/tests/connect/test_connect_plan.py

@@ -54,6 +54,7 @@
         WriteOperation,
     )
     from pyspark.sql.connect.readwriter import DataFrameReader
+    from pyspark.sql.connect.streaming.readwriter import DataStreamReader
     from pyspark.sql.connect.types import (
         proto_schema_to_pyspark_data_type,
         pyspark_types_to_proto_types,
@@ -800,6 +801,26 @@ def test_datasource_read(self):
         self.assertEqual(len(data_source.paths), 1)
         self.assertEqual(data_source.paths[0], "test_path")
 
+    def test_streaming_table_source_name(self):
+        for source_name in (None, "orders_source"):
+            with self.subTest(source_name=source_name):
+                reader = DataStreamReader(self.connect).option("maxFilesPerTrigger", "1")
+                if source_name is not None:
+                    reader.name(source_name)
+                plan = reader.table("db.orders")._plan.to_proto(self.connect)
+                self.assertTrue(plan.root.read.is_streaming)
+                named_table = plan.root.read.named_table
+                self.assertEqual(named_table.unparsed_identifier, "db.orders")
+                self.assertEqual(dict(named_table.options), {"maxFilesPerTrigger": "1"})
+                self.assertEqual(named_table.HasField("source_name"), source_name is not None)
+                if source_name is not None:
+                    self.assertEqual(named_table.source_name, source_name)
+                self.assertEqual(proto.Plan.FromString(plan.SerializeToString()), plan)
+
+        batch_plan = DataFrameReader(self.connect).table("db.orders")._plan.to_proto(self.connect)
+        self.assertFalse(batch_plan.root.read.is_streaming)
+        self.assertFalse(batch_plan.root.read.named_table.HasField("source_name"))
+
     def test_reader_options_case_insensitive(self):
         reader = DataFrameReader(self.connect)
         df = (

python/pyspark/sql/tests/streaming/test_streaming.py

@@ -19,8 +19,9 @@
 import shutil
 import tempfile
 import time
+import uuid
 
-from pyspark.errors import PySparkTypeError, PySparkValueError
+from pyspark.errors import AnalysisException, PySparkTypeError, PySparkValueError
 from pyspark.sql import Row
 from pyspark.sql.functions import lit
 from pyspark.sql.types import IntegerType, StringType, StructField, StructType, TimestampType
@@ -539,6 +540,56 @@ def test_name_with_valid_names(self):
                         df.isStreaming, f"DataFrame should be streaming for name: {name}"
                     )
 
+    def test_name_with_table(self):
+        table_name = f"test_streaming_source_name_{uuid.uuid4().hex}"
+        config = "spark.sql.streaming.queryEvolution.enableSourceEvolution"
+        with self.table(table_name), tempfile.TemporaryDirectory() as tmpdir:
+            self.spark.range(1).write.format("parquet").saveAsTable(table_name)
+            with self.sql_conf({config: "true"}):
+                frame = self.spark.readStream.name("orders_source").table(table_name)
+                self.assertTrue(frame.isStreaming)
+                output_path = os.path.join(tmpdir, "output")
+                query = (
+                    frame.writeStream.format("parquet")
+                    .option("checkpointLocation", os.path.join(tmpdir, "checkpoint"))
+                    .trigger(availableNow=True)
+                    .start(output_path)
+                )
+                try:
+                    self.assertTrue(query.awaitTermination(60))
+                    self.assertEqual(self.spark.read.parquet(output_path).collect(), [Row(id=0)])
+                    with open(os.path.join(tmpdir, "checkpoint", "offsets", "0")) as offset_file:
+                        offset_lines = offset_file.read().splitlines()
+                    self.assertEqual(offset_lines[0], "v2")
+                    self.assertTrue(any(line.startswith("orders_source:") for line in offset_lines))
+                finally:
+                    query.stop()
+                    self.spark.streams.resetTerminated()
+
+                with self.assertRaises(AnalysisException) as context:
+                    self.spark.readStream.table(table_name).schema
+                self.check_error(
+                    context.exception,
+                    errorClass=(
+                        "STREAMING_QUERY_EVOLUTION_ERROR.UNNAMED_STREAMING_SOURCES_WITH_ENFORCEMENT"
+                    ),
+                    messageParameters={"sourceInfo": "Leaf position 0 - .*"},
+                    matchPVals=True,
+                )
+
+            with self.sql_conf({config: "false"}):
+                self.assertEqual(
+                    self.spark.readStream.table(table_name).schema,
+                    self.spark.table(table_name).schema,
+                )
+                with self.assertRaises(AnalysisException) as context:
+                    self.spark.readStream.name("orders_source").table(table_name).schema
+                self.check_error(
+                    context.exception,
+                    errorClass="STREAMING_QUERY_EVOLUTION_ERROR.SOURCE_NAMING_NOT_SUPPORTED",
+                    messageParameters={"name": "orders_source"},
+                )
+
     def test_name_method_chaining(self):
         """Test that name() returns the reader for method chaining."""
         with self.sql_conf(

sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/streaming/ClientStreamingQuerySuite.scala

@@ -770,6 +770,19 @@ class ClientStreamingQuerySuite extends QueryTest with RemoteSparkSession with L
   }
 
   // Tests for DataStreamReader.name() method
+  test("stream reader table source name is serialized") {
+    Seq(None, Some("orders_source")).foreach { sourceName =>
+      val reader = spark.readStream.option("maxFilesPerTrigger", "1")
+      sourceName.foreach(reader.name)
+      val read = reader.table("db.orders").plan.getRoot.getRead
+      assert(read.getIsStreaming)
+      assert(read.getNamedTable.getUnparsedIdentifier == "db.orders")
+      assert(read.getNamedTable.getOptionsMap.get("maxFilesPerTrigger") == "1")
+      assert(read.getNamedTable.hasSourceName == sourceName.isDefined)
+      sourceName.foreach(name => assert(read.getNamedTable.getSourceName == name))
+    }
+  }
+
   testWithSourceEvolution("stream reader name() with valid source names") {
     Seq("mySource", "my_source", "MySource123", "_private", "source_123_test", "123source")
       .foreach { name =>

sql/connect/common/src/main/protobuf/spark/connect/relations.proto

@@ -234,6 +234,8 @@ message Read {
 
     // Options for the named table. The map key is case insensitive.
     map<string, string> options = 2;
+
+    optional string source_name = 3;
   }
 
   message DataSource {

sql/connect/common/src/main/scala/org/apache/spark/sql/connect/DataStreamReader.scala

@@ -103,11 +103,14 @@ final class DataStreamReader private[sql] (sparkSession: SparkSession)
   def table(tableName: String): DataFrame = {
     require(tableName != null, "The table name can't be null")
     sparkSession.newDataFrame { builder =>
-      builder.getReadBuilder
+      val tableBuilder = builder.getReadBuilder
         .setIsStreaming(true)
         .getNamedTableBuilder
         .setUnparsedIdentifier(tableName)
         .putAllOptions(sourceBuilder.getOptionsMap)
+      if (sourceBuilder.hasSourceName) {
+        tableBuilder.setSourceName(sourceBuilder.getSourceName)
+      }
     }
   }
 

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala

@@ -44,7 +44,7 @@ import org.apache.spark.internal.LogKeys.{DATAFRAME_ID, SESSION_ID}
 import org.apache.spark.resource.{ExecutorResourceRequest, ResourceProfile, TaskResourceProfile, TaskResourceRequest}
 import org.apache.spark.sql.{AnalysisException, Column, Encoders, ForeachWriter, Row}
 import org.apache.spark.sql.catalyst.{expressions, AliasIdentifier, FunctionIdentifier, InternalRow, QueryPlanningTracker}
-import org.apache.spark.sql.catalyst.analysis.{ChangelogContextUtils, FunctionRegistry, GlobalTempView, LocalTempView, MultiAlias, RelationChanges, UnresolvedAlias, UnresolvedAttribute, UnresolvedDataFrameStar, UnresolvedDeduplicate, UnresolvedDeserializer, UnresolvedExtractValue, UnresolvedFunction, UnresolvedOrdinal, UnresolvedPlanId, UnresolvedRegex, UnresolvedRelation, UnresolvedStar, UnresolvedStarWithColumns, UnresolvedStarWithColumnsRenames, UnresolvedSubqueryColumnAliases, UnresolvedTableValuedFunction, UnresolvedTranspose}
+import org.apache.spark.sql.catalyst.analysis.{ChangelogContextUtils, FunctionRegistry, GlobalTempView, LocalTempView, MultiAlias, NamedStreamingRelation, RelationChanges, UnresolvedAlias, UnresolvedAttribute, UnresolvedDataFrameStar, UnresolvedDeduplicate, UnresolvedDeserializer, UnresolvedExtractValue, UnresolvedFunction, UnresolvedOrdinal, UnresolvedPlanId, UnresolvedRegex, UnresolvedRelation, UnresolvedStar, UnresolvedStarWithColumns, UnresolvedStarWithColumnsRenames, UnresolvedSubqueryColumnAliases, UnresolvedTableValuedFunction, UnresolvedTranspose}
 import org.apache.spark.sql.catalyst.encoders.{encoderFor, AgnosticEncoder, ExpressionEncoder, RowEncoder}
 import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{ProductEncoder, RowEncoder => AgnosticRowEncoder, StringEncoder, UnboundRowEncoder}
 import org.apache.spark.sql.catalyst.expressions._
@@ -1730,7 +1730,18 @@ class SparkConnectPlanner(
           temporalIdent.nameParts,
           new CaseInsensitiveStringMap(rel.getNamedTable.getOptionsMap),
           isStreaming = rel.getIsStreaming)
-        temporalIdent.wrapTimeTravel(relation)
+        if (rel.getIsStreaming) {
+          val sourceName = if (rel.getNamedTable.hasSourceName) {
+            val name = rel.getNamedTable.getSourceName
+            session.readStream.validateSourceName(name)
+            Some(name)
+          } else {
+            None
+          }
+          NamedStreamingRelation.withUserProvidedName(relation, sourceName)
+        } else {
+          temporalIdent.wrapTimeTravel(relation)
+        }
 
       case proto.Read.ReadTypeCase.DATA_SOURCE if !rel.getIsStreaming =>
         val reader = session.read

sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectPlannerSuite.scala

@@ -26,17 +26,19 @@ import org.apache.spark.connect.proto
 import org.apache.spark.connect.proto.Expression.{Alias, ExpressionString, UnresolvedStar}
 import org.apache.spark.sql.{AnalysisException, Row}
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.analysis.{RelationTimeTravel, UnresolvedAlias, UnresolvedFunction, UnresolvedRelation}
+import org.apache.spark.sql.catalyst.analysis.{NamedStreamingRelation, NameStreamingSources, RelationTimeTravel, UnresolvedAlias, UnresolvedFunction, UnresolvedRelation}
 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, UnsafeProjection}
 import org.apache.spark.sql.catalyst.plans.logical
 import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
+import org.apache.spark.sql.catalyst.streaming.{Unassigned, UserProvided}
 import org.apache.spark.sql.catalyst.trees.{CurrentOrigin, TreePattern}
 import org.apache.spark.sql.catalyst.types.DataTypeUtils
 import org.apache.spark.sql.classic.Dataset
 import org.apache.spark.sql.connect.SparkConnectTestUtils
 import org.apache.spark.sql.connect.common.InvalidPlanInput
 import org.apache.spark.sql.connect.common.LiteralValueProtoConverter.toLiteralProto
 import org.apache.spark.sql.execution.arrow.ArrowConverters
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
@@ -159,6 +161,71 @@ class SparkConnectPlannerSuite extends SparkFunSuite with SparkConnectPlanTest {
     assert(res.nodeName == "UnresolvedRelation")
   }
 
+  test("streaming table source name") {
+    for {
+      isStreaming <- Seq(false, true)
+      sourceName <- Seq(None, Some("orders_source"))
+    } {
+      val table = proto.Read.NamedTable
+        .newBuilder()
+        .setUnparsedIdentifier("db.orders")
+        .putOptions("maxFilesPerTrigger", "1")
+      sourceName.foreach(table.setSourceName)
+      val read = proto.Read.newBuilder().setNamedTable(table).setIsStreaming(isStreaming)
+      val plan = transform(proto.Relation.newBuilder().setRead(read).build())
+      val relation = plan match {
+        case NamedStreamingRelation(child: UnresolvedRelation, UserProvided(name)) =>
+          assert(isStreaming && sourceName.contains(name))
+          child
+        case NamedStreamingRelation(child: UnresolvedRelation, Unassigned) =>
+          assert(isStreaming && sourceName.isEmpty)
+          child
+        case child: UnresolvedRelation =>
+          assert(!isStreaming)
+          child
+        case other => fail(s"Unexpected table read: $other")
+      }
+      assert(relation.multipartIdentifier == Seq("db", "orders"))
+      assert(relation.options.get("maxFilesPerTrigger") == "1")
+      assert(relation.isStreaming == isStreaming)
+    }
+  }
+
+  test("unnamed streaming table sources obey source naming enforcement") {
+    val read = proto.Read
+      .newBuilder()
+      .setNamedTable(proto.Read.NamedTable.newBuilder().setUnparsedIdentifier("db.orders"))
+      .setIsStreaming(true)
+    val plan = transform(proto.Relation.newBuilder().setRead(read).build())
+    withSQLConf(SQLConf.ENABLE_STREAMING_SOURCE_EVOLUTION.key -> "true") {
+      checkError(
+        exception = intercept[AnalysisException] {
+          NameStreamingSources(plan)
+        },
+        condition = "STREAMING_QUERY_EVOLUTION_ERROR.UNNAMED_STREAMING_SOURCES_WITH_ENFORCEMENT",
+        parameters = Map("sourceInfo" -> "Leaf position 0 - UnresolvedRelation"))
+    }
+  }
+
+  test("streaming table source names are validated") {
+    val table = proto.Read.NamedTable
+      .newBuilder()
+      .setUnparsedIdentifier("db.orders")
+      .setSourceName("orders-source")
+    val read = proto.Read.newBuilder().setNamedTable(table).setIsStreaming(true)
+    checkError(
+      exception = intercept[AnalysisException] {
+        transform(proto.Relation.newBuilder().setRead(read).build())
+      },
+      condition = "STREAMING_QUERY_EVOLUTION_ERROR.INVALID_SOURCE_NAME",
+      parameters = Map("sourceName" -> "orders-source"))
+
+    read.setNamedTable(table.setSourceName(""))
+    intercept[IllegalArgumentException] {
+      transform(proto.Relation.newBuilder().setRead(read).build())
+    }
+  }
+
   test("Read with at syntax time travel") {
     // Version time travel
     val read = proto.Read

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala

@@ -343,16 +343,16 @@ class FindDataSourceTable(sparkSession: SparkSession) extends Rule[LogicalPlan]
     // is processed first but doesn't match the case above due to the !isStreaming guard,
     // so the NamedStreamingRelation case here can match.
     // We pass the sourceIdentifyingName directly to getStreamingRelation.
-    case NamedStreamingRelation(u: UnresolvedCatalogRelation, sourceIdentifyingName) =>
-      getStreamingRelation(u.tableMeta, u.options, sourceIdentifyingName)
+    case named @ NamedStreamingRelation(u: UnresolvedCatalogRelation, sourceIdentifyingName) =>
+      named.copy(child = getStreamingRelation(u.tableMeta, u.options, sourceIdentifyingName))
 
     // Handle NamedStreamingRelation wrapping SubqueryAlias(UnresolvedCatalogRelation)
     // - this happens when resolving streaming tables from catalogs where the table lookup
     // creates a SubqueryAlias wrapper around the UnresolvedCatalogRelation.
-    case NamedStreamingRelation(
+    case named @ NamedStreamingRelation(
         SubqueryAlias(alias, u: UnresolvedCatalogRelation), sourceIdentifyingName) =>
       val resolved = getStreamingRelation(u.tableMeta, u.options, sourceIdentifyingName)
-      SubqueryAlias(alias, resolved)
+      named.copy(child = SubqueryAlias(alias, resolved))
 
     // Fallback for streaming UnresolvedCatalogRelation that is NOT wrapped in
     // NamedStreamingRelation (e.g., from .readStream.table() API path).