sqlframe 3.21.1__py3-none-any.whl → 3.22.0__py3-none-any.whl

sqlframe/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.21.1'
-__version_tuple__ = version_tuple = (3, 21, 1)
+__version__ = version = '3.22.0'
+__version_tuple__ = version_tuple = (3, 22, 0)

sqlframe/base/dataframe.py

@@ -342,7 +342,7 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
             return types.StructType(
                 [
                     types.StructField(
-                        c.name,
+                        self.display_name_mapping.get(c.name, c.name),
                         sqlglot_to_spark(
                             exp.DataType.build(c.dataType, dialect=self.session.output_dialect)
                         ),
@@ -1898,7 +1898,7 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         print("root")
         for column in self._typed_columns:
             print_schema(
-                column.name,
+                self.display_name_mapping.get(column.name, column.name),
                 exp.DataType.build(column.dataType, dialect=self.session.output_dialect),
                 column.nullable,
                 0,

sqlframe/base/functions.py

@@ -4504,7 +4504,7 @@ def median(col: ColumnOrName) -> Column:
     return Column.invoke_expression_over_column(col, expression.Median)
 
 
-@meta(unsupported_engines="*")
+@meta(unsupported_engines=["bigquery", "postgres"])
 def mode(col: ColumnOrName) -> Column:
     """
     Returns the most frequent value in a group.
@@ -4540,6 +4540,7 @@ def mode(col: ColumnOrName) -> Column:
     |dotNET|      2012|
     +------+----------+
     """
+
     return Column.invoke_anonymous_function(col, "mode")
 
 

sqlframe/base/mixins/readwriter_mixins.py

@@ -82,6 +82,10 @@ class PandasLoaderMixin(_BaseDataFrameReader, t.Generic[SESSION, DF]):
         elif format == "parquet":
             df = pd.read_parquet(path, **kwargs)  # type: ignore
         elif format == "csv":
+            kwargs.pop("inferSchema", None)
+            if "header" in kwargs:
+                if isinstance(kwargs["header"], bool) and kwargs["header"]:
+                    kwargs["header"] = "infer"
             df = pd.read_csv(path, **kwargs)  # type: ignore
         else:
             raise UnsupportedOperationError(f"Unsupported format: {format}")

sqlframe/base/readerwriter.py

@@ -393,10 +393,12 @@ class _BaseDataFrameWriter(t.Generic[SESSION, DF]):
         df: DF,
         mode: t.Optional[str] = None,
         by_name: bool = False,
+        state_format_to_write: t.Optional[str] = None,
     ):
         self._df = df
         self._mode = mode
         self._by_name = by_name
+        self._state_format_to_write = state_format_to_write
 
     @property
     def _session(self) -> SESSION:
@@ -484,6 +486,44 @@ class _BaseDataFrameWriter(t.Generic[SESSION, DF]):
     def _write(self, path: str, mode: t.Optional[str], format: str, **options) -> None:
         raise NotImplementedError
 
+    def format(self, source: str) -> "Self":
+        """Specifies the input data source format.
+
+        .. versionadded:: 1.4.0
+
+        .. versionchanged:: 3.4.0
+            Supports Spark Connect.
+
+        Parameters
+        ----------
+        source : str
+            string, name of the data source, e.g. 'json', 'parquet'.
+
+        Examples
+        --------
+        >>> spark.read.format('json')
+        <...readwriter.DataFrameReader object ...>
+
+        Write a DataFrame into a JSON file and read it back.
+
+        >>> import tempfile
+        >>> with tempfile.TemporaryDirectory() as d:
+        ...     # Write a DataFrame into a JSON file
+        ...     spark.createDataFrame(
+        ...         [{"age": 100, "name": "Hyukjin Kwon"}]
+        ...     ).write.mode("overwrite").format("json").save(d)
+        ...
+        ...     # Read the JSON file as a DataFrame.
+        ...     spark.read.format('json').load(d).show()
+        +---+------------+
+        |age|        name|
+        +---+------------+
+        |100|Hyukjin Kwon|
+        +---+------------+
+        """
+        self._state_format_to_write = source
+        return self
+
     def json(
         self,
         path: str,

sqlframe/base/util.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import importlib
+import random
+import string
 import typing as t
 import unicodedata
 
@@ -427,3 +429,20 @@ def normalize_string(
         for pos in star_positions:
             normalized_value = normalized_value[:pos] + "*" + normalized_value[pos:]
     return normalized_value
+
+
+def generate_random_identifier(size=6, chars=string.ascii_uppercase + string.digits):
+    return "_" + "".join(random.choice(chars) for _ in range(size))
+
+
+def split_filepath(filepath: str) -> tuple[str, str]:
+    if filepath.startswith("dbfs:") or filepath.startswith("/dbfs"):
+        prefix = "dbfs:"
+        return prefix, filepath[len(prefix) :]
+    if filepath.startswith("file://"):
+        prefix = "file://"
+        return "", filepath[len(prefix) :]
+    split_ = str(filepath).split("://", 1)
+    if len(split_) == 2:  # noqa: PLR2004
+        return split_[0] + "://", split_[1]
+    return "", split_[0]

sqlframe/duckdb/readwriter.py

@@ -92,6 +92,7 @@ class DuckDBDataFrameReader(
         if format == "delta":
             from_clause = f"delta_scan('{path}')"
         elif format:
+            options.pop("inferSchema", None)
             paths = ",".join([f"'{path}'" for path in ensure_list(path)])
             from_clause = f"read_{format}([{paths}], {to_csv(options)})"
         else:

sqlframe/spark/readwriter.py

@@ -4,27 +4,160 @@ from __future__ import annotations
 
 import typing as t
 
-from sqlframe.base.mixins.readwriter_mixins import PandasLoaderMixin, PandasWriterMixin
+from sqlglot import exp
+from sqlglot.helper import ensure_list
+
 from sqlframe.base.readerwriter import (
     _BaseDataFrameReader,
     _BaseDataFrameWriter,
+    _infer_format,
 )
+from sqlframe.base.util import ensure_column_mapping, generate_random_identifier, to_csv
 
 if t.TYPE_CHECKING:
+    from sqlframe.base._typing import OptionalPrimitiveType, PathOrPaths
+    from sqlframe.base.types import StructType
     from sqlframe.spark.dataframe import SparkDataFrame
     from sqlframe.spark.session import SparkSession
     from sqlframe.spark.table import SparkTable
 
 
 class SparkDataFrameReader(
-    PandasLoaderMixin["SparkSession", "SparkDataFrame"],
     _BaseDataFrameReader["SparkSession", "SparkDataFrame", "SparkTable"],
 ):
-    pass
+    def load(
+        self,
+        path: t.Optional[PathOrPaths] = None,
+        format: t.Optional[str] = None,
+        schema: t.Optional[t.Union[StructType, str]] = None,
+        **options: OptionalPrimitiveType,
+    ) -> SparkDataFrame:
+        """Loads data from a data source and returns it as a :class:`DataFrame`.
+
+        .. versionadded:: 1.4.0
+
+        .. versionchanged:: 3.4.0
+            Supports Spark Connect.
+
+        Parameters
+        ----------
+        path : str or list, t.Optional
+            t.Optional string or a list of string for file-system backed data sources.
+        format : str, t.Optional
+            t.Optional string for format of the data source. Default to 'parquet'.
+        schema : :class:`pyspark.sql.types.StructType` or str, t.Optional
+            t.Optional :class:`pyspark.sql.types.StructType` for the input schema
+            or a DDL-formatted string (For example ``col0 INT, col1 DOUBLE``).
+        **options : dict
+            all other string options
+
+        Examples
+        --------
+        Load a CSV file with format, schema and options specified.
+
+        >>> import tempfile
+        >>> with tempfile.TemporaryDirectory() as d:
+        ...     # Write a DataFrame into a CSV file with a header
+        ...     df = spark.createDataFrame([{"age": 100, "name": "Hyukjin Kwon"}])
+        ...     df.write.option("header", True).mode("overwrite").format("csv").save(d)
+        ...
+        ...     # Read the CSV file as a DataFrame with 'nullValue' option set to 'Hyukjin Kwon',
+        ...     # and 'header' option set to `True`.
+        ...     df = spark.read.load(
+        ...         d, schema=df.schema, format="csv", nullValue="Hyukjin Kwon", header=True)
+        ...     df.printSchema()
+        ...     df.show()
+        root
+         |-- age: long (nullable = true)
+         |-- name: string (nullable = true)
+        +---+----+
+        |age|name|
+        +---+----+
+        |100|NULL|
+        +---+----+
+        """
+        assert path is not None, "path is required"
+        assert isinstance(path, str), "path must be a string"
+        format = format or self.state_format_to_read or _infer_format(path)
+        if schema:
+            column_mapping = ensure_column_mapping(schema)
+            select_column_mapping = column_mapping.copy()
+            select_columns = [x.expression for x in self._to_casted_columns(select_column_mapping)]
+
+            if hasattr(schema, "simpleString"):
+                schema = schema.simpleString()
+        else:
+            select_columns = [exp.Star()]
+
+        if format == "delta":
+            from_clause = f"delta.`{path}`"
+        elif format:
+            paths = ",".join([f"{path}" for path in ensure_list(path)])
+            tmp_view_key = options.get("_tmp_view_key_", f"{generate_random_identifier()}_vw")
+            options["_tmp_view_key_"] = tmp_view_key
+
+            format_options: dict[str, OptionalPrimitiveType] = {
+                k: v for k, v in options.items() if v is not None
+            }
+            format_options.pop("_tmp_view_key_")
+            format_options["path"] = paths
+            if schema:
+                format_options["schema"] = f"{schema}"
+                format_options.pop("inferSchema", None)
+            format_options = {key: f"'{val}'" for key, val in format_options.items()}
+            format_options_str = to_csv(format_options, " ")
+
+            tmp_view = f"CREATE OR REPLACE TEMPORARY VIEW {tmp_view_key} USING {format}" + (
+                f" OPTIONS ({format_options_str})" if format_options_str else ""
+            )
+            self.session.spark_session.sql(tmp_view).collect()
+
+            from_clause = f"{tmp_view_key}"
+        else:
+            from_clause = f"'{path}'"
+
+        df = self.session.sql(
+            exp.select(*select_columns).from_(from_clause, dialect=self.session.input_dialect),
+            qualify=False,
+        )
+        if select_columns == [exp.Star()] and df.schema:
+            return self.load(path=path, format=format, schema=df.schema, **options)
+        self.session._last_loaded_file = path  # type: ignore
+        return df
 
 
 class SparkDataFrameWriter(
-    PandasWriterMixin["SparkSession", "SparkDataFrame"],
     _BaseDataFrameWriter["SparkSession", "SparkDataFrame"],
 ):
-    pass
+    def save(
+        self,
+        path: str,
+        mode: t.Optional[str] = None,
+        format: t.Optional[str] = None,
+        partitionBy: t.Optional[t.Union[str, t.List[str]]] = None,
+        **options,
+    ):
+        format = str(format or self._state_format_to_write)
+        self._write(path, mode, format, partitionBy=partitionBy, **options)
+
+    def _write(self, path: str, mode: t.Optional[str], format: str, **options):
+        spark_df = None
+        expressions = self._df._get_expressions()
+        for i, expression in enumerate(expressions):
+            if i < len(expressions) - 1:
+                self._df.session._collect(expressions)
+            else:
+                sql = self._df.session._to_sql(expression)
+                spark_df = self._session.spark_session.sql(sql)
+        if spark_df is not None:
+            options = {k: v for k, v in options.items() if v is not None}
+            mode = str(mode or self._mode or "default")
+            spark_writer = spark_df.write.format(format).mode(mode)
+            partition_columns = options.pop("partitionBy", None)
+            compression = options.pop("compression", None)
+            if partition_columns:
+                partition_columns = options.pop("partitionBy")
+                spark_writer = spark_writer.partitionBy(*partition_columns)
+            if compression:
+                spark_writer = spark_writer.option("compression", compression)
+            spark_writer.save(path=path, **options)

{sqlframe-3.21.1.dist-info → sqlframe-3.22.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlframe
-Version: 3.21.1
+Version: 3.22.0
 Summary: Turning PySpark Into a Universal DataFrame API
 Home-page: https://github.com/eakmanrq/sqlframe
 Author: Ryan Eakman

{sqlframe-3.21.1.dist-info → sqlframe-3.22.0.dist-info}/RECORD

@@ -1,29 +1,29 @@
 sqlframe/__init__.py,sha256=wfqm98eLoLid9oV_FzzpG5loKC6LxOhj2lXpfN7SARo,3138
-sqlframe/_version.py,sha256=fmhKf9XPZdwZdKpQ-ESJ_LGssm7Q8K_NJEGVKwXLGQM,413
+sqlframe/_version.py,sha256=z_95mpOa--N235aXeiPpCYjkuWuYr-3uQDssTEpMtnc,413
 sqlframe/base/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sqlframe/base/_typing.py,sha256=b2clI5HI1zEZKB_3Msx3FeAJQyft44ubUifJwQRVXyQ,1298
 sqlframe/base/catalog.py,sha256=SzFQalTWdhWzxUY-4ut1f9TfOECp_JmJEgNPfrRKCe0,38457
 sqlframe/base/column.py,sha256=oHVwkSWABO3ZlAbgBShsxSSlgbI06BOup5XJrRhgqJI,18097
-sqlframe/base/dataframe.py,sha256=FOgLdCpscLsBntkRvutcgSVqXqMgXo9DYa892mXu00E,83907
+sqlframe/base/dataframe.py,sha256=9wcN5I5bSiGKs6m-mSCwlDUSRfrZ_ymfRUtJa8hggd4,83990
 sqlframe/base/decorators.py,sha256=ms-CvDOIW3T8IVB9VqDmLwAiaEsqXLYRXEqVQaxktiM,1890
 sqlframe/base/exceptions.py,sha256=9Uwvqn2eAkDpqm4BrRgbL61qM-GMCbJEMAW8otxO46s,370
 sqlframe/base/function_alternatives.py,sha256=KFkEm0aIHzajvQmiPZnzTLh-Ud9wjeg4lJ4Rk0vk-YU,53674
-sqlframe/base/functions.py,sha256=jfLgboldiTB9CPkoZMtKUAwx6XSvFnEOIpCZQfoEJJU,223060
+sqlframe/base/functions.py,sha256=0lLefa959J1z7Ea6QBnQrru06lKCLurEX1xFYxGJbTc,223082
 sqlframe/base/group.py,sha256=fsyG5990_Pd7gFPjTFrH9IEoAquL_wEkVpIlBAIkZJU,4091
 sqlframe/base/normalize.py,sha256=nXAJ5CwxVf4DV0GsH-q1w0p8gmjSMlv96k_ez1eVul8,3880
 sqlframe/base/operations.py,sha256=xSPw74e59wYvNd6U1AlwziNCTG6Aftrbl4SybN9u9VE,3450
-sqlframe/base/readerwriter.py,sha256=w8926cqIrXF7NGHiINw5UHzP_3xpjsqbijTBTzycBRM,26605
+sqlframe/base/readerwriter.py,sha256=g6h9ldC4ImYdoUwZJxwfd5Iu2p9aqZ39bEFhwrHVlvI,27868
 sqlframe/base/session.py,sha256=G5_bI_z1iJtAGm2SgEdjkKiyJmS0yOUopx9P5TEGdR4,27273
 sqlframe/base/table.py,sha256=rCeh1W5SWbtEVfkLAUiexzrZwNgmZeptLEmLcM1ABkE,6961
 sqlframe/base/transforms.py,sha256=y0j3SGDz3XCmNGrvassk1S-owllUWfkHyMgZlY6SFO4,467
 sqlframe/base/types.py,sha256=iBNk9bpFtb2NBIogYS8i7OlQZMRvpR6XxqzBebsjQDU,12280
 sqlframe/base/udf.py,sha256=O6hMhBUy9NVv-mhJRtfFhXTIa_-Z8Y_FkmmuOHu0l90,1117
-sqlframe/base/util.py,sha256=_s2M-qHzTLgyGu1v8laRHJorUpUO6-fr3kk7CsvcuXw,15161
+sqlframe/base/util.py,sha256=RmeAqsZEb1zlUEvlRXJnOcDEWJC7OOYEw229s3LzpFA,15805
 sqlframe/base/window.py,sha256=8hOv-ignPPIsZA9FzvYzcLE9J_glalVaYjIAUdRUX3o,4943
 sqlframe/base/mixins/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sqlframe/base/mixins/catalog_mixins.py,sha256=9tn0mK8oPoqIIjNItystD5tdBMdK9YpkxTG7G9KQl8k,18619
 sqlframe/base/mixins/dataframe_mixins.py,sha256=9U556vWjhmAwnbqkQ4mDkKHxQRs-FkncwEEY9RWmm6U,1408
-sqlframe/base/mixins/readwriter_mixins.py,sha256=ap8j_g7PoUGHaHKCPMnRPbXofOsUhUzlaF7Loxy2m-I,4752
+sqlframe/base/mixins/readwriter_mixins.py,sha256=ItQ_0jZ5RljgmLjGDIzLMRP_NQdy3wAyKwJ6K5NjaqA,4954
 sqlframe/base/mixins/table_mixins.py,sha256=2TnGFpbDSGw_NswpZwLACqvdD4zCA7hXekQ9IEkoTOk,13784
 sqlframe/bigquery/__init__.py,sha256=kbaomhYAANPdxeDQhajv8IHfMg_ENKivtYK-rPwaV08,939
 sqlframe/bigquery/catalog.py,sha256=w6cxAV6OWa2icVt5vwgNXy08EP9WJh3DkQJzQkn2kIo,11661
@@ -58,7 +58,7 @@ sqlframe/duckdb/dataframe.py,sha256=Z8_K69UQGZVeBfVGXVwIJP8OMuIvNBB3DPKTP3Lfu4w,
 sqlframe/duckdb/functions.py,sha256=ix2efGGD4HLaY1rtCtEd3IrsicGEVGiBAeKOo5OD8rA,424
 sqlframe/duckdb/functions.pyi,sha256=P0ky6k-J7LdCDrQ0OjfRC3ARIYNHPmAmmaB_jBEO5L0,12383
 sqlframe/duckdb/group.py,sha256=IkhbW42Ng1U5YT3FkIdiB4zBqRkW4QyTb-1detY1e_4,383
-sqlframe/duckdb/readwriter.py,sha256=5EP8DEoX3N_xYavWpetsZdzvtYF-oCrAz3n-siNE8yY,4938
+sqlframe/duckdb/readwriter.py,sha256=-_Ama7evadIa3PYvynKDK6RcTMTDBHpHJzfANTine7g,4983
 sqlframe/duckdb/session.py,sha256=Uf7im6eBbBYRvIhVGV0VCTCF76FQ00A5FbKPCdNllzw,2898
 sqlframe/duckdb/table.py,sha256=AmEKoH2TZo98loS5NbNaTuqv0eg76SY_OckVBMmQ6Co,410
 sqlframe/duckdb/types.py,sha256=KwNyuXIo-2xVVd4bZED3YrQOobKCtemlxGrJL7DrTC8,34
@@ -109,7 +109,7 @@ sqlframe/spark/dataframe.py,sha256=WyXHWsH8Ldu2cWTNmsLy5hEFrjJvQh_Aqv3JJcbDy6k,1
 sqlframe/spark/functions.py,sha256=MYCgHsjRQWylT-rezWRBuLV6BivcaVarbaQtP4T0toQ,331
 sqlframe/spark/functions.pyi,sha256=GyOdUzv2Z7Qt99JAKEPKgV2t2Rn274OuqwAfcoAXlN0,24259
 sqlframe/spark/group.py,sha256=MrvV_v-YkBc6T1zz882WrEqtWjlooWIyHBCmTQg3fCA,379
-sqlframe/spark/readwriter.py,sha256=zXZcCPWpQMMN90wdIx8AD4Y5tWBcpRSL4-yKX2aZyik,874
+sqlframe/spark/readwriter.py,sha256=YVGgkYpmuQj8wGHxJx6ivAAKSqyGes-0GhCezvd7kYU,6565
 sqlframe/spark/session.py,sha256=irlsTky06pKRKAyPLwVzUtLGe4O8mALSgxIqLvqJNF8,5675
 sqlframe/spark/table.py,sha256=puWV8h_CqA64zwpzq0ydY9LoygMAvprkODyxyzZeF9M,186
 sqlframe/spark/types.py,sha256=KwNyuXIo-2xVVd4bZED3YrQOobKCtemlxGrJL7DrTC8,34
@@ -129,8 +129,8 @@ sqlframe/standalone/udf.py,sha256=azmgtUjHNIPs0WMVNId05SHwiYn41MKVBhKXsQJ5dmY,27
 sqlframe/standalone/window.py,sha256=6GKPzuxeSapJakBaKBeT9VpED1ACdjggDv9JRILDyV0,35
 sqlframe/testing/__init__.py,sha256=VVCosQhitU74A3NnE52O4mNtGZONapuEXcc20QmSlnQ,132
 sqlframe/testing/utils.py,sha256=PFsGZpwNUE_4-g_f43_vstTqsK0AQ2lBneb5Eb6NkFo,13008
-sqlframe-3.21.1.dist-info/LICENSE,sha256=VZu79YgW780qxaFJMr0t5ZgbOYEh04xWoxaWOaqIGWk,1068
-sqlframe-3.21.1.dist-info/METADATA,sha256=AauznGD-zSbh2cqT63w2MIrg_-0SlewyyRMNElL5O2I,8970
-sqlframe-3.21.1.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-sqlframe-3.21.1.dist-info/top_level.txt,sha256=T0_RpoygaZSF6heeWwIDQgaP0varUdSK1pzjeJZRjM8,9
-sqlframe-3.21.1.dist-info/RECORD,,
+sqlframe-3.22.0.dist-info/LICENSE,sha256=VZu79YgW780qxaFJMr0t5ZgbOYEh04xWoxaWOaqIGWk,1068
+sqlframe-3.22.0.dist-info/METADATA,sha256=uN5N8oTC6EkNmVWzBgxkT_Ri1eIgTN3XIaC7VH3_zL4,8970
+sqlframe-3.22.0.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+sqlframe-3.22.0.dist-info/top_level.txt,sha256=T0_RpoygaZSF6heeWwIDQgaP0varUdSK1pzjeJZRjM8,9
+sqlframe-3.22.0.dist-info/RECORD,,