sqlframe 1.11.0__tar.gz → 1.12.0__tar.gz

{sqlframe-1.11.0 → sqlframe-1.12.0}/PKG-INFO

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

{sqlframe-1.11.0 → sqlframe-1.12.0}/sqlframe/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.11.0'
-__version_tuple__ = version_tuple = (1, 11, 0)
+__version__ = version = '1.12.0'
+__version_tuple__ = version_tuple = (1, 12, 0)

{sqlframe-1.11.0 → sqlframe-1.12.0}/sqlframe/base/column.py

@@ -407,3 +407,44 @@ class Column:
         window_expression = window.expression.copy()
         window_expression.set("this", self.column_expression)
         return Column(window_expression)
+
+    def getItem(self, key: t.Any) -> Column:
+        """
+        An expression that gets an item at position ``ordinal`` out of a list,
+        or gets an item by key out of a dict.
+
+        .. versionadded:: 1.3.0
+
+        .. versionchanged:: 3.4.0
+            Supports Spark Connect.
+
+        Parameters
+        ----------
+        key
+            a literal value, or a :class:`Column` expression.
+            The result will only be true at a location if the item matches in the column.
+
+             .. deprecated:: 3.0.0
+                 :class:`Column` as a parameter is deprecated.
+
+        Returns
+        -------
+        :class:`Column`
+            Column representing the item(s) got at position out of a list or by key out of a dict.
+
+        Examples
+        --------
+        >>> df = spark.createDataFrame([([1, 2], {"key": "value"})], ["l", "d"])
+        >>> df.select(df.l.getItem(0), df.d.getItem("key")).show()
+        +----+------+
+        |l[0]|d[key]|
+        +----+------+
+        |   1| value|
+        +----+------+
+        """
+        element_at = get_func_from_session("element_at")
+        lit = get_func_from_session("lit")
+        key = lit(key) if not isinstance(key, Column) else key
+        if isinstance(key.expression, exp.Literal) and key.expression.is_number:
+            key = key + lit(1)
+        return element_at(self, key)

{sqlframe-1.11.0 → sqlframe-1.12.0}/sqlframe/base/dataframe.py

@@ -361,7 +361,7 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
 
         cols = self._ensure_list_of_columns(cols)
         normalize(self.session, expression or self.expression, cols)
-        return cols
+        return list(flatten([self._expand_star(col) for col in cols]))
 
     def _ensure_and_normalize_col(self, col):
         from sqlframe.base.column import Column
@@ -514,6 +514,27 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         select_expressions.append(expression_select_pair)  # type: ignore
         return select_expressions
 
