sqlframe 1.9.0__py3-none-any.whl → 1.11.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.9.0'
-__version_tuple__ = version_tuple = (1, 9, 0)
+__version__ = version = '1.11.0'
+__version_tuple__ = version_tuple = (1, 11, 0)

sqlframe/base/dataframe.py

@@ -22,6 +22,7 @@ from sqlglot.optimizer.pushdown_projections import pushdown_projections
 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
@@ -29,6 +30,7 @@ from sqlframe.base.util import (
     get_func_from_session,
     get_tables_from_expression_with_join,
     quote_preserving_alias_or_name,
+    sqlglot_to_spark,
     verify_openai_installed,
 )
 
@@ -231,6 +233,10 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
     def __copy__(self):
         return self.copy()
 
+    @property
+    def _typed_columns(self) -> t.List[CatalogColumn]:
+        raise NotImplementedError
+
     @property
     def write(self) -> WRITER:
         return self.session._writer(self)
@@ -293,7 +299,24 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         StructType([StructField('age', LongType(), True),
                     StructField('name', StringType(), True)])
         """
-        raise NotImplementedError
+        from sqlframe.base import types
+
+        try:
+            return types.StructType(
+                [
+                    types.StructField(
+                        c.name,
+                        sqlglot_to_spark(
+                            exp.DataType.build(c.dataType, dialect=self.session.output_dialect)
+                        ),
+                    )
+                    for c in self._typed_columns
+                ]
+            )
+        except NotImplementedError as e:
+            raise NotImplementedError(
+                "This engine does not support schema inference likely since it does not have an active connection."
+            ) from e
 
     def _replace_cte_names_with_hashes(self, expression: exp.Select):
         replacement_mapping = {}
@@ -1537,6 +1560,36 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
                 table.add_row(list(row))
         print(table)
 
+    def printSchema(self, level: t.Optional[int] = None) -> None:
+        def print_schema(
+            column_name: str, column_type: exp.DataType, nullable: bool, current_level: int
+        ):
+            if level and current_level >= level:
+                return
+            if current_level > 0:
+                print(" |   " * current_level, end="")
+            print(
+                f" |-- {column_name}: {column_type.sql(self.session.output_dialect).lower()} (nullable = {str(nullable).lower()})"
+            )
+            if column_type.this in (exp.DataType.Type.STRUCT, exp.DataType.Type.OBJECT):
+                for column_def in column_type.expressions:
+                    print_schema(column_def.name, column_def.args["kind"], True, current_level + 1)
+            if column_type.this == exp.DataType.Type.ARRAY:
+                for data_type in column_type.expressions:
+                    print_schema("element", data_type, True, current_level + 1)
+            if column_type.this == exp.DataType.Type.MAP:
+                print_schema("key", column_type.expressions[0], True, current_level + 1)
+                print_schema("value", column_type.expressions[1], True, current_level + 1)
+
+        print("root")
+        for column in self._typed_columns:
+            print_schema(
+                column.name,
+                exp.DataType.build(column.dataType, dialect=self.session.output_dialect),
+                column.nullable,
+                0,
+            )
+
     def toPandas(self) -> pd.DataFrame:
         sql_kwargs = dict(
             pretty=False, optimize=False, dialect=self.session.output_dialect, as_list=True

sqlframe/base/exceptions.py

@@ -12,3 +12,15 @@ class RowError(SQLFrameException):
 
 class TableSchemaError(SQLFrameException):
     pass
+
+
+class PandasDiffError(SQLFrameException):
+    pass
+
+
+class DataFrameDiffError(SQLFrameException):
+    pass
+
+
+class SchemaDiffError(SQLFrameException):
+    pass

sqlframe/base/function_alternatives.py

@@ -1424,3 +1424,99 @@ def bit_length_from_length(col: ColumnOrName) -> Column:
     col_func = get_func_from_session("col")
 
     return Column(expression.Length(this=col_func(col).expression)) * lit(8)
+
+
+def any_value_always_ignore_nulls(
+    col: ColumnOrName, ignoreNulls: t.Optional[t.Union[bool, Column]] = None
+) -> Column:
+    from sqlframe.base.functions import any_value
+
+    if not ignoreNulls:
+        logger.warning("Nulls are always ignored when using `ANY_VALUE` on this engine")
+    return any_value(col)
+
+
+def any_value_ignore_nulls_not_supported(
+    col: ColumnOrName, ignoreNulls: t.Optional[t.Union[bool, Column]] = None
+) -> Column:
+    from sqlframe.base.functions import any_value
+
+    if ignoreNulls:
+        logger.warning("Ignoring nulls is not supported in this dialect")
+    return any_value(col)
+
+
+def current_user_from_session_user() -> Column:
+    return Column(expression.Anonymous(this="SESSION_USER"))
+
+
+def extract_convert_to_var(field: ColumnOrName, source: ColumnOrName) -> Column:
+    from sqlframe.base.functions import extract
+
+    field = expression.Var(this=Column.ensure_col(field).alias_or_name)  # type: ignore
+    return extract(field, source)  # type: ignore
+
+
+def left_cast_len(str: ColumnOrName, len: ColumnOrName) -> Column:
+    from sqlframe.base.functions import left
+
+    len = Column.ensure_col(len).cast("integer")
+    return left(str, len)
+
+
+def right_cast_len(str: ColumnOrName, len: ColumnOrName) -> Column:
+    from sqlframe.base.functions import right
+
+    len = Column.ensure_col(len).cast("integer")
+    return right(str, len)
+
+
+def position_cast_start(
+    substr: ColumnOrName, str: ColumnOrName, start: t.Optional[ColumnOrName] = None
+) -> Column:
+    from sqlframe.base.functions import position
+
+    start = Column.ensure_col(start).cast("integer") if start else None
+    return position(substr, str, start)
+
+
+def position_as_strpos(
+    substr: ColumnOrName, str: ColumnOrName, start: t.Optional[ColumnOrName] = None
+) -> Column:
+    substr_func = get_func_from_session("substr")
+    lit = get_func_from_session("lit")
+
+    if start:
+        str = substr_func(str, start)
+    column = Column.invoke_anonymous_function(str, "STRPOS", substr)
+    if start:
+        return column + start - lit(1)
+    return column
+
+
+def to_number_using_to_double(col: ColumnOrName, format: ColumnOrName) -> Column:
+    return Column.invoke_anonymous_function(col, "TO_DOUBLE", format)
+
+
+def try_element_at_zero_based(col: ColumnOrName, extraction: ColumnOrName) -> Column:
+    from sqlframe.base.functions import try_element_at
+
+    lit = get_func_from_session("lit")
+    index = Column.ensure_col(extraction)
+    if isinstance(index.expression, expression.Literal) and index.expression.is_number:
+        index = index - lit(1)
+    return try_element_at(col, index)
+
+
+def to_unix_timestamp_include_default_format(
+    timestamp: ColumnOrName,
+    format: t.Optional[ColumnOrName] = None,
+) -> Column:
+    from sqlframe.base.functions import to_unix_timestamp
+
+    lit = get_func_from_session("lit")
+
+    if not format:
+        format = lit("%Y-%m-%d %H:%M:%S")
+
+    return to_unix_timestamp(timestamp, format)