sqlframe 1.14.0__py3-none-any.whl → 2.1.0__py3-none-any.whl

sqlframe/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.14.0'
-__version_tuple__ = version_tuple = (1, 14, 0)
+__version__ = version = '2.1.0'
+__version_tuple__ = version_tuple = (2, 1, 0)

sqlframe/base/_typing.py

@@ -24,6 +24,7 @@ OutputExpressionContainer = t.Union[exp.Select, exp.Create, exp.Insert]
 StorageLevel = str
 PathOrPaths = t.Union[str, t.List[str]]
 OptionalPrimitiveType = t.Optional[PrimitiveType]
+DataTypeOrString = t.Union[DataType, str]
 
 
 class UserDefinedFunctionLike(t.Protocol):

sqlframe/base/catalog.py

@@ -3,12 +3,12 @@
 from __future__ import annotations
 
 import typing as t
+from collections import defaultdict
 
 from sqlglot import MappingSchema, exp
 
-from sqlframe.base.decorators import normalize
 from sqlframe.base.exceptions import TableSchemaError
-from sqlframe.base.util import ensure_column_mapping, to_schema
+from sqlframe.base.util import ensure_column_mapping, normalize_string, to_schema
 
 if t.TYPE_CHECKING:
     from sqlglot.schema import ColumnMapping
@@ -33,6 +33,7 @@ class _BaseCatalog(t.Generic[SESSION, DF]):
         """Create a new Catalog that wraps the underlying JVM object."""
         self.session = sparkSession
         self._schema = schema or MappingSchema()
+        self._quoted_columns: t.Dict[exp.Table, t.List[str]] = defaultdict(list)
 
     @property
     def spark(self) -> SESSION:
