sqlframe 1.1.3__py3-none-any.whl

sqlframe/bigquery/functions.py

@@ -0,0 +1,378 @@
+from __future__ import annotations
+
+import inspect
+import sys
+import typing as t
+
+from sqlglot import exp as sqlglot_expression
+
+import sqlframe.base.functions
+from sqlframe.base.util import get_func_from_session
+from sqlframe.bigquery.column import Column
+
+if t.TYPE_CHECKING:
+    from sqlframe.base._typing import ColumnOrLiteral, ColumnOrName
+
+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 "bigquery" 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,
+    factorial_from_case_statement as factorial,
+    log1p_from_log as log1p,
+    rint_from_round as rint,
+    collect_set_from_list_distinct as collect_set,
+    isnull_using_equal as isnull,
+    nanvl_as_case as nanvl,
+    percentile_approx_without_accuracy_and_plural as percentile_approx,
+    rand_no_seed as rand,
+    year_from_extract as year,
+    quarter_from_extract as quarter,
+    month_from_extract as month,
+    dayofweek_from_extract as dayofweek,
+    dayofmonth_from_extract_with_day as dayofmonth,
+    dayofyear_from_extract as dayofyear,
+    hour_from_extract as hour,
+    minute_from_extract as minute,
+    second_from_extract as second,
+    weekofyear_from_extract_as_isoweek as weekofyear,
+    make_date_from_date_func as make_date,
+    to_date_from_timestamp as to_date,
+    last_day_with_cast as last_day,
+    sha1_force_sha1_and_to_hex as sha1,
+    hash_from_farm_fingerprint as hash,
+    base64_from_blob as base64,
+    concat_ws_from_array_to_string as concat_ws,
+    format_string_with_format as format_string,
+    instr_using_strpos as instr,
+    overlay_from_substr as overlay,
+    split_with_split as split,
+    regexp_extract_only_one_group as regexp_extract,
+    hex_casted_as_bytes as hex,
+    bit_length_from_length as bit_length,
+    element_at_using_brackets as element_at,
+    array_union_using_array_concat as array_union,
+    sequence_from_generate_array as sequence,
+)
+
+
+def typeof(col: ColumnOrName) -> Column:
+    return Column(
+        sqlglot_expression.Anonymous(
+            this="bqutil.fn.typeof", expressions=[Column.ensure_col(col).expression]
+        )
+    )
+
+
+def degrees(col: ColumnOrName) -> Column:
+    return Column(
+        sqlglot_expression.Anonymous(
+            this="bqutil.fn.degrees", expressions=[Column.ensure_col(col).expression]
+        )
+    )
+
+
+def radians(col: ColumnOrName) -> Column:
+    return Column(
+        sqlglot_expression.Anonymous(
+            this="bqutil.fn.radians", expressions=[Column.ensure_col(col).expression]
+        )
+    )
+
+
+def bround(col: ColumnOrName, scale: t.Optional[int] = None) -> Column:
+    from sqlframe.base.session import _BaseSession
+
+    lit = get_func_from_session("lit", _BaseSession())
+
+    expressions = [Column.ensure_col(col).cast("bignumeric").expression]
+    if scale is not None:
+        expressions.append(lit(scale).expression)
+    return Column(
+        sqlglot_expression.Anonymous(
+            this="bqutil.fn.cw_round_half_even",
+            expressions=expressions,
+        )
+    )
+
+
+def months_between(
+    date1: ColumnOrName, date2: ColumnOrName, roundOff: t.Optional[bool] = None
+) -> Column:
+    roundOff = True if roundOff is None else roundOff
+    round = get_func_from_session("round")
+    lit = get_func_from_session("lit")
+
+    value = Column(
+        sqlglot_expression.Anonymous(
+            this="bqutil.fn.cw_months_between",
+            expressions=[
+                Column.ensure_col(date1).cast("datetime").expression,
+                Column.ensure_col(date2).cast("datetime").expression,
+            ],
+        )
+    )
+    if roundOff:
+        value = round(value, lit(8))
+    return value
+
+
+def next_day(col: ColumnOrName, dayOfWeek: str) -> Column:
+    lit = get_func_from_session("lit")
+
+    return Column(
+        sqlglot_expression.Anonymous(
+            this="bqutil.fn.cw_next_day",
+            expressions=[Column.ensure_col(col).cast("date").expression, lit(dayOfWeek).expression],
+        )
+    )
+
+
+def from_unixtime(col: ColumnOrName, format: t.Optional[str] = None) -> Column:
+    from sqlframe.base.session import _BaseSession
+
+    session: _BaseSession = _BaseSession()
+    lit = get_func_from_session("lit")
+    to_timestamp = get_func_from_session("to_timestamp")
+
+    expressions = [Column.ensure_col(col).expression]
+    if format is not None:
+        expressions.append(lit(format).expression)
+    return Column(
+        sqlglot_expression.Anonymous(
+            this="FORMAT_TIMESTAMP",
+            expressions=[
+                lit(session.DEFAULT_TIME_FORMAT).expression,
+                to_timestamp(
+                    Column(
+                        sqlglot_expression.Anonymous(
+                            this="TIMESTAMP_SECONDS", expressions=expressions
+                        )
+                    ),
+                    format,
+                ).expression,
+            ],
+        )
+    )
+
+
+def unix_timestamp(
+    timestamp: t.Optional[ColumnOrName] = None, format: t.Optional[str] = None
+) -> Column:
+    from sqlframe.base.session import _BaseSession
+
+    lit = get_func_from_session("lit")
+
+    if format is None:
+        format = _BaseSession().DEFAULT_TIME_FORMAT
+    return Column(
+        sqlglot_expression.Anonymous(
+            this="UNIX_SECONDS",
+            expressions=[
+                sqlglot_expression.Anonymous(
+                    this="PARSE_TIMESTAMP",
+                    expressions=[
+                        lit(format).expression,
+                        Column.ensure_col(timestamp).expression,
+                        lit("UTC").expression,
+                    ],
+                )
+            ],
+        )
+    )
+
+
+def format_number(col: ColumnOrName, d: int) -> Column:
+    round = get_func_from_session("round")
+    lit = get_func_from_session("lit")
+
+    return Column(
+        sqlglot_expression.Anonymous(
+            this="FORMAT",
+            expressions=[
+                lit(f"%'.{d}f").expression,
+                round(Column.ensure_col(col).cast("float"), d).expression,
+            ],
+        )
+    )
+
+
+def substring_index(str: ColumnOrName, delim: str, count: int) -> Column:
+    lit = get_func_from_session("lit")
+
+    return Column(
+        sqlglot_expression.Anonymous(
+            this="bqutil.fn.cw_substring_index",
+            expressions=[
+                Column.ensure_col(str).expression,
+                lit(delim).expression,
+                lit(count).expression,
+            ],
+        )
+    )
+
+
+def bin(col: ColumnOrName) -> Column:
+    return (
+        Column(
+            sqlglot_expression.Anonymous(
+                this="bqutil.fn.to_binary",
+                expressions=[Column.ensure_col(col).expression],
+            )
+        )
+        .cast("int")
+        .cast("string")
+    )
+
+
+def slice(
+    x: ColumnOrName, start: t.Union[ColumnOrName, int], length: t.Union[ColumnOrName, int]
+) -> Column:
+    lit = get_func_from_session("lit")
+
+    start_col = start if isinstance(start, Column) else lit(start)
+    length_col = length if isinstance(length, Column) else lit(length)
+
+    subquery = (
+        sqlglot_expression.select(
+            sqlglot_expression.column("x"),
+        )
+        .from_(
+            sqlglot_expression.Unnest(
+                expressions=[Column.ensure_col(x).expression],
+                alias=sqlglot_expression.TableAlias(
+                    columns=[sqlglot_expression.to_identifier("x")],
+                ),
+                offset=sqlglot_expression.to_identifier("offset"),
+            )
+        )
+        .where(
+            sqlglot_expression.Between(
+                this=sqlglot_expression.column("offset"),
+                low=(start_col - lit(1)).expression,
+                high=(start_col + length_col).expression,
+            )
+        )
+    )
+
+    return Column(
+        sqlglot_expression.Anonymous(
+            this="ARRAY",
+            expressions=[subquery],
+        )
+    )
+
+
+def array_position(col: ColumnOrName, value: ColumnOrLiteral) -> Column:
+    lit = get_func_from_session("lit")
+
+    value_col = value if isinstance(value, Column) else lit(value)
+
+    return Column(
+        sqlglot_expression.Coalesce(
+            this=sqlglot_expression.Anonymous(
+                this="bqutil.fn.find_in_set",
+                expressions=[
+                    value_col.expression,
+                    sqlglot_expression.Anonymous(
+                        this="ARRAY_TO_STRING",
+                        expressions=[Column.ensure_col(col).expression, lit(",").expression],
+                    ),
+                ],
+            ),
+            expressions=[lit(0).expression],
+        )
+    )
+
+
+def array_remove(col: ColumnOrName, value: ColumnOrLiteral) -> Column:
+    lit = get_func_from_session("lit")
+
+    value_col = value if isinstance(value, Column) else lit(value)
+
+    filter_subquery = sqlglot_expression.select(
+        "*",
+    ).from_(
+        sqlglot_expression.Unnest(
+            expressions=[Column.ensure_col(col).expression],
+            alias=sqlglot_expression.TableAlias(
+                columns=[sqlglot_expression.to_identifier("x")],
+            ),
+        )
+    )
+
+    agg_subquery = (
+        sqlglot_expression.select(
+            sqlglot_expression.Anonymous(
+                this="ARRAY_AGG",
+                expressions=[sqlglot_expression.column("x")],
+            ),
+        )
+        .from_(filter_subquery.subquery("t"))
+        .where(
+            sqlglot_expression.NEQ(
+                this=sqlglot_expression.column("x", "t"),
+                expression=value_col.expression,
+            )
+        )
+    )
+
+    return Column(agg_subquery.subquery())
+
+
+def array_distinct(col: ColumnOrName) -> Column:
+    return Column(
+        sqlglot_expression.Anonymous(
+            this="bqutil.fn.cw_array_distinct",
+            expressions=[Column.ensure_col(col).expression],
+        )
+    )
+
+
+def array_min(col: ColumnOrName) -> Column:
+    return Column(
+        sqlglot_expression.Anonymous(
+            this="bqutil.fn.cw_array_min",
+            expressions=[Column.ensure_col(col).expression],
+        )
+    )
+
+
+def array_max(col: ColumnOrName) -> Column:
+    return Column(
+        sqlglot_expression.Anonymous(
+            this="bqutil.fn.cw_array_max",
+            expressions=[Column.ensure_col(col).expression],
+        )
+    )
+
+
+def sort_array(col: ColumnOrName, asc: t.Optional[bool] = None) -> Column:
+    order = "ASC" if asc or asc is None else "DESC"
+    subquery = (
+        sqlglot_expression.select("x")
+        .from_(
+            sqlglot_expression.Unnest(
+                expressions=[Column.ensure_col(col).expression],
+                alias=sqlglot_expression.TableAlias(
+                    columns=[sqlglot_expression.to_identifier("x")],
+                ),
+            )
+        )
+        .order_by(f"x {order}")
+    )
+
+    return Column(sqlglot_expression.Anonymous(this="ARRAY", expressions=[subquery]))
+
+
+array_sort = sort_array

