sqlframe 3.31.4__tar.gz → 3.32.1__tar.gz

{sqlframe-3.31.4 → sqlframe-3.32.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlframe
-Version: 3.31.4
+Version: 3.32.1
 Summary: Turning PySpark Into a Universal DataFrame API
 Home-page: https://github.com/eakmanrq/sqlframe
 Author: Ryan Eakman

{sqlframe-3.31.4 → sqlframe-3.32.1}/sqlframe/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.31.4'
-__version_tuple__ = version_tuple = (3, 31, 4)
+__version__ = version = '3.32.1'
+__version_tuple__ = version_tuple = (3, 32, 1)

{sqlframe-3.31.4 → sqlframe-3.32.1}/sqlframe/base/catalog.py

@@ -6,9 +6,16 @@ import typing as t
 from collections import defaultdict
 
 from sqlglot import MappingSchema, exp
+from sqlglot.helper import seq_get
 
+from sqlframe.base import types
 from sqlframe.base.exceptions import TableSchemaError
-from sqlframe.base.util import ensure_column_mapping, normalize_string, to_schema
+from sqlframe.base.util import (
+    ensure_column_mapping,
+    normalize_string,
+    spark_to_sqlglot,
+    to_schema,
+)
 
 if t.TYPE_CHECKING:
     from sqlglot.schema import ColumnMapping
@@ -99,6 +106,10 @@ class _BaseCatalog(t.Generic[SESSION, DF, TABLE]):
                     "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
+        if isinstance(column_mapping, dict) and isinstance(
+            seq_get(list(column_mapping.values()), 0), types.DataType
+        ):
+            column_mapping = {k: spark_to_sqlglot(v) for k, v in column_mapping.items()}
         for column_name in column_mapping:
             column = exp.to_column(column_name, dialect=self.session.input_dialect)
             if column.this.quoted:

{sqlframe-3.31.4 → sqlframe-3.32.1}/sqlframe/base/dataframe.py

@@ -260,6 +260,9 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
     def __copy__(self):
         return self.copy()
 
+    def _display_(self) -> str:
+        return self.__repr__()
+
     @property
     def _typed_columns(self) -> t.List[CatalogColumn]:
         raise NotImplementedError

{sqlframe-3.31.4 → sqlframe-3.32.1}/sqlframe/base/functions.py

@@ -494,21 +494,21 @@ def skewness(col: ColumnOrName) -> Column:
         func_name = "SKEW"
 
     if session._is_duckdb or session._is_snowflake:
+        col = Column.ensure_col(col)
         when_func = get_func_from_session("when")
         count_func = get_func_from_session("count")
-        count_star = count_func("*")
+        count_col = count_func(col)
         lit_func = get_func_from_session("lit")
         sqrt_func = get_func_from_session("sqrt")
-        col = Column.ensure_col(col)
         full_calc = (
             Column.invoke_anonymous_function(col, func_name)
-            * (count_star - lit_func(2))
-            / (sqrt_func(count_star * (count_star - lit_func(1))))
+            * (count_col - lit_func(2))
+            / (sqrt_func(count_col * (count_col - lit_func(1))))
         )
         return (
-            when_func(count_star == lit_func(0), lit_func(None))
-            .when(count_star == lit_func(1), lit_func(float("nan")))
-            .when(count_star == lit_func(2), lit_func(None))
+            when_func(count_col == lit_func(0), lit_func(None))
+            .when(count_col == lit_func(1), lit_func(None))
+            .when(count_col == lit_func(2), lit_func(0.0))
             .otherwise(full_calc)
         )
 

{sqlframe-3.31.4 → sqlframe-3.32.1}/sqlframe/base/util.py

@@ -347,6 +347,93 @@ def sqlglot_to_spark(sqlglot_dtype: exp.DataType) -> types.DataType:
     raise NotImplementedError(f"Unsupported data type: {sqlglot_dtype}")
 
 
+def spark_to_sqlglot(spark_dtype: types.DataType) -> exp.DataType:
+    """
+    Convert a Spark data type to a SQLGlot data type.
+
+    This function is the opposite of sqlglot_to_spark.
+
+    Args:
+        spark_dtype: A Spark data type
+
+    Returns:
+        The equivalent SQLGlot data type
+    """
+    from sqlframe.base import types
+
+    # Handle primitive types
+    if isinstance(spark_dtype, types.StringType):
+        return exp.DataType(this=exp.DataType.Type.TEXT)
+    elif isinstance(spark_dtype, types.VarcharType):
+        return exp.DataType(
+            this=exp.DataType.Type.VARCHAR,
+            expressions=[exp.DataTypeParam(this=exp.Literal.number(spark_dtype.length))],
+        )
+    elif isinstance(spark_dtype, types.CharType):
+        return exp.DataType(
+            this=exp.DataType.Type.CHAR,
+            expressions=[exp.DataTypeParam(this=exp.Literal.number(spark_dtype.length))],
+        )
+    elif isinstance(spark_dtype, types.BinaryType):
+        return exp.DataType(this=exp.DataType.Type.BINARY)
+    elif isinstance(spark_dtype, types.BooleanType):
+        return exp.DataType(this=exp.DataType.Type.BOOLEAN)
+    elif isinstance(spark_dtype, types.IntegerType):
+        return exp.DataType(this=exp.DataType.Type.INT)
+    elif isinstance(spark_dtype, types.LongType):
+        return exp.DataType(this=exp.DataType.Type.BIGINT)
+    elif isinstance(spark_dtype, types.ShortType):
+        return exp.DataType(this=exp.DataType.Type.SMALLINT)
+    elif isinstance(spark_dtype, types.ByteType):
+        return exp.DataType(this=exp.DataType.Type.TINYINT)
+    elif isinstance(spark_dtype, types.FloatType):
+        return exp.DataType(this=exp.DataType.Type.FLOAT)
+    elif isinstance(spark_dtype, types.DoubleType):
+        return exp.DataType(this=exp.DataType.Type.DOUBLE)
+    elif isinstance(spark_dtype, types.DecimalType):
+        if spark_dtype.precision is not None and spark_dtype.scale is not None:
+            return exp.DataType(
+                this=exp.DataType.Type.DECIMAL,
+                expressions=[
+                    exp.DataTypeParam(this=exp.Literal.number(spark_dtype.precision)),
+                    exp.DataTypeParam(this=exp.Literal.number(spark_dtype.scale)),
+                ],
+            )
+        return exp.DataType(this=exp.DataType.Type.DECIMAL)
+    elif isinstance(spark_dtype, types.TimestampType):
+        return exp.DataType(this=exp.DataType.Type.TIMESTAMP)
+    elif isinstance(spark_dtype, types.TimestampNTZType):
+        return exp.DataType(this=exp.DataType.Type.TIMESTAMPNTZ)
+    elif isinstance(spark_dtype, types.DateType):
+        return exp.DataType(this=exp.DataType.Type.DATE)
+
+    # Handle complex types
+    elif isinstance(spark_dtype, types.ArrayType):
+        return exp.DataType(
+            this=exp.DataType.Type.ARRAY, expressions=[spark_to_sqlglot(spark_dtype.elementType)]
+        )
+    elif isinstance(spark_dtype, types.MapType):
+        return exp.DataType(
+            this=exp.DataType.Type.MAP,
+            expressions=[
+                spark_to_sqlglot(spark_dtype.keyType),
+                spark_to_sqlglot(spark_dtype.valueType),
+            ],
+        )
+    elif isinstance(spark_dtype, types.StructType):
+        return exp.DataType(
+            this=exp.DataType.Type.STRUCT,
+            expressions=[
+                exp.ColumnDef(
+                    this=exp.to_identifier(field.name), kind=spark_to_sqlglot(field.dataType)
+                )
+                for field in spark_dtype
+            ],
+        )
+
+    raise NotImplementedError(f"Unsupported data type: {spark_dtype}")
+
+
 def normalize_string(
     value: t.Union[str, exp.Expression],
     from_dialect: DialectType = None,

{sqlframe-3.31.4 → sqlframe-3.32.1}/sqlframe.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlframe
-Version: 3.31.4
+Version: 3.32.1
 Summary: Turning PySpark Into a Universal DataFrame API
 Home-page: https://github.com/eakmanrq/sqlframe
 Author: Ryan Eakman

{sqlframe-3.31.4 → sqlframe-3.32.1}/sqlframe.egg-info/SOURCES.txt

@@ -356,6 +356,7 @@ tests/unit/__init__.py
 tests/unit/conftest.py
 tests/unit/test_activate.py
 tests/unit/test_base_reader_options.py
+tests/unit/test_catalog.py
 tests/unit/test_util.py
 tests/unit/bigquery/__init__.py
 tests/unit/bigquery/test_activate.py

{sqlframe-3.31.4 → sqlframe-3.32.1}/tests/integration/engines/test_int_functions.py

@@ -694,6 +694,8 @@ def test_skewness(get_session_and_func):
     assert math.isclose(df.select(skewness(df.c)).first()[0], 0.7071067811865475, rel_tol=1e-5)
     df = session.createDataFrame([{"a": 2.0}, {"a": None}])
     assert df.agg(skewness(df.a)).collect() == [Row(value=None)]
+    df = session.createDataFrame([{"a": 1}, {"a": 2}])
+    assert df.agg(skewness(df.a)).collect() == [Row(value=0.0)]
 
 
 def test_kurtosis(get_session_and_func):

sqlframe-3.32.1/tests/unit/test_catalog.py

@@ -0,0 +1,57 @@
+import pytest
+from sqlglot import exp
+
+from sqlframe.base import types
+from sqlframe.base.util import spark_to_sqlglot
+from sqlframe.standalone.session import StandaloneSession
+
+
+@pytest.fixture(scope="function")
+def standalone_session() -> StandaloneSession:
+    return StandaloneSession()
+
+
+def test_add_table_with_spark_types(standalone_session: StandaloneSession):
+    """Test that add_table properly converts spark types to sqlglot types."""
+    # Create a dictionary with column names as keys and spark types as values
+    column_mapping = {
+        "col_string": types.StringType(),
+        "col_int": types.IntegerType(),
+        "col_long": types.LongType(),
+        "col_float": types.FloatType(),
+        "col_double": types.DoubleType(),
+        "col_boolean": types.BooleanType(),
+        "col_timestamp": types.TimestampType(),
+        "col_date": types.DateType(),
+        "col_decimal": types.DecimalType(10, 2),
+        "col_binary": types.BinaryType(),
+        "col_array": types.ArrayType(types.StringType()),
+        "col_map": types.MapType(types.StringType(), types.IntegerType()),
+        "col_struct": types.StructType(
+            [
+                types.StructField("nested_string", types.StringType()),
+                types.StructField("nested_int", types.IntegerType()),
+            ]
+        ),
+    }
+
+    # Call add_table with the dictionary
+    table_name = "test_table"
+    standalone_session.catalog.add_table(table_name, column_mapping)
+
+    # Get the schema from the catalog
+    table = exp.to_table(table_name, dialect=standalone_session.input_dialect)
+    schema = standalone_session.catalog._schema.find(table)
+
+    # Verify that the schema has been properly updated with the expected sqlglot types
+    assert schema is not None
+
+    # Check each column type
+    for col_name, spark_type in column_mapping.items():
+        expected_sqlglot_type = spark_to_sqlglot(spark_type)
+        actual_sqlglot_type = schema[col_name]
+
+        # Compare the SQL representation of the types
+        assert actual_sqlglot_type.sql() == expected_sqlglot_type.sql(), (
+            f"Column {col_name}: expected {expected_sqlglot_type.sql()}, got {actual_sqlglot_type.sql()}"
+        )

sqlframe-3.32.1/tests/unit/test_util.py

@@ -0,0 +1,216 @@
+import typing as t
+
+import pytest
+from sqlglot import exp, parse_one
+from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+
+from sqlframe.base import types
+from sqlframe.base.util import (
+    quote_preserving_alias_or_name,
+    spark_to_sqlglot,
+    sqlglot_to_spark,
+)
+
+
+@pytest.mark.parametrize(
+    "expression, expected",
+    [
+        ("a", "a"),
+        ("a AS b", "b"),
+        ("`a`", "`a`"),
+        ("`a` AS b", "b"),
+        ("`a` AS `b`", "`b`"),
+        ("`aB`", "`aB`"),
+        ("`aB` AS c", "c"),
+        ("`aB` AS `c`", "`c`"),
+        ("`aB` AS `Cd`", "`Cd`"),
+        # We assume inputs have been normalized so `Cd` is returned as is instead of normalized `cd`
+        ("`aB` AS Cd", "Cd"),
+    ],
+)
+def test_quote_preserving_alias_or_name(expression: t.Union[exp.Column, exp.Alias], expected: str):
+    assert quote_preserving_alias_or_name(parse_one(expression, dialect="bigquery")) == expected  # type: ignore
+
+
+@pytest.mark.parametrize(
+    "dtype, expected",
+    [
+        ("STRING", types.StringType()),
+        ("VARCHAR(100)", types.VarcharType(100)),
+        ("CHAR(100)", types.CharType(100)),
+        ("DECIMAL(10, 2)", types.DecimalType(10, 2)),
+        ("STRING", types.StringType()),
+        ("INTEGER", types.IntegerType()),
+        ("BIGINT", types.LongType()),
+        ("SMALLINT", types.ShortType()),
+        ("TINYINT", types.ByteType()),
+        ("FLOAT", types.FloatType()),
+        ("DOUBLE", types.DoubleType()),
+        ("BOOLEAN", types.BooleanType()),
+        ("TIMESTAMP", types.TimestampType()),
+        ("DATE", types.DateType()),
+        ("DECIMAL", types.DecimalType()),
+        ("BINARY", types.BinaryType()),
+        ("ARRAY<STRING>", types.ArrayType(types.StringType())),
+        ("MAP<STRING, INTEGER>", types.MapType(types.StringType(), types.IntegerType())),
+        (
+            "STRUCT<a STRING, b INTEGER>",
+            types.StructType(
+                [
+                    types.StructField("a", types.StringType()),
+                    types.StructField("b", types.IntegerType()),
+                ]
+            ),
+        ),
+        (
+            "ARRAY<STRUCT<a STRING, b INTEGER>>",
+            types.ArrayType(
+                types.StructType(
+                    [
+                        types.StructField("a", types.StringType()),
+                        types.StructField("b", types.IntegerType()),
+                    ]
+                )
+            ),
+        ),
+    ],
+)
+def test_sqlglot_to_spark(dtype: str, expected: types.DataType):
+    assert sqlglot_to_spark(exp.DataType.build(dtype)) == expected
+
+
+@pytest.mark.parametrize(
+    "spark_dtype, expected_str",
+    [
+        (types.StringType(), "TEXT"),
+        (types.VarcharType(100), "VARCHAR(100)"),
+        (types.CharType(100), "CHAR(100)"),
+        (types.DecimalType(11, 2), "DECIMAL(11, 2)"),
+        (types.IntegerType(), "INT"),
+        (types.LongType(), "BIGINT"),
+        (types.ShortType(), "SMALLINT"),
+        (types.ByteType(), "TINYINT"),
+        (types.FloatType(), "FLOAT"),
+        (types.DoubleType(), "DOUBLE"),
+        (types.BooleanType(), "BOOLEAN"),
+        (types.TimestampType(), "TIMESTAMP"),
+        (types.TimestampNTZType(), "TIMESTAMPNTZ"),
+        (types.DateType(), "DATE"),
+        (types.DecimalType(), "DECIMAL(10, 0)"),
+        (types.BinaryType(), "BINARY"),
+        (types.ArrayType(types.StringType()), "ARRAY(TEXT)"),
+        (types.MapType(types.StringType(), types.IntegerType()), "MAP(TEXT, INT)"),
+        (
+            types.StructType(
+                [
+                    types.StructField("a", types.StringType()),
+                    types.StructField("b", types.IntegerType()),
+                ]
+            ),
+            "STRUCT(a TEXT, b INT)",
+        ),
+        (
+            types.ArrayType(
+                types.StructType(
+                    [
+                        types.StructField("a", types.StringType()),
+                        types.StructField("b", types.IntegerType()),
+                    ]
+                )
+            ),
+            "ARRAY(STRUCT(a TEXT, b INT))",
+        ),
+    ],
+)
+def test_spark_to_sqlglot(spark_dtype: types.DataType, expected_str: str):
+    sqlglot_dtype = spark_to_sqlglot(spark_dtype)
+    assert sqlglot_dtype.sql() == expected_str
+
+
+@pytest.mark.parametrize(
+    "dtype, expected_str",
+    [
+        ("STRING", "TEXT"),
+        ("VARCHAR(100)", "VARCHAR(100)"),
+        ("CHAR(100)", "CHAR(100)"),
+        ("DECIMAL(10, 2)", "DECIMAL(10, 2)"),
+        ("INTEGER", "INT"),
+        ("BIGINT", "BIGINT"),
+        ("SMALLINT", "SMALLINT"),
+        ("TINYINT", "TINYINT"),
+        ("FLOAT", "FLOAT"),
+        ("DOUBLE", "DOUBLE"),
+        ("BOOLEAN", "BOOLEAN"),
+        ("TIMESTAMP", "TIMESTAMP"),
+        ("DATE", "DATE"),
+        ("DECIMAL", "DECIMAL(10, 0)"),
+        ("BINARY", "BINARY"),
+        ("ARRAY<STRING>", "ARRAY(TEXT)"),
+        ("MAP<STRING, INTEGER>", "MAP(TEXT, INT)"),
+        (
+            "STRUCT<a STRING, b INTEGER>",
+            "STRUCT(a TEXT, b INT)",
+        ),
+        (
+            "ARRAY<STRUCT<a STRING, b INTEGER>>",
+            "ARRAY(STRUCT(a TEXT, b INT))",
+        ),
+    ],
+)
+def test_sqlglot_to_spark_to_sqlglot(dtype: str, expected_str: str):
+    """Test round trip conversion from SQLGlot to Spark and back to SQLGlot."""
+    spark_dtype = sqlglot_to_spark(exp.DataType.build(dtype))
+    sqlglot_dtype = spark_to_sqlglot(spark_dtype)
+
+    assert sqlglot_dtype.sql() == expected_str
+
+
+@pytest.mark.parametrize(
+    "spark_dtype, expected_str",
+    [
+        (types.StringType(), "TEXT"),
+        (types.VarcharType(100), "VARCHAR(100)"),
+        (types.CharType(100), "CHAR(100)"),
+        (types.DecimalType(11, 2), "DECIMAL(11, 2)"),
+        (types.IntegerType(), "INT"),
+        (types.LongType(), "BIGINT"),
+        (types.ShortType(), "SMALLINT"),
+        (types.ByteType(), "TINYINT"),
+        (types.FloatType(), "FLOAT"),
+        (types.DoubleType(), "DOUBLE"),
+        (types.BooleanType(), "BOOLEAN"),
+        (types.TimestampType(), "TIMESTAMP"),
+        (types.TimestampNTZType(), "TIMESTAMPNTZ"),
+        (types.DateType(), "DATE"),
+        (types.DecimalType(), "DECIMAL(10, 0)"),
+        (types.BinaryType(), "BINARY"),
+        (types.ArrayType(types.StringType()), "ARRAY(TEXT)"),
+        (types.MapType(types.StringType(), types.IntegerType()), "MAP(TEXT, INT)"),
+        (
+            types.StructType(
+                [
+                    types.StructField("a", types.StringType()),
+                    types.StructField("b", types.IntegerType()),
+                ]
+            ),
+            "STRUCT(a TEXT, b INT)",
+        ),
+        (
+            types.ArrayType(
+                types.StructType(
+                    [
+                        types.StructField("a", types.StringType()),
+                        types.StructField("b", types.IntegerType()),
+                    ]
+                )
+            ),
+            "ARRAY(STRUCT(a TEXT, b INT))",
+        ),
+    ],
+)
+def test_spark_to_sqlglot_to_spark(spark_dtype: types.DataType, expected_str: str):
+    """Test round trip conversion from Spark to SQLGlot and back to Spark."""
+    sqlglot_dtype = spark_to_sqlglot(spark_dtype)
+    final_spark_dtype = sqlglot_to_spark(sqlglot_dtype)
+
+    assert str(final_spark_dtype) == str(spark_dtype)

sqlframe-3.31.4/tests/unit/test_util.py

@@ -1,75 +0,0 @@
-import typing as t
-
-import pytest
-from sqlglot import exp, parse_one
-from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
-
-from sqlframe.base import types
-from sqlframe.base.util import quote_preserving_alias_or_name, sqlglot_to_spark
-
-
-@pytest.mark.parametrize(
-    "expression, expected",
-    [
-        ("a", "a"),
-        ("a AS b", "b"),
-        ("`a`", "`a`"),
-        ("`a` AS b", "b"),
-        ("`a` AS `b`", "`b`"),
-        ("`aB`", "`aB`"),
-        ("`aB` AS c", "c"),
-        ("`aB` AS `c`", "`c`"),
-        ("`aB` AS `Cd`", "`Cd`"),
-        # We assume inputs have been normalized so `Cd` is returned as is instead of normalized `cd`
-        ("`aB` AS Cd", "Cd"),
-    ],
-)
-def test_quote_preserving_alias_or_name(expression: t.Union[exp.Column, exp.Alias], expected: str):
-    assert quote_preserving_alias_or_name(parse_one(expression, dialect="bigquery")) == expected  # type: ignore
-
-
-@pytest.mark.parametrize(
-    "dtype, expected",
-    [
-        ("STRING", types.StringType()),
-        ("VARCHAR(100)", types.VarcharType(100)),
-        ("CHAR(100)", types.CharType(100)),
-        ("DECIMAL(10, 2)", types.DecimalType(10, 2)),
-        ("STRING", types.StringType()),
-        ("INTEGER", types.IntegerType()),
-        ("BIGINT", types.LongType()),
-        ("SMALLINT", types.ShortType()),
-        ("TINYINT", types.ByteType()),
-        ("FLOAT", types.FloatType()),
-        ("DOUBLE", types.DoubleType()),
-        ("BOOLEAN", types.BooleanType()),
-        ("TIMESTAMP", types.TimestampType()),
-        ("DATE", types.DateType()),
-        ("DECIMAL", types.DecimalType()),
-        ("BINARY", types.BinaryType()),
-        ("ARRAY<STRING>", types.ArrayType(types.StringType())),
-        ("MAP<STRING, INTEGER>", types.MapType(types.StringType(), types.IntegerType())),
-        (
-            "STRUCT<a STRING, b INTEGER>",
-            types.StructType(
-                [
-                    types.StructField("a", types.StringType()),
-                    types.StructField("b", types.IntegerType()),
-                ]
-            ),
-        ),
-        (
-            "ARRAY<STRUCT<a STRING, b INTEGER>>",
-            types.ArrayType(
-                types.StructType(
-                    [
-                        types.StructField("a", types.StringType()),
-                        types.StructField("b", types.IntegerType()),
-                    ]
-                )
-            ),
-        ),
-    ],
-)
-def test_sqlglot_to_spark(dtype: str, expected: types.DataType):
-    assert sqlglot_to_spark(exp.DataType.build(dtype)) == expected