@@ -52,7 +53,7 @@ class _BaseCatalog(t.Generic[SESSION, DF]):
     def get_columns_from_schema(self, table: exp.Table | str) -> t.Dict[str, exp.DataType]:
         table = self.ensure_table(table)
         return {
-            exp.column(name, quoted=True).sql(
+            exp.column(name, quoted=name in self._quoted_columns[table]).sql(
                 dialect=self.session.input_dialect
             ): exp.DataType.build(dtype, dialect=self.session.input_dialect)
             for name, dtype in self._schema.find(table, raise_on_missing=True).items()  # type: ignore
@@ -64,9 +65,7 @@ class _BaseCatalog(t.Generic[SESSION, DF]):
         if not columns:
             return {}
         return {
-            exp.column(c.name, quoted=True).sql(
-                dialect=self.session.input_dialect
-            ): exp.DataType.build(c.dataType, dialect=self.session.input_dialect)
+            c.name: exp.DataType.build(c.dataType, dialect=self.session.output_dialect)
             for c in columns
         }
 
@@ -79,16 +78,30 @@ class _BaseCatalog(t.Generic[SESSION, DF]):
             return
         if not column_mapping:
             try:
-                column_mapping = self.get_columns(table)
+                column_mapping = {
+                    normalize_string(
+                        k, from_dialect="output", to_dialect="input", is_column=True
+                    ): normalize_string(
+                        v.sql(dialect=self.session.output_dialect),
+                        from_dialect="output",
+                        to_dialect="input",
+                        is_datatype=True,
+                    )
+                    for k, v in self.get_columns(table).items()
+                }
             except NotImplementedError:
                 # TODO: Add doc link
                 raise TableSchemaError(
                     "This session does not have access to a catalog that can lookup column information. See docs for explicitly defining columns or using a session that can automatically determine this."
                 )
         column_mapping = ensure_column_mapping(column_mapping)  # type: ignore
+        for column_name in column_mapping:
+            column = exp.to_column(column_name, dialect=self.session.input_dialect)
+            if column.this.quoted:
+                self._quoted_columns[table].append(column.this.name)
+
         self._schema.add_table(table, column_mapping, dialect=self.session.input_dialect)
 
-    @normalize(["dbName"])
     def getDatabase(self, dbName: str) -> Database:
         """Get the database with the specified name.
         This throws an :class:`AnalysisException` when the database cannot be found.
@@ -115,6 +128,7 @@ class _BaseCatalog(t.Generic[SESSION, DF]):
         >>> spark.catalog.getDatabase("spark_catalog.default")
         Database(name='default', catalog='spark_catalog', description='default database', ...
         """
+        dbName = normalize_string(dbName, from_dialect="input", is_schema=True)
         schema = to_schema(dbName, dialect=self.session.input_dialect)
         database_name = schema.db
         databases = self.listDatabases(pattern=database_name)
@@ -122,12 +136,16 @@ class _BaseCatalog(t.Generic[SESSION, DF]):
             raise ValueError(f"Database '{dbName}' not found")
         if len(databases) > 1:
             if schema.catalog is not None:
-                filtered_databases = [db for db in databases if db.catalog == schema.catalog]
+                filtered_databases = [
+                    db
+                    for db in databases
+                    if normalize_string(db.catalog, from_dialect="output", to_dialect="input")  # type: ignore
+                    == schema.catalog
+                ]
                 if filtered_databases:
                     return filtered_databases[0]
         return databases[0]
 
-    @normalize(["dbName"])
     def databaseExists(self, dbName: str) -> bool:
         """Check if the database with the specified name exists.
 
@@ -168,7 +186,6 @@ class _BaseCatalog(t.Generic[SESSION, DF]):
         except ValueError:
             return False
 
-    @normalize(["tableName"])
     def getTable(self, tableName: str) -> Table:
         """Get the table or view with the specified name. This table can be a temporary view or a
         table/view. This throws an :class:`AnalysisException` when no Table can be found.
@@ -210,13 +227,18 @@ class _BaseCatalog(t.Generic[SESSION, DF]):
             ...
         AnalysisException: ...
         """
+        tableName = normalize_string(tableName, from_dialect="input", is_table=True)
         table = exp.to_table(tableName, dialect=self.session.input_dialect)
         schema = table.copy()
         schema.set("this", None)
         tables = self.listTables(
             schema.sql(dialect=self.session.input_dialect) if schema.db else None
         )
-        matching_tables = [t for t in tables if t.name == table.name]
+        matching_tables = [
+            t
+            for t in tables
+            if normalize_string(t.name, from_dialect="output", to_dialect="input") == table.name
+        ]
         if not matching_tables:
             raise ValueError(f"Table '{tableName}' not found")
         return matching_tables[0]
@@ -315,7 +337,6 @@ class _BaseCatalog(t.Generic[SESSION, DF]):
             raise ValueError(f"Function '{functionName}' not found")
         return matching_functions[0]
 
-    @normalize(["tableName", "dbName"])
     def tableExists(self, tableName: str, dbName: t.Optional[str] = None) -> bool:
         """Check if the table or view with the specified name exists.
         This can either be a temporary view or a table/view.
@@ -389,6 +410,8 @@ class _BaseCatalog(t.Generic[SESSION, DF]):
         >>> spark.catalog.tableExists("view1")
         False
         """
+        tableName = normalize_string(tableName, from_dialect="input", is_table=True)
+        dbName = normalize_string(dbName, from_dialect="input", is_schema=True) if dbName else None
         table = exp.to_table(tableName, dialect=self.session.input_dialect)
         schema_arg = to_schema(dbName, dialect=self.session.input_dialect) if dbName else None
         if not table.db:

sqlframe/base/column.py

@@ -7,11 +7,11 @@ import math
 import typing as t
 
 import sqlglot
+from sqlglot import Dialect
 from sqlglot import expressions as exp
 from sqlglot.helper import flatten, is_iterable
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 
-from sqlframe.base.decorators import normalize
 from sqlframe.base.exceptions import UnsupportedOperationError
 from sqlframe.base.types import DataType
 from sqlframe.base.util import get_func_from_session, quote_preserving_alias_or_name
@@ -211,9 +211,8 @@ class Column:
     def binary_op(
         self, klass: t.Callable, other: ColumnOrLiteral, paren: bool = False, **kwargs
     ) -> Column:
-        op = klass(
-            this=self.column_expression, expression=Column(other).column_expression, **kwargs
-        )
+        other = self._lit(other) if isinstance(other, str) else Column(other)
+        op = klass(this=self.column_expression, expression=other.column_expression, **kwargs)
         if paren:
             return Column(exp.Paren(this=op))
         return Column(op)
@@ -221,9 +220,8 @@ class Column:
     def inverse_binary_op(
         self, klass: t.Callable, other: ColumnOrLiteral, paren: bool = False, **kwargs
     ) -> Column:
-        op = klass(
-            this=Column(other).column_expression, expression=self.column_expression, **kwargs
-        )
+        other = self._lit(other) if isinstance(other, str) else Column(other)
+        op = klass(this=other.column_expression, expression=self.column_expression, **kwargs)
         if paren:
             return Column(exp.Paren(this=op))
         return Column(op)
@@ -340,13 +338,17 @@ class Column:
         new_expression = exp.Not(this=exp.Is(this=self.column_expression, expression=exp.Null()))
         return Column(new_expression)
 
-    def cast(self, dataType: t.Union[str, DataType]) -> Column:
+    def cast(
+        self, dataType: t.Union[str, DataType], dialect: t.Optional[t.Union[str, Dialect]] = None
+    ) -> Column:
         from sqlframe.base.session import _BaseSession
 
         if isinstance(dataType, DataType):
             dataType = dataType.simpleString()
         return Column(
-            exp.cast(self.column_expression, dataType, dialect=_BaseSession().input_dialect)
+            exp.cast(
+                self.column_expression, dataType, dialect=dialect or _BaseSession().input_dialect
+            )
         )
 
     def startswith(self, value: t.Union[str, Column]) -> Column:

sqlframe/base/dataframe.py

@@ -23,12 +23,12 @@ from sqlglot.optimizer.qualify import qualify
 from sqlglot.optimizer.qualify_columns import quote_identifiers
 
 from sqlframe.base.catalog import Column as CatalogColumn
-from sqlframe.base.decorators import normalize
 from sqlframe.base.operations import Operation, operation
 from sqlframe.base.transforms import replace_id_value
 from sqlframe.base.util import (
     get_func_from_session,
     get_tables_from_expression_with_join,
+    normalize_string,
     quote_preserving_alias_or_name,
     sqlglot_to_spark,
     verify_openai_installed,
@@ -41,6 +41,7 @@ else:
 
 if t.TYPE_CHECKING:
     import pandas as pd
+    from pyarrow import Table as ArrowTable
     from sqlglot.dialects.dialect import DialectType
 
     from sqlframe.base._typing import (
@@ -535,39 +536,11 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
             )
         return [col]
 
-    @t.overload
-    def sql(
-        self,
-        dialect: DialectType = ...,
-        optimize: bool = ...,
-        pretty: bool = ...,
-        *,
-        as_list: t.Literal[False],
-        **kwargs: t.Any,
-    ) -> str: ...
-
-    @t.overload
-    def sql(
-        self,
-        dialect: DialectType = ...,
-        optimize: bool = ...,
-        pretty: bool = ...,
-        *,
-        as_list: t.Literal[True],
-        **kwargs: t.Any,
-    ) -> t.List[str]: ...
-
-    def sql(
+    def _get_expressions(
         self,
-        dialect: DialectType = None,
         optimize: bool = True,
-        pretty: bool = True,
         openai_config: t.Optional[t.Union[t.Dict[str, t.Any], OpenAIConfig]] = None,
-        as_list: bool = False,
-        **kwargs,
-    ) -> t.Union[str, t.List[str]]:
-        dialect = Dialect.get_or_raise(dialect or self.session.output_dialect)
-
+    ) -> t.List[exp.Expression]:
         df = self._resolve_pending_hints()
         select_expressions = df._get_select_expressions()
         output_expressions: t.List[t.Union[exp.Select, exp.Cache, exp.Drop]] = []
@@ -583,12 +556,16 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
                 replace_id_value, replacement_mapping
             ).assert_is(exp.Select)
             if optimize:
-                quote_identifiers(select_expression, dialect=dialect)
                 select_expression = t.cast(
-                    exp.Select, self.session._optimize(select_expression, dialect=dialect)
+                    exp.Select,
+                    self.session._optimize(select_expression),
                 )
             elif openai_config:
-                qualify(select_expression, dialect=dialect, schema=self.session.catalog._schema)
+                qualify(
+                    select_expression,
+                    dialect=self.session.input_dialect,
+                    schema=self.session.catalog._schema,
+                )
                 pushdown_projections(select_expression, schema=self.session.catalog._schema)
 
             select_expression = df._replace_cte_names_with_hashes(select_expression)
@@ -606,7 +583,7 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
                     cache_table_name,
                     {
                         quote_preserving_alias_or_name(expression): expression.type.sql(
-                            dialect=dialect
+                            dialect=self.session.input_dialect
                         )
                         if expression.type
                         else "UNKNOWN"
@@ -642,10 +619,43 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
                 raise ValueError(f"Invalid expression type: {expression_type}")
 
             output_expressions.append(expression)
+        return output_expressions  # type: ignore
+
+    @t.overload
+    def sql(
+        self,
+        dialect: DialectType = ...,
+        optimize: bool = ...,
+        pretty: bool = ...,
+        *,
+        as_list: t.Literal[False],
+        **kwargs: t.Any,
+    ) -> str: ...
+
+    @t.overload
+    def sql(
+        self,
+        dialect: DialectType = ...,
+        optimize: bool = ...,
+        pretty: bool = ...,
+        *,
+        as_list: t.Literal[True],
+        **kwargs: t.Any,
+    ) -> t.List[str]: ...
 
+    def sql(
+        self,
+        dialect: DialectType = None,
+        optimize: bool = True,
+        pretty: bool = True,
+        openai_config: t.Optional[t.Union[t.Dict[str, t.Any], OpenAIConfig]] = None,
+        as_list: bool = False,
+        **kwargs,
+    ) -> t.Union[str, t.List[str]]:
+        dialect = Dialect.get_or_raise(dialect) if dialect else self.session.output_dialect
         results = []
-        for expression in output_expressions:
-            sql = expression.sql(dialect=dialect, pretty=pretty, **kwargs)
+        for expression in self._get_expressions(optimize=optimize, openai_config=openai_config):
+            sql = self.session._to_sql(expression, dialect=dialect, pretty=pretty, **kwargs)
             if openai_config:
                 assert isinstance(openai_config, OpenAIConfig)
                 verify_openai_installed()
@@ -828,7 +838,7 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         from sqlframe.base.functions import coalesce
 
         if on is None:
-            logger.warning("Got no value for on. This appears change the join to a cross join.")
+            logger.warning("Got no value for on. This appears to change the join to a cross join.")
             how = "cross"
         other_df = other_df._convert_leaf_to_cte()
         # We will determine actual "join on" expression later so we don't provide it at first
@@ -871,13 +881,12 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
                             )
                             join_column_pairs.append((left_column, right_column))
                             num_matching_ctes += 1
-                    if num_matching_ctes > 1:
-                        raise ValueError(
-                            f"Column {join_column.alias_or_name} is ambiguous. Please specify the table name."
-                        )
-                    elif num_matching_ctes == 0:
+                            # We only want to match one table to the column and that should be matched left -> right
+                            # so we break after the first match
+                            break
+                    if num_matching_ctes == 0:
                         raise ValueError(
-                            f"Column {join_column.alias_or_name} does not exist in any of the tables."
+                            f"Column `{join_column.alias_or_name}` does not exist in any of the tables."
                         )
                 join_clause = functools.reduce(
                     lambda x, y: x & y,
@@ -1154,7 +1163,7 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         if len(sql_queries) > 1:
             raise ValueError("Cannot explain a DataFrame with multiple queries")
         sql_query = "EXPLAIN " + sql_queries[0]
-        self.session._execute(sql_query, quote_identifiers=False)
+        self.session._execute(sql_query)
 
     @operation(Operation.FROM)
     def fillna(
@@ -1373,23 +1382,20 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
             self._ensure_and_normalize_col(k).alias_or_name: self._ensure_and_normalize_col(v)
             for k, v in colsMap[0].items()
         }
-        existing_col_names = [
-            x.alias_or_name for x in self._get_outer_select_columns(self.expression)
-        ]
-        updated_expression = self.expression.copy()
-        select_columns = []
+        existing_cols = self._get_outer_select_columns(self.expression)
+        existing_col_names = [x.alias_or_name for x in existing_cols]
+        select_columns = existing_cols
         for column_name, col_value in col_map.items():
             existing_col_index = (
                 existing_col_names.index(column_name) if column_name in existing_col_names else None
             )
-            if existing_col_index:
-                updated_expression.expressions[existing_col_index] = col_value.alias(
+            if existing_col_index is not None:
+                select_columns[existing_col_index] = col_value.alias(  # type: ignore
                     column_name
                 ).expression
             else:
                 select_columns.append(col_value.alias(column_name))
-        df = self.copy(expression=updated_expression)
-        return df.select.__wrapped__(df, *select_columns, append=True)  # type: ignore
+        return self.select.__wrapped__(self, *select_columns)  # type: ignore
 
     @operation(Operation.SELECT)
     def drop(self, *cols: t.Union[str, Column]) -> Self:
@@ -1539,10 +1545,10 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         return self._group_data(self, grouping_columns, self.last_op)
 
     def collect(self) -> t.List[Row]:
-        result = []
-        for sql in self.sql(pretty=False, optimize=False, as_list=True):
-            result = self.session._fetch_rows(sql)
-        return result
+        return self._collect()
+
+    def _collect(self, **kwargs) -> t.List[Row]:
+        return self.session._collect(self._get_expressions(optimize=False), **kwargs)
 
     @t.overload
     def head(self) -> t.Optional[Row]: ...
@@ -1569,11 +1575,7 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
             logger.warning("Truncate is ignored so full results will be displayed")
         # Make sure that the limit we add doesn't affect the results
         df = self._convert_leaf_to_cte()
-        sql = df.limit(n).sql(
-            pretty=False, optimize=False, dialect=self.session.output_dialect, as_list=True
-        )
-        for sql in ensure_list(sql):
-            result = self.session._fetch_rows(sql)
+        result = df.limit(n).collect()
         table = PrettyTable()
         if row := seq_get(result, 0):
             table.field_names = row._unique_field_names
@@ -1612,18 +1614,10 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
             )
 
     def toPandas(self) -> pd.DataFrame:
-        sql_kwargs = dict(
-            pretty=False, optimize=False, dialect=self.session.output_dialect, as_list=True
-        )
-        sqls = [None] + self.sql(**sql_kwargs)  # type: ignore
-        for sql in self.sql(**sql_kwargs)[:-1]:  # type: ignore
-            if sql:
-                self.session._execute(sql)
-        assert sqls[-1] is not None
-        return self.session._fetchdf(sqls[-1])
-
-    @normalize("name")
+        return self.session._fetchdf(self._get_expressions(optimize=False))
+
     def createOrReplaceTempView(self, name: str) -> None:
+        name = normalize_string(name, from_dialect="input")
         self.session.temp_views[name] = self.copy()._convert_leaf_to_cte()
 
     def count(self) -> int:
@@ -1632,11 +1626,7 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
 
         df = self._convert_leaf_to_cte()
         df = self.copy(expression=df.expression.select("count(*)", append=False))
-        for sql in df.sql(
-            dialect=self.session.output_dialect, pretty=False, optimize=False, as_list=True
-        ):
-            result = self.session._fetch_rows(sql)
-        return result[0][0]
+        return df.collect()[0][0]
 
     def createGlobalTempView(self, name: str) -> None:
         raise NotImplementedError("Global temp views are not yet supported")
@@ -1816,3 +1806,6 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         col_func = get_func_from_session("col")
 
         return self.select(covar_samp(col_func(col1), col_func(col2))).collect()[0][0]
+
+    def toArrow(self) -> ArrowTable:
+        raise NotImplementedError("Arrow conversion is not supported by this engine")

sqlframe/base/decorators.py

@@ -1,50 +1,12 @@
 from __future__ import annotations
 
-import functools
 import typing as t
 
-from sqlglot import parse_one
 from sqlglot.helper import ensure_list
-from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
-
-if t.TYPE_CHECKING:
-    from sqlframe.base.catalog import _BaseCatalog
 
 CALLING_CLASS = t.TypeVar("CALLING_CLASS")
 
 
-def normalize(normalize_kwargs: t.Union[str, t.List[str]]) -> t.Callable[[t.Callable], t.Callable]:
-    """
-    Decorator used to normalize identifiers in the kwargs of a method.
-    """
-
-    def decorator(func: t.Callable) -> t.Callable:
-        @functools.wraps(func)
-        def wrapper(self: CALLING_CLASS, *args, **kwargs) -> CALLING_CLASS:
-            from sqlframe.base.session import _BaseSession
-
-            input_dialect = _BaseSession().input_dialect
-            kwargs.update(dict(zip(func.__code__.co_varnames[1:], args)))
-            for kwarg in ensure_list(normalize_kwargs):
-                if kwarg in kwargs:
-                    value = kwargs.get(kwarg)
-                    if value:
-                        expression = (
-                            parse_one(value, dialect=input_dialect)
-                            if isinstance(value, str)
-                            else value
-                        )
-                        kwargs[kwarg] = normalize_identifiers(expression, input_dialect).sql(
-                            dialect=input_dialect
-                        )
-            return func(self, **kwargs)
-
-        wrapper.__wrapped__ = func  # type: ignore
-        return wrapper
-
-    return decorator
-
-
 def func_metadata(unsupported_engines: t.Optional[t.Union[str, t.List[str]]] = None) -> t.Callable:
     def _metadata(func: t.Callable) -> t.Callable:
         func.unsupported_engines = ensure_list(unsupported_engines) if unsupported_engines else []  # type: ignore