Fix!: Handle maps during test generation (#2364)

* Fix: Handle maps during test generation

* Bump sqlglot to 23.3.0

* Refactor generate_test tests

* Fix typo in comment

* Remove unused variables

* Add more tests

* Update sqlmesh/core/test/definition.py

Co-authored-by: Jo <46752250+georgesittas@users.noreply.github.com>

---------

Co-authored-by: Jo <46752250+georgesittas@users.noreply.github.com>

Author: vchan

setup.py

@@ -46,7 +46,7 @@
         "requests",
         "rich[jupyter]",
         "ruamel.yaml",
-        "sqlglot[rs]~=23.2.0",
+        "sqlglot[rs]~=23.3.0",
     ],
     extras_require={
         "bigquery": [

sqlmesh/core/test/definition.py

@@ -592,4 +592,11 @@ def _normalize_dataframe(value: t.Any) -> t.Any:
     """Normalize data in a pandas dataframe so ruamel and sqlglot can deal with it."""
     if isinstance(value, np.ndarray):
         return [_normalize_dataframe(v) for v in value]
+    if isinstance(value, dict):
+        if "key" in value and "value" in value:
+            # Maps returned by DuckDB have the following structure: {'key': ['key1', 'key2', 'key3'], 'value': [10, 20, 30]}
+            # so we convert to {'key1': 10, 'key2': 20, 'key3': 30}
+            # TODO: handle more dialects here
+            return {k: _normalize_dataframe(v) for k, v in zip(value["key"], value["value"])}
+        return {k: _normalize_dataframe(v) for k, v in value.items()}
     return value

tests/core/test_test.py

@@ -921,7 +921,15 @@ def test_successes(sushi_context: Context) -> None:
     assert "test_customer_revenue_by_day" in successful_tests
 
 
-def test_test_generation_with_array(tmp_path: Path) -> None:
+def test_test_generation_with_data_structures(tmp_path: Path) -> None:
+    def create_test(query: str):
+        context.create_test(
+            "sqlmesh_example.foo",
+            input_queries={"sqlmesh_example.bar": query},
+            overwrite=True,
+        )
+        return load_yaml(context.path / c.TESTS / "test_foo.yaml")
+
     init_example_project(tmp_path, dialect="duckdb")
 
     config = Config(
@@ -930,57 +938,92 @@ def test_test_generation_with_array(tmp_path: Path) -> None:
     )
     foo_sql_file = tmp_path / "models" / "foo.sql"
     foo_sql_file.write_text(
-        "MODEL (name sqlmesh_example.foo); SELECT array_col FROM sqlmesh_example.bar;"
+        "MODEL (name sqlmesh_example.foo); SELECT col FROM sqlmesh_example.bar;"
     )
     bar_sql_file = tmp_path / "models" / "bar.sql"
-    bar_sql_file.write_text(
-        "MODEL (name sqlmesh_example.bar); SELECT array_col FROM external_table;"
-    )
+    bar_sql_file.write_text("MODEL (name sqlmesh_example.bar); SELECT col FROM external_table;")
 
     context = Context(paths=tmp_path, config=config)
 
-    input_queries = {"sqlmesh_example.bar": "SELECT ['value1', 'value2'] AS array_col"}
-
-    context.create_test(
-        "sqlmesh_example.foo",
-        input_queries=input_queries,
-        overwrite=True,
-        variables={"start": "2020-01-01", "end": "2024-01-01"},
-    )
-
-    test = load_yaml(context.path / c.TESTS / "test_foo.yaml")
-
-    assert len(test) == 1
-    assert "test_foo" in test
-    assert "vars" in test["test_foo"]
-    assert test["test_foo"]["inputs"] == {
-        "sqlmesh_example.bar": [{"array_col": ["value1", "value2"]}]
-    }
-    assert test["test_foo"]["outputs"] == {"query": [{"array_col": ["value1", "value2"]}]}
+    # Array of strings
+    test = create_test("SELECT ['value1', 'value2'] AS col")
+    expected_value: t.Any = [{"col": ["value1", "value2"]}]
+    assert test["test_foo"]["inputs"] == {"sqlmesh_example.bar": expected_value}
+    assert test["test_foo"]["outputs"] == {"query": expected_value}
 
     # Array of arrays
-    input_queries = {
-        "sqlmesh_example.bar": "SELECT [['value1'], ['value2', 'value3']] AS array_col"
-    }
+    test = create_test("SELECT [['value1'], ['value2', 'value3']] AS col")
+    expected_value = [{"col": [["value1"], ["value2", "value3"]]}]
+    assert test["test_foo"]["inputs"] == {"sqlmesh_example.bar": expected_value}
+    assert test["test_foo"]["outputs"] == {"query": expected_value}
+
+    # Array of maps
+    test = create_test("SELECT [MAP {'key': 'value1'}, MAP {'key': 'value2'}] AS col")
+    expected_value = [{"col": [{"key": "value1"}, {"key": "value2"}]}]
+    assert test["test_foo"]["inputs"] == {"sqlmesh_example.bar": expected_value}
+    assert test["test_foo"]["outputs"] == {"query": expected_value}
+
+    # Array of structs
+    test = create_test("SELECT [{'key': 'value1'}, {'key': 'value2'}] AS col")
+    expected_value = [{"col": [{"key": "value1"}, {"key": "value2"}]}]
+    assert test["test_foo"]["inputs"] == {"sqlmesh_example.bar": expected_value}
+    assert test["test_foo"]["outputs"] == {"query": expected_value}
+
+    # Map of strings
+    test = create_test("SELECT MAP {'key1': 'value1', 'key2': 'value2'} AS col")
+    expected_value = [{"col": {"key1": "value1", "key2": "value2"}}]
+    assert test["test_foo"]["inputs"] == {"sqlmesh_example.bar": expected_value}
+    assert test["test_foo"]["outputs"] == {"query": expected_value}
+
+    # Struct of strings
+    test = create_test("SELECT {'key1': 'value1', 'key2': 'value2'} AS col")
+    expected_value = [{"col": {"key1": "value1", "key2": "value2"}}]
+    assert test["test_foo"]["inputs"] == {"sqlmesh_example.bar": expected_value}
+    assert test["test_foo"]["outputs"] == {"query": expected_value}
+
+    # Map of arrays
+    test = create_test("SELECT MAP {'key1': ['value1'], 'key2': ['value2']} AS col")
+    expected_value = [{"col": {"key1": ["value1"], "key2": ["value2"]}}]
+    assert test["test_foo"]["inputs"] == {"sqlmesh_example.bar": expected_value}
+    assert test["test_foo"]["outputs"] == {"query": expected_value}
+
+    # Struct of arrays
+    test = create_test("SELECT {'key1': ['value1'], 'key2': ['value2']} AS col")
+    expected_value = [{"col": {"key1": ["value1"], "key2": ["value2"]}}]
+    assert test["test_foo"]["inputs"] == {"sqlmesh_example.bar": expected_value}
+    assert test["test_foo"]["outputs"] == {"query": expected_value}
+
+    # Map of maps
+    test = create_test(
+        "SELECT MAP {'key1': MAP {'subkey1': 'value1'}, 'key2': MAP {'subkey2': 'value2'}} AS col"
+    )
+    expected_value = [{"col": {"key1": {"subkey1": "value1"}, "key2": {"subkey2": "value2"}}}]
+    assert test["test_foo"]["inputs"] == {"sqlmesh_example.bar": expected_value}
+    assert test["test_foo"]["outputs"] == {"query": expected_value}
 
-    context.create_test(
-        "sqlmesh_example.foo",
-        input_queries=input_queries,
-        overwrite=True,
-        variables={"start": "2020-01-01", "end": "2024-01-01"},
+    # Map of structs
+    test = create_test(
+        "SELECT MAP {'key1': {'subkey': 'value1'}, 'key2': {'subkey': 'value2'}} AS col"
     )
+    expected_value = [{"col": {"key1": {"subkey": "value1"}, "key2": {"subkey": "value2"}}}]
+    assert test["test_foo"]["inputs"] == {"sqlmesh_example.bar": expected_value}
+    assert test["test_foo"]["outputs"] == {"query": expected_value}
 
-    test = load_yaml(context.path / c.TESTS / "test_foo.yaml")
+    # Struct of structs
+    test = create_test(
+        "SELECT {'key1': {'subkey1': 'value1'}, 'key2': {'subkey2': 'value2'}} AS col"
+    )
+    expected_value = [{"col": {"key1": {"subkey1": "value1"}, "key2": {"subkey2": "value2"}}}]
+    assert test["test_foo"]["inputs"] == {"sqlmesh_example.bar": expected_value}
+    assert test["test_foo"]["outputs"] == {"query": expected_value}
 
-    assert len(test) == 1
-    assert "test_foo" in test
-    assert "vars" in test["test_foo"]
-    assert test["test_foo"]["inputs"] == {
-        "sqlmesh_example.bar": [{"array_col": [["value1"], ["value2", "value3"]]}]
-    }
-    assert test["test_foo"]["outputs"] == {
-        "query": [{"array_col": [["value1"], ["value2", "value3"]]}]
-    }
+    # Struct of maps
+    test = create_test(
+        "SELECT {'key1': MAP {'subkey1': 'value1'}, 'key2': MAP {'subkey2': 'value2'}} AS col"
+    )
+    expected_value = [{"col": {"key1": {"subkey1": "value1"}, "key2": {"subkey2": "value2"}}}]
+    assert test["test_foo"]["inputs"] == {"sqlmesh_example.bar": expected_value}
+    assert test["test_foo"]["outputs"] == {"query": expected_value}
 
 
 def test_test_generation_with_timestamp(tmp_path: Path) -> None: