feat(risingwave): properly support structs

Author: cpcloud

ci/schema/risingwave.sql

@@ -188,3 +188,17 @@ DROP TABLE IF EXISTS "topk";
 
 CREATE TABLE "topk" ("x" BIGINT);
 INSERT INTO "topk" VALUES (1), (1), (NULL);
+
+DROP TABLE IF EXISTS "struct";
+CREATE TABLE "struct" (
+    "abc" STRUCT<"a" DOUBLE, "b" STRING, "c" BIGINT>
+);
+
+INSERT INTO "struct" VALUES
+    (ROW(1.0, 'banana',  2)),
+    (ROW(2.0, 'apple', 3)),
+    (ROW(3.0, 'orange',  4)),
+    (ROW(NULL, 'banana', 2)),
+    (ROW(2.0, NULL, 3)),
+    (NULL),
+    (ROW(3.0, 'orange', NULL));

ibis/backends/risingwave/__init__.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import contextlib
+import itertools
 from operator import itemgetter
 from typing import TYPE_CHECKING, Any
 from urllib.parse import unquote_plus
@@ -16,39 +17,66 @@
 import ibis.backends.sql.compilers as sc
 import ibis.common.exceptions as com
 import ibis.common.exceptions as exc
+import ibis.expr.datatypes as dt
 import ibis.expr.operations as ops
 import ibis.expr.schema as sch
 import ibis.expr.types as ir
 from ibis import util
 from ibis.backends import CanCreateDatabase, CanListCatalog, NoExampleLoader
 from ibis.backends.sql import SQLBackend
-from ibis.backends.sql.compilers.base import TRUE, C, ColGen
+from ibis.backends.sql.compilers.base import TRUE, C
 from ibis.util import experimental
 
 if TYPE_CHECKING:
+    from collections.abc import Iterable, Mapping
     from urllib.parse import ParseResult
 
     import pandas as pd
     import polars as pl
     import pyarrow as pa
 
 
+def dict_to_struct(struct):
+    return dt.Struct(
+        {
+            field: dtype if isinstance(dtype, dt.DataType) else dict_to_struct(dtype)
+            for field, dtype in struct.items()
+        }
+    )
+
+
+def string_to_struct(
+    dtypes: Iterable[tuple[str, dt.DataType, int]], top: str
+) -> Mapping[str, dt.DataType]:
+    result = {}
+    for field, dtype, _ in dtypes:
+        field_top, *components, bottom = field.split(".")
+        assert top == field_top, f"{top} != {field_top}"
+        child = result.setdefault(top, {})
+        for component in components:
+            child = child.setdefault(component, {})
+        child[bottom] = dtype
+    return result[top]
+
+
 def data_and_encode_format(data_format, encode_format, encode_properties):
-    res = ""
+    res = []
     if data_format is not None:
-        res = res + " FORMAT " + data_format.upper()
+        res.append("FORMAT")
+        res.append(data_format.upper())
     if encode_format is not None:
-        res = res + " ENCODE " + encode_format.upper()
+        res.append("ENCODE")
+        res.append(encode_format.upper())
         if encode_properties is not None:
-            res = res + " " + format_properties(encode_properties)
-    return res
+            res.append(format_properties(encode_properties))
+    return " ".join(res)
 
 
 def format_properties(props):
     tokens = []
     for k, v in props.items():
         tokens.append(f"{k}='{v}'")
-    return "( {} ) ".format(", ".join(tokens))
+    return "({}) ".format(", ".join(tokens))
 
 
 class Backend(SQLBackend, CanListCatalog, CanCreateDatabase, NoExampleLoader):
@@ -264,12 +292,6 @@ def get_schema(
         catalog: str | None = None,
         database: str | None = None,
     ):
-        a = ColGen(table="a")
-        c = ColGen(table="c")
-        n = ColGen(table="n")
-
-        format_type = self.compiler.f["pg_catalog.format_type"]
-
         # If no database is specified, assume the current database
         db = database or self.current_database
 
