feat(postgres): implement maps in terms of JSONB instead of HSTORE

Author: cpcloud

ci/schema/postgres.sql

@@ -1,4 +1,3 @@
-CREATE EXTENSION IF NOT EXISTS hstore;
 CREATE EXTENSION IF NOT EXISTS postgis;
 CREATE EXTENSION IF NOT EXISTS plpython3u;
 CREATE EXTENSION IF NOT EXISTS fuzzystrmatch;
@@ -298,10 +297,10 @@ INSERT INTO win VALUES
     ('a', 4, 1);
 
 DROP TABLE IF EXISTS map CASCADE;
-CREATE TABLE map (idx BIGINT, kv HSTORE);
+CREATE TABLE map (idx BIGINT, kv JSONB);
 INSERT INTO map VALUES
-    (1, 'a=>1,b=>2,c=>3'),
-    (2, 'd=>4,e=>5,c=>6');
+    (1, '{"a": 1, "b": 2, "c": 3}'),
+    (2, '{"d": 4, "e": 5, "f": 6}');
 
 DROP TABLE IF EXISTS topk;
 

ibis/backends/postgres/__init__.py

@@ -4,7 +4,6 @@
 
 import contextlib
 import inspect
-import warnings
 from operator import itemgetter
 from typing import TYPE_CHECKING, Any
 from urllib.parse import unquote_plus
@@ -208,7 +207,6 @@ def do_connect(
         database: str | None = None,
         schema: str | None = None,
         autocommit: bool = True,
-        enable_map_support: bool = True,
         **kwargs: Any,
     ) -> None:
         """Create an Ibis client connected to PostgreSQL database.
@@ -229,9 +227,6 @@ def do_connect(
             PostgreSQL schema to use. If `None`, use the default `search_path`.
         autocommit
             Whether or not to autocommit
-        enable_map_support
-            Whether or not to enable map support. If `True`, the HSTORE
-            extension will be loaded to support maps of string -> string.
         kwargs
             Additional keyword arguments to pass to the backend client connection.
 
@@ -265,9 +260,6 @@ def do_connect(
         """
         import pandas as pd
         import psycopg
