sqlframe 1.1.3__py3-none-any.whl

sqlframe/base/group.py

@@ -0,0 +1,102 @@
+# 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.operations import Operation, group_operation, operation
+
+if t.TYPE_CHECKING:
+    from sqlframe.base.column import Column
+    from sqlframe.base.session import DF
+else:
+    DF = t.TypeVar("DF")
+
+
+# https://spark.apache.org/docs/latest/sql-ref-syntax-qry-select-groupby.html
+# https://stackoverflow.com/questions/37975227/what-is-the-difference-between-cube-rollup-and-groupby-operators
+class _BaseGroupedData(t.Generic[DF]):
+    def __init__(
+        self,
+        df: DF,
+        group_by_cols: t.Union[t.List[Column], t.List[t.List[Column]]],
+        last_op: Operation,
+    ):
+        self._df = df.copy()
+        self.session = df.session
+        self.last_op = last_op
+        self.group_by_cols = group_by_cols
+
+    def _get_function_applied_columns(
+        self, func_name: str, cols: t.Tuple[str, ...]
+    ) -> t.List[Column]:
+        from sqlframe.base import functions as F
+
+        func_name = func_name.lower()
+        return [
+            getattr(F, func_name)(name).alias(
+                self.session._sanitize_column_name(f"{func_name}({name})")
+            )
+            for name in cols
+        ]
+
+    @group_operation(Operation.SELECT)
+    def agg(self, *exprs: t.Union[Column, t.Dict[str, str]]) -> DF:
+        from sqlframe.base.column import Column
+
+        columns = (
+            [Column(f"{agg_func}({column_name})") for column_name, agg_func in exprs[0].items()]
+            if isinstance(exprs[0], dict)
+            else exprs
+        )
+        cols = self._df._ensure_and_normalize_cols(columns)
+
+        if not self.group_by_cols or not isinstance(self.group_by_cols[0], (list, tuple, set)):
+            expression = self._df.expression.group_by(
+                *[x.expression for x in self.group_by_cols]  # type: ignore
+            ).select(*[x.expression for x in self.group_by_cols + cols], append=False)  # type: ignore
+            group_by_cols = self.group_by_cols
+        else:
+            from sqlglot import exp
+
+            expression = self._df.expression
+            all_grouping_sets = []
+            group_by_cols = []
+            for grouping_set in self.group_by_cols:
+                all_grouping_sets.append(
+                    exp.Tuple(expressions=[x.expression for x in grouping_set])  # type: ignore
+                )
+                group_by_cols.extend(grouping_set)  # type: ignore
+            group_by_cols = list(dict.fromkeys(group_by_cols))
+            group_by = exp.Group(grouping_sets=all_grouping_sets)
+            expression.set("group", group_by)
+        for col in cols:
+            # Spark supports having an empty grouping_id which means all of the columns but other dialects
+            # like duckdb don't support this so we expand the grouping_id to include all of the columns
+            if col.column_expression.this == "GROUPING_ID":
+                col.column_expression.set("expressions", [x.expression for x in group_by_cols])  # type: ignore
+        expression = expression.select(*[x.expression for x in group_by_cols + cols], append=False)  # type: ignore
+        return self._df.copy(expression=expression)
+
+    def count(self) -> DF:
+        from sqlframe.base import functions as F
+
+        return self.agg(F.count("*").alias("count"))
+
+    def mean(self, *cols: str) -> DF:
+        return self.avg(*cols)
+
+    def avg(self, *cols: str) -> DF:
+        return self.agg(*self._get_function_applied_columns("avg", cols))
+
+    def max(self, *cols: str) -> DF:
+        return self.agg(*self._get_function_applied_columns("max", cols))
+
+    def min(self, *cols: str) -> DF:
+        return self.agg(*self._get_function_applied_columns("min", cols))
+
+    def sum(self, *cols: str) -> DF:
+        return self.agg(*self._get_function_applied_columns("sum", cols))
+
+    def pivot(self, *cols: str) -> DF:
+        raise NotImplementedError("Sum distinct is not currently implemented")

sqlframe/base/mixins/catalog_mixins.py

