improve array function

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "splink_data_normalisation"
-version = "0.1.1"
+version = "0.1.2"
 description = ""
 authors = ["Robin Linacre <robin.linacre@digital.justice.gov.uk>"]
 license = "MIT"

splink_data_normalisation/arrays.py

@@ -15,7 +15,7 @@ def fix_zero_length_arrays(df:DataFrame):
 
     stmt = """
     case 
-    when size({c}) > 0 then {c}
+    when size(filter({c}, x -> x is not null and trim(x) != '')) > 0 then filter({c}, x -> x is not null and trim(x) != '')
     else null
     end
     """

tests/test_array.py

@@ -10,19 +10,26 @@ def test_fix_1(spark):
     names_list = [
         {"id": 1, "my_arr1": ["a", "b", "c"], "other_arr": [ ],"my_str": "a"},
         {"id": 2, "my_arr1": [             ], "other_arr": [1],"my_str": "a"},
-
-            ]
+        {"id": 3, "my_arr1": [None, '', 'c'], "other_arr": [1],"my_str": "a"},
+        {"id": 4, "my_arr1": [None, ''     ], "other_arr": [1],"my_str": "a"},
+        {"id": 5, "my_arr1": [''           ], "other_arr": [1],"my_str": "a"},
+        {"id": 6, "my_arr1": [None, None   ], "other_arr": [1],"my_str": "a"},
+        ]
 
     df = spark.createDataFrame(Row(**x) for x in names_list)
     df = df.select(list(names_list[0].keys()))
-
     df = fix_zero_length_arrays(df)
 
     df_result = df.toPandas()
 
     df_expected = [
         {"id": 1, "my_arr1": ["a", "b", "c"], "other_arr": None,"my_str": "a"},
         {"id": 2, "my_arr1": None,            "other_arr": [1] ,"my_str": "a"},
+        {"id": 3, "my_arr1": ['c'],           "other_arr": [1] ,"my_str": "a"},
+        {"id": 4, "my_arr1": None,            "other_arr": [1] ,"my_str": "a"},
+        {"id": 5, "my_arr1": None,            "other_arr": [1] ,"my_str": "a"},
+        {"id": 6, "my_arr1": None,            "other_arr": [1] ,"my_str": "a"},
+
     ]
 
     df_expected = pd.DataFrame(df_expected)