Add csv & parquet write functions and toPandas to experimental PySpark API

tools/pythonpkg/duckdb/experimental/spark/sql/dataframe.py

@@ -10,6 +10,7 @@
 from .column import Column
 import duckdb
 from functools import reduce
+from pandas import DataFrame as PandasDataFrame
 
 if TYPE_CHECKING:
     from .session import SparkSession
@@ -27,6 +28,9 @@ def __init__(self, relation: duckdb.DuckDBPyRelation, session: "SparkSession"):
 
     def show(self, **kwargs) -> None:
         self.relation.show()
+
+    def toPandas(self) -> PandasDataFrame:
+        return self.relation.df()
 
     def createOrReplaceTempView(self, name: str) -> None:
         """Creates or replaces a local temporary view with this :class:`DataFrame`.

tools/pythonpkg/duckdb/experimental/spark/sql/readwriter.py

@@ -20,6 +20,47 @@ def saveAsTable(self, table_name: str) -> None:
         relation = self.dataframe.relation
         relation.create(table_name)
 
+    def parquet(self, path: str, mode: Optional[str] = None, partitionBy: Union[str, List[str], None] = None, compression: Optional[str] = None)  -> None:
+        relation = self.dataframe.relation
+        if mode:
+            raise NotImplementedError
+        if partitionBy:
+            raise NotImplementedError
+
+        relation.write_parquet(path, compression=compression)
+
+    def csv(self, path, mode=None, compression=None, sep=None, quote=None, escape=None,
+            header=None, nullValue=None, escapeQuotes=None, quoteAll=None, dateFormat=None,
+            timestampFormat=None, ignoreLeadingWhiteSpace=None, ignoreTrailingWhiteSpace=None,
+            charToEscapeQuoteEscaping=None, encoding=None, emptyValue=None, lineSep=None):
+        if mode not in (None, "overwrite"):
+            raise NotImplementedError
+        if escapeQuotes:
+            raise NotImplementedError
+        if ignoreLeadingWhiteSpace:
+            raise NotImplementedError
+        if ignoreTrailingWhiteSpace:
+            raise NotImplementedError
+        if charToEscapeQuoteEscaping:
+            raise NotImplementedError
+        if emptyValue:
+            raise NotImplementedError
+        if lineSep:
+            raise NotImplementedError
+        relation = self.dataframe.relation
+        relation.write_csv(path,
+                        sep=sep, 
+                        na_rep=nullValue,
+                        quotechar=quote,
+                        compression=compression, 
+                        escapechar=escape,
+                        header=header if isinstance(header, bool) else header == "True",
+                        encoding=encoding,
+                        quoting=quoteAll, # ~ check this
+                        date_format=dateFormat,
+                        timestamp_format = timestampFormat
+                        )
+
 
 class DataFrameReader:
     def __init__(self, session: "SparkSession"):

tools/pythonpkg/tests/fast/spark/test_spark_pandas_dataframe.py

@@ -18,6 +18,7 @@
 from duckdb.experimental.spark.sql.functions import col, struct, when
 import duckdb
 import re
+from pandas.testing import assert_frame_equal
 
 
 @pytest.fixture
@@ -48,3 +49,8 @@ def test_pd_conversion_schema(self, spark, pandasDF):
         res = sparkDF.collect()
         expected = "[Row(First Name='Scott', Age=50), Row(First Name='Jeff', Age=45), Row(First Name='Thomas', Age=54), Row(First Name='Ann', Age=34)]"
         assert str(res) == expected
+
+    def test_spark_to_pandas_dataframe(self, spark, pandasDF):
+        sparkDF = spark.createDataFrame(pandasDF)
+        res = sparkDF.toPandas()
+        assert_frame_equal(res, pandasDF)

tools/pythonpkg/tests/fast/spark/test_spark_to_csv.py

@@ -0,0 +1,199 @@
+import pytest
+import tempfile
+
+import os
+
+_ = pytest.importorskip("duckdb.experimental.spark")
+
+from duckdb.experimental.spark.sql import SparkSession as session
+from duckdb import connect, InvalidInputException, read_csv
+from conftest import NumpyPandas, ArrowPandas
+import pandas._testing as tm
+import datetime
+import csv
+
+
+@pytest.fixture
+def df(spark):
+    simpleData = (
+        ("Java", 4000, 5),
+        ("Python", 4600, 10),
+        ("Scala", 4100, 15),
+        ("Scala", 4500, 15),
+        ("PHP", 3000, 20),
+    )
+    columns = ["CourseName", "fee", "discount"]
+    dataframe = spark.createDataFrame(data=simpleData, schema=columns)
+    yield dataframe
+
+
+class TestSparkToCSV(object):
+    @pytest.mark.parametrize('pandas', [NumpyPandas(), ArrowPandas()])
+    def test_basic_to_csv(self, pandas, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+
+        pandas_df = pandas.DataFrame({'a': [5, 3, 23, 2], 'b': [45, 234, 234, 2]})
+        df = spark.createDataFrame(pandas_df)
+
+        df.write.csv(temp_file_name, header=False)
+
+        csv_rel = spark.read.csv(temp_file_name, header=False)
+
+        assert df.collect() == csv_rel.collect()
+
+    @pytest.mark.parametrize('pandas', [NumpyPandas(), ArrowPandas()])
+    def test_to_csv_sep(self, pandas, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+        pandas_df = pandas.DataFrame({'a': [5, 3, 23, 2], 'b': [45, 234, 234, 2]})
+
+        df = spark.createDataFrame(pandas_df)
+
+        df.write.csv(temp_file_name, sep=',', header=False)
+
+        csv_rel = spark.read.csv(temp_file_name, header=False, sep=',')
+        assert df.collect() == csv_rel.collect()
+
+    @pytest.mark.parametrize('pandas', [NumpyPandas(), ArrowPandas()])
+    def test_to_csv_na_rep(self, pandas, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+        pandas_df = pandas.DataFrame({'a': [5, None, 23, 2], 'b': [45, 234, 234, 2]})
+
+        df = spark.createDataFrame(pandas_df)
+
+        df.write.csv(temp_file_name, nullValue="test", header=False)
+
+        csv_rel = spark.read.csv(temp_file_name, nullValue="test")
+        assert df.collect() == csv_rel.collect()
+
+    @pytest.mark.parametrize('pandas', [NumpyPandas(), ArrowPandas()])
+    def test_to_csv_header(self, pandas, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+        pandas_df = pandas.DataFrame({'a': [5, None, 23, 2], 'b': [45, 234, 234, 2]})
+
+        df = spark.createDataFrame(pandas_df)
+
+        df.write.csv(temp_file_name, header=True)
+
+        csv_rel = spark.read.csv(temp_file_name, header=True)
+        assert df.collect() == csv_rel.collect()
+
+    @pytest.mark.parametrize('pandas', [NumpyPandas(), ArrowPandas()])
+    def test_to_csv_quotechar(self, pandas, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+
+        pandas_df = pandas.DataFrame({'a': ["\'a,b,c\'", None, "hello", "bye"], 'b': [45, 234, 234, 2]})
+
+        df = spark.createDataFrame(pandas_df)
+
+        df.write.csv(temp_file_name, quote='\'', sep=',', header=False)
+
+        csv_rel = spark.read.csv(temp_file_name, sep=',', quote='\'')
+        assert df.collect() == csv_rel.collect()
+
+    @pytest.mark.parametrize('pandas', [NumpyPandas(), ArrowPandas()])
+    def test_to_csv_escapechar(self, pandas, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+        pandas_df = pandas.DataFrame(
+            {
+                "c_bool": [True, False],
+                "c_float": [1.0, 3.2],
+                "c_int": [42, None],
+                "c_string": ["a", "b,c"],
+            }
+        )
+
+        df = spark.createDataFrame(pandas_df)
+
+        df.write.csv(temp_file_name, header=True, quote='"', escape='!')
+        csv_rel = spark.read.csv(temp_file_name, quote='"', escape='!', header=True)
+        assert df.collect() == csv_rel.collect()
+
+    @pytest.mark.parametrize('pandas', [NumpyPandas(), ArrowPandas()])
+    def test_to_csv_date_format(self, pandas, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+        pandas_df = pandas.DataFrame(tm.getTimeSeriesData())
+        dt_index = pandas_df.index
+        pandas_df = pandas.DataFrame({"A": dt_index, "B": dt_index.shift(1)}, index=dt_index)
+
+        df = spark.createDataFrame(pandas_df)
+
+        df.write.csv(temp_file_name, dateFormat="%Y%m%d", header=False)
+
+        csv_rel = spark.read.csv(temp_file_name, dateFormat="%Y%m%d")
+
+        assert df.collect() == csv_rel.collect()
+
+    @pytest.mark.parametrize('pandas', [NumpyPandas(), ArrowPandas()])
+    def test_to_csv_timestamp_format(self, pandas, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+        data = [datetime.time(hour=23, minute=1, second=34, microsecond=234345)]
+        pandas_df = pandas.DataFrame({'0': pandas.Series(data=data, dtype='object')})
+
+        df = spark.createDataFrame(pandas_df)
+
+        df.write.csv(temp_file_name, timestampFormat='%m/%d/%Y', header=False)
+
+        csv_rel = spark.read.csv(temp_file_name, timestampFormat='%m/%d/%Y')
+
+        assert df.collect() == csv_rel.collect()
+
+    @pytest.mark.parametrize('pandas', [NumpyPandas(), ArrowPandas()])
+    def test_to_csv_quoting_off(self, pandas, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+        pandas_df = pandas.DataFrame({'a': ['string1', 'string2', 'string3']})
+
+        df = spark.createDataFrame(pandas_df)
+        df.write.csv(temp_file_name, quoteAll=None, header=False)
+
+        csv_rel = spark.read.csv(temp_file_name)
+        assert df.collect() == csv_rel.collect()
+
+    @pytest.mark.parametrize('pandas', [NumpyPandas(), ArrowPandas()])
+    def test_to_csv_quoting_on(self, pandas, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+        pandas_df = pandas.DataFrame({'a': ['string1', 'string2', 'string3']})
+        df = spark.createDataFrame(pandas_df)
+        df.write.csv(temp_file_name, quoteAll="force", header=False)
+
+        csv_rel = spark.read.csv(temp_file_name)
+        assert df.collect() == csv_rel.collect()
+
+    @pytest.mark.parametrize('pandas', [NumpyPandas(), ArrowPandas()])
+    def test_to_csv_quoting_quote_all(self, pandas, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+        pandas_df = pandas.DataFrame({'a': ['string1', 'string2', 'string3']})
+        df = spark.createDataFrame(pandas_df)
+        df.write.csv(temp_file_name, quoteAll=csv.QUOTE_ALL, header=False)
+
+        csv_rel = spark.read.csv(temp_file_name)
+        assert df.collect() == csv_rel.collect()
+
+    @pytest.mark.parametrize('pandas', [NumpyPandas(), ArrowPandas()])
+    def test_to_csv_encoding_incorrect(self, pandas, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+        pandas_df = pandas.DataFrame({'a': ['string1', 'string2', 'string3']})
+        df = spark.createDataFrame(pandas_df)
+        with pytest.raises(
+            InvalidInputException, match="Invalid Input Error: The only supported encoding option is 'UTF8"
+        ):
+            df.write.csv(temp_file_name, encoding="nope", header=False)
+
+    @pytest.mark.parametrize('pandas', [NumpyPandas(), ArrowPandas()])
+    def test_to_csv_encoding_correct(self, pandas, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+        pandas_df = pandas.DataFrame({'a': ['string1', 'string2', 'string3']})
+        df = spark.createDataFrame(pandas_df)
+        df.write.csv(temp_file_name, encoding="UTF-8", header=False)
+        csv_rel = spark.read.csv(temp_file_name)
+        assert df.collect() == csv_rel.collect()
+
+    @pytest.mark.parametrize('pandas', [NumpyPandas(), ArrowPandas()])
+    def test_compression_gzip(self, pandas, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+        pandas_df = pandas.DataFrame({'a': ['string1', 'string2', 'string3']})
+        df = spark.createDataFrame(pandas_df)
+        df.write.csv(temp_file_name, compression="gzip", header=False)
+
+        # slightly convoluted - pyspark .read.csv does not take a compression argument
+        csv_rel = spark.createDataFrame(read_csv(temp_file_name, compression="gzip").df())
+        assert df.collect() == csv_rel.collect()

tools/pythonpkg/tests/fast/spark/test_spark_to_parquet.py

@@ -0,0 +1,43 @@
+import pytest
+import tempfile
+
+import os
+
+_ = pytest.importorskip("duckdb.experimental.spark")
+
+from duckdb.experimental.spark.sql import SparkSession as session
+from duckdb import connect
+
+
+@pytest.fixture
+def df(spark):
+    simpleData = (
+        ("Java", 4000, 5),
+        ("Python", 4600, 10),
+        ("Scala", 4100, 15),
+        ("Scala", 4500, 15),
+        ("PHP", 3000, 20),
+    )
+    columns = ["CourseName", "fee", "discount"]
+    dataframe = spark.createDataFrame(data=simpleData, schema=columns)
+    yield dataframe
+
+
+class TestSparkToParquet(object):
+    def test_basic_to_parquet(self, df, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+
+        df.write.parquet(temp_file_name)
+
+        csv_rel = spark.read.parquet(temp_file_name)
+
+        assert df.collect() == csv_rel.collect()
+
+    def test_compressed_to_parquet(self, df, spark):
+        temp_file_name = os.path.join(tempfile.mkdtemp(), next(tempfile._get_candidate_names()))
+
+        df.write.parquet(temp_file_name, compression="ZSTD")
+
+        csv_rel = spark.read.parquet(temp_file_name)
+
+        assert df.collect() == csv_rel.collect()