#
# 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 os
import platform
import tempfile
import unittest
from datetime import datetime
from decimal import Decimal
from typing import Callable, Iterable, List, Union

from pyspark.errors import AnalysisException, PythonException
from pyspark.sql.datasource import (
    CaseInsensitiveDict,
    DataSource,
    DataSourceArrowWriter,
    DataSourceReader,
    DataSourceWriter,
    EqualNullSafe,
    EqualTo,
    Filter,
    GreaterThan,
    GreaterThanOrEqual,
    In,
    InputPartition,
    IsNotNull,
    IsNull,
    LessThan,
    LessThanOrEqual,
    Not,
    StringContains,
    StringEndsWith,
    StringStartsWith,
    WriterCommitMessage,
)
from pyspark.sql.functions import spark_partition_id
from pyspark.sql.session import SparkSession
from pyspark.sql.types import Row, StructType, VariantVal
from pyspark.testing import assertDataFrameEqual
from pyspark.testing.sqlutils import (
    SPARK_HOME,
    ReusedSQLTestCase,
    have_pyarrow,
    pyarrow_requirement_message,
)


@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message)
class BasePythonDataSourceTestsMixin:
    spark: SparkSession

    def test_basic_data_source_class(self):
        class MyDataSource(DataSource):
            ...

        options = dict(a=1, b=2)
        ds = MyDataSource(options=options)
        self.assertEqual(ds.options, options)
        self.assertEqual(ds.name(), "MyDataSource")
        with self.assertRaises(NotImplementedError):
            ds.schema()
        with self.assertRaises(NotImplementedError):
            ds.reader(None)
        with self.assertRaises(NotImplementedError):
            ds.writer(None, None)

    def test_basic_data_source_reader_class(self):
        class MyDataSourceReader(DataSourceReader):
            def read(self, partition):
                yield None,

        reader = MyDataSourceReader()
        self.assertEqual(list(reader.read(None)), [(None,)])

    def test_data_source_register(self):
        class TestReader(DataSourceReader):
            def read(self, partition):
                yield (
                    0,
                    1,
                    VariantVal.parseJson('{"c":1}'),
                    {"v": VariantVal.parseJson('{"d":2}')},
                    [VariantVal.parseJson('{"e":3}')],
                    {"v1": VariantVal.parseJson('{"f":4}'), "v2": VariantVal.parseJson('{"g":5}')},
                )

        class TestDataSource(DataSource):
            def schema(self):
                return (
                    "a INT, b INT, c VARIANT, d STRUCT<v VARIANT>, e ARRAY<VARIANT>,"
                    "f MAP<STRING, VARIANT>"
                )

            def reader(self, schema):
                return TestReader()

        self.spark.dataSource.register(TestDataSource)
        df = self.spark.read.format("TestDataSource").load()
        assertDataFrameEqual(
            df.selectExpr(
                "a", "b", "to_json(c) c", "to_json(d.v) d", "to_json(e[0]) e", "to_json(f['v2']) f"
            ),
            [Row(a=0, b=1, c='{"c":1}', d='{"d":2}', e='{"e":3}', f='{"g":5}')],
        )

        class MyDataSource(TestDataSource):
            @classmethod
            def name(cls):
                return "TestDataSource"

            def schema(self):
                return (
                    "c INT, d INT, e VARIANT, f STRUCT<v VARIANT>, g ARRAY<VARIANT>,"
                    "h MAP<STRING, VARIANT>"
                )

        # Should be able to register the data source with the same name.
        self.spark.dataSource.register(MyDataSource)

        df = self.spark.read.format("TestDataSource").load()
        assertDataFrameEqual(
            df.selectExpr(
                "c", "d", "to_json(e) e", "to_json(f.v) f", "to_json(g[0]) g", "to_json(h['v2']) h"
            ),
            [Row(c=0, d=1, e='{"c":1}', f='{"d":2}', g='{"e":3}', h='{"g":5}')],
        )

    def register_data_source(
        self,
        read_func: Callable,
        partition_func: Callable = None,
        output: Union[str, StructType] = "i int, j int",
        name: str = "test",
    ):
        class TestDataSourceReader(DataSourceReader):
            def __init__(self, schema):
                self.schema = schema

            def partitions(self):
                if partition_func is not None:
                    return partition_func()
                else:
                    raise NotImplementedError

            def read(self, partition):
                return read_func(self.schema, partition)

        class TestDataSource(DataSource):
            @classmethod
            def name(cls):
                return name

            def schema(self):
                return output

            def reader(self, schema) -> "DataSourceReader":
                return TestDataSourceReader(schema)

        self.spark.dataSource.register(TestDataSource)

    def test_data_source_read_output_tuple(self):
        self.register_data_source(read_func=lambda schema, partition: iter([(0, 1)]))
        df = self.spark.read.format("test").load()
        assertDataFrameEqual(df, [Row(0, 1)])

    def test_data_source_read_output_list(self):
        self.register_data_source(read_func=lambda schema, partition: iter([[0, 1]]))
        df = self.spark.read.format("test").load()
        assertDataFrameEqual(df, [Row(0, 1)])

    def test_data_source_read_output_row(self):
        self.register_data_source(read_func=lambda schema, partition: iter([Row(0, 1)]))
        df = self.spark.read.format("test").load()
        assertDataFrameEqual(df, [Row(0, 1)])

    def test_data_source_read_output_named_row(self):
        self.register_data_source(
            read_func=lambda schema, partition: iter([Row(j=1, i=0), Row(i=1, j=2)])
        )
        df = self.spark.read.format("test").load()
        assertDataFrameEqual(df, [Row(0, 1), Row(1, 2)])

    def test_data_source_read_output_named_row_with_wrong_schema(self):
        self.register_data_source(
            read_func=lambda schema, partition: iter([Row(i=1, j=2), Row(j=3, k=4)])
        )
        with self.assertRaisesRegex(
            PythonException,
            r"\[DATA_SOURCE_RETURN_SCHEMA_MISMATCH\] Return schema mismatch in the result",
        ):
            self.spark.read.format("test").load().show()

    def test_data_source_read_output_none(self):
        self.register_data_source(read_func=lambda schema, partition: None)
        df = self.spark.read.format("test").load()
        with self.assertRaisesRegex(PythonException, "DATA_SOURCE_INVALID_RETURN_TYPE"):
            assertDataFrameEqual(df, [])

    def test_data_source_read_output_empty_iter(self):
        self.register_data_source(read_func=lambda schema, partition: iter([]))
        df = self.spark.read.format("test").load()
        assertDataFrameEqual(df, [])

    def test_data_source_read_cast_output_schema(self):
        self.register_data_source(
            read_func=lambda schema, partition: iter([(0, 1)]), output="i long, j string"
        )
        df = self.spark.read.format("test").load()
        assertDataFrameEqual(df, [Row(i=0, j="1")])

    def test_data_source_read_output_with_partition(self):
        def partition_func():
            return [InputPartition(0), InputPartition(1)]

        def read_func(schema, partition):
            if partition.value == 0:
                return iter([])
            elif partition.value == 1:
                yield (0, 1)

        self.register_data_source(read_func=read_func, partition_func=partition_func)
        df = self.spark.read.format("test").load()
        assertDataFrameEqual(df, [Row(0, 1)])

    def test_data_source_read_output_with_schema_mismatch(self):
        self.register_data_source(read_func=lambda schema, partition: iter([(0, 1)]))
        df = self.spark.read.format("test").schema("i int").load()
        with self.assertRaisesRegex(PythonException, "DATA_SOURCE_RETURN_SCHEMA_MISMATCH"):
            df.collect()
        self.register_data_source(
            read_func=lambda schema, partition: iter([(0, 1)]), output="i int, j int, k int"
        )
        with self.assertRaisesRegex(PythonException, "DATA_SOURCE_RETURN_SCHEMA_MISMATCH"):
            df.collect()

    def test_read_with_invalid_return_row_type(self):
        self.register_data_source(read_func=lambda schema, partition: iter([1]))
        df = self.spark.read.format("test").load()
        with self.assertRaisesRegex(PythonException, "DATA_SOURCE_INVALID_RETURN_TYPE"):
            df.collect()

    def test_in_memory_data_source(self):
        class InMemDataSourceReader(DataSourceReader):
            DEFAULT_NUM_PARTITIONS: int = 3

            def __init__(self, options):
                self.options = options

            def partitions(self):
                if "num_partitions" in self.options:
                    num_partitions = int(self.options["num_partitions"])
                else:
                    num_partitions = self.DEFAULT_NUM_PARTITIONS
                return [InputPartition(i) for i in range(num_partitions)]

            def read(self, partition):
                yield partition.value, str(partition.value)

        class InMemoryDataSource(DataSource):
            @classmethod
            def name(cls):
                return "memory"

            def schema(self):
                return "x INT, y STRING"

            def reader(self, schema) -> "DataSourceReader":
                return InMemDataSourceReader(self.options)

        self.spark.dataSource.register(InMemoryDataSource)
        df = self.spark.read.format("memory").load()
        self.assertEqual(df.select(spark_partition_id()).distinct().count(), 3)
        assertDataFrameEqual(df, [Row(x=0, y="0"), Row(x=1, y="1"), Row(x=2, y="2")])

        df = self.spark.read.format("memory").option("num_partitions", 2).load()
        assertDataFrameEqual(df, [Row(x=0, y="0"), Row(x=1, y="1")])
        self.assertEqual(df.select(spark_partition_id()).distinct().count(), 2)

    def test_filter_pushdown(self):
        class TestDataSourceReader(DataSourceReader):
            def __init__(self):
                self.has_filter = False

            def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
                assert set(filters) == {
                    IsNotNull(("x",)),
                    IsNotNull(("y",)),
                    EqualTo(("x",), 1),
                    EqualTo(("y",), 2),
                }, filters
                self.has_filter = True
                # pretend we support x = 1 filter but in fact we don't
                # so we only return y = 2 filter
                yield filters[filters.index(EqualTo(("y",), 2))]

            def partitions(self):
                assert self.has_filter
                return super().partitions()

            def read(self, partition):
                assert self.has_filter
                yield [1, 1]
                yield [1, 2]
                yield [2, 2]

        class TestDataSource(DataSource):
            @classmethod
            def name(cls):
                return "test"

            def schema(self):
                return "x int, y int"

            def reader(self, schema) -> "DataSourceReader":
                return TestDataSourceReader()

        with self.sql_conf({"spark.sql.python.filterPushdown.enabled": True}):
            self.spark.dataSource.register(TestDataSource)
            df = self.spark.read.format("test").load().filter("x = 1 and y = 2")
            # only the y = 2 filter is applied post scan
            assertDataFrameEqual(df, [Row(x=1, y=2), Row(x=2, y=2)])

    def test_extraneous_filter(self):
        class TestDataSourceReader(DataSourceReader):
            def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
                yield EqualTo(("x",), 1)

            def partitions(self):
                assert False

            def read(self, partition):
                assert False

        class TestDataSource(DataSource):
            @classmethod
            def name(cls):
                return "test"

            def schema(self):
                return "x int"

            def reader(self, schema) -> "DataSourceReader":
                return TestDataSourceReader()

        with self.sql_conf({"spark.sql.python.filterPushdown.enabled": True}):
            self.spark.dataSource.register(TestDataSource)
            with self.assertRaisesRegex(Exception, "DATA_SOURCE_EXTRANEOUS_FILTERS"):
                self.spark.read.format("test").load().filter("x = 1").show()

    def test_filter_pushdown_error(self):
        error_str = "dummy error"

        class TestDataSourceReader(DataSourceReader):
            def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
                raise Exception(error_str)

            def read(self, partition):
                yield [1]

        class TestDataSource(DataSource):
            def schema(self):
                return "x int"

            def reader(self, schema) -> "DataSourceReader":
                return TestDataSourceReader()

        with self.sql_conf({"spark.sql.python.filterPushdown.enabled": True}):
            self.spark.dataSource.register(TestDataSource)
            df = self.spark.read.format("TestDataSource").load().filter("x = 1 or x is null")
            assertDataFrameEqual(df, [Row(x=1)])  # works when not pushing down filters
            with self.assertRaisesRegex(Exception, error_str):
                df.filter("x = 1").explain()

    def test_filter_pushdown_disabled(self):
        class TestDataSourceReader(DataSourceReader):
            def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
                assert False

            def read(self, partition):
                assert False

        class TestDataSource(DataSource):
            def reader(self, schema) -> "DataSourceReader":
                return TestDataSourceReader()

        with self.sql_conf({"spark.sql.python.filterPushdown.enabled": False}):
            self.spark.dataSource.register(TestDataSource)
            df = self.spark.read.format("TestDataSource").schema("x int").load()
            with self.assertRaisesRegex(Exception, "DATA_SOURCE_PUSHDOWN_DISABLED"):
                df.show()

    def _check_filters(self, sql_type, sql_filter, python_filters):
        """
        Parameters
        ----------
        sql_type: str
            The SQL type of the column x.
        sql_filter: str
            A SQL filter using the column x.
        python_filters: List[Filter]
            The expected python filters to be pushed down.
        """

        class TestDataSourceReader(DataSourceReader):
            def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
                actual = [f for f in filters if not isinstance(f, IsNotNull)]
                expected = python_filters
                assert actual == expected, (actual, expected)
                return filters

            def read(self, partition):
                yield from []

        class TestDataSource(DataSource):
            def schema(self):
                return f"x {sql_type}"

            def reader(self, schema) -> "DataSourceReader":
                return TestDataSourceReader()

        with self.sql_conf({"spark.sql.python.filterPushdown.enabled": True}):
            self.spark.dataSource.register(TestDataSource)
            df = self.spark.read.format("TestDataSource").load().filter(sql_filter)
            df.count()

    def test_unsupported_filter(self):
        self._check_filters(
            "struct<a:int, b:int, c:int>", "x.a = 1 and x.b = x.c", [EqualTo(("x", "a"), 1)]
        )
        self._check_filters("int", "x = 1 or x > 2", [])
        self._check_filters("int", "(0 < x and x < 1) or x = 2", [])
        self._check_filters("int", "x % 5 = 1", [])
        self._check_filters("array<int>", "x[0] = 1", [])
        self._check_filters("string", "x like 'a%a%'", [])
        self._check_filters("string", "x ilike 'a'", [])
        self._check_filters("string", "x = 'a' collate zh", [])

    def test_filter_value_type(self):
        self._check_filters("int", "x = 1", [EqualTo(("x",), 1)])
        self._check_filters("int", "x = null", [EqualTo(("x",), None)])
        self._check_filters("float", "x = 3 / 2", [EqualTo(("x",), 1.5)])
        self._check_filters("string", "x = '1'", [EqualTo(("x",), "1")])
        self._check_filters("array<int>", "x = array(1, 2)", [EqualTo(("x",), [1, 2])])
        self._check_filters(
            "struct<x:int>", "x = named_struct('x', 1)", [EqualTo(("x",), {"x": 1})]
        )
        self._check_filters(
            "decimal", "x in (1.1, 2.1)", [In(("x",), [Decimal(1.1), Decimal(2.1)])]
        )
        self._check_filters(
            "timestamp_ntz",
            "x = timestamp_ntz '2020-01-01 00:00:00'",
            [EqualTo(("x",), datetime.strptime("2020-01-01 00:00:00", "%Y-%m-%d %H:%M:%S"))],
        )
        self._check_filters(
            "interval second",
            "x = interval '2' second",
            [],  # intervals are not supported
        )

    def test_filter_type(self):
        self._check_filters("boolean", "x", [EqualTo(("x",), True)])
        self._check_filters("boolean", "not x", [Not(EqualTo(("x",), True))])
        self._check_filters("int", "x is null", [IsNull(("x",))])
        self._check_filters("int", "x <> 0", [Not(EqualTo(("x",), 0))])
        self._check_filters("int", "x <=> 1", [EqualNullSafe(("x",), 1)])
        self._check_filters("int", "1 < x", [GreaterThan(("x",), 1)])
        self._check_filters("int", "1 <= x", [GreaterThanOrEqual(("x",), 1)])
        self._check_filters("int", "x < 1", [LessThan(("x",), 1)])
        self._check_filters("int", "x <= 1", [LessThanOrEqual(("x",), 1)])
        self._check_filters("string", "x like 'a%'", [StringStartsWith(("x",), "a")])
        self._check_filters("string", "x like '%a'", [StringEndsWith(("x",), "a")])
        self._check_filters("string", "x like '%a%'", [StringContains(("x",), "a")])
        self._check_filters(
            "string", "x like 'a%b'", [StringStartsWith(("x",), "a"), StringEndsWith(("x",), "b")]
        )
        self._check_filters("int", "x in (1, 2)", [In(("x",), [1, 2])])

    def test_filter_nested_column(self):
        self._check_filters("struct<y:int>", "x.y = 1", [EqualTo(("x", "y"), 1)])

    def _get_test_json_data_source(self):
        import json
        import os
        from dataclasses import dataclass

        class TestJsonReader(DataSourceReader):
            def __init__(self, options):
                self.options = options

            def read(self, partition):
                path = self.options.get("path")
                if path is None:
                    raise Exception("path is not specified")
                with open(path, "r") as file:
                    for line in file.readlines():
                        if line.strip():
                            data = json.loads(line)
                            yield data.get("name"), data.get("age")

        @dataclass
        class TestCommitMessage(WriterCommitMessage):
            count: int

        class TestJsonWriter(DataSourceWriter):
            def __init__(self, options):
                self.options = options
                self.path = self.options.get("path")

            def write(self, iterator):
                from pyspark import TaskContext

                context = TaskContext.get()
                output_path = os.path.join(self.path, f"{context.partitionId()}.json")
                count = 0
                with open(output_path, "w") as file:
                    for row in iterator:
                        count += 1
                        if "id" in row and row.id > 5:
                            raise Exception("id > 5")
                        file.write(json.dumps(row.asDict()) + "\n")
                return TestCommitMessage(count=count)

            def commit(self, messages):
                total_count = sum(message.count for message in messages)
                with open(os.path.join(self.path, "_success.txt"), "a") as file:
                    file.write(f"count: {total_count}\n")

            def abort(self, messages):
                with open(os.path.join(self.path, "_failed.txt"), "a") as file:
                    file.write("failed")

        class TestJsonDataSource(DataSource):
            @classmethod
            def name(cls):
                return "my-json"

            def schema(self):
                return "name STRING, age INT"

            def reader(self, schema) -> "DataSourceReader":
                return TestJsonReader(self.options)

            def writer(self, schema, overwrite):
                return TestJsonWriter(self.options)

        return TestJsonDataSource

    def test_custom_json_data_source_read(self):
        data_source = self._get_test_json_data_source()
        self.spark.dataSource.register(data_source)
        path1 = os.path.join(SPARK_HOME, "python/test_support/sql/people.json")
        path2 = os.path.join(SPARK_HOME, "python/test_support/sql/people1.json")
        assertDataFrameEqual(
            self.spark.read.format("my-json").load(path1),
            [Row(name="Michael", age=None), Row(name="Andy", age=30), Row(name="Justin", age=19)],
        )
        assertDataFrameEqual(
            self.spark.read.format("my-json").load(path2),
            [Row(name="Jonathan", age=None)],
        )

    def test_custom_json_data_source_write(self):
        data_source = self._get_test_json_data_source()
        self.spark.dataSource.register(data_source)
        input_path = os.path.join(SPARK_HOME, "python/test_support/sql/people.json")
        df = self.spark.read.json(input_path)
        with tempfile.TemporaryDirectory(prefix="test_custom_json_data_source_write") as d:
            df.write.format("my-json").mode("append").save(d)
            assertDataFrameEqual(self.spark.read.json(d), self.spark.read.json(input_path))

    def test_custom_json_data_source_commit(self):
        data_source = self._get_test_json_data_source()
        self.spark.dataSource.register(data_source)
        with tempfile.TemporaryDirectory(prefix="test_custom_json_data_source_commit") as d:
            self.spark.range(0, 5, 1, 3).write.format("my-json").mode("append").save(d)
            with open(os.path.join(d, "_success.txt"), "r") as file:
                text = file.read()
            assert text == "count: 5\n"

    def test_custom_json_data_source_abort(self):
        data_source = self._get_test_json_data_source()
        self.spark.dataSource.register(data_source)
        with tempfile.TemporaryDirectory(prefix="test_custom_json_data_source_abort") as d:
            with self.assertRaises(PythonException):
                self.spark.range(0, 8, 1, 3).write.format("my-json").mode("append").save(d)
            with open(os.path.join(d, "_failed.txt"), "r") as file:
                text = file.read()
            assert text == "failed"

    def test_case_insensitive_dict(self):
        d = CaseInsensitiveDict({"foo": 1, "Bar": 2})
        self.assertEqual(d["foo"], d["FOO"])
        self.assertEqual(d["bar"], d["BAR"])
        self.assertTrue("baR" in d)
        d["BAR"] = 3
        self.assertEqual(d["BAR"], 3)
        # Test update
        d.update({"BaZ": 3})
        self.assertEqual(d["BAZ"], 3)
        d.update({"FOO": 4})
        self.assertEqual(d["foo"], 4)
        # Test delete
        del d["FoO"]
        self.assertFalse("FOO" in d)
        # Test copy
        d2 = d.copy()
        self.assertEqual(d2["BaR"], 3)
        self.assertEqual(d2["baz"], 3)

    def test_arrow_batch_data_source(self):
        import pyarrow as pa

        class ArrowBatchDataSource(DataSource):
            """
            A data source for testing Arrow Batch Serialization
            """

            @classmethod
            def name(cls):
                return "arrowbatch"

            def schema(self):
                return "key int, value string"

            def reader(self, schema: str):
                return ArrowBatchDataSourceReader(schema, self.options)

        class ArrowBatchDataSourceReader(DataSourceReader):
            def __init__(self, schema, options):
                self.schema: str = schema
                self.options = options

            def read(self, partition):
                # Create Arrow Record Batch
                keys = pa.array([1, 2, 3, 4, 5], type=pa.int32())
                values = pa.array(["one", "two", "three", "four", "five"], type=pa.string())
                schema = pa.schema([("key", pa.int32()), ("value", pa.string())])
                record_batch = pa.RecordBatch.from_arrays([keys, values], schema=schema)
                yield record_batch

            def partitions(self):
                # hardcoded number of partitions
                num_part = 1
                return [InputPartition(i) for i in range(num_part)]

        self.spark.dataSource.register(ArrowBatchDataSource)
        df = self.spark.read.format("arrowbatch").load()
        expected_data = [
            Row(key=1, value="one"),
            Row(key=2, value="two"),
            Row(key=3, value="three"),
            Row(key=4, value="four"),
            Row(key=5, value="five"),
        ]
        assertDataFrameEqual(df, expected_data)

        with self.assertRaisesRegex(
            PythonException,
            "PySparkRuntimeError: \\[DATA_SOURCE_RETURN_SCHEMA_MISMATCH\\] Return schema"
            " mismatch in the result from 'read' method\\. Expected: 1 columns, Found: 2 columns",
        ):
            self.spark.read.format("arrowbatch").schema("dummy int").load().show()

        with self.assertRaisesRegex(
            PythonException,
            "PySparkRuntimeError: \\[DATA_SOURCE_RETURN_SCHEMA_MISMATCH\\] Return schema mismatch"
            " in the result from 'read' method\\. Expected: \\['key', 'dummy'\\] columns, Found:"
            " \\['key', 'value'\\] columns",
        ):
            self.spark.read.format("arrowbatch").schema("key int, dummy string").load().show()

    def test_arrow_batch_sink(self):
        class TestDataSource(DataSource):
            @classmethod
            def name(cls):
                return "arrow_sink"

            def writer(self, schema, overwrite):
                return TestArrowWriter(self.options["path"])

        class TestArrowWriter(DataSourceArrowWriter):
            def __init__(self, path):
                self.path = path

            def write(self, iterator):
                from pyspark import TaskContext

                context = TaskContext.get()
                output_path = os.path.join(self.path, f"{context.partitionId()}.json")
                with open(output_path, "w") as file:
                    for batch in iterator:
                        df = batch.to_pandas()
                        df.to_json(file, orient="records", lines=True)
                return WriterCommitMessage()

        self.spark.dataSource.register(TestDataSource)
        df = self.spark.range(3)
        with tempfile.TemporaryDirectory(prefix="test_arrow_batch_sink") as d:
            df.write.format("arrow_sink").mode("append").save(d)
            df2 = self.spark.read.format("json").load(d)
            assertDataFrameEqual(df2, df)

    def test_data_source_type_mismatch(self):
        class TestDataSource(DataSource):
            @classmethod
            def name(cls):
                return "test"

            def schema(self):
                return "id int"

            def reader(self, schema):
                return TestReader()

            def writer(self, schema, overwrite):
                return TestWriter()

        class TestReader:
            def partitions(self):
                return []

            def read(self, partition):
                yield (0,)

        class TestWriter:
            def write(self, iterator):
                return WriterCommitMessage()

        self.spark.dataSource.register(TestDataSource)

        with self.assertRaisesRegex(
            AnalysisException,
            r"\[DATA_SOURCE_TYPE_MISMATCH\] Expected an instance of DataSourceReader",
        ):
            self.spark.read.format("test").load().show()

        df = self.spark.range(10)
        with self.assertRaisesRegex(
            AnalysisException,
            r"\[DATA_SOURCE_TYPE_MISMATCH\] Expected an instance of DataSourceWriter",
        ):
            df.write.format("test").mode("append").saveAsTable("test_table")

    @unittest.skipIf(
        "pypy" in platform.python_implementation().lower(), "cannot run in environment pypy"
    )
    def test_data_source_segfault(self):
        import ctypes

        for enabled, expected in [
            (True, "Segmentation fault"),
            (False, "Consider setting .* for the better Python traceback."),
        ]:
            with self.subTest(enabled=enabled), self.sql_conf(
                {"spark.sql.execution.pyspark.udf.faulthandler.enabled": enabled}
            ):
                with self.subTest(worker="pyspark.sql.worker.create_data_source"):

                    class TestDataSource(DataSource):
                        @classmethod
                        def name(cls):
                            return "test"

                        def schema(self):
                            return ctypes.string_at(0)

                    self.spark.dataSource.register(TestDataSource)

                    with self.assertRaisesRegex(Exception, expected):
                        self.spark.read.format("test").load().show()

                with self.subTest(worker="pyspark.sql.worker.plan_data_source_read"):

                    class TestDataSource(DataSource):
                        @classmethod
                        def name(cls):
                            return "test"

                        def schema(self):
                            return "x string"

                        def reader(self, schema):
                            return TestReader()

                    class TestReader(DataSourceReader):
                        def partitions(self):
                            ctypes.string_at(0)
                            return []

                        def read(self, partition):
                            return []

                    self.spark.dataSource.register(TestDataSource)

                    with self.assertRaisesRegex(Exception, expected):
                        self.spark.read.format("test").load().show()

                with self.subTest(worker="pyspark.worker"):

                    class TestDataSource(DataSource):
                        @classmethod
                        def name(cls):
                            return "test"

                        def schema(self):
                            return "x string"

                        def reader(self, schema):
                            return TestReader()

                    class TestReader(DataSourceReader):
                        def read(self, partition):
                            ctypes.string_at(0)
                            yield "x",

                    self.spark.dataSource.register(TestDataSource)

                    with self.assertRaisesRegex(Exception, expected):
                        self.spark.read.format("test").load().show()

                with self.subTest(worker="pyspark.sql.worker.write_into_data_source"):

                    class TestDataSource(DataSource):
                        @classmethod
                        def name(cls):
                            return "test"

                        def writer(self, schema, overwrite):
                            return TestWriter()

                    class TestWriter(DataSourceWriter):
                        def write(self, iterator):
                            ctypes.string_at(0)
                            return WriterCommitMessage()

                    self.spark.dataSource.register(TestDataSource)

                    with self.assertRaisesRegex(Exception, expected):
                        self.spark.range(10).write.format("test").mode("append").saveAsTable(
                            "test_table"
                        )

                with self.subTest(worker="pyspark.sql.worker.commit_data_source_write"):

                    class TestDataSource(DataSource):
                        @classmethod
                        def name(cls):
                            return "test"

                        def writer(self, schema, overwrite):
                            return TestWriter()

                    class TestWriter(DataSourceWriter):
                        def write(self, iterator):
                            return WriterCommitMessage()

                        def commit(self, messages):
                            ctypes.string_at(0)

                    self.spark.dataSource.register(TestDataSource)

                    with self.assertRaisesRegex(Exception, expected):
                        self.spark.range(10).write.format("test").mode("append").saveAsTable(
                            "test_table"
                        )


class PythonDataSourceTests(BasePythonDataSourceTestsMixin, ReusedSQLTestCase):
    ...


if __name__ == "__main__":
    from pyspark.sql.tests.test_python_datasource import *  # noqa: F401

    try:
        import xmlrunner  # type: ignore

        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
    except ImportError:
        testRunner = None
    unittest.main(testRunner=testRunner, verbosity=2)