sqlframe 1.1.3__py3-none-any.whl

sqlframe/duckdb/functions.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import inspect
+import sys
+
+import sqlframe.base.functions  # noqa
+
+module = sys.modules["sqlframe.base.functions"]
+globals().update(
+    {
+        name: func
+        for name, func in inspect.getmembers(module, inspect.isfunction)
+        if hasattr(func, "unsupported_engines")
+        and "duckdb" not in func.unsupported_engines
+        and "*" not in func.unsupported_engines
+    }
+)
+
+
+from sqlframe.base.function_alternatives import (  # noqa
+    e_literal as e,
+    expm1_from_exp as expm1,
+    log1p_from_log as log1p,
+    rint_from_round as rint,
+    kurtosis_from_kurtosis_pop as kurtosis,
+    collect_set_from_list_distinct as collect_set,
+    first_always_ignore_nulls as first,
+    factorial_ensure_int as factorial,
+    isnull_using_equal as isnull,
+    nanvl_as_case as nanvl,
+    percentile_approx_without_accuracy as percentile_approx,
+    rand_no_seed as rand,
+    base64_from_blob as base64,
+    decode_from_blob as decode,
+    format_string_with_pipes as format_string,
+    overlay_from_substr as overlay,
+    split_no_limit as split,
+    arrays_overlap_using_intersect as arrays_overlap,
+    slice_as_list_slice as slice,
+    array_join_null_replacement_with_transform as array_join,
+    element_at_using_brackets as element_at,
+    array_remove_using_filter as array_remove,
+    array_union_using_list_concat as array_union,
+    array_min_from_sort as array_min,
+    array_max_from_sort as array_max,
+    sequence_from_generate_series as sequence,
+)

sqlframe/duckdb/group.py

@@ -0,0 +1,14 @@
+# This code is based on code from Apache Spark under the license found in the LICENSE file located in the 'sqlframe' folder.
+
+from __future__ import annotations
+
+import typing as t
+
+from sqlframe.base.group import _BaseGroupedData
+
+if t.TYPE_CHECKING:
+    from sqlframe.duckdb.dataframe import DuckDBDataFrame
+
+
+class DuckDBGroupedData(_BaseGroupedData["DuckDBDataFrame"]):
+    pass

sqlframe/duckdb/readwriter.py

@@ -0,0 +1,111 @@
+# This code is based on code from Apache Spark under the license found in the LICENSE file located in the 'sqlframe' folder.
+
+from __future__ import annotations
+
+import logging
+import typing as t
+
+from sqlglot import exp
+from sqlglot.helper import ensure_list
+
+from sqlframe.base.readerwriter import _BaseDataFrameReader, _BaseDataFrameWriter
+from sqlframe.base.util import ensure_column_mapping, to_csv
+
+if t.TYPE_CHECKING:
+    from sqlframe.base._typing import OptionalPrimitiveType, PathOrPaths
+    from sqlframe.base.types import StructType
+    from sqlframe.duckdb.dataframe import DuckDBDataFrame
+    from sqlframe.duckdb.session import DuckDBSession  # noqa
+
+logger = logging.getLogger(__name__)
+
+
+class DuckDBDataFrameReader(_BaseDataFrameReader["DuckDBSession", "DuckDBDataFrame"]):
+    def load(
+        self,
+        path: t.Optional[PathOrPaths] = None,
+        format: t.Optional[str] = None,
+        schema: t.Optional[t.Union[StructType, str]] = None,
+        **options: OptionalPrimitiveType,
+    ) -> DuckDBDataFrame:
+        """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|
+        +---+----+
+        """
+        if schema:
+            column_mapping = ensure_column_mapping(schema)
+            select_column_mapping = column_mapping.copy()
+            if options.get("filename"):
+                select_column_mapping["filename"] = "VARCHAR"
+            select_columns = [x.expression for x in self._to_casted_columns(select_column_mapping)]
+            if format == "csv":
+                duckdb_columns = ", ".join(
+                    [f"'{column}': '{dtype}'" for column, dtype in column_mapping.items()]
+                )
+                options["columns"] = "{" + duckdb_columns + "}"
+        else:
+            select_columns = [exp.Star()]
+        if format:
+            paths = ",".join([f"'{path}'" for path in ensure_list(path)])
+            from_clause = f"read_{format}([{paths}], {to_csv(options)})"
+        else:
+            from_clause = f"'{path}'"
+        df = self.session.sql(exp.select(*select_columns).from_(from_clause), optimize=False)
+        self.session._last_loaded_file = path  # type: ignore
+        return df
+
+
+class DuckDBDataFrameWriter(_BaseDataFrameWriter["DuckDBSession", "DuckDBDataFrame"]):
+    def _write(self, path: str, mode: t.Optional[str], **options):  # type: ignore
+        mode, skip = self._validate_mode(path, mode)
+        if skip:
+            return
+        if mode == "append":
+            raise NotImplementedError("Append mode not supported")
+        options = to_csv(options, equality_char=" ")  # type: ignore
+        sqls = self._df.sql(pretty=False, optimize=False, as_list=True)
+        for i, sql in enumerate(sqls):
+            if i < len(sqls) - 1:
+                self._df.session._fetch_rows(sql)
+            else:
+                self._df.session._fetch_rows(f"COPY ({sqls[0]}) TO '{path}' ({options})")