+    def _expand_star(self, col: Column) -> t.List[Column]:
+        from sqlframe.base.column import Column
+
+        if isinstance(col.column_expression, exp.Star):
+            return self._get_outer_select_columns(self.expression)
+        elif (
+            isinstance(col.column_expression, exp.Column)
+            and isinstance(col.column_expression.this, exp.Star)
+            and col.column_expression.args.get("table")
+        ):
+            for cte in self.expression.ctes:
+                if cte.alias_or_name == col.column_expression.args["table"].this:
+                    return [
+                        Column.ensure_col(exp.column(x.column_alias_or_name, cte.alias_or_name))
+                        for x in self._get_outer_select_columns(cte)
+                    ]
+            raise ValueError(
+                f"Could not find table to expand star: {col.column_expression.args['table']}"
+            )
+        return [col]
+
     @t.overload
     def sql(
         self,
@@ -1555,7 +1576,7 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
             result = self.session._fetch_rows(sql)
         table = PrettyTable()
         if row := seq_get(result, 0):
-            table.field_names = list(row.asDict().keys())
+            table.field_names = row._unique_field_names
             for row in result:
                 table.add_row(list(row))
         print(table)

{sqlframe-1.11.0 → sqlframe-1.12.0}/sqlframe/base/function_alternatives.py

@@ -1135,13 +1135,11 @@ def array_intersect_using_intersection(col1: ColumnOrName, col2: ColumnOrName) -
 def element_at_using_brackets(col: ColumnOrName, value: ColumnOrLiteral) -> Column:
     col_func = get_func_from_session("col")
     lit = get_func_from_session("lit")
-    #  SQLGlot will auto add 1 to whatever we pass in for the brackets even though the value is already 1 based.
-    if not isinstance(value, int):
-        raise ValueError("This dialect requires the value must be an integer")
-    value_lit = lit(value - 1)
-    return Column(
-        expression.Bracket(this=col_func(col).expression, expressions=[value_lit.expression])
-    )
+    # SQLGlot will auto add 1 to whatever we pass in for the brackets even though the value is already 1 based.
+    value = value if isinstance(value, Column) else lit(value)
+    if [x for x in value.expression.find_all(expression.Literal) if x.is_number]:
+        value = value - lit(1)
+    return Column(expression.Bracket(this=col_func(col).expression, expressions=[value.expression]))  # type: ignore
 
 
 def array_remove_using_filter(col: ColumnOrName, value: ColumnOrLiteral) -> Column:

{sqlframe-1.11.0 → sqlframe-1.12.0}/sqlframe/base/functions.py

@@ -1923,7 +1923,9 @@ def call_function(funcName: str, *cols: ColumnOrName) -> Column:
     cols = ensure_list(cols)  # type: ignore
     if len(cols) > 1:
         return Column.invoke_anonymous_function(cols[0], funcName, *cols[1:])
-    return Column.invoke_anonymous_function(cols[0], funcName)
+    elif len(cols) == 1:
+        return Column.invoke_anonymous_function(cols[0], funcName)
+    return Column.invoke_anonymous_function(None, funcName)
 
 
 # @meta(unsupported_engines="*")
@@ -2028,7 +2030,7 @@ def character_length(str: ColumnOrName) -> Column:
     return Column.invoke_anonymous_function(str, "character_length")
 
 
-@meta(unsupported_engines="*")
+@meta()
 def contains(left: ColumnOrName, right: ColumnOrName) -> Column:
     return Column.invoke_anonymous_function(left, "contains", right)
 

{sqlframe-1.11.0 → sqlframe-1.12.0}/sqlframe/base/types.py

@@ -416,3 +416,13 @@ class Row(tuple):
             )
         else:
             return "<Row(%s)>" % ", ".join(repr(field) for field in self)
+
+    # SQLFrame Specific
+    @property
+    def _unique_field_names(self) -> t.List[str]:
+        fields = []
+        for i, field in enumerate(self.__fields__):
+            if field in fields:
+                field = field + "_" + str(i)
+            fields.append(field)
+        return fields

{sqlframe-1.11.0 → sqlframe-1.12.0}/sqlframe/bigquery/functions.py