@@ -280,46 +302,44 @@ def get_schema(
         if database is None and (temp_table_db := self._session_temp_db) is not None:
             dbs.append(sge.convert(temp_table_db))
 
-        type_info = (
-            sg.select(
-                a.attname.as_("column_name"),
-                format_type(a.atttypid, a.atttypmod).as_("data_type"),
-                sg.not_(a.attnotnull).as_("nullable"),
-            )
-            .from_(sg.table("pg_attribute", db="pg_catalog").as_("a"))
-            .join(
-                sg.table("pg_class", db="pg_catalog").as_("c"),
-                on=c.oid.eq(a.attrelid),
-                join_type="INNER",
-            )
-            .join(
-                sg.table("pg_namespace", db="pg_catalog").as_("n"),
-                on=n.oid.eq(c.relnamespace),
-                join_type="INNER",
-            )
-            .where(
-                a.attnum > 0,
-                sg.not_(a.attisdropped),
-                n.nspname.isin(*dbs),
-                c.relname.eq(sge.convert(name)),
-            )
-            .order_by(a.attnum)
-        )
+        ident = sg.table(name, catalog=catalog, db=database, quoted=True)
+        try:
+            with self._safe_raw_sql(sge.Describe(this=ident)) as cur:
+                raw_rows = cur.fetchall()
+        except psycopg2.InternalError as exc:
+            raise com.TableNotFound(name) from exc
 
         type_mapper = self.compiler.type_mapper
 
-        with self._safe_raw_sql(type_info) as cur:
-            rows = cur.fetchall()
+        rows = []
+        field_number = 0
+        for raw_name, dtype, hidden, *_ in raw_rows:
+            if hidden == "false":
+                field_number += (not dtype) or "." not in raw_name
+                rows.append(
+                    (
+                        raw_name,
+                        None
+                        if not dtype
+                        else type_mapper.from_string(dtype, nullable=True),
+                        field_number,
+                    )
+                )
 
-        if not rows:
-            raise com.TableNotFound(name)
+        schema = {}
 
-        return sch.Schema(
-            {
-                col: type_mapper.from_string(typestr, nullable=nullable)
-                for col, typestr, nullable in rows
-            }
-        )
+        for _, values in itertools.groupby(rows, key=lambda x: x[-1]):
+            vals = list(values)
+            assert vals, "vals is empty"
+
+            if len(vals) == 1:
+                name, dtype, _ = vals.pop()
+                schema[name] = dtype
+            else:
+                name, _, _ = vals[0]
+                schema[name] = dict_to_struct(string_to_struct(vals[1:], name))
+
+        return sch.Schema(schema)
 
     def _get_schema_using_query(self, query: str) -> sch.Schema:
         name = util.gen_name(f"{self.name}_metadata")
@@ -583,7 +603,9 @@ def create_table(
             create_stmt = sge.Create(
                 kind="TABLE",
                 this=target,
-                properties=sge.Properties.from_dict(connector_properties),
+                properties=sge.Properties(
+                    expressions=sge.Properties.from_dict(connector_properties)
+                ),
             )
             create_stmt = create_stmt.sql(self.dialect) + data_and_encode_format(
                 data_format, encode_format, encode_properties
@@ -742,7 +764,6 @@ def create_source(
         data_format: str,
         encode_format: str,
         encode_properties: dict | None = None,
-        includes: dict[str, str] | None = None,
     ) -> ir.Table:
         """Creating a source.
 
@@ -763,32 +784,23 @@ def create_source(
             The encode format for the new source, e.g., "JSON". data_format and encode_format must be specified at the same time.
         encode_properties
             The properties of encode format, providing information like schema registry url. Refer https://docs.risingwave.com/docs/current/sql-create-source/ for more details.
-        includes
-            A dict of `INCLUDE` clauses of the form `{field: alias, ...}`.
-            Set value(s) to `None` if no alias is needed. Refer to https://docs.risingwave.com/docs/current/sql-create-source/ for more details.
 
         Returns
         -------
         Table
             Table expression
         """
-        quoted = self.compiler.quoted
-        table = sg.table(name, db=database, quoted=quoted)
+        table = sg.table(name, db=database, quoted=self.compiler.quoted)
         target = sge.Schema(this=table, expressions=schema.to_sqlglot(self.dialect))
 
-        properties = sge.Properties.from_dict(connector_properties)
-        properties.expressions.extend(
-            sge.IncludeProperty(
-                this=sg.to_identifier(include_type),
-                alias=sg.to_identifier(column_name, quoted=quoted)
-                if column_name
-                else None,
-            )
-            for include_type, column_name in (includes or {}).items()
+        create_stmt = sge.Create(
+            kind="SOURCE",
+            this=target,
+            properties=sge.Properties(
+                expressions=sge.Properties.from_dict(connector_properties)
+            ),
         )
 
-        create_stmt = sge.Create(kind="SOURCE", this=target, properties=properties)
-
         create_stmt = create_stmt.sql(self.dialect) + data_and_encode_format(
             data_format, encode_format, encode_properties
         )

ibis/backends/risingwave/tests/conftest.py

@@ -30,7 +30,7 @@ class TestConf(ServiceBackendTest):
     # for numeric and decimal
 
     returned_timestamp_unit = "s"
-    supports_structs = False
+    supports_structs = True
     rounding_method = "half_to_even"
     service_name = "risingwave"
     deps = ("psycopg2",)

ibis/backends/risingwave/tests/test_client.py

@@ -10,6 +10,7 @@
 import ibis
 import ibis.expr.datatypes as dt
 import ibis.expr.types as ir
+from ibis.backends.risingwave import string_to_struct
 from ibis.util import gen_name
 
 pytest.importorskip("psycopg2")
@@ -132,3 +133,17 @@ def test_insert_with_cte(con):
     assert Y.execute().empty
     con.drop_table("Y")
     con.drop_table("X")
+
+
+@pytest.mark.parametrize(
+    ("expected", "shredded", "top"),
+    [
+        (
+            {"b": {"c": "int", "d": "string"}},
+            [("a.b.c", "int", 1), ("a.b.d", "string", 1)],
+            "a",
+        ),
+    ],
+)
+def test_string_to_struct(expected, shredded, top):
+    assert string_to_struct(shredded, top) == expected

ibis/backends/sql/compilers/risingwave.py

@@ -76,6 +76,13 @@ def to_sqlglot(
             for col, typ in schema.items()
             if typ.is_map()
         )
+        # similar to maps are structs (spelled ROW in RisingWave), which come
+        # back as stringified tuples without first converting
+        conversions.update(
+            (col, table_expr[col].cast(dt.JSON(binary=True)))
+            for col, typ in schema.items()
+            if typ.is_struct()
+        )
 
         if conversions:
             table_expr = table_expr.mutate(**conversions)
@@ -84,6 +91,9 @@ def to_sqlglot(
     def visit_DateNow(self, op):
         return self.cast(sge.CurrentTimestamp(), dt.date)
 
+    def visit_Array(self, op, *, exprs):
+        return self.cast(self.f.array(*exprs), op.dtype)
+
     def visit_Cast(self, op, *, arg, to):
         if to.is_json():
             return self.f.to_jsonb(arg)
@@ -154,6 +164,9 @@ def visit_TimestampTruncate(self, op, *, arg, unit):
     def _make_interval(self, arg, unit):
         return arg * sge.Interval(this=sge.convert(1), unit=self.v[unit.name])
 
+    def visit_StructColumn(self, op, *, names, values):
+        return self.cast(self.f.row(*values), op.dtype)
+
     def visit_NonNullLiteral(self, op, *, value, dtype):
         if dtype.is_binary():
             return self.cast("".join(map(r"\x{:0>2x}".format, value)), dt.binary)
@@ -165,6 +178,8 @@ def visit_NonNullLiteral(self, op, *, value, dtype):
             return self.f.map_from_key_values(
                 self.f.array(*value.keys()), self.f.array(*value.values())
             )
+        elif dtype.is_struct():
+            return self.cast(self.f.row(*value.values()), dtype)
         return None
 
     def visit_MapGet(self, op, *, arg, key, default):
@@ -193,5 +208,8 @@ def visit_MapContains(self, op, *, arg, key):
             self.cast(arg, op.arg.dtype), self.cast(key, op.key.dtype)
         )
 
+    def visit_StructField(self, op, *, arg, field):
+        return sge.Dot(this=sge.Paren(this=arg), expression=sge.to_identifier(field))
+
 
 compiler = RisingWaveCompiler()

ibis/backends/tests/test_array.py

@@ -283,11 +283,6 @@ def test_array_index(con, idx):
     reason="backend does not support nullable nested types",
     raises=AssertionError,
 )
-@pytest.mark.notimpl(
-    ["risingwave"],
-    raises=AssertionError,
-    reason="Do not nest ARRAY types; ARRAY(basetype) handles multi-dimensional arrays of basetype",
-)
 @pytest.mark.never(
     ["bigquery"], reason="doesn't support arrays of arrays", raises=AssertionError
 )

ibis/backends/tests/test_generic.py

@@ -1741,7 +1741,6 @@ def hash_256(col):
                 pytest.mark.notimpl(["druid"], raises=PyDruidProgrammingError),
                 pytest.mark.notimpl(["oracle"], raises=OracleDatabaseError),
                 pytest.mark.notimpl(["postgres"], raises=PsycoPgSyntaxError),
-                pytest.mark.notimpl(["risingwave"], raises=PsycoPg2InternalError),
                 pytest.mark.notimpl(["snowflake"], raises=AssertionError),
                 pytest.mark.never(
                     ["datafusion", "exasol", "impala", "mssql", "mysql", "sqlite"],

ibis/backends/tests/test_map.py

@@ -427,7 +427,6 @@ def test_literal_map_getitem_broadcast(backend, alltypes, df):
             marks=[
                 pytest.mark.notyet("clickhouse", reason="nested types can't be null"),
                 mark_notyet_postgres,
-                pytest.mark.notimpl(["risingwave"], raises=PsycoPg2InternalError),
             ],
             id="struct",
         ),

ibis/backends/tests/test_param.py

@@ -71,7 +71,7 @@ def test_scalar_param_array(con):
     assert result == len(value)
 
 
-@pytest.mark.notimpl(["impala", "postgres", "risingwave", "druid", "oracle", "exasol"])
+@pytest.mark.notimpl(["impala", "postgres", "druid", "oracle", "exasol"])
 @pytest.mark.never(
     ["mysql", "sqlite", "mssql"],
     reason="mysql and sqlite will never implement struct types",