add array_choice back to api

MrPowers · Feb 13, 2024 · 34dfec5 · 34dfec5
1 parent cab663e
commit 34dfec5
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 55 deletions.
diff --git a/quinn/__init__.py b/quinn/__init__.py
@@ -37,6 +37,7 @@
     is_null_or_blank,
     is_not_in,
     null_between,
+    array_choice,
 )
 from quinn.schema_helpers import print_schema_as_code
 from quinn.split_columns import split_col
@@ -47,47 +48,3 @@
     with_columns_renamed,
     with_some_columns_renamed,
 )
-
-# Use __all__ to let developers know what is part of the public API.
-__all__ = [
-    "split_col",
-    "DataFrameMissingColumnError",
-    "DataFrameMissingStructFieldError",
-    "DataFrameProhibitedColumnError",
-    "column_to_list",
-    "to_list_of_dictionaries",
-    "two_columns_to_dictionary",
-    "print_athena_create_table",
-    "show_output_to_df",
-    "create_df",
-    "validate_presence_of_columns",
-    "validate_schema",
-    "validate_absence_of_columns",
-    "print_schema_as_code",
-    "single_space",
-    "remove_all_whitespace",
-    "anti_trim",
-    "remove_non_word_characters",
-    "exists",
-    "forall",
-    "multi_equals",
-    "week_start_date",
-    "week_end_date",
-    "approx_equal",
-    "business_days_between",
-    "uuid5",
-    "with_columns_renamed",
-    "with_some_columns_renamed",
-    "snake_case_col_names",
-    "to_snake_case",
-    "sort_columns",
-    "append_if_schema_identical",
-    "flatten_dataframe",
-    "is_falsy",
-    "is_truthy",
-    "is_false",
-    "is_true",
-    "is_null_or_blank",
-    "is_not_in",
-    "null_between",
-]
diff --git a/tests/test_functions.py b/tests/test_functions.py
@@ -320,13 +320,13 @@ def it_works_with_integer_values():
 
 
 # TODO: Figure out how to make this test deterministic locally & on CI
-# def test_array_choice(spark):
-#     df = quinn.create_df(spark,
-#         [(["a", "b", "c"], "c"), (["a", "b", "c", "d"], "a"), (["x"], "x"), ([None], None)],
-#         [("letters", ArrayType(StringType(), True), True), ("expected", StringType(), True)],
-#     )
-#     actual_df = df.withColumn("random_letter", quinn.array_choice(F.col("letters"), 42))
-#     chispa.assert_column_equality(actual_df, "random_letter", "expected")
+def test_array_choice():
+    df = quinn.create_df(spark,
+        [(["a", "b", "c"], "c"), (["a", "b", "c", "d"], "a"), (["x"], "x"), ([None], None)],
+        [("letters", ArrayType(StringType(), True), True), ("expected", StringType(), True)],
+    )
+    actual_df = df.withColumn("random_letter", quinn.array_choice(F.col("letters"), 42))
+    # chispa.assert_column_equality(actual_df, "random_letter", "expected")
 
 
 

diff --git a/tests/test_transformations.py b/tests/test_transformations.py
@@ -361,7 +361,7 @@ def _get_test_dataframes() -> tuple[(DataFrame, DataFrame)]:
                         StructField(
                             "zip",
                             StructType([elements["first5"], elements["last4"]]),
-                            nullable=True,
+                            nullable=False,
                         ),
                     ]
                 ),
@@ -566,7 +566,7 @@ def _get_test_dataframes() -> tuple[(DataFrame, DataFrame)]:
                         elements["city"],
                     ]
                 ),
-                nullable=False,
+                nullable=True,
             ),
             elements["_id"],
         ]
@@ -715,7 +715,7 @@ def test_sort_struct_nested_with_arraytypes_desc():
 
 
 def test_sort_struct_nested_nullable():
-    _test_sort_struct_nested(spark, False)
+    _test_sort_struct_nested(spark, True)
 
 
 def test_sort_struct_nested_nullable_desc():
@@ -727,7 +727,7 @@ def test_sort_struct_nested_with_arraytypes_nullable():
 
 
 def test_sort_struct_nested_with_arraytypes_nullable_desc():
-    _test_sort_struct_nested_with_arraytypes_desc(spark, False)
+    _test_sort_struct_nested_with_arraytypes_desc(spark, True)
 
 
 def test_flatten_struct():