@@ -52,6 +52,7 @@ from sqlframe.base.function_alternatives import (  # noqa
     make_date_from_date_func as make_date,
     to_date_from_timestamp as to_date,
     last_day_with_cast as last_day,
+    sha1_force_sha1_and_to_hex as sha,
     sha1_force_sha1_and_to_hex as sha1,
     hash_from_farm_fingerprint as hash,
     base64_from_blob as base64,

{sqlframe-1.11.0 → sqlframe-1.12.0}/sqlframe.egg-info/PKG-INFO

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

{sqlframe-1.11.0 → sqlframe-1.12.0}/sqlframe.egg-info/SOURCES.txt

@@ -170,6 +170,7 @@ tests/integration/test_int_dataframe_stats.py
 tests/integration/test_int_grouped_data.py
 tests/integration/test_int_session.py
 tests/integration/engines/__init__.py
+tests/integration/engines/test_engine_column.py
 tests/integration/engines/test_engine_dataframe.py
 tests/integration/engines/test_engine_reader.py
 tests/integration/engines/test_engine_session.py

sqlframe-1.12.0/tests/integration/engines/test_engine_column.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import typing as t
+
+import pytest
+
+from sqlframe.base import types
+from sqlframe.bigquery import BigQuerySession
+from sqlframe.postgres import PostgresSession
+
+if t.TYPE_CHECKING:
+    from sqlframe.base.session import _BaseSession
+
+pytest_plugins = ["tests.integration.fixtures"]
+
+
+def test_column_get_item_array(get_session: t.Callable[[], _BaseSession], get_func):
+    session = get_session()
+    lit = get_func("lit", session)
+    assert session.range(1).select(lit(["a", "b", "c"]).getItem(0).alias("value")).first()[0] == "a"
+
+
+def test_column_get_item_map(get_session: t.Callable[[], _BaseSession], get_func):
+    session = get_session()
+    lit = get_func("lit", session)
+    if not isinstance(session, (PostgresSession, BigQuerySession)):
+        assert session.range(1).select(lit({"key": "value"}).getItem("key")).first()[0] == "value"

{sqlframe-1.11.0 → sqlframe-1.12.0}/tests/integration/engines/test_engine_dataframe.py

@@ -24,36 +24,42 @@ def test_collect(get_engine_df: t.Callable[[str], _BaseDataFrame], get_func):
 
 
 def test_show(
-    get_engine_df: t.Callable[[str], _BaseDataFrame], capsys, caplog, is_snowflake: t.Callable
+    get_engine_df: t.Callable[[str], _BaseDataFrame],
+    get_func,
+    capsys,
+    caplog,
+    is_snowflake: t.Callable,
 ):
     employee = get_engine_df("employee")
+    lit = get_func("lit", employee.session)
+    employee = employee.select("*", lit(1).alias("one"))
     employee.show()
     captured = capsys.readouterr()
     if is_snowflake():
         assert (
             captured.out
-            == """+-------------+--------+-----------+-----+----------+
-| EMPLOYEE_ID | FNAME  |   LNAME   | AGE | STORE_ID |
-+-------------+--------+-----------+-----+----------+
-|      1      |  Jack  |  Shephard |  37 |    1     |
-|      2      |  John  |   Locke   |  65 |    1     |
-|      3      |  Kate  |   Austen  |  37 |    2     |
-|      4      | Claire | Littleton |  27 |    2     |
-|      5      |  Hugo  |   Reyes   |  29 |   100    |
-+-------------+--------+-----------+-----+----------+\n"""
+            == """+-------------+--------+-----------+-----+----------+-----+
+| EMPLOYEE_ID | FNAME  |   LNAME   | AGE | STORE_ID | ONE |
++-------------+--------+-----------+-----+----------+-----+
+|      1      |  Jack  |  Shephard |  37 |    1     |  1  |
+|      2      |  John  |   Locke   |  65 |    1     |  1  |
+|      3      |  Kate  |   Austen  |  37 |    2     |  1  |
+|      4      | Claire | Littleton |  27 |    2     |  1  |
+|      5      |  Hugo  |   Reyes   |  29 |   100    |  1  |
++-------------+--------+-----------+-----+----------+-----+\n"""
         )
     else:
         assert (
             captured.out
-            == """+-------------+--------+-----------+-----+----------+
-| employee_id | fname  |   lname   | age | store_id |
-+-------------+--------+-----------+-----+----------+
-|      1      |  Jack  |  Shephard |  37 |    1     |
-|      2      |  John  |   Locke   |  65 |    1     |
-|      3      |  Kate  |   Austen  |  37 |    2     |
-|      4      | Claire | Littleton |  27 |    2     |
-|      5      |  Hugo  |   Reyes   |  29 |   100    |
-+-------------+--------+-----------+-----+----------+\n"""
+            == """+-------------+--------+-----------+-----+----------+-----+
+| employee_id | fname  |   lname   | age | store_id | one |
++-------------+--------+-----------+-----+----------+-----+
+|      1      |  Jack  |  Shephard |  37 |    1     |  1  |
+|      2      |  John  |   Locke   |  65 |    1     |  1  |
+|      3      |  Kate  |   Austen  |  37 |    2     |  1  |
+|      4      | Claire | Littleton |  27 |    2     |  1  |
+|      5      |  Hugo  |   Reyes   |  29 |   100    |  1  |
++-------------+--------+-----------+-----+----------+-----+\n"""
         )
     assert "Truncate is ignored so full results will be displayed" not in caplog.text
     employee.show(truncate=True)

{sqlframe-1.11.0 → sqlframe-1.12.0}/tests/integration/test_int_dataframe.py

@@ -51,6 +51,16 @@ def test_simple_select_from_table(
     compare_frames(df, dfs)
 
 
+def test_select_star_from_table(
+    pyspark_employee: PySparkDataFrame,
+    get_df: t.Callable[[str], _BaseDataFrame],
+    compare_frames: t.Callable,
+):
+    df = pyspark_employee
+    dfs = get_df("employee").session.read.table("employee")
+    compare_frames(df, dfs)
+
+
 def test_simple_select_df_attribute(
     pyspark_employee: PySparkDataFrame,
     get_df: t.Callable[[str], _BaseDataFrame],

{sqlframe-1.11.0 → sqlframe-1.12.0}/tests/unit/standalone/test_column.py

@@ -220,3 +220,7 @@ def test_over():
         ).sql()
         == "SUM(cola) OVER (PARTITION BY colb ORDER BY colc RANGE BETWEEN 1 PRECEDING AND UNBOUNDED FOLLOWING)"
     )
+
+
+def test_get_item():
+    assert F.col("cola").getItem(1).sql() == "ELEMENT_AT(cola, (1 + 1))"

{sqlframe-1.11.0 → sqlframe-1.12.0}/tests/unit/standalone/test_dataframe.py

@@ -123,3 +123,17 @@ def test_missing_method(standalone_employee: StandaloneDataFrame):
         UnsupportedOperationError, match="Tried to call a column which is unexpected.*"
     ):
         standalone_employee.missing_method("blah")
