[SPARK-42570][CONNECT][PYTHON] Fix DataFrameReader to use the default source

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

@@ -122,8 +122,10 @@ message Read {
   }
 
   message DataSource {
-    // (Required) Supported formats include: parquet, orc, text, json, parquet, csv, avro.
-    string format = 1;
+    // (Optional) Supported formats include: parquet, orc, text, json, parquet, csv, avro.
+    //
+    // If not set, the value from SQL conf 'spark.sql.sources.default' will be used.
+    optional string format = 1;
 
     // (Optional) If not set, Spark will infer the schema.
     //

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

@@ -667,12 +667,11 @@ class SparkConnectPlanner(val session: SparkSession) {
         UnresolvedRelation(multipartIdentifier)
 
       case proto.Read.ReadTypeCase.DATA_SOURCE =>
-        if (rel.getDataSource.getFormat == "") {
-          throw InvalidPlanInput("DataSource requires a format")
-        }
         val localMap = CaseInsensitiveMap[String](rel.getDataSource.getOptionsMap.asScala.toMap)
         val reader = session.read
-        reader.format(rel.getDataSource.getFormat)
+        if (rel.getDataSource.hasFormat) {
+          reader.format(rel.getDataSource.getFormat)
+        }
         localMap.foreach { case (key, value) => reader.option(key, value) }
         if (rel.getDataSource.hasSchema && rel.getDataSource.getSchema.nonEmpty) {
 

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

@@ -332,18 +332,6 @@ class SparkConnectPlannerSuite extends SparkFunSuite with SparkConnectPlanTest {
     assert(res.nodeName == "Aggregate")
   }
 
-  test("Invalid DataSource") {
-    val dataSource = proto.Read.DataSource.newBuilder()
-
-    val e = intercept[InvalidPlanInput](
-      transform(
-        proto.Relation
-          .newBuilder()
-          .setRead(proto.Read.newBuilder().setDataSource(dataSource))
-          .build()))
-    assert(e.getMessage.contains("DataSource requires a format"))
-  }
-
   test("Test invalid deduplicate") {
     val deduplicate = proto.Deduplicate
       .newBuilder()

python/pyspark/sql/connect/plan.py

@@ -255,15 +255,14 @@ class DataSource(LogicalPlan):
 
     def __init__(
         self,
-        format: str,
+        format: Optional[str] = None,
         schema: Optional[str] = None,
         options: Optional[Mapping[str, str]] = None,
         paths: Optional[List[str]] = None,
     ) -> None:
         super().__init__(None)
 
-        assert isinstance(format, str) and format != ""
-
+        assert format is None or isinstance(format, str)
         assert schema is None or isinstance(schema, str)
 
         if options is not None:
@@ -282,7 +281,8 @@ def __init__(
 
     def plan(self, session: "SparkConnectClient") -> proto.Relation:
         plan = self._create_proto_relation()
-        plan.read.data_source.format = self._format
+        if self._format is not None:
+            plan.read.data_source.format = self._format
         if self._schema is not None:
             plan.read.data_source.schema = self._schema
         if self._options is not None and len(self._options) > 0:

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

@@ -602,7 +602,10 @@ class Read(google.protobuf.message.Message):
         OPTIONS_FIELD_NUMBER: builtins.int
         PATHS_FIELD_NUMBER: builtins.int
         format: builtins.str
-        """(Required) Supported formats include: parquet, orc, text, json, parquet, csv, avro."""
+        """(Optional) Supported formats include: parquet, orc, text, json, parquet, csv, avro.
+
+        If not set, the value from SQL conf 'spark.sql.sources.default' will be used.
+        """
         schema: builtins.str
         """(Optional) If not set, Spark will infer the schema.
 
@@ -624,17 +627,29 @@ class Read(google.protobuf.message.Message):
         def __init__(
             self,
             *,
-            format: builtins.str = ...,
+            format: builtins.str | None = ...,
             schema: builtins.str | None = ...,
             options: collections.abc.Mapping[builtins.str, builtins.str] | None = ...,
             paths: collections.abc.Iterable[builtins.str] | None = ...,
         ) -> None: ...
         def HasField(
-            self, field_name: typing_extensions.Literal["_schema", b"_schema", "schema", b"schema"]
+            self,
+            field_name: typing_extensions.Literal[
+                "_format",
+                b"_format",
+                "_schema",
+                b"_schema",
+                "format",
+                b"format",
+                "schema",
+                b"schema",
+            ],
         ) -> builtins.bool: ...
         def ClearField(
             self,
             field_name: typing_extensions.Literal[
+                "_format",
+                b"_format",
                 "_schema",
                 b"_schema",
                 "format",
@@ -647,6 +662,11 @@ class Read(google.protobuf.message.Message):
                 b"schema",
             ],
         ) -> None: ...
+        @typing.overload
+        def WhichOneof(
+            self, oneof_group: typing_extensions.Literal["_format", b"_format"]
+        ) -> typing_extensions.Literal["format"] | None: ...
+        @typing.overload
         def WhichOneof(
             self, oneof_group: typing_extensions.Literal["_schema", b"_schema"]
         ) -> typing_extensions.Literal["schema"] | None: ...

python/pyspark/sql/connect/readwriter.py

@@ -63,7 +63,7 @@ class DataFrameReader(OptionUtils):
 
     def __init__(self, client: "SparkSession"):
         self._client = client
-        self._format = ""
+        self._format: Optional[str] = None
         self._schema = ""
         self._options: Dict[str, str] = {}
 

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

@@ -22,15 +22,7 @@
 
 
 class ReadwriterParityTests(ReadwriterTestsMixin, ReusedConnectTestCase):
-    # TODO(SPARK-41834): Implement SparkSession.conf
-    @unittest.skip("Fails in Spark Connect, should enable.")
-    def test_save_and_load(self):
-        super().test_save_and_load()
-
-    # TODO(SPARK-41834): Implement SparkSession.conf
-    @unittest.skip("Fails in Spark Connect, should enable.")
-    def test_save_and_load_builder(self):
-        super().test_save_and_load_builder()
+    pass
 
 
 class ReadwriterV2ParityTests(ReadwriterV2TestsMixin, ReusedConnectTestCase):

python/pyspark/sql/tests/test_readwriter.py

@@ -31,75 +31,77 @@ def test_save_and_load(self):
         df = self.df
         tmpPath = tempfile.mkdtemp()
         shutil.rmtree(tmpPath)
-        df.write.json(tmpPath)
-        actual = self.spark.read.json(tmpPath)
-        self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
-
-        schema = StructType([StructField("value", StringType(), True)])
-        actual = self.spark.read.json(tmpPath, schema)
-        self.assertEqual(sorted(df.select("value").collect()), sorted(actual.collect()))
-
-        df.write.json(tmpPath, "overwrite")
-        actual = self.spark.read.json(tmpPath)
-        self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
-
-        df.write.save(
-            format="json",
-            mode="overwrite",
-            path=tmpPath,
-            noUse="this options will not be used in save.",
-        )
-        actual = self.spark.read.load(
-            format="json", path=tmpPath, noUse="this options will not be used in load."
-        )
-        self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
-
-        defaultDataSourceName = self.spark.conf.get(
-            "spark.sql.sources.default", "org.apache.spark.sql.parquet"
-        )
-        self.spark.sql("SET spark.sql.sources.default=org.apache.spark.sql.json")
-        actual = self.spark.read.load(path=tmpPath)
-        self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
-        self.spark.sql("SET spark.sql.sources.default=" + defaultDataSourceName)
+        try:
+            df.write.json(tmpPath)
+            actual = self.spark.read.json(tmpPath)
+            self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
+
+            schema = StructType([StructField("value", StringType(), True)])
+            actual = self.spark.read.json(tmpPath, schema)
+            self.assertEqual(sorted(df.select("value").collect()), sorted(actual.collect()))
+
+            df.write.json(tmpPath, "overwrite")
+            actual = self.spark.read.json(tmpPath)
+            self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
+
+            df.write.save(
+                format="json",
+                mode="overwrite",
+                path=tmpPath,
+                noUse="this options will not be used in save.",
+            )
+            actual = self.spark.read.load(
+                format="json", path=tmpPath, noUse="this options will not be used in load."
+            )
+            self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
 
-        csvpath = os.path.join(tempfile.mkdtemp(), "data")
-        df.write.option("quote", None).format("csv").save(csvpath)
+            try:
+                self.spark.sql("SET spark.sql.sources.default=org.apache.spark.sql.json").collect()
+                actual = self.spark.read.load(path=tmpPath)
+                self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
+            finally:
+                self.spark.sql("RESET spark.sql.sources.default").collect()
 
-        shutil.rmtree(tmpPath)
+            csvpath = os.path.join(tempfile.mkdtemp(), "data")
+            df.write.option("quote", None).format("csv").save(csvpath)
+        finally:
+            shutil.rmtree(tmpPath)
 
     def test_save_and_load_builder(self):
         df = self.df
         tmpPath = tempfile.mkdtemp()
         shutil.rmtree(tmpPath)
-        df.write.json(tmpPath)
-        actual = self.spark.read.json(tmpPath)
-        self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
-
-        schema = StructType([StructField("value", StringType(), True)])
-        actual = self.spark.read.json(tmpPath, schema)
-        self.assertEqual(sorted(df.select("value").collect()), sorted(actual.collect()))
-
-        df.write.mode("overwrite").json(tmpPath)
-        actual = self.spark.read.json(tmpPath)
-        self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
-
-        df.write.mode("overwrite").options(noUse="this options will not be used in save.").option(
-            "noUse", "this option will not be used in save."
-        ).format("json").save(path=tmpPath)
-        actual = self.spark.read.format("json").load(
-            path=tmpPath, noUse="this options will not be used in load."
-        )
-        self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
-
-        defaultDataSourceName = self.spark.conf.get(
-            "spark.sql.sources.default", "org.apache.spark.sql.parquet"
-        )
-        self.spark.sql("SET spark.sql.sources.default=org.apache.spark.sql.json")
-        actual = self.spark.read.load(path=tmpPath)
-        self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
-        self.spark.sql("SET spark.sql.sources.default=" + defaultDataSourceName)
-
-        shutil.rmtree(tmpPath)
+        try:
+            df.write.json(tmpPath)
+            actual = self.spark.read.json(tmpPath)
+            self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
+
+            schema = StructType([StructField("value", StringType(), True)])
+            actual = self.spark.read.json(tmpPath, schema)
+            self.assertEqual(sorted(df.select("value").collect()), sorted(actual.collect()))
+
+            df.write.mode("overwrite").json(tmpPath)
+            actual = self.spark.read.json(tmpPath)
+            self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
+
+            df.write.mode("overwrite").options(
+                noUse="this options will not be used in save."
+            ).option("noUse", "this option will not be used in save.").format("json").save(
+                path=tmpPath
+            )
+            actual = self.spark.read.format("json").load(
+                path=tmpPath, noUse="this options will not be used in load."
+            )
+            self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
+
+            try:
+                self.spark.sql("SET spark.sql.sources.default=org.apache.spark.sql.json").collect()
+                actual = self.spark.read.load(path=tmpPath)
+                self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
+            finally:
+                self.spark.sql("RESET spark.sql.sources.default").collect()
+        finally:
+            shutil.rmtree(tmpPath)
 
     def test_bucketed_write(self):
         data = [