@@ -0,0 +1,419 @@
+import fnmatch
+import typing as t
+
+from sqlglot import exp
+
+from sqlframe.base.catalog import (
+    DF,
+    SESSION,
+    CatalogMetadata,
+    Column,
+    Database,
+    Table,
+    _BaseCatalog,
+)
+from sqlframe.base.decorators import normalize
+from sqlframe.base.util import decoded_str, schema_, to_schema
+
+
+class _BaseInfoSchemaMixin(_BaseCatalog, t.Generic[SESSION, DF]):
+    QUALIFY_INFO_SCHEMA_WITH_DATABASE = False
+    UPPERCASE_INFO_SCHEMA = False
+
+    def _get_info_schema_table(
+        self,
+        table_name: str,
+        database: t.Optional[str] = None,
+        qualify_override: t.Optional[bool] = None,
+    ) -> exp.Table:
+        table = f"information_schema.{table_name}"
+        if self.UPPERCASE_INFO_SCHEMA:
+            table = table.upper()
+        qualify = (
+            qualify_override
+            if qualify_override is not None
+            else self.QUALIFY_INFO_SCHEMA_WITH_DATABASE
+        )
+        if qualify:
+            db = database or self.currentDatabase()
+            if not db:
+                raise ValueError("Table name must be qualified with a database.")
+            table = f"{db}.{table}"
+        return exp.to_table(table)
+
+
+class GetCurrentCatalogFromFunctionMixin(_BaseCatalog, t.Generic[SESSION, DF]):
+    CURRENT_CATALOG_EXPRESSION: exp.Expression = exp.func("current_catalog")
+
+    def currentCatalog(self) -> str:
+        """Returns the current default catalog in this session.
+
+        .. versionadded:: 3.4.0
+
+        Examples
+        --------
+        >>> spark.catalog.currentCatalog()
+        'spark_catalog'
+        """
+        return self.session._fetch_rows(
+            exp.select(self.CURRENT_CATALOG_EXPRESSION), quote_identifiers=False
+        )[0][0]
+
+
+class GetCurrentDatabaseFromFunctionMixin(_BaseCatalog, t.Generic[SESSION, DF]):
+    CURRENT_DATABASE_EXPRESSION: exp.Expression = exp.func("current_schema")
+
+    def currentDatabase(self) -> str:
+        """Returns the current default schema in this session.
+
+        .. versionadded:: 3.4.0
+
+        Examples
+        --------
+        >>> spark.catalog.currentDatabase()
+        'default'
+        """
+        return self.session._fetch_rows(exp.select(self.CURRENT_DATABASE_EXPRESSION))[0][0]
+
+
+class SetCurrentCatalogFromUseMixin(_BaseCatalog, t.Generic[SESSION, DF]):
+    def setCurrentCatalog(self, catalogName: str) -> None:
+        """Sets the current default catalog in this session.
+
+        .. versionadded:: 3.4.0
+
+        Parameters
+        ----------
+        catalogName : str
+            name of the catalog to set
+
+        Examples
+        --------
+        >>> spark.catalog.setCurrentCatalog("spark_catalog")
+        """
+        self.session._execute(
+            exp.Use(this=exp.parse_identifier(catalogName, dialect=self.session.input_dialect))
+        )
+
+
+class ListDatabasesFromInfoSchemaMixin(_BaseInfoSchemaMixin, t.Generic[SESSION, DF]):
+    def listDatabases(self, pattern: t.Optional[str] = None) -> t.List[Database]:
+        """
+        Returns a t.List of databases available across all sessions.
+
+        .. versionadded:: 2.0.0
+
+        Parameters
+        ----------
+        pattern : str
+            The pattern that the database name needs to match.
+
+            .. versionchanged: 3.5.0
+                Adds ``pattern`` argument.
+
+        Returns
+        -------
+        t.List
+            A t.List of :class:`Database`.
+
+        Examples
+        --------
+        >>> spark.catalog.t.listDatabases()
+        [Database(name='default', catalog='spark_catalog', description='default database', ...
+
+        >>> spark.catalog.t.listDatabases("def*")
+        [Database(name='default', catalog='spark_catalog', description='default database', ...
+
+        >>> spark.catalog.t.listDatabases("def2*")
+        []
+        """
+        table = self._get_info_schema_table("schemata", qualify_override=False)
+        results = self.session._fetch_rows(
+            exp.Select().select("schema_name", "catalog_name").from_(table)
+        )
+        databases = [
+            Database(name=x[0], catalog=x[1], description=None, locationUri="") for x in results
+        ]
+        if pattern:
+            databases = [db for db in databases if fnmatch.fnmatch(db.name, pattern)]
+        return databases
+
+
+class ListCatalogsFromInfoSchemaMixin(_BaseInfoSchemaMixin, t.Generic[SESSION, DF]):
+    def listCatalogs(self, pattern: t.Optional[str] = None) -> t.List[CatalogMetadata]:
+        """
+        Returns a t.List of databases available across all sessions.
+
+        .. versionadded:: 2.0.0
+
+        Parameters
+        ----------
+        pattern : str
+            The pattern that the database name needs to match.
+
+            .. versionchanged: 3.5.0
+                Adds ``pattern`` argument.
+
+        Returns
+        -------
+        t.List
+            A t.List of :class:`Database`.
+
+        Examples
+        --------
+        >>> spark.catalog.t.listDatabases()
+        [Database(name='default', catalog='spark_catalog', description='default database', ...
+
+        >>> spark.catalog.t.listDatabases("def*")
+        [Database(name='default', catalog='spark_catalog', description='default database', ...
+
+        >>> spark.catalog.t.listDatabases("def2*")
+        []
+        """
+        table = self._get_info_schema_table("schemata")
+        results = self.session._fetch_rows(
+            exp.Select().select("catalog_name").from_(table).distinct()
+        )
+        catalogs = [CatalogMetadata(name=x[0], description=None) for x in results]
+        if pattern:
+            catalogs = [catalog for catalog in catalogs if fnmatch.fnmatch(catalog.name, pattern)]
+        return catalogs
+
+
+class SetCurrentDatabaseFromSearchPathMixin(_BaseCatalog, t.Generic[SESSION, DF]):
+    def setCurrentDatabase(self, dbName: str) -> None:
+        """
+        Sets the current default database in this session.
+
+        .. versionadded:: 2.0.0
+
+        Examples
+        --------
+        >>> spark.catalog.setCurrentDatabase("default")
+        """
+        self.session._execute(f'SET search_path TO "{dbName}"')
+
+
+class SetCurrentDatabaseFromUseMixin(_BaseCatalog, t.Generic[SESSION, DF]):
+    def setCurrentDatabase(self, dbName: str) -> None:
+        """
+        Sets the current default database in this session.
+
+        .. versionadded:: 2.0.0
+
+        Examples
+        --------
+        >>> spark.catalog.setCurrentDatabase("default")
+        """
+        schema = to_schema(dbName, dialect=self.session.input_dialect)
+        if not schema.catalog:
+            schema.set(
+                "catalog",
+                exp.parse_identifier(self.currentCatalog(), dialect=self.session.input_dialect),
+            )
+        self.session._execute(exp.Use(this=schema))
+
+
+class ListTablesFromInfoSchemaMixin(_BaseInfoSchemaMixin, t.Generic[SESSION, DF]):
+    @normalize(["dbName"])
+    def listTables(
+        self, dbName: t.Optional[str] = None, pattern: t.Optional[str] = None
+    ) -> t.List[Table]:
+        """Returns a t.List of tables/views in the specified database.
+
+        .. versionadded:: 2.0.0
+
+        Parameters
+        ----------
+        dbName : str
+            name of the database to t.List the tables.
+
+            .. versionchanged:: 3.4.0
+               Allow ``dbName`` to be qualified with catalog name.
+
+        pattern : str
+            The pattern that the database name needs to match.
+
+            .. versionchanged: 3.5.0
+                Adds ``pattern`` argument.
+
+        Returns
+        -------
+        t.List
+            A t.List of :class:`Table`.
+
+        Notes
+        -----
+        If no database is specified, the current database and catalog
+        are used. This API includes all temporary views.
+
+        Examples
+        --------
+        >>> spark.range(1).createTempView("test_view")
+        >>> spark.catalog.t.listTables()
+        [Table(name='test_view', catalog=None, namespace=[], description=None, ...
+
+        >>> spark.catalog.t.listTables(pattern="test*")
+        [Table(name='test_view', catalog=None, namespace=[], description=None, ...
+
+        >>> spark.catalog.t.listTables(pattern="table*")
+        []
+
+        >>> _ = spark.catalog.dropTempView("test_view")
+        >>> spark.catalog.t.listTables()
+        []
+        """
+        if dbName is None and pattern is None:
+            schema = schema_(
+                db=exp.parse_identifier(self.currentDatabase(), dialect=self.session.input_dialect),
+                catalog=exp.parse_identifier(
+                    self.currentCatalog(), dialect=self.session.input_dialect
+                ),
+            )
+        elif dbName:
+            schema = to_schema(dbName, dialect=self.session.input_dialect)
+        else:
+            schema = None
+        table = self._get_info_schema_table("tables", database=schema.db if schema else None)
+        select = exp.select(
+            'table_name AS "table_name"',
+            'table_schema AS "table_schema"',
+            'table_catalog AS "table_catalog"',
+            'table_type AS "table_type"',
+        ).from_(table)
+        if schema and schema.db:
+            select = select.where(exp.column("table_schema").eq(schema.db))
+        if schema and schema.catalog:
+            select = select.where(exp.column("table_catalog").eq(schema.catalog))
+        results = self.session._fetch_rows(select)
+        tables = [
+            Table(
+                name=x["table_name"],
+                catalog=x["table_catalog"],
+                namespace=[x["table_schema"]],
+                description=None,
+                tableType="VIEW" if x["table_type"] == "VIEW" else "MANAGED",
+                isTemporary=False,
+            )
+            for x in results
+        ]
+        for table in self.session.temp_views.keys():
+            tables.append(
+                Table(
+                    name=table,  # type: ignore
+                    catalog=None,
+                    namespace=[],
+                    description=None,
+                    tableType="VIEW",
+                    isTemporary=True,
+                )
+            )
+        if pattern:
+            tables = [x for x in tables if fnmatch.fnmatch(x.name, pattern)]
+        return tables
+
+
+class ListColumnsFromInfoSchemaMixin(_BaseInfoSchemaMixin, t.Generic[SESSION, DF]):
+    @normalize(["tableName", "dbName"])
+    def listColumns(self, tableName: str, dbName: t.Optional[str] = None) -> t.List[Column]:
+        """Returns a t.List of columns for the given table/view in the specified database.
+
+        .. versionadded:: 2.0.0
+
+        Parameters
+        ----------
+        tableName : str
+            name of the table to t.List columns.
+
+            .. versionchanged:: 3.4.0
+               Allow ``tableName`` to be qualified with catalog name when ``dbName`` is None.
+
+        dbName : str, t.Optional
+            name of the database to find the table to t.List columns.
+
+        Returns
+        -------
+        t.List
+            A t.List of :class:`Column`.
+
+        Notes
+        -----
+        The order of arguments here is different from that of its JVM counterpart
+        because Python does not support method overloading.
+
+        If no database is specified, the current database and catalog
+        are used. This API includes all temporary views.
+
+        Examples
+        --------
+        >>> _ = spark.sql("DROP TABLE IF EXISTS tbl1")
+        >>> _ = spark.sql("CREATE TABLE tblA (name STRING, age INT) USING parquet")
+        >>> spark.catalog.t.listColumns("tblA")
+        [Column(name='name', description=None, dataType='string', nullable=True, ...
+        >>> _ = spark.sql("DROP TABLE tblA")
+        """
+        if df := self.session.temp_views.get(tableName):
+            return [
+                Column(
+                    name=x,
+                    description=None,
+                    dataType="",
+                    nullable=True,
+                    isPartition=False,
+                    isBucket=False,
+                )
+                for x in df.columns
+            ]
+
+        table = exp.to_table(tableName, dialect=self.session.input_dialect)
+        schema = to_schema(dbName, dialect=self.session.input_dialect) if dbName else None
+        if not table.db:
+            if schema and schema.db:
+                table.set("db", schema.args["db"])
+            else:
+                table.set(
+                    "db",
+                    exp.parse_identifier(
+                        self.currentDatabase(), dialect=self.session.input_dialect
+                    ),
+                )
+        if not table.catalog:
+            if schema and schema.catalog:
+                table.set("catalog", schema.args["catalog"])
+            else:
+                table.set(
+                    "catalog",
+                    exp.parse_identifier(self.currentCatalog(), dialect=self.session.input_dialect),
+                )
+        # if self.QUALIFY_INFO_SCHEMA_WITH_DATABASE:
+        #     if not table.db:
+        #         raise ValueError("dbName must be specified when listing columns from INFORMATION_SCHEMA")
+        #     source_table = f"{table.db}.INFORMATION_SCHEMA.COLUMNS"
+        # else:
+        #     source_table = "INFORMATION_SCHEMA.COLUMNS"
+        source_table = self._get_info_schema_table("columns", database=table.db)
+        select = (
+            exp.select(
+                'column_name AS "column_name"',
+                'data_type AS "data_type"',
+                'is_nullable AS "is_nullable"',
+            )
+            .from_(source_table)
+            .where(exp.column("table_name").eq(table.name))
+        )
+        if table.db:
+            select = select.where(exp.column("table_schema").eq(table.db))
+        if table.catalog:
+            select = select.where(exp.column("table_catalog").eq(table.catalog))
+        results = self.session._fetch_rows(select)
+        return [
+            Column(
+                name=x["column_name"],
+                description=None,
+                dataType=x["data_type"],
+                nullable=x["is_nullable"] == "YES",
+                isPartition=False,
+                isBucket=False,
+            )
+            for x in results
+        ]