+
+
+def test_expand_star(standalone_employee: StandaloneDataFrame):
+    assert (
+        standalone_employee.select("*").sql(pretty=False, optimize=False)
+        == "WITH t51718876 AS (SELECT CAST(employee_id AS INT) AS employee_id, CAST(fname AS STRING) AS fname, CAST(lname AS STRING) AS lname, CAST(age AS INT) AS age, CAST(store_id AS INT) AS store_id FROM VALUES (1, 'Jack', 'Shephard', 37, 1), (2, 'John', 'Locke', 65, 1), (3, 'Kate', 'Austen', 37, 2), (4, 'Claire', 'Littleton', 27, 2), (5, 'Hugo', 'Reyes', 29, 100) AS a1(employee_id, fname, lname, age, store_id)) SELECT employee_id, fname, lname, age, store_id FROM t51718876"
+    )
+
+
+def test_expand_star_table_alias(standalone_employee: StandaloneDataFrame):
+    assert (
+        standalone_employee.alias("blah").select("blah.*").sql(pretty=False, optimize=False)
+        == "WITH t51718876 AS (SELECT CAST(employee_id AS INT) AS employee_id, CAST(fname AS STRING) AS fname, CAST(lname AS STRING) AS lname, CAST(age AS INT) AS age, CAST(store_id AS INT) AS store_id FROM VALUES (1, 'Jack', 'Shephard', 37, 1), (2, 'John', 'Locke', 65, 1), (3, 'Kate', 'Austen', 37, 2), (4, 'Claire', 'Littleton', 27, 2), (5, 'Hugo', 'Reyes', 29, 100) AS a1(employee_id, fname, lname, age, store_id)), t37842204 AS (SELECT employee_id, fname, lname, age, store_id FROM t51718876) SELECT t37842204.employee_id, t37842204.fname, t37842204.lname, t37842204.age, t37842204.store_id FROM t37842204"
+    )

{sqlframe-1.11.0 → sqlframe-1.12.0}/tests/unit/standalone/test_functions.py

@@ -3244,6 +3244,7 @@ def test_bucket(expression, expected):
         (SF.call_function("blah", "cola"), "BLAH(cola)"),
         (SF.call_function("blah", SF.col("cola")), "BLAH(cola)"),
         (SF.call_function("blah", "cola", "colb"), "BLAH(cola, colb)"),
+        (SF.call_function("blah"), "BLAH()"),
     ],
 )
 def test_call_function(expression, expected):