feat: implement concat function (#33)

voltrondata · Jun 19, 2024 · c88854e · c88854e
1 parent 51d8f01
commit c88854e
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/src/gateway/tests/test_dataframe_api.py b/src/gateway/tests/test_dataframe_api.py
@@ -13,6 +13,7 @@
     broadcast,
     coalesce,
     col,
+    concat,
     expr,
     greatest,
     isnan,
@@ -1026,3 +1027,16 @@ def test_named_struct(self, spark_session):
             outcome = df.select(named_struct(lit('a'), col('s'), lit('b'), col('i'))).collect()
 
         assertDataFrameEqual(outcome, expected)
+
+    def test_concat(self, users_dataframe):
+        expected = [
+            Row(a='user669344115Joshua Browntrue'),
+            Row(a='user849118289Brooke Jonesfalse'),
+            Row(a='user954079192Collin Frankfalse'),
+        ]
+
+        with utilizes_valid_plans(users_dataframe):
+            outcome = users_dataframe.select(
+                concat('user_id', 'name', 'paid_for_service')).limit(3).collect()
+
+        assertDataFrameEqual(outcome, expected)