sqlframe/duckdb/session.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import typing as t
+from functools import cached_property
+
+from sqlframe.base.session import _BaseSession
+from sqlframe.base.util import soundex
+from sqlframe.duckdb.catalog import DuckDBCatalog
+from sqlframe.duckdb.dataframe import DuckDBDataFrame
+from sqlframe.duckdb.readwriter import (
+    DuckDBDataFrameReader,
+    DuckDBDataFrameWriter,
+)
+
+if t.TYPE_CHECKING:
+    from duckdb import DuckDBPyConnection
+
+else:
+    DuckDBPyConnection = t.Any
+
+
+class DuckDBSession(
+    _BaseSession[  # type: ignore
+        DuckDBCatalog,
+        DuckDBDataFrameReader,
+        DuckDBDataFrameWriter,
+        DuckDBDataFrame,
+        DuckDBPyConnection,
+    ]
+):
+    _catalog = DuckDBCatalog
+    _reader = DuckDBDataFrameReader
+    _writer = DuckDBDataFrameWriter
+    _df = DuckDBDataFrame
+
+    DEFAULT_TIME_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+    def __init__(self, conn: t.Optional[DuckDBPyConnection] = None, *args, **kwargs):
+        import duckdb
+        from duckdb.typing import VARCHAR
+
+        if not hasattr(self, "_conn"):
+            conn = conn or duckdb.connect()
+            conn.create_function("SOUNDEX", lambda x: soundex(x), return_type=VARCHAR)
+            super().__init__(conn, *args, **kwargs)
+
+    @classmethod
+    def _try_get_map(cls, value: t.Any) -> t.Optional[t.Dict[str, t.Any]]:
+        if value and isinstance(value, dict) and "key" in value and "value" in value:
+            return dict(zip(value["key"], value["value"]))
+        return None
+
+    class Builder(_BaseSession.Builder):
+        DEFAULT_INPUT_DIALECT = "duckdb"
+        DEFAULT_OUTPUT_DIALECT = "duckdb"
+
+        @cached_property
+        def session(self) -> DuckDBSession:
+            return DuckDBSession(**self._session_kwargs)
+
+        def getOrCreate(self) -> DuckDBSession:
+            self._set_session_properties()
+            return self.session
+
+    builder = Builder()

sqlframe/duckdb/types.py

@@ -0,0 +1 @@
+from sqlframe.base.types import *

sqlframe/duckdb/window.py

@@ -0,0 +1 @@
+from sqlframe.base.window import *

sqlframe/postgres/__init__.py

@@ -0,0 +1,23 @@
+from sqlframe.postgres.catalog import PostgresCatalog
+from sqlframe.postgres.column import Column
+from sqlframe.postgres.dataframe import PostgresDataFrame, PostgresDataFrameNaFunctions
+from sqlframe.postgres.group import PostgresGroupedData
+from sqlframe.postgres.readwriter import (
+    PostgresDataFrameReader,
+    PostgresDataFrameWriter,
+)
+from sqlframe.postgres.session import PostgresSession
+from sqlframe.postgres.window import Window, WindowSpec
+
+__all__ = [
+    "PostgresCatalog",
+    "Column",
+    "PostgresDataFrame",
+    "PostgresDataFrameNaFunctions",
+    "PostgresGroupedData",
+    "PostgresDataFrameReader",
+    "PostgresDataFrameWriter",
+    "PostgresSession",
+    "Window",
+    "WindowSpec",
+]

sqlframe/postgres/catalog.py

