python/pyspark/sql/tests/test_datasources.py - master - githubmirror/spark

python/pyspark/sql/tests/test_datasources.py
464 строки18 KB
Yicong-Huang
[SPARK-56254][PYTHON][CONNECT] Make spark.read.xml accept DataFrame input
15 апр 2026, 01:36
15 апр 2026, 01:36190e3b8
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import shutil
import tempfile
import uuid
import os

from pyspark.sql import Row
from pyspark.sql.datasource import InputPartition, DataSource
from pyspark.sql.types import IntegerType, StructField, StructType, LongType, StringType
from pyspark.errors import PySparkNotImplementedError
from pyspark.testing.sqlutils import ReusedSQLTestCase


class DataSourcesTestsMixin:
    def test_linesep_text(self):
        df = self.spark.read.text("python/test_support/sql/ages_newlines.csv", lineSep=",")
        expected = [
            Row(value="Joe"),
            Row(value="20"),
            Row(value='"Hi'),
            Row(value='\nI am Jeo"\nTom'),
            Row(value="30"),
            Row(value='"My name is Tom"\nHyukjin'),
            Row(value="25"),
            Row(value='"I am Hyukjin\n\nI love Spark!"\n'),
        ]
        self.assertEqual(df.collect(), expected)

        tpath = tempfile.mkdtemp()
        shutil.rmtree(tpath)
        try:
            df.write.text(tpath, lineSep="!")
            expected = [
                Row(value='Joe!20!"Hi!'),
                Row(value='I am Jeo"'),
                Row(value='Tom!30!"My name is Tom"'),
                Row(value='Hyukjin!25!"I am Hyukjin'),
                Row(value=""),
                Row(value='I love Spark!"'),
                Row(value="!"),
            ]
            readback = self.spark.read.text(tpath)
            self.assertEqual(readback.collect(), expected)
        finally:
            shutil.rmtree(tpath)

    def test_multiline_json(self):
        people1 = self.spark.read.json("python/test_support/sql/people.json")
        people_array = self.spark.read.json(
            "python/test_support/sql/people_array.json", multiLine=True
        )
        self.assertEqual(people1.collect(), people_array.collect())

    def test_encoding_json(self):
        people_array = self.spark.read.json(
            "python/test_support/sql/people_array_utf16le.json", multiLine=True, encoding="UTF-16LE"
        )
        expected = [Row(age=30, name="Andy"), Row(age=19, name="Justin")]
        self.assertEqual(people_array.collect(), expected)

    def test_linesep_json(self):
        df = self.spark.read.json("python/test_support/sql/people.json", lineSep=",")
        expected = [
            Row(_corrupt_record=None, name="Michael"),
            Row(_corrupt_record=' "age":30}\n{"name":"Justin"', name=None),
            Row(_corrupt_record=' "age":19}\n', name=None),
        ]
        self.assertEqual(df.collect(), expected)

        tpath = tempfile.mkdtemp()
        shutil.rmtree(tpath)
        try:
            df = self.spark.read.json("python/test_support/sql/people.json")
            df.write.json(tpath, lineSep="!!")
            readback = self.spark.read.json(tpath, lineSep="!!")
            self.assertEqual(readback.collect(), df.collect())
        finally:
            shutil.rmtree(tpath)

    def test_json_with_dataframe_input(self):
        json_df = self.spark.createDataFrame(
            [('{"name": "Alice", "age": 25}',), ('{"name": "Bob", "age": 30}',)],
            schema="value STRING",
        )
        result = self.spark.read.json(json_df)
        expected = [Row(age=25, name="Alice"), Row(age=30, name="Bob")]
        self.assertEqual(sorted(result.collect(), key=lambda r: r.name), expected)

    def test_json_with_dataframe_input_and_schema(self):
        json_df = self.spark.createDataFrame(
            [('{"name": "Alice", "age": 25}',), ('{"name": "Bob", "age": 30}',)],
            schema="value STRING",
        )
        result = self.spark.read.json(json_df, schema="name STRING, age INT")
        expected = [Row(name="Alice", age=25), Row(name="Bob", age=30)]
        self.assertEqual(sorted(result.collect(), key=lambda r: r.name), expected)

    def test_json_with_dataframe_input_non_string_column(self):
        int_df = self.spark.createDataFrame([(1,), (2,)], schema="value INT")
        with self.assertRaisesRegex(Exception, "DATAFRAME_INPUT_NOT_STRING_TYPE"):
            self.spark.read.json(int_df).collect()

    def test_json_with_dataframe_input_multiple_columns(self):
        multi_df = self.spark.createDataFrame(
            [('{"name": "Alice"}', "extra"), ('{"name": "Bob"}', "extra")],
            schema="value STRING, other STRING",
        )
        with self.assertRaisesRegex(Exception, "DATAFRAME_INPUT_NOT_SINGLE_COLUMN"):
            self.spark.read.json(multi_df).collect()

    def test_json_with_dataframe_input_zero_columns(self):
        empty_schema_df = self.spark.range(1).select()
        with self.assertRaisesRegex(Exception, "DATAFRAME_INPUT_NOT_SINGLE_COLUMN"):
            self.spark.read.json(empty_schema_df).collect()

    def test_multiline_csv(self):
        ages_newlines = self.spark.read.csv(
            "python/test_support/sql/ages_newlines.csv", multiLine=True
        )
        expected = [
            Row(_c0="Joe", _c1="20", _c2="Hi,\nI am Jeo"),
            Row(_c0="Tom", _c1="30", _c2="My name is Tom"),
            Row(_c0="Hyukjin", _c1="25", _c2="I am Hyukjin\n\nI love Spark!"),
        ]
        self.assertEqual(ages_newlines.collect(), expected)

    def test_ignorewhitespace_csv(self):
        tmpPath = tempfile.mkdtemp()
        shutil.rmtree(tmpPath)
        self.spark.createDataFrame([[" a", "b  ", " c "]]).write.csv(
            tmpPath, ignoreLeadingWhiteSpace=False, ignoreTrailingWhiteSpace=False
        )

        expected = [Row(value=" a,b  , c ")]
        readback = self.spark.read.text(tmpPath)
        self.assertEqual(readback.collect(), expected)
        shutil.rmtree(tmpPath)

    def test_csv_with_dataframe_input(self):
        csv_df = self.spark.createDataFrame(
            [("Alice,25",), ("Bob,30",)],
            schema="value STRING",
        )
        result = self.spark.read.csv(csv_df)
        expected = [Row(_c0="Alice", _c1="25"), Row(_c0="Bob", _c1="30")]
        self.assertEqual(sorted(result.collect(), key=lambda r: r._c0), expected)

    def test_csv_with_dataframe_input_and_schema(self):
        csv_df = self.spark.createDataFrame(
            [("Alice,25",), ("Bob,30",)],
            schema="value STRING",
        )
        result = self.spark.read.csv(csv_df, schema="name STRING, age INT")
        expected = [Row(name="Alice", age=25), Row(name="Bob", age=30)]
        self.assertEqual(sorted(result.collect(), key=lambda r: r.name), expected)

    def test_csv_with_dataframe_input_non_string_column(self):
        int_df = self.spark.createDataFrame([(1,), (2,)], schema="value INT")
        with self.assertRaisesRegex(Exception, "DATAFRAME_INPUT_NOT_STRING_TYPE"):
            self.spark.read.csv(int_df).collect()

    def test_csv_with_dataframe_input_multiple_columns(self):
        multi_df = self.spark.createDataFrame(
            [("Alice,25", "extra"), ("Bob,30", "extra")],
            schema="value STRING, other STRING",
        )
        with self.assertRaisesRegex(Exception, "DATAFRAME_INPUT_NOT_SINGLE_COLUMN"):
            self.spark.read.csv(multi_df).collect()

    def test_csv_with_dataframe_input_zero_columns(self):
        empty_schema_df = self.spark.range(1).select()
        with self.assertRaisesRegex(Exception, "DATAFRAME_INPUT_NOT_SINGLE_COLUMN"):
            self.spark.read.csv(empty_schema_df).collect()

    def test_xml(self):
        tmpPath = tempfile.mkdtemp()
        shutil.rmtree(tmpPath)
        xsdPath = tempfile.mkdtemp()
        xsdString = """<?xml version="1.0" encoding="UTF-8" ?>
          <xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema">
            <xs:element name="person">
              <xs:complexType>
                <xs:sequence>
                  <xs:element name="name" type="xs:string" />
                  <xs:element name="age" type="xs:long" />
                </xs:sequence>
              </xs:complexType>
            </xs:element>
          </xs:schema>"""
        try:
            with open(os.path.join(xsdPath, "people.xsd"), "w") as f:
                _ = f.write(xsdString)
            df = self.spark.createDataFrame([("Hyukjin", 100), ("Aria", 101), ("Arin", 102)]).toDF(
                "name", "age"
            )
            df.write.xml(tmpPath, rootTag="people", rowTag="person")
            people = self.spark.read.xml(
                tmpPath, rowTag="person", rowValidationXSDPath=os.path.join(xsdPath, "people.xsd")
            )
            expected = [
                Row(age=100, name="Hyukjin"),
                Row(age=101, name="Aria"),
                Row(age=102, name="Arin"),
            ]
            self.assertEqual(people.sort("age").collect(), expected)
            self.assertEqual(
                people.schema,
                StructType(
                    [StructField("age", LongType(), True), StructField("name", StringType(), True)]
                ),
            )
        finally:
            shutil.rmtree(tmpPath)
            shutil.rmtree(xsdPath)

    def test_xml_sampling_ratio(self):
        rdd = self.spark.sparkContext.range(0, 100, 1, 1).map(
            lambda x: "<p><a>0.1</a></p>" if x == 1 else "<p><a>%s</a></p>" % str(x)
        )
        schema = self.spark.read.option("samplingRatio", 0.5).xml(rdd).schema
        self.assertEqual(schema, StructType([StructField("a", LongType(), True)]))

    def test_xml_with_dataframe_input(self):
        xml_df = self.spark.createDataFrame(
            [
                ("<person><name>Alice</name><age>25</age></person>",),
                ("<person><name>Bob</name><age>30</age></person>",),
            ],
            schema="value STRING",
        )
        result = self.spark.read.option("rowTag", "person").xml(xml_df)
        expected = [Row(age=25, name="Alice"), Row(age=30, name="Bob")]
        self.assertEqual(sorted(result.collect(), key=lambda r: r.name), expected)

    def test_xml_with_dataframe_input_and_schema(self):
        xml_df = self.spark.createDataFrame(
            [
                ("<person><name>Alice</name><age>25</age></person>",),
                ("<person><name>Bob</name><age>30</age></person>",),
            ],
            schema="value STRING",
        )
        result = self.spark.read.option("rowTag", "person").xml(
            xml_df, schema="name STRING, age INT"
        )
        expected = [Row(name="Alice", age=25), Row(name="Bob", age=30)]
        self.assertEqual(sorted(result.collect(), key=lambda r: r.name), expected)

    def test_xml_with_dataframe_input_non_string_column(self):
        int_df = self.spark.createDataFrame([(1,), (2,)], schema="value INT")
        with self.assertRaisesRegex(Exception, "DATAFRAME_INPUT_NOT_STRING_TYPE"):
            self.spark.read.option("rowTag", "person").xml(int_df).collect()

    def test_xml_with_dataframe_input_multiple_columns(self):
        multi_df = self.spark.createDataFrame(
            [
                ("<person><name>Alice</name></person>", "extra"),
                ("<person><name>Bob</name></person>", "extra"),
            ],
            schema="value STRING, other STRING",
        )
        with self.assertRaisesRegex(Exception, "DATAFRAME_INPUT_NOT_SINGLE_COLUMN"):
            self.spark.read.option("rowTag", "person").xml(multi_df).collect()

    def test_xml_with_dataframe_input_zero_columns(self):
        empty_schema_df = self.spark.range(1).select()
        with self.assertRaisesRegex(Exception, "DATAFRAME_INPUT_NOT_SINGLE_COLUMN"):
            self.spark.read.option("rowTag", "person").xml(empty_schema_df).collect()

    def test_read_multiple_orc_file(self):
        df = self.spark.read.orc(
            [
                "python/test_support/sql/orc_partitioned/b=0/c=0",
                "python/test_support/sql/orc_partitioned/b=1/c=1",
            ]
        )
        self.assertEqual(2, df.count())

    def test_read_text_file_list(self):
        df = self.spark.read.text(
            ["python/test_support/sql/text-test.txt", "python/test_support/sql/text-test.txt"]
        )
        count = df.count()
        self.assertEqual(count, 4)

    def test_json_sampling_ratio(self):
        rdd = self.spark.sparkContext.range(0, 100, 1, 1).map(
            lambda x: '{"a":0.1}' if x == 1 else '{"a":%s}' % str(x)
        )
        schema = (
            self.spark.read.option("inferSchema", True)
            .option("samplingRatio", 0.5)
            .json(rdd)
            .schema
        )
        self.assertEqual(schema, StructType([StructField("a", LongType(), True)]))

    def test_csv_sampling_ratio(self):
        rdd = self.spark.sparkContext.range(0, 100, 1, 1).map(lambda x: "0.1" if x == 1 else str(x))
        schema = self.spark.read.option("inferSchema", True).csv(rdd, samplingRatio=0.5).schema
        self.assertEqual(schema, StructType([StructField("_c0", IntegerType(), True)]))

    def test_checking_csv_header(self):
        path = tempfile.mkdtemp()
        shutil.rmtree(path)
        try:
            self.spark.createDataFrame([[1, 1000], [2000, 2]]).toDF("f1", "f2").write.option(
                "header", "true"
            ).csv(path)
            schema = StructType(
                [
                    StructField("f2", IntegerType(), nullable=True),
                    StructField("f1", IntegerType(), nullable=True),
                ]
            )
            df = (
                self.spark.read.option("header", "true")
                .option("quote", None)
                .schema(schema)
                .csv(path, enforceSchema=False)
            )
            self.assertRaisesRegex(
                Exception, "CSV header does not conform to the schema", lambda: df.collect()
            )
        finally:
            shutil.rmtree(path)

    def test_ignore_column_of_all_nulls(self):
        path = tempfile.mkdtemp()
        shutil.rmtree(path)
        try:
            df = self.spark.createDataFrame(
                [
                    ["""{"a":null, "b":1, "c":3.0}"""],
                    ["""{"a":null, "b":null, "c":"string"}"""],
                    ["""{"a":null, "b":null, "c":null}"""],
                ]
            )
            df.write.text(path)
            schema = StructType(
                [
                    StructField("b", LongType(), nullable=True),
                    StructField("c", StringType(), nullable=True),
                ]
            )
            readback = self.spark.read.json(path, dropFieldIfAllNull=True)
            self.assertEqual(readback.schema, schema)
        finally:
            shutil.rmtree(path)

    def test_jdbc(self):
        db = f"memory:{uuid.uuid4()}"
        url = f"jdbc:derby:{db}"
        dbtable = "test_table"

        try:
            df = self.spark.range(10)
            df.write.jdbc(url=f"{url};create=true", table=dbtable)

            readback = self.spark.read.jdbc(url=url, table=dbtable)
            self.assertEqual(sorted(df.collect()), sorted(readback.collect()))

            additional_arguments = dict(column="id", lowerBound=3, upperBound=8, numPartitions=10)
            readback = self.spark.read.jdbc(url=url, table=dbtable, **additional_arguments)
            self.assertEqual(sorted(df.collect()), sorted(readback.collect()))

            additional_arguments = dict(predicates=['"id" < 5'])
            readback = self.spark.read.jdbc(url=url, table=dbtable, **additional_arguments)
            self.assertEqual(sorted(df.filter("id < 5").collect()), sorted(readback.collect()))
        finally:
            # Clean up.
            with self.assertRaisesRegex(Exception, f"Database '{db}' dropped."):
                self.spark.read.jdbc(url=f"{url};drop=true", table=dbtable).collect()

    def test_jdbc_format(self):
        db = f"memory:{uuid.uuid4()}"
        url = f"jdbc:derby:{db}"
        dbtable = "test_table"

        try:
            df = self.spark.range(10)
            df.write.format("jdbc").options(url=f"{url};create=true", dbtable=dbtable).save()
            readback = self.spark.read.format("jdbc").options(url=url, dbtable=dbtable).load()
            self.assertEqual(sorted(df.collect()), sorted(readback.collect()))
        finally:
            # Clean up.
            with self.assertRaisesRegex(Exception, f"Database '{db}' dropped."):
                self.spark.read.format("jdbc").options(
                    url=f"{url};drop=true", dbtable=dbtable
                ).load().collect()

    def test_custom_data_source(self):
        class MyCustomDataSource(DataSource):
            pass

        custom_data_source = MyCustomDataSource(options={"path": "/path/to/custom/data"})

        with self.assertRaises(PySparkNotImplementedError) as pe:
            custom_data_source.schema()

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_IMPLEMENTED",
            messageParameters={"feature": "schema"},
        )

        with self.assertRaises(PySparkNotImplementedError) as pe:
            custom_data_source.reader(schema=None)

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_IMPLEMENTED",
            messageParameters={"feature": "reader"},
        )

        with self.assertRaises(PySparkNotImplementedError) as pe:
            custom_data_source.writer(schema=None, overwrite=False)

        self.check_error(
            exception=pe.exception,
            errorClass="NOT_IMPLEMENTED",
            messageParameters={"feature": "writer"},
        )

    def test_input_partition(self):
        partition = InputPartition(1)
        expected_repr = "InputPartition(value=1)"
        actual_repr = repr(partition)
        self.assertEqual(expected_repr, actual_repr)

        class RangeInputPartition(InputPartition):
            def __init__(self, start, end):
                super().__init__((start, end))

        partition = RangeInputPartition(1, 3)
        expected_repr = "RangeInputPartition(value=(1, 3))"
        actual_repr = repr(partition)
        self.assertEqual(expected_repr, actual_repr)


class DataSourcesTests(DataSourcesTestsMixin, ReusedSQLTestCase):
    pass


if __name__ == "__main__":
    from pyspark.testing import main

    main()