sqlframe/base/mixins/readwriter_mixins.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import pathlib
+import typing as t
+
+import pandas as pd
+
+from sqlframe.base.exceptions import UnsupportedOperationError
+from sqlframe.base.readerwriter import (
+    DF,
+    SESSION,
+    _BaseDataFrameReader,
+    _BaseDataFrameWriter,
+    _infer_format,
+)
+from sqlframe.base.util import pandas_to_spark_schema
+
+if t.TYPE_CHECKING:
+    from sqlframe.base._typing import OptionalPrimitiveType, PathOrPaths
+    from sqlframe.base.types import StructType
+
+
+class PandasLoaderMixin(_BaseDataFrameReader, t.Generic[SESSION, DF]):
+    def load(
+        self,
+        path: t.Optional[PathOrPaths] = None,
+        format: t.Optional[str] = None,
+        schema: t.Optional[t.Union[StructType, str]] = None,
+        **options: OptionalPrimitiveType,
+    ) -> DF:
+        """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 _infer_format(path)
+        kwargs = {k: v for k, v in options.items() if v is not None}
+        if format == "json":
+            df = pd.read_json(path, lines=True, **kwargs)  # type: ignore
+        elif format == "parquet":
+            df = pd.read_parquet(path, **kwargs)  # type: ignore
+        elif format == "csv":
+            df = pd.read_csv(path, **kwargs)  # type: ignore
+        else:
+            raise UnsupportedOperationError(f"Unsupported format: {format}")
+        schema = schema or pandas_to_spark_schema(df)
+        self.session._last_loaded_file = path
+        return self._session.createDataFrame(list(df.itertuples(index=False)), schema)
+
+
+class PandasWriterMixin(_BaseDataFrameWriter, t.Generic[SESSION, DF]):
+    def _write(self, path: str, mode: t.Optional[str], format: str, **options):  # type: ignore
+        mode, skip = self._validate_mode(path, mode)
+        if skip:
+            return
+        pandas_df = self._df.toPandas()
+        mode = self._mode_to_pandas_mode(mode)
+        kwargs = {k: v for k, v in options.items() if v is not None}
+        kwargs["index"] = False
+        if format == "csv":
+            kwargs["mode"] = mode
+            if mode == "a" and pathlib.Path(path).exists():
+                kwargs["header"] = False
+            pandas_df.to_csv(path, **kwargs)
+        elif format == "parquet":
+            if mode == "a":
+                raise NotImplementedError("Append mode is not supported for parquet.")
+            pandas_df.to_parquet(path, **kwargs)
+        elif format == "json":
+            # Pandas versions are inconsistent on how to handle True/False index so we just remove it
+            # since in all versions it will not result in an index column in the output.
+            del kwargs["index"]
+            kwargs["mode"] = mode
+            kwargs["orient"] = "records"
+            pandas_df.to_json(path, lines=True, **kwargs)
+        else:
+            raise NotImplementedError(f"Unsupported format: {format}")