@@ -0,0 +1,106 @@
+# This code is based on code from Apache Spark under the license found in the LICENSE file located in the 'sqlframe' folder.
+
+from __future__ import annotations
+
+import fnmatch
+import typing as t
+
+from sqlglot import exp, parse_one
+
+from sqlframe.base.catalog import Function, _BaseCatalog
+from sqlframe.base.mixins.catalog_mixins import (
+    GetCurrentCatalogFromFunctionMixin,
+    GetCurrentDatabaseFromFunctionMixin,
+    ListCatalogsFromInfoSchemaMixin,
+    ListColumnsFromInfoSchemaMixin,
+    ListDatabasesFromInfoSchemaMixin,
+    ListTablesFromInfoSchemaMixin,
+    SetCurrentDatabaseFromSearchPathMixin,
+)
+
+if t.TYPE_CHECKING:
+    from sqlframe.postgres.session import PostgresSession  # noqa
+    from sqlframe.postgres.dataframe import PostgresDataFrame  # noqa
+
+
+class PostgresCatalog(
+    GetCurrentCatalogFromFunctionMixin["PostgresSession", "PostgresDataFrame"],
+    GetCurrentDatabaseFromFunctionMixin["PostgresSession", "PostgresDataFrame"],
+    ListDatabasesFromInfoSchemaMixin["PostgresSession", "PostgresDataFrame"],
+    ListCatalogsFromInfoSchemaMixin["PostgresSession", "PostgresDataFrame"],
+    SetCurrentDatabaseFromSearchPathMixin["PostgresSession", "PostgresDataFrame"],
+    ListTablesFromInfoSchemaMixin["PostgresSession", "PostgresDataFrame"],
+    ListColumnsFromInfoSchemaMixin["PostgresSession", "PostgresDataFrame"],
+    _BaseCatalog["PostgresSession", "PostgresDataFrame"],
+):
+    CURRENT_CATALOG_EXPRESSION: exp.Expression = exp.column("current_catalog")
+
+    def listFunctions(
+        self, dbName: t.Optional[str] = None, pattern: t.Optional[str] = None
+    ) -> t.List[Function]:
+        """
+        Returns a t.List of functions registered in the specified database.
+
+        .. versionadded:: 3.4.0
+
+        Parameters
+        ----------
+        dbName : str
+            name of the database to t.List the functions.
+            ``dbName`` can be qualified with catalog name.
+        pattern : str
+            The pattern that the function name needs to match.
+
+            .. versionchanged: 3.5.0
+                Adds ``pattern`` argument.
+
+        Returns
+        -------
+        t.List
+            A t.List of :class:`Function`.
+
+        Notes
+        -----
+        If no database is specified, the current database and catalog
+        are used. This API includes all temporary functions.
+
+        Examples
+        --------
+        >>> spark.catalog.t.listFunctions()
+        [Function(name=...
+
+        >>> spark.catalog.t.listFunctions(pattern="to_*")
+        [Function(name=...
+
+        >>> spark.catalog.t.listFunctions(pattern="*not_existing_func*")
+        []
+        """
+        # SO: https://stackoverflow.com/questions/44143816/any-way-to-list-all-user-defined-postgresql-functions
+        query = parse_one(
+            """SELECT n.nspname as "namespace",
+  p.proname as "name"
+FROM pg_catalog.pg_proc p
+     LEFT JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace
+WHERE pg_catalog.pg_function_is_visible(p.oid)
+      AND n.nspname <> 'pg_catalog'
+      AND n.nspname <> 'information_schema'
+ORDER BY 1, 2;
+        """,
+            dialect=self.session.input_dialect,
+        )
+        functions = self.session._fetch_rows(query)
+        catalog = self.currentCatalog()
+        results = [
+            Function(
+                name=x["name"],
+                catalog=catalog,
+                namespace=[x["namespace"]],
+                description=None,
+                className="",
+                isTemporary=False,
+            )
+            for x in functions
+        ]
+        if pattern:
+            results = [x for x in results if fnmatch.fnmatch(x.name, pattern)]
+        return results

sqlframe/postgres/column.py

@@ -0,0 +1 @@
+from sqlframe.base.column import Column

sqlframe/postgres/dataframe.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import logging
+import sys
+import typing as t
+
+from sqlframe.base.dataframe import (
+    _BaseDataFrame,
+    _BaseDataFrameNaFunctions,
+    _BaseDataFrameStatFunctions,
+)
+from sqlframe.postgres.group import PostgresGroupedData
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+if t.TYPE_CHECKING:
+    from sqlframe.postgres.readwriter import PostgresDataFrameWriter
+    from sqlframe.postgres.session import PostgresSession
+
+
+logger = logging.getLogger(__name__)
+
+
+class PostgresDataFrameNaFunctions(_BaseDataFrameNaFunctions["PostgresDataFrame"]):
+    pass
+
+
+class PostgresDataFrameStatFunctions(_BaseDataFrameStatFunctions["PostgresDataFrame"]):
+    pass
+
+
+class PostgresDataFrame(
+    _BaseDataFrame[
+        "PostgresSession",
+        "PostgresDataFrameWriter",
+        "PostgresDataFrameNaFunctions",
+        "PostgresDataFrameStatFunctions",
+        "PostgresGroupedData",
+    ]
+):
+    _na = PostgresDataFrameNaFunctions
+    _stat = PostgresDataFrameStatFunctions
+    _group_data = PostgresGroupedData
+
+    def cache(self) -> Self:
+        logger.warning("Postgres does not support caching. Ignoring cache() call.")
+        return self
+
+    def persist(self) -> Self:
+        logger.warning("Postgres does not support persist. Ignoring persist() call.")
+        return self

sqlframe/postgres/functions.py

@@ -0,0 +1,61 @@
+import inspect
+import sys
+
+import sqlframe.base.functions
+
+module = sys.modules["sqlframe.base.functions"]
+globals().update(
+    {
+        name: func
+        for name, func in inspect.getmembers(module, inspect.isfunction)
+        if hasattr(func, "unsupported_engines")
+        and "postgres" not in func.unsupported_engines
+        and "*" not in func.unsupported_engines
+    }
+)
+
+
+from sqlframe.base.function_alternatives import (  # noqa
+    e_literal as e,
+    expm1_from_exp as expm1,
+    log1p_from_log as log1p,
+    rint_from_round as rint,
+    collect_set_from_list_distinct as collect_set,
+    isnan_using_equal as isnan,
+    isnull_using_equal as isnull,
+    nanvl_as_case as nanvl,
+    rand_no_seed as rand,
+    round_cast_as_numeric as round,
+    year_from_extract as year,
+    quarter_from_extract as quarter,
+    month_from_extract as month,
+    dayofweek_from_extract_with_isodow as dayofweek,
+    dayofmonth_from_extract_with_day as dayofmonth,
+    dayofyear_from_extract_doy as dayofyear,
+    hour_from_extract as hour,
+    minute_from_extract as minute,
+    second_from_extract as second,
+    weekofyear_from_extract_as_week as weekofyear,
+    make_date_casted_as_integer as make_date,
+    date_add_by_multiplication as date_add,
+    date_sub_by_multiplication as date_sub,
+    date_diff_with_subtraction as date_diff,
+    add_months_by_multiplication as add_months,
+    months_between_from_age_and_extract as months_between,
+    from_unixtime_from_timestamp as from_unixtime,
+    unix_timestamp_from_extract as unix_timestamp,
+    base64_from_blob as base64,
+    bas64_from_encode as base64,
+    unbase64_from_decode as unbase64,
+    decode_from_convert_from as decode,
+    encode_from_convert_to as encode,
+    format_number_from_to_char as format_number,
+    format_string_with_format as format_string,
+    split_from_regex_split_to_array as split,
+    array_contains_any as array_contains,
+    slice_with_brackets as slice,
+    element_at_using_brackets as element_at,
+    get_json_object_using_arrow_op as get_json_object,
+    array_min_from_subquery as array_min,
+    array_max_from_subquery as array_max,
+)

sqlframe/postgres/group.py

@@ -0,0 +1,14 @@
+# This code is based on code from Apache Spark under the license found in the LICENSE file located in the 'sqlframe' folder.
+
+from __future__ import annotations
+
+import typing as t
+
+from sqlframe.base.group import _BaseGroupedData
+
+if t.TYPE_CHECKING:
+    from sqlframe.postgres.dataframe import PostgresDataFrame
+
+
+class PostgresGroupedData(_BaseGroupedData["PostgresDataFrame"]):
+    pass

sqlframe/postgres/readwriter.py

@@ -0,0 +1,29 @@
+# This code is based on code from Apache Spark under the license found in the LICENSE file located in the 'sqlframe' folder.
+
+from __future__ import annotations
+
+import typing as t
+
+from sqlframe.base.mixins.readwriter_mixins import PandasLoaderMixin, PandasWriterMixin
+from sqlframe.base.readerwriter import (
+    _BaseDataFrameReader,
+    _BaseDataFrameWriter,
+)
+
+if t.TYPE_CHECKING:
+    from sqlframe.postgres.session import PostgresSession  # noqa
+    from sqlframe.postgres.dataframe import PostgresDataFrame  # noqa
+
+
+class PostgresDataFrameReader(
+    PandasLoaderMixin["PostgresSession", "PostgresDataFrame"],
+    _BaseDataFrameReader["PostgresSession", "PostgresDataFrame"],
+):
+    pass
+
+
+class PostgresDataFrameWriter(
+    PandasWriterMixin["PostgresSession", "PostgresDataFrame"],
+    _BaseDataFrameWriter["PostgresSession", "PostgresDataFrame"],
+):
+    pass

sqlframe/postgres/session.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot import exp
+
+from sqlframe.base.session import _BaseSession
+from sqlframe.postgres.catalog import PostgresCatalog
+from sqlframe.postgres.dataframe import PostgresDataFrame
+from sqlframe.postgres.readwriter import (
+    PostgresDataFrameReader,
+    PostgresDataFrameWriter,
+)
+
+if t.TYPE_CHECKING:
+    from psycopg2.extensions import connection as psycopg2_connection
+
+    from sqlframe.base.types import Row
+else:
+    psycopg2_connection = t.Any
+
+
+class PostgresSession(
+    _BaseSession[  # type: ignore
+        PostgresCatalog,
+        PostgresDataFrameReader,
+        PostgresDataFrameWriter,
+        PostgresDataFrame,
+        psycopg2_connection,
+    ],
+):
+    _catalog = PostgresCatalog
+    _reader = PostgresDataFrameReader
+    _writer = PostgresDataFrameWriter
+    _df = PostgresDataFrame
+
+    DEFAULT_TIME_FORMAT = "yyyy-MM-dd HH:MI:SS"
+
+    def __init__(self, conn: t.Optional[psycopg2_connection] = None):
+        if not hasattr(self, "_conn"):
+            super().__init__(conn)
+            self._execute("CREATE EXTENSION IF NOT EXISTS fuzzystrmatch")
+
+    def _fetch_rows(
+        self, sql: t.Union[str, exp.Expression], *, quote_identifiers: bool = True
+    ) -> t.List[Row]:
+        from psycopg2 import ProgrammingError
+
+        try:
+            return super()._fetch_rows(sql, quote_identifiers=quote_identifiers)
+        except ProgrammingError as e:
+            if "no results to fetch" in str(e):
+                return []
+            raise e
+
+    class Builder(_BaseSession.Builder):
+        DEFAULT_INPUT_DIALECT = "postgres"
+        DEFAULT_OUTPUT_DIALECT = "postgres"
+
+        @property
+        def session(self) -> PostgresSession:
+            return PostgresSession(**self._session_kwargs)
+
+        def getOrCreate(self) -> PostgresSession:
+            self._set_session_properties()
+            return self.session
+
+    builder = Builder()

sqlframe/postgres/types.py

@@ -0,0 +1 @@
+from sqlframe.base.types import *

sqlframe/postgres/window.py

@@ -0,0 +1 @@
+from sqlframe.base.window import *

sqlframe/redshift/__init__.py

@@ -0,0 +1,23 @@
+from sqlframe.redshift.catalog import RedshiftCatalog
+from sqlframe.redshift.column import Column
+from sqlframe.redshift.dataframe import RedshiftDataFrame, RedshiftDataFrameNaFunctions
+from sqlframe.redshift.group import RedshiftGroupedData
+from sqlframe.redshift.readwriter import (
+    RedshiftDataFrameReader,
+    RedshiftDataFrameWriter,
+)
+from sqlframe.redshift.session import RedshiftSession
+from sqlframe.redshift.window import Window, WindowSpec
+
+__all__ = [
+    "RedshiftCatalog",
+    "Column",
+    "RedshiftDataFrame",
+    "RedshiftDataFrameNaFunctions",
+    "RedshiftGroupedData",
+    "RedshiftDataFrameReader",
+    "RedshiftDataFrameWriter",
+    "RedshiftSession",
+    "Window",
+    "WindowSpec",
+]