sqlframe 1.8.0__tar.gz → 1.9.0__tar.gz

{sqlframe-1.8.0 → sqlframe-1.9.0}/PKG-INFO

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

{sqlframe-1.8.0 → sqlframe-1.9.0}/docs/spark.md

@@ -183,6 +183,8 @@ df.show(5)
 * [acos](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.acos.html)
 * [acosh](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.acosh.html)
 * [add_months](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.add_months.html)
+* [aes_encrypt](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.aes_encrypt.html)
+* [aes_decrypt](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.aes_decrypt.html)
 * [aggregate](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.aggregate.html)
 * [approxCountDistinct](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.approxCountDistinct.html)
 * [approx_count_distinct](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.approx_count_distinct.html)
@@ -398,6 +400,7 @@ df.show(5)
 * [timestamp_seconds](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.timestamp_seconds.html)
 * [toDegrees](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.toDegrees.html)
 * [toRadians](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.toRadians.html)
+* [to_binary](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.to_binary.html)
 * [to_csv](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.to_csv.html)
 * [to_date](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.to_date.html)
 * [to_json](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.to_json.html)

{sqlframe-1.8.0 → sqlframe-1.9.0}/setup.py

@@ -20,7 +20,7 @@ setup(
     python_requires=">=3.8",
     install_requires=[
         "prettytable<3.11.0",
-        "sqlglot>=24.0.0,<25.1",
+        "sqlglot>=24.0.0,<25.4",
         "typing_extensions>=4.8,<5",
     ],
     extras_require={
@@ -31,7 +31,7 @@ setup(
         "dev": [
             "duckdb>=0.9,<1.1",
             "mypy>=1.10.0,<1.11",
-            "openai>=1.30,<1.34",
+            "openai>=1.30,<1.36",
             "pandas>=2,<3",
             "pandas-stubs>=2,<3",
             "psycopg>=3.1,<4",
@@ -57,7 +57,7 @@ setup(
             "pandas>=2,<3",
         ],
         "openai": [
-            "openai>=1.30,<1.34",
+            "openai>=1.30,<1.36",
         ],
         "pandas": [
             "pandas>=2,<3",
@@ -69,7 +69,7 @@ setup(
             "redshift_connector>=2.1.1,<2.2.0",
         ],
         "snowflake": [
-            "snowflake-connector-python[secure-local-storage]>=3.10.0,<3.11",
+            "snowflake-connector-python[secure-local-storage]>=3.10.0,<3.12",
         ],
         "spark": [
             "pyspark>=2,<3.6",

{sqlframe-1.8.0 → sqlframe-1.9.0}/sqlframe/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.8.0'
-__version_tuple__ = version_tuple = (1, 8, 0)
+__version__ = version = '1.9.0'
+__version_tuple__ = version_tuple = (1, 9, 0)

{sqlframe-1.8.0 → sqlframe-1.9.0}/sqlframe/base/column.py

@@ -229,7 +229,7 @@ class Column:
         return Column(op)
 
     def unary_op(self, klass: t.Callable, **kwargs) -> Column:
-        return Column(klass(this=self.column_expression, **kwargs))
+        return Column(klass(this=exp.Paren(this=self.column_expression), **kwargs))
 
     @property
     def is_alias(self):

{sqlframe-1.8.0 → sqlframe-1.9.0}/sqlframe/base/functions.py

@@ -1692,14 +1692,9 @@ def make_interval(
     mins: t.Optional[ColumnOrName] = None,
     secs: t.Optional[ColumnOrName] = None,
 ) -> Column:
-    values = [years, months, weeks, days, hours, mins, secs]
-    for value in reversed(values.copy()):
-        if value is not None:
-            break
-        values = values[:-1]
-    else:
+    columns = _ensure_column_of_optionals([years, months, weeks, days, hours, mins, secs])
+    if not columns:
         raise ValueError("At least one value must be provided")
-    columns = [Column.ensure_col(x) if x is not None else lit(None) for x in values]
     return Column.invoke_anonymous_function(columns[0], "MAKE_INTERVAL", *columns[1:])
 
 
@@ -1747,6 +1742,38 @@ def try_to_number(col: ColumnOrName, format: t.Optional[ColumnOrName] = None) ->
     return Column.invoke_anonymous_function(col, "TRY_TO_NUMBER")
 
 
+@meta(unsupported_engines="*")
+def aes_decrypt(
+    input: ColumnOrName,
+    key: ColumnOrName,
+    mode: t.Optional[ColumnOrName] = None,
+    padding: t.Optional[ColumnOrName] = None,
+    aad: t.Optional[ColumnOrName] = None,
+) -> Column:
+    columns = _ensure_column_of_optionals([key, mode, padding, aad])
+    return Column.invoke_anonymous_function(input, "AES_DECRYPT", *columns)
+
+
+@meta(unsupported_engines="*")
+def aes_encrypt(
+    input: ColumnOrName,
+    key: ColumnOrName,
+    mode: t.Optional[ColumnOrName] = None,
+    padding: t.Optional[ColumnOrName] = None,
+    iv: t.Optional[ColumnOrName] = None,
+    aad: t.Optional[ColumnOrName] = None,
+) -> Column:
+    columns = _ensure_column_of_optionals([key, mode, padding, iv, aad])
+    return Column.invoke_anonymous_function(input, "AES_ENCRYPT", *columns)
+
+
+@meta(unsupported_engines="*")
+def to_binary(col: ColumnOrName, format: t.Optional[ColumnOrName] = None) -> Column:
+    if format is not None:
+        return Column.invoke_anonymous_function(col, "TO_BINARY", format)
+    return Column.invoke_anonymous_function(col, "TO_BINARY")
+
+
 @meta()
 def _lambda_quoted(value: str) -> t.Optional[bool]:
     return False if value == "_" else None
@@ -1762,3 +1789,11 @@ def _get_lambda_from_func(lambda_expression: t.Callable):
         this=lambda_expression(*[Column(x) for x in variables]).expression,
         expressions=variables,
     )
+
+
+def _ensure_column_of_optionals(optionals: t.List[t.Optional[ColumnOrName]]) -> t.List[Column]:
+    for value in reversed(optionals.copy()):
+        if value is not None:
+            break
+        optionals = optionals[:-1]
+    return [Column.ensure_col(x) if x is not None else lit(None) for x in optionals]

{sqlframe-1.8.0 → sqlframe-1.9.0}/sqlframe/spark/functions.pyi

@@ -7,6 +7,8 @@ from sqlframe.base.functions import (
     abs as abs,
     acos as acos,
     acosh as acosh,
+    aes_encrypt as aes_encrypt,
+    aes_decrypt as aes_decrypt,
     aggregate as aggregate,
     approxCountDistinct as approxCountDistinct,
     approx_count_distinct as approx_count_distinct,
@@ -218,6 +220,7 @@ from sqlframe.base.functions import (
     timestamp_seconds as timestamp_seconds,
     toDegrees as toDegrees,
     toRadians as toRadians,
+    to_binary as to_binary,
     to_csv as to_csv,
     to_date as to_date,
     to_json as to_json,

{sqlframe-1.8.0 → sqlframe-1.9.0}/sqlframe.egg-info/PKG-INFO

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

{sqlframe-1.8.0 → sqlframe-1.9.0}/sqlframe.egg-info/requires.txt

@@ -1,5 +1,5 @@
 prettytable<3.11.0
-sqlglot<25.1,>=24.0.0
+sqlglot<25.4,>=24.0.0
 typing_extensions<5,>=4.8
 
 [bigquery]
@@ -9,7 +9,7 @@ google-cloud-bigquery[pandas]<4,>=3
 [dev]
 duckdb<1.1,>=0.9
 mypy<1.11,>=1.10.0
-openai<1.34,>=1.30
+openai<1.36,>=1.30
 pandas-stubs<3,>=2
 pandas<3,>=2
 psycopg<4,>=3.1
@@ -39,7 +39,7 @@ duckdb<1.1,>=0.9
 pandas<3,>=2
 
 [openai]
-openai<1.34,>=1.30
+openai<1.36,>=1.30
 
 [pandas]
 pandas<3,>=2
@@ -51,7 +51,7 @@ psycopg2<3,>=2.8
 redshift_connector<2.2.0,>=2.1.1
 
 [snowflake]
-snowflake-connector-python[secure-local-storage]<3.11,>=3.10.0
+snowflake-connector-python[secure-local-storage]<3.12,>=3.10.0
 
 [spark]
 pyspark<3.6,>=2

{sqlframe-1.8.0 → sqlframe-1.9.0}/tests/conftest.py

@@ -7,7 +7,7 @@ import pytest
 
 def pytest_collection_modifyitems(items, *args, **kwargs):
     for item in items:
-        if not list(item.iter_markers()):
+        if not [x for x in item.iter_markers() if x.name != "parametrize"]:
             item.add_marker("fast")
 
 

{sqlframe-1.8.0 → sqlframe-1.9.0}/tests/integration/engines/duck/test_duckdb_session.py

@@ -10,4 +10,4 @@ def test_session_from_config():
     conn.execute("CREATE TABLE test_table (cola INT, colb STRING)")
     session = DuckDBSession.builder.config("sqlframe.conn", conn).getOrCreate()
     columns = session.catalog.get_columns("test_table")
-    assert columns == {'"cola"': exp.DataType.build("INT"), '"colb"': exp.DataType.build("VARCHAR")}
+    assert columns == {'"cola"': exp.DataType.build("INT"), '"colb"': exp.DataType.build("TEXT")}

{sqlframe-1.8.0 → sqlframe-1.9.0}/tests/integration/engines/test_int_functions.py

@@ -194,7 +194,7 @@ def test_typeof(get_session_and_func, get_types, arg, expected):
     if isinstance(session, DuckDBSession):
         if expected == "binary":
             pytest.skip("DuckDB doesn't support binary")
-        expected = expected.replace("string", "varchar")
+        expected = expected.replace("string", "varchar").replace("struct<a:", "struct<a ")
     if isinstance(session, BigQuerySession):
         if expected.startswith("map"):
             pytest.skip("BigQuery doesn't support map types")
@@ -202,8 +202,12 @@ def test_typeof(get_session_and_func, get_types, arg, expected):
             expected = expected.split("<")[0]
         if expected == "binary":
             pytest.skip("BigQuery doesn't support binary")
+        if expected == "timestamp":
+            expected = "datetime"
     result = df.select(typeof("col").alias("test")).first()[0]
-    assert exp.DataType.build(result, dialect=dialect) == exp.DataType.build(expected)
+    assert exp.DataType.build(result, dialect=dialect) == exp.DataType.build(
+        expected, dialect=dialect
+    )
 
 
 def test_alias(get_session_and_func):
@@ -3071,3 +3075,148 @@ def test_try_to_number(get_session_and_func, get_func):
     else:
         expected = Decimal("78.12")
     assert actual == expected
+
+
+def test_to_binary(get_session_and_func, get_func):
+    session, to_binary = get_session_and_func("to_binary")
+    lit = get_func("lit", session)
+    df = session.createDataFrame([("abc",)], ["e"])
+    assert df.select(to_binary(df.e, lit("utf-8")).alias("r")).first()[0] == bytearray(b"abc")
+    df = session.createDataFrame([("414243",)], ["e"])
+    assert df.select(to_binary(df.e).alias("r")).first()[0] == bytearray(b"ABC")
+
+
+def test_aes_decrypt(get_session_and_func, get_func):
+    session, aes_decrypt = get_session_and_func("aes_decrypt")
+    unhex = get_func("unhex", session)
+    unbase64 = get_func("unbase64", session)
+    df = session.createDataFrame(
+        [
+            (
+                "AAAAAAAAAAAAAAAAQiYi+sTLm7KD9UcZ2nlRdYDe/PX4",
+                "abcdefghijklmnop12345678ABCDEFGH",
+                "GCM",
+                "DEFAULT",
+                "This is an AAD mixed into the input",
+            )
+        ],
+        ["input", "key", "mode", "padding", "aad"],
+    )
+    assert df.select(
+        aes_decrypt(unbase64(df.input), df.key, df.mode, df.padding, df.aad).alias("r")
+    ).first()[0] == bytearray(b"Spark")
+    df = session.createDataFrame(
+        [
+            (
+                "AAAAAAAAAAAAAAAAAAAAAPSd4mWyMZ5mhvjiAPQJnfg=",
+                "abcdefghijklmnop12345678ABCDEFGH",
+                "CBC",
+                "DEFAULT",
+            )
+        ],
+        ["input", "key", "mode", "padding"],
+    )
+    assert df.select(
+        aes_decrypt(unbase64(df.input), df.key, df.mode, df.padding).alias("r")
+    ).first()[0] == bytearray(b"Spark")
+    assert df.select(aes_decrypt(unbase64(df.input), df.key, df.mode).alias("r")).first()[
+        0
+    ] == bytearray(b"Spark")
+    df = session.createDataFrame(
+        [
+            (
+                "83F16B2AA704794132802D248E6BFD4E380078182D1544813898AC97E709B28A94",
+                "0000111122223333",
+            )
+        ],
+        ["input", "key"],
+    )
+    assert df.select(aes_decrypt(unhex(df.input), df.key).alias("r")).first()[0] == bytearray(
+        b"Spark"
+    )
+
+
+def test_aes_encrypt(get_session_and_func, get_func):
+    session, aes_encrypt = get_session_and_func("aes_encrypt")
+    to_binary = get_func("to_binary", session)
+    base64 = get_func("base64", session)
+    unbase64 = get_func("unbase64", session)
+    lit = get_func("lit", session)
+    aes_decrypt = get_func("aes_decrypt", session)
+    df = session.createDataFrame(
+        [
+            (
+                "Spark",
+                "abcdefghijklmnop12345678ABCDEFGH",
+                "GCM",
+                "DEFAULT",
+                "000000000000000000000000",
+                "This is an AAD mixed into the input",
+            )
+        ],
+        ["input", "key", "mode", "padding", "iv", "aad"],
+    )
+    assert (
+        df.select(
+            base64(
+                aes_encrypt(
+                    df.input, df.key, df.mode, df.padding, to_binary(df.iv, lit("hex")), df.aad
+                )
+            ).alias("r")
+        ).first()[0]
+        == "AAAAAAAAAAAAAAAAQiYi+sTLm7KD9UcZ2nlRdYDe/PX4"
+    )
+    assert (
+        df.select(
+            base64(
+                aes_encrypt(df.input, df.key, df.mode, df.padding, to_binary(df.iv, lit("hex")))
+            ).alias("r")
+        ).first()[0]
+        == "AAAAAAAAAAAAAAAAQiYi+sRNYDAOTjdSEcYBFsAWPL1f"
+    )
+    df = session.createDataFrame(
+        [
+            (
+                "Spark SQL",
+                "1234567890abcdef",
+                "ECB",
+                "PKCS",
+            )
+        ],
+        ["input", "key", "mode", "padding"],
+    )
+    assert df.select(
+        aes_decrypt(
+            aes_encrypt(df.input, df.key, df.mode, df.padding), df.key, df.mode, df.padding
+        ).alias("r")
+    ).first()[0] == bytearray(b"Spark SQL")
+    df = session.createDataFrame(
+        [
+            (
+                "Spark SQL",
+                "0000111122223333",
+                "ECB",
+            )
+        ],
+        ["input", "key", "mode"],
+    )
+    assert df.select(
+        aes_decrypt(aes_encrypt(df.input, df.key, df.mode), df.key, df.mode).alias("r")
+    ).first()[0] == bytearray(b"Spark SQL")
+    df = session.createDataFrame(
+        [
+            (
+                "Spark SQL",
+                "abcdefghijklmnop",
+            )
+        ],
+        ["input", "key"],
+    )
+    assert (
+        df.select(
+            aes_decrypt(unbase64(base64(aes_encrypt(df.input, df.key))), df.key)
+            .cast("STRING")
+            .alias("r")
+        ).first()[0]
+        == "Spark SQL"
+    )

{sqlframe-1.8.0 → sqlframe-1.9.0}/tests/unit/standalone/test_column.py

@@ -85,7 +85,11 @@ def test_rpow():
 
 
 def test_invert():
-    assert (~F.col("cola")).sql() == "NOT cola"
+    assert (~F.col("cola")).sql() == "NOT (cola)"
+
+
+def test_invert_conjuction():
+    assert (~(F.col("cola") | F.col("colb"))).sql() == "NOT (cola OR colb)"
 
 
 def test_paren():

{sqlframe-1.8.0 → sqlframe-1.9.0}/tests/unit/standalone/test_functions.py

@@ -2348,14 +2348,14 @@ def test_to_json(expression, expected):
     "expression, expected",
     [
         (
-            SF.schema_of_json("cola", dict(timestampFormat="dd/MM/yyyy")),
+            SF.schema_of_json(SF.col("cola"), dict(timestampFormat="dd/MM/yyyy")),
             "SCHEMA_OF_JSON(cola, MAP('timestampFormat', 'dd/MM/yyyy'))",
         ),
         (
             SF.schema_of_json(SF.col("cola"), dict(timestampFormat="dd/MM/yyyy")),
             "SCHEMA_OF_JSON(cola, MAP('timestampFormat', 'dd/MM/yyyy'))",
         ),
-        (SF.schema_of_json("cola"), "SCHEMA_OF_JSON(cola)"),
+        (SF.schema_of_json(SF.col("cola")), "SCHEMA_OF_JSON(cola)"),
     ],
 )
 def test_schema_of_json(expression, expected):
@@ -2366,14 +2366,14 @@ def test_schema_of_json(expression, expected):
     "expression, expected",
     [
         (
-            SF.schema_of_csv("cola", dict(timestampFormat="dd/MM/yyyy")),
+            SF.schema_of_csv(SF.col("cola"), dict(timestampFormat="dd/MM/yyyy")),
             "SCHEMA_OF_CSV(cola, MAP('timestampFormat', 'dd/MM/yyyy'))",
         ),
         (
             SF.schema_of_csv(SF.col("cola"), dict(timestampFormat="dd/MM/yyyy")),
             "SCHEMA_OF_CSV(cola, MAP('timestampFormat', 'dd/MM/yyyy'))",
         ),
-        (SF.schema_of_csv("cola"), "SCHEMA_OF_CSV(cola)"),
+        (SF.schema_of_csv(SF.col("cola")), "SCHEMA_OF_CSV(cola)"),
     ],
 )
 def test_schema_of_csv(expression, expected):
@@ -2903,3 +2903,68 @@ def test_try_to_binary(expression, expected):
 )
 def test_try_to_number(expression, expected):
     assert expression.sql() == expected
+
+
+@pytest.mark.parametrize(
+    "expression, expected",
+    [
+        (
+            SF.aes_decrypt("cola", "colb", "colc", "cold", "cole"),
+            "AES_DECRYPT(cola, colb, colc, cold, cole)",
+        ),
+        (
+            SF.aes_decrypt(
+                SF.col("cola"), SF.col("colb"), SF.col("colc"), SF.col("cold"), SF.col("cole")
+            ),
+            "AES_DECRYPT(cola, colb, colc, cold, cole)",
+        ),
+        (SF.aes_decrypt("cola", SF.col("colb")), "AES_DECRYPT(cola, colb)"),
+        (
+            SF.aes_decrypt(SF.col("cola"), SF.col("colb"), padding="colc"),
+            "AES_DECRYPT(cola, colb, NULL, colc)",
+        ),
+    ],
+)
+def test_aes_decrypt(expression, expected):
+    assert expression.sql() == expected
+
+
+@pytest.mark.parametrize(
+    "expression, expected",
+    [
+        (
+            SF.aes_encrypt("cola", "colb", "colc", "cold", "cole", "colf"),
+            "AES_ENCRYPT(cola, colb, colc, cold, cole, colf)",
+        ),
+        (
+            SF.aes_encrypt(
+                SF.col("cola"),
+                SF.col("colb"),
+                SF.col("colc"),
+                SF.col("cold"),
+                SF.col("cole"),
+                SF.col("colf"),
+            ),
+            "AES_ENCRYPT(cola, colb, colc, cold, cole, colf)",
+        ),
+        (SF.aes_encrypt("cola", SF.col("colb")), "AES_ENCRYPT(cola, colb)"),
+        (
+            SF.aes_encrypt(SF.col("cola"), SF.col("colb"), padding="colc"),
+            "AES_ENCRYPT(cola, colb, NULL, colc)",
+        ),
+    ],
+)
+def test_aes_encrypt(expression, expected):
+    assert expression.sql() == expected
+
+
+@pytest.mark.parametrize(
+    "expression, expected",
+    [
+        (SF.to_binary("cola"), "TO_BINARY(cola)"),
+        (SF.to_binary(SF.col("cola")), "TO_BINARY(cola)"),
+        (SF.to_binary("cola", SF.lit("UTF-8")), "TO_BINARY(cola, 'UTF-8')"),
+    ],
+)
+def test_to_binary(expression, expected):
+    assert expression.sql() == expected