feat: add array_append support all engines (#106)

eakmanrq · Jun 29, 2024 · eeffd41 · eeffd41
1 parent b8dea22
commit eeffd41
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 3 deletions.
diff --git a/sqlframe/base/function_alternatives.py b/sqlframe/base/function_alternatives.py
@@ -1529,6 +1529,19 @@ def to_unix_timestamp_include_default_format(
     return to_unix_timestamp(timestamp, format)
 
 
+def array_append_list_append(col: ColumnOrName, value: ColumnOrLiteral) -> Column:
+    lit = get_func_from_session("lit")
+    value = value if isinstance(value, Column) else lit(value)
+    return Column.invoke_anonymous_function(col, "LIST_APPEND", value)
+
+
+def array_append_using_array_cat(col: ColumnOrName, value: ColumnOrLiteral) -> Column:
+    lit = get_func_from_session("lit")
+    array = get_func_from_session("array")
+    value = value if isinstance(value, Column) else lit(value)
+    return Column.invoke_anonymous_function(col, "ARRAY_CONCAT", array(value))
+
+
 def day_with_try_to_timestamp(col: ColumnOrName) -> Column:
     from sqlframe.base.functions import day
 

diff --git a/sqlframe/base/functions.py b/sqlframe/base/functions.py
@@ -1300,7 +1300,7 @@ def array_agg(col: ColumnOrName) -> Column:
     return Column.invoke_expression_over_column(col, expression.ArrayAgg)
 
 
-@meta(unsupported_engines="*")
+@meta()
 def array_append(col: ColumnOrName, value: ColumnOrLiteral) -> Column:
     value = value if isinstance(value, Column) else lit(value)
     return Column.invoke_anonymous_function(col, "ARRAY_APPEND", value)

diff --git a/sqlframe/bigquery/functions.py b/sqlframe/bigquery/functions.py
@@ -74,6 +74,7 @@
     position_as_strpos as position,
     try_to_timestamp_safe as try_to_timestamp,
     _is_string_using_typeof_string as _is_string,
+    array_append_using_array_cat as array_append,
 )
 
 

diff --git a/sqlframe/duckdb/functions.py b/sqlframe/duckdb/functions.py
@@ -49,4 +49,5 @@
     day_with_try_to_timestamp as day,
     try_to_timestamp_strptime as try_to_timestamp,
     _is_string_using_typeof_varchar as _is_string,
+    array_append_list_append as array_append,
 )
diff --git a/tests/integration/engines/test_int_functions.py b/tests/integration/engines/test_int_functions.py
@@ -2067,9 +2067,10 @@ def test_array_agg(get_session_and_func):
     ]
 
 
-def test_array_append(get_session_and_func):
+def test_array_append(get_session_and_func, get_func):
     session, array_append = get_session_and_func("array_append")
-    df = session.createDataFrame([Row(c1=["b", "a", "c"], c2="c")])
+    lit = get_func("lit", session)
+    df = session.range(1).select(lit(["b", "a", "c"]).alias("c1"), lit("c").alias("c2"))
     assert df.select(array_append(df.c1, df.c2)).collect() == [
         Row(value=["b", "a", "c", "c"]),
     ]