fix(presto): expand data with null item (#17470)

* fix(presto): expand data with null item * Fixed pre-commit check
apache · Nov 21, 2021 · 8da982b · 8da982b
1 parent 19f1c40
commit 8da982b
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 2 deletions.
diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py
@@ -871,8 +871,9 @@ def expand_data(  # pylint: disable=too-many-locals
                 for row in data:
                     values = row.get(name) or []
                     if isinstance(values, str):
-                        row[name] = values = cast(List[Any], destringify(values))
-                    for value, col in zip(values, expanded):
+                        values = cast(Optional[List[Any]], destringify(values))
+                        row[name] = values
+                    for value, col in zip(values or [], expanded):
                         row[col["name"]] = value
 
         data = [

diff --git a/tests/integration_tests/db_engine_specs/presto_tests.py b/tests/integration_tests/db_engine_specs/presto_tests.py
@@ -286,6 +286,60 @@ def test_presto_expand_data_with_complex_row_columns(self):
         self.assertEqual(actual_data, expected_data)
         self.assertEqual(actual_expanded_cols, expected_expanded_cols)
 
+    @mock.patch.dict(
+        "superset.extensions.feature_flag_manager._feature_flags",
+        {"PRESTO_EXPAND_DATA": True},
+        clear=True,
+    )
+    def test_presto_expand_data_with_complex_row_columns_and_null_values(self):
+        cols = [
+            {"name": "row_column", "type": "ROW(NESTED_ROW ROW(NESTED_OBJ VARCHAR))",}
+        ]
+        data = [
+            {"row_column": '[["a"]]'},
+            {"row_column": "[[null]]"},
+            {"row_column": "[null]"},
+            {"row_column": "null"},
+        ]
+        actual_cols, actual_data, actual_expanded_cols = PrestoEngineSpec.expand_data(
+            cols, data
+        )
+        expected_cols = [
+            {"name": "row_column", "type": "ROW(NESTED_ROW ROW(NESTED_OBJ VARCHAR))",},
+            {"name": "row_column.nested_row", "type": "ROW(NESTED_OBJ VARCHAR)"},
+            {"name": "row_column.nested_row.nested_obj", "type": "VARCHAR"},
+        ]
+        expected_data = [
+            {
+                "row_column": [["a"]],
+                "row_column.nested_row": ["a"],
+                "row_column.nested_row.nested_obj": "a",
+            },
+            {
+                "row_column": [[None]],
+                "row_column.nested_row": [None],
+                "row_column.nested_row.nested_obj": None,
+            },
+            {
+                "row_column": [None],
+                "row_column.nested_row": None,
+                "row_column.nested_row.nested_obj": "",
+            },
+            {
+                "row_column": None,
+                "row_column.nested_row": "",
+                "row_column.nested_row.nested_obj": "",
+            },
+        ]
+
+        expected_expanded_cols = [
+            {"name": "row_column.nested_row", "type": "ROW(NESTED_OBJ VARCHAR)"},
+            {"name": "row_column.nested_row.nested_obj", "type": "VARCHAR"},
+        ]
+        self.assertEqual(actual_cols, expected_cols)
+        self.assertEqual(actual_data, expected_data)
+        self.assertEqual(actual_expanded_cols, expected_expanded_cols)
+
     @mock.patch.dict(
         "superset.extensions.feature_flag_manager._feature_flags",
         {"PRESTO_EXPAND_DATA": True},