-        import psycopg.types.json
-
-        psycopg.types.json.set_json_loads(loads=lambda x: x)
 
         self.con = psycopg.connect(
             host=host,
@@ -282,7 +274,7 @@ def do_connect(
 
         self.con.adapters.register_dumper(type(pd.NaT), NatDumper)
 
-        self._post_connect(enable_map_support)
+        self._post_connect()
 
     @util.experimental
     @classmethod
@@ -300,26 +292,8 @@ def from_connection(cls, con: psycopg.Connection, /) -> Backend:
         new_backend._post_connect()
         return new_backend
 
-    def _post_connect(self, enable_map_support: bool = True) -> None:
-        import psycopg.types
-        import psycopg.types.hstore
-
-        con = self.con
-
-        try:
-            # try to load hstore
-            if enable_map_support:
-                with con.cursor() as cursor, con.transaction():
-                    cursor.execute("CREATE EXTENSION IF NOT EXISTS hstore")
-                psycopg.types.hstore.register_hstore(
-                    psycopg.types.TypeInfo.fetch(self.con, "hstore"), self.con
-                )
-        except psycopg.Error as e:
-            warnings.warn(f"Failed to load hstore extension: {e}")
-        except TypeError:
-            pass
-
-        with con.cursor() as cursor, con.transaction():
+    def _post_connect(self) -> None:
+        with (con := self.con).cursor() as cursor, con.transaction():
             cursor.execute("SET TIMEZONE = UTC")
 
     @property

ibis/backends/postgres/converter.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import json
+
 from ibis.formats.pandas import PandasData
 
 
@@ -18,3 +20,19 @@ def convert_GeoSpatial(cls, s, dtype, pandas_type):
     @classmethod
     def convert_Binary(cls, s, dtype, pandas_type):
         return s.map(bytes, na_action="ignore")
+
+    @classmethod
+    def convert_Map_element(cls, dtype):
+        convert_key = cls.get_element_converter(dtype.key_type)
+        convert_value = cls.get_element_converter(dtype.value_type)
+
+        def convert(raw_row):
+            if raw_row is None:
+                return raw_row
+
+            row = json.loads(raw_row)
+            return dict(
+                zip(map(convert_key, row.keys()), map(convert_value, row.values()))
+            )
+
+        return convert

ibis/backends/postgres/tests/conftest.py

@@ -25,6 +25,8 @@
     from collections.abc import Iterable
     from pathlib import Path
 
+    import ibis.expr.types as ir
+
 PG_USER = os.environ.get(
     "IBIS_TEST_POSTGRES_USER", os.environ.get("PGUSER", "postgres")
 )
@@ -46,12 +48,17 @@ class TestConf(ServiceBackendTest):
 
     returned_timestamp_unit = "s"
     supports_structs = False
+    supports_map = True
     rounding_method = "half_to_even"
     service_name = "postgres"
     deps = ("psycopg",)
 
     driver_supports_multiple_statements = True
 
+    @property
+    def map(self) -> ir.Table | None:
+        return self.connection.table("map").cast({"kv": "map<string, int>"})
+
     @property
     def test_files(self) -> Iterable[Path]:
         return self.data_dir.joinpath("csv").glob("*.csv")

ibis/backends/sql/compilers/postgres.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import inspect
+import json
 import string
 import textwrap
 from functools import partial, reduce
@@ -114,9 +115,7 @@ class PostgresCompiler(SQLGlotCompiler):
         ops.GeoWithin: "st_within",
         ops.GeoX: "st_x",
         ops.GeoY: "st_y",
-        ops.MapContains: "exist",
-        ops.MapKeys: "akeys",
-        ops.MapValues: "avals",
+        ops.MapContains: "jsonb_contains",
         ops.RegexSearch: "regexp_like",
         ops.TimeFromHMS: "make_time",
         ops.RandomUUID: "gen_random_uuid",
@@ -130,8 +129,14 @@ def to_sqlglot(
         params: Mapping[ir.Expr, Any] | None = None,
     ):
         table_expr = expr.as_table()
-        geocols = table_expr.schema().geospatial
-        conversions = {name: table_expr[name].as_ewkb() for name in geocols}
+        schema = table_expr.schema()
+
+        conversions = {name: table_expr[name].as_ewkb() for name in schema.geospatial}
+        conversions.update(
+            (col, table_expr[col].cast(dt.string))
+            for col, typ in schema.items()
+            if typ.is_map() or typ.is_json()
+        )
 
         if conversions:
             table_expr = table_expr.mutate(**conversions)
@@ -160,25 +165,44 @@ def _compile_python_udf(self, udf_node: ops.ScalarUDF):
 
         type_mapper = self.type_mapper
         argnames = udf_node.argnames
-        return """\
+        args = ", ".join(argnames)
+        name = type(udf_node).__name__
+        argsig = ", ".join(argnames)
+        raw_args = [
+            f"json.loads({argname})" if arg.dtype.is_map() else argname
+            for argname, arg in zip(argnames, udf_node.args)
+        ]
+        args = ", ".join(raw_args)
+        call = f"{name}({args})"
+        defn = """\
 CREATE OR REPLACE FUNCTION {ident}({signature})
 RETURNS {return_type}
 LANGUAGE {language}
 AS $$
+{json_import}
+def {name}({argsig}):
 {source}
-return {name}({args})
+return {call}
 $$""".format(
-            name=type(udf_node).__name__,
             ident=self.__sql_name__(udf_node),
             signature=", ".join(
                 f"{argname} {type_mapper.to_string(arg.dtype)}"
                 for argname, arg in zip(argnames, udf_node.args)
             ),
             return_type=type_mapper.to_string(udf_node.dtype),
             language=config.get("language", "plpython3u"),
-            source=source,
-            args=", ".join(argnames),
+            json_import=(
+                "import json"
+                if udf_node.dtype.is_map()
+                or any(arg.dtype.is_map() for arg in udf_node.args)
+                else ""
+            ),
+            name=name,
+            argsig=argsig,
+            source=textwrap.indent(source, " " * 4),
+            call=call if not udf_node.dtype.is_map() else f"json.dumps({call})",
         )
+        return defn
 
     def visit_Mode(self, op, *, arg, where):
         expr = self.f.mode()
@@ -513,22 +537,61 @@ def visit_ToJSONArray(self, op, *, arg):
     def visit_Map(self, op, *, keys, values):
         # map(["a", "b"], NULL) results in {"a": NULL, "b": NULL} in regular postgres,
         # so we need to modify it to return NULL instead
-        regular = self.f.map(keys, values)
-        return self.if_(values.is_(NULL), NULL, regular)
+        k, v = map(sg.to_identifier, "kv")
+        regular = (
+            sg.select(self.f.jsonb_object_agg(k, v))
+            .from_(
+                sg.select(
+                    self.f.unnest(keys).as_(k), self.f.unnest(values).as_(v)
+                ).subquery()
+            )
+            .subquery()
+        )
+        return self.if_(keys.is_(NULL).or_(values.is_(NULL)), NULL, regular)
 
     def visit_MapLength(self, op, *, arg):
-        return self.f.cardinality(self.f.akeys(arg))
+        return (
+            sg.select(self.f.count(sge.Star()))
+            .from_(self.f.jsonb_object_keys(arg))
+            .subquery()
+        )
 
     def visit_MapGet(self, op, *, arg, key, default):
-        return self.if_(
-            self.f.exist(arg, key),
-            self.f.jsonb_extract_path_text(self.f.to_jsonb(arg), key),
-            default,
-        )
+        if op.dtype.is_null():
+            return NULL
+        else:
+            return self.cast(
+                self.if_(
+                    self.f.jsonb_contains(arg, key),
+                    self.f.jsonb_extract_path_text(arg, key),
+                    default,
+                ),
+                op.dtype,
+            )
 
     def visit_MapMerge(self, op, *, left, right):
         return sge.DPipe(this=left, expression=right)
 
+    def visit_MapKeys(self, op, *, arg):
+        return self.if_(
+            arg.is_(NULL), NULL, self.f.array(sg.select(self.f.jsonb_object_keys(arg)))
+        )
+
+    def visit_MapValues(self, op, *, arg):
+        col = gen_name("json_each_col")
+        return self.if_(
+            arg.is_(NULL),
+            NULL,
+            self.f.array(
+                sg.select(
+                    sge.Dot(
+                        this=sg.to_identifier(col),
+                        expression=sg.to_identifier("value", quoted=True),
+                    )
+                ).from_(self.f.jsonb_each(arg).as_(col))
+            ),
+        )
+
     def visit_TypeOf(self, op, *, arg):
         typ = self.cast(self.f.pg_typeof(arg), dt.string)
         return self.if_(
@@ -591,6 +654,11 @@ def visit_NonNullLiteral(self, op, *, value, dtype):
             )
         elif dtype.is_json():
             return self.cast(value, dt.json)
+        elif dtype.is_map():
+            return sge.Cast(
+                this=sge.convert(json.dumps(value)),
+                to=sge.DataType(this=sge.DataType.Type.JSONB),
+            )
         return None
 
     def visit_TimestampFromYMDHMS(

ibis/backends/sql/datatypes.py

@@ -519,9 +519,7 @@ class PostgresType(SqlglotType):
     def _from_ibis_Map(cls, dtype: dt.Map) -> sge.DataType:
         if not dtype.key_type.is_string():
             raise com.IbisTypeError("Postgres only supports string keys in maps")
-        if not dtype.value_type.is_string():
-            raise com.IbisTypeError("Postgres only supports string values in maps")
-        return sge.DataType(this=typecode.HSTORE)
+        return sge.DataType(this=typecode.JSONB)
 
     @classmethod
     def from_string(cls, text: str, nullable: bool | None = None) -> dt.DataType:

ibis/backends/sql/dialects.py

@@ -435,7 +435,6 @@ class Polars(Postgres):
 
 
 Postgres.Generator.TRANSFORMS |= {
-    sge.Map: rename_func("hstore"),
     sge.Split: rename_func("string_to_array"),
     sge.RegexpSplit: rename_func("regexp_split_to_array"),
     sge.DateFromParts: rename_func("make_date"),