sqlframe/bigquery/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.bigquery.dataframe import BigQueryDataFrame
+
+
+class BigQueryGroupedData(_BaseGroupedData["BigQueryDataFrame"]):
+    pass

sqlframe/bigquery/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.bigquery.session import BigQuerySession  # noqa
+    from sqlframe.bigquery.dataframe import BigQueryDataFrame  # noqa
+
+
+class BigQueryDataFrameReader(
+    PandasLoaderMixin["BigQuerySession", "BigQueryDataFrame"],
+    _BaseDataFrameReader["BigQuerySession", "BigQueryDataFrame"],
+):
+    pass
+
+
+class BigQueryDataFrameWriter(
+    PandasWriterMixin["BigQuerySession", "BigQueryDataFrame"],
+    _BaseDataFrameWriter["BigQuerySession", "BigQueryDataFrame"],
+):
+    pass

sqlframe/bigquery/session.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlframe.base.session import _BaseSession
+from sqlframe.bigquery.catalog import BigQueryCatalog
+from sqlframe.bigquery.dataframe import BigQueryDataFrame
+from sqlframe.bigquery.readwriter import (
+    BigQueryDataFrameReader,
+    BigQueryDataFrameWriter,
+)
+
+if t.TYPE_CHECKING:
+    from google.cloud.bigquery.client import Client as BigQueryClient
+    from google.cloud.bigquery.dbapi.connection import Connection as BigQueryConnection
+else:
+    BigQueryClient = t.Any
+    BigQueryConnection = t.Any
+
+
+class BigQuerySession(
+    _BaseSession[  # type: ignore
+        BigQueryCatalog,
+        BigQueryDataFrameReader,
+        BigQueryDataFrameWriter,
+        BigQueryDataFrame,
+        BigQueryConnection,
+    ],
+):
+    _catalog = BigQueryCatalog
+    _reader = BigQueryDataFrameReader
+    _writer = BigQueryDataFrameWriter
+    _df = BigQueryDataFrame
+
+    DEFAULT_TIME_FORMAT = "%Y-%m-%d %H:%M:%S"
+    QUALIFY_INFO_SCHEMA_WITH_DATABASE = True
+    SANITIZE_COLUMN_NAMES = True
+
+    def __init__(
+        self, conn: t.Optional[BigQueryConnection] = None, default_dataset: t.Optional[str] = None
+    ):
+        from google.cloud import bigquery
+        from google.cloud.bigquery.dbapi import connect
+
+        if not hasattr(self, "_conn"):
+            super().__init__(conn or connect())
+            if self._client.default_query_job_config is None:
+                self._client.default_query_job_config = bigquery.QueryJobConfig()
+            self.default_dataset = default_dataset
+
+    @property
+    def _client(self) -> BigQueryClient:
+        assert self._connection
+        return self._connection._client
+
+    @property
+    def default_dataset(self) -> t.Optional[str]:
+        return self._default_dataset
+
+    @default_dataset.setter
+    def default_dataset(self, dataset: str) -> None:
+        self._default_dataset = dataset
+        self._client.default_query_job_config.default_dataset = dataset  # type: ignore
+
+    @property
+    def default_project(self) -> str:
+        return self._client.project
+
+    @default_project.setter
+    def default_project(self, project: str) -> None:
+        self._client.project = project
+
+    @classmethod
+    def _try_get_map(cls, value: t.Any) -> t.Optional[t.Dict[str, t.Any]]:
+        return None
+
+    class Builder(_BaseSession.Builder):
+        DEFAULT_INPUT_DIALECT = "bigquery"
+        DEFAULT_OUTPUT_DIALECT = "bigquery"
+
+        @property
+        def session(self) -> BigQuerySession:
+            return BigQuerySession(**self._session_kwargs)
+
+        def getOrCreate(self) -> BigQuerySession:
+            self._set_session_properties()
+            return self.session
+
+    builder = Builder()

sqlframe/bigquery/types.py

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

sqlframe/bigquery/window.py

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

sqlframe/duckdb/__init__.py

@@ -0,0 +1,20 @@
+from sqlframe.duckdb.catalog import DuckDBCatalog
+from sqlframe.duckdb.column import DuckDBColumn
+from sqlframe.duckdb.dataframe import DuckDBDataFrame, DuckDBDataFrameNaFunctions
+from sqlframe.duckdb.group import DuckDBGroupedData
+from sqlframe.duckdb.readwriter import DuckDBDataFrameReader, DuckDBDataFrameWriter
+from sqlframe.duckdb.session import DuckDBSession
+from sqlframe.duckdb.window import Window, WindowSpec
+
+__all__ = [
+    "DuckDBCatalog",
+    "DuckDBColumn",
+    "DuckDBDataFrame",
+    "DuckDBDataFrameNaFunctions",
+    "DuckDBGroupedData",
+    "DuckDBDataFrameReader",
+    "DuckDBDataFrameWriter",
+    "DuckDBSession",
+    "Window",
+    "WindowSpec",
+]

sqlframe/duckdb/catalog.py

@@ -0,0 +1,108 @@
+# 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
+
+from sqlframe.base.catalog import Function, _BaseCatalog
+from sqlframe.base.mixins.catalog_mixins import (
+    GetCurrentCatalogFromFunctionMixin,
+    GetCurrentDatabaseFromFunctionMixin,
+    ListCatalogsFromInfoSchemaMixin,
+    ListColumnsFromInfoSchemaMixin,
+    ListDatabasesFromInfoSchemaMixin,
+    ListTablesFromInfoSchemaMixin,
+    SetCurrentCatalogFromUseMixin,
+    SetCurrentDatabaseFromUseMixin,
+)
+from sqlframe.base.util import schema_, to_schema
+
+if t.TYPE_CHECKING:
+    from sqlframe.duckdb.session import DuckDBSession  # noqa
+    from sqlframe.duckdb.dataframe import DuckDBDataFrame  # noqa
+
+
+class DuckDBCatalog(
+    GetCurrentCatalogFromFunctionMixin["DuckDBSession", "DuckDBDataFrame"],
+    SetCurrentCatalogFromUseMixin["DuckDBSession", "DuckDBDataFrame"],
+    GetCurrentDatabaseFromFunctionMixin["DuckDBSession", "DuckDBDataFrame"],
+    ListDatabasesFromInfoSchemaMixin["DuckDBSession", "DuckDBDataFrame"],
+    ListCatalogsFromInfoSchemaMixin["DuckDBSession", "DuckDBDataFrame"],
+    SetCurrentDatabaseFromUseMixin["DuckDBSession", "DuckDBDataFrame"],
+    ListTablesFromInfoSchemaMixin["DuckDBSession", "DuckDBDataFrame"],
+    ListColumnsFromInfoSchemaMixin["DuckDBSession", "DuckDBDataFrame"],
+    _BaseCatalog["DuckDBSession", "DuckDBDataFrame"],
+):
+    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*")
+        []
+        """
+        if not dbName:
+            schema = schema_(
+                db=exp.parse_identifier(self.currentDatabase(), dialect=self.session.input_dialect),
+                catalog=exp.parse_identifier(
+                    self.currentCatalog(), dialect=self.session.input_dialect
+                ),
+            )
+        else:
+            schema = to_schema(dbName, dialect=self.session.input_dialect)
+        select = (
+            exp.select("function_name", "schema_name", "database_name")
+            .from_("duckdb_functions()")
+            .where(exp.column("schema_name").eq(schema.db))
+        )
+        if schema.catalog:
+            select = select.where(exp.column("database_name").eq(schema.catalog))
+        functions = self.session._fetch_rows(select)
+        if pattern:
+            functions = [x for x in functions if fnmatch.fnmatch(x["function_name"], pattern)]
+        return [
+            Function(
+                name=x["function_name"],
+                catalog=x["database_name"],
+                namespace=[x["schema_name"]],
+                description=None,
+                className="",
+                isTemporary=False,
+            )
+            for x in functions
+        ]

sqlframe/duckdb/column.py

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

sqlframe/duckdb/dataframe.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+import logging
+import sys
+import typing as t
+
+from sqlframe.base.dataframe import (
+    _BaseDataFrame,
+    _BaseDataFrameNaFunctions,
+    _BaseDataFrameStatFunctions,
+)
+from sqlframe.duckdb.group import DuckDBGroupedData
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+if t.TYPE_CHECKING:
+    from sqlframe.duckdb.session import DuckDBSession  # noqa
+    from sqlframe.duckdb.readwriter import DuckDBDataFrameWriter  # noqa
+    from sqlframe.duckdb.group import DuckDBGroupedData  # noqa
+
+
+logger = logging.getLogger(__name__)
+
+
+class DuckDBDataFrameNaFunctions(_BaseDataFrameNaFunctions["DuckDBDataFrame"]):
+    pass
+
+
+class DuckDBDataFrameStatFunctions(_BaseDataFrameStatFunctions["DuckDBDataFrame"]):
+    pass
+
+
+class DuckDBDataFrame(
+    _BaseDataFrame[
+        "DuckDBSession",
+        "DuckDBDataFrameWriter",
+        "DuckDBDataFrameNaFunctions",
+        "DuckDBDataFrameStatFunctions",
+        "DuckDBGroupedData",
+    ]
+):
+    _na = DuckDBDataFrameNaFunctions
+    _stat = DuckDBDataFrameStatFunctions
+    _group_data = DuckDBGroupedData
+
+    def cache(self) -> Self:
+        logger.warning("DuckDB does not support caching. Ignoring cache() call.")
+        return self
+
+    def persist(self) -> Self:
+        logger.warning("DuckDB does not support persist. Ignoring persist() call.")
+        return self