Support subscripting on duckdb.map (#773)

In #689 the subscript operator was not implemented for the duckdb.map
type. This worked fine for MAPs returned by `duckdb.query()` because
that returned duckdb.unresolved_type. But not when getting a map from a
`duckdb` table column, i.e. a MotherDuck table. So this implements the
subscript operator directly on the `duckdb.map` type.

Fixes #732

Author: JelteF

sql/pg_duckdb--0.3.0--1.0.0.sql

@@ -686,10 +686,12 @@ LANGUAGE C AS 'MODULE_PATHNAME', 'pgduckdb_enable_motherduck';
 CREATE TYPE duckdb.map;
 CREATE FUNCTION duckdb.map_in(cstring) RETURNS duckdb.map AS 'MODULE_PATHNAME', 'duckdb_map_in' LANGUAGE C IMMUTABLE STRICT;
 CREATE FUNCTION duckdb.map_out(duckdb.map) RETURNS cstring AS 'MODULE_PATHNAME', 'duckdb_map_out' LANGUAGE C IMMUTABLE STRICT;
+CREATE FUNCTION duckdb.map_subscript(internal) RETURNS internal AS 'MODULE_PATHNAME', 'duckdb_map_subscript' LANGUAGE C IMMUTABLE STRICT;
 CREATE TYPE duckdb.map(
     INTERNALLENGTH = VARIABLE,
     INPUT = duckdb.map_in,
-    OUTPUT = duckdb.map_out
+    OUTPUT = duckdb.map_out,
+    SUBSCRIPT = duckdb.map_subscript
 );
 
 -- Drop legacy secret objects

src/pgduckdb_options.cpp

@@ -422,17 +422,6 @@ CoerceSubscriptToText(struct ParseState *pstate, A_Indices *subscript, const cha
 	return coerced_expr;
 }
 
-Node *
-CoerceRowSubscriptToText(struct ParseState *pstate, A_Indices *subscript) {
-	return CoerceSubscriptToText(pstate, subscript, "duckdb.row");
-}
-
-// Cloned implementation from CoerceRowSubscriptToText
-Node *
-CoerceStructSubscriptToText(struct ParseState *pstate, A_Indices *subscript) {
-	return CoerceSubscriptToText(pstate, subscript, "duckdb.struct");
-}
-
 /*
  * In Postgres all index operations in a row ar all slices or all plain
  * index operations. If you mix them, all are converted to slices.
@@ -473,9 +462,9 @@ AddSubscriptExpressions(SubscriptingRef *sbsref, struct ParseState *pstate, A_In
 
 /*
  * DuckdbSubscriptTransform is called by the parser when a subscripting
- * operation is performed on a duckdb.row. It has two main puprposes:
- * 1. Ensure that the row is being indexed using a string literal
- * 2. Ensure that the return type of this index operation is duckdb.unresolved_type
+ * operation is performed on a duckdb type that can be indexed by arbitrary
+ * expressions. All this does is parse those expressions and make sure the
+ * subscript returns an an duckdb.unresolved_type again.
  */
 void
 DuckdbSubscriptTransform(SubscriptingRef *sbsref, List *indirection, struct ParseState *pstate, bool isSlice,
@@ -497,6 +486,46 @@ DuckdbSubscriptTransform(SubscriptingRef *sbsref, List *indirection, struct Pars
 		elog(ERROR, "Subscripting %s with an empty subscript is not supported", type_name);
 	}
 
+	// Transform each subscript expression
+	foreach_node(A_Indices, subscript, indirection) {
+		AddSubscriptExpressions(sbsref, pstate, subscript, isSlice);
+	}
+
+	// Set the result type of the subscripting operation
+	sbsref->refrestype = pgduckdb::DuckdbUnresolvedTypeOid();
+	sbsref->reftypmod = -1;
+}
+
+/*
+ * DuckdbTextSubscriptTransform is called by the parser when a subscripting
+ * operation is performed on type that can only be indexed by string literals.
+ * It has two main puprposes:
+ * 1. Ensure that the row is being indexed using a string literal
+ * 2. Ensure that the return type of this index operation is
+ *    duckdb.unresolved_type
+ *
+ * Currently this is used for duckdb.row and duckdb.struct types.
+ */
+void
+DuckdbTextSubscriptTransform(SubscriptingRef *sbsref, List *indirection, struct ParseState *pstate, bool isSlice,
+                             bool isAssignment, const char *type_name) {
+	/*
+	 * We need to populate our cache for some of the code below. Normally this
+	 * cache is populated at the start of our planner hook, but this function
+	 * is being called from the parser.
+	 */
+	if (!pgduckdb::IsExtensionRegistered()) {
+		elog(ERROR, "BUG: Using %s but the pg_duckdb extension is not installed", type_name);
+	}
+
+	if (isAssignment) {
+		elog(ERROR, "Assignment to %s is not supported", type_name);
+	}
+
+	if (indirection == NIL) {
+		elog(ERROR, "Subscripting %s with an empty subscript is not supported", type_name);
+	}
+
 	bool first = true;
 
 	// Transform each subscript expression
@@ -505,7 +534,8 @@ DuckdbSubscriptTransform(SubscriptingRef *sbsref, List *indirection, struct Pars
 		 * a column reference. But the subscripts after that can be anything,
 		 * DuckDB should interpret those. */
 		if (first) {
-			sbsref->refupperindexpr = lappend(sbsref->refupperindexpr, CoerceRowSubscriptToText(pstate, subscript));
+			sbsref->refupperindexpr =
+			    lappend(sbsref->refupperindexpr, CoerceSubscriptToText(pstate, subscript, type_name));
 			if (isSlice) {
 				sbsref->reflowerindexpr = lappend(sbsref->reflowerindexpr, NULL);
 			}
@@ -524,13 +554,13 @@ DuckdbSubscriptTransform(SubscriptingRef *sbsref, List *indirection, struct Pars
 void
 DuckdbRowSubscriptTransform(SubscriptingRef *sbsref, List *indirection, struct ParseState *pstate, bool isSlice,
                             bool isAssignment) {
-	DuckdbSubscriptTransform(sbsref, indirection, pstate, isSlice, isAssignment, "duckdb.row");
+	DuckdbTextSubscriptTransform(sbsref, indirection, pstate, isSlice, isAssignment, "duckdb.row");
 }
 
 void
 DuckdbStructSubscriptTransform(SubscriptingRef *sbsref, List *indirection, struct ParseState *pstate, bool isSlice,
                                bool isAssignment) {
-	DuckdbSubscriptTransform(sbsref, indirection, pstate, isSlice, isAssignment, "duckdb.struct");
+	DuckdbTextSubscriptTransform(sbsref, indirection, pstate, isSlice, isAssignment, "duckdb.struct");
 }
 
 static bool
@@ -617,52 +647,40 @@ DECLARE_PG_FUNCTION(duckdb_struct_subscript) {
 	PG_RETURN_POINTER(&duckdb_struct_subscript_routines);
 }
 
-/*
- * DuckdbUnresolvedTypeSubscriptTransform is called by the parser when a
- * subscripting operation is performed on a duckdb.unresolved_type. All this
- * does is parse ensre that any subscript on duckdb.unresolved_type returns an
- * unrsolved type again.
- */
 void
-DuckdbUnresolvedTypeSubscriptTransform(SubscriptingRef *sbsref, List *indirection, struct ParseState *pstate,
-                                       bool isSlice, bool isAssignment) {
-	/*
-	 * We need to populate our cache for some of the code below. Normally this
-	 * cache is populated at the start of our planner hook, but this function
-	 * is being called from the parser.
-	 */
-	if (!pgduckdb::IsExtensionRegistered()) {
-		elog(ERROR, "BUG: Using duckdb.unresolved_type but the pg_duckdb extension is not installed");
-	}
+DuckdbMapSubscriptExecSetup(const SubscriptingRef *sbsref, SubscriptingRefState *sbsrefstate,
+                            SubscriptExecSteps *methods) {
+	DuckdbSubscriptExecSetup(sbsref, sbsrefstate, methods, "duckdb.map");
+}
 
-	if (isAssignment) {
-		elog(ERROR, "Assignment to duckdb.unresolved_type is not supported");
-	}
+void
+DuckdbMapSubscriptTransform(SubscriptingRef *sbsref, List *indirection, struct ParseState *pstate, bool isSlice,
+                            bool isAssignment) {
+	DuckdbSubscriptTransform(sbsref, indirection, pstate, isSlice, isAssignment, "duckdb.map");
+}
 
-	if (indirection == NIL) {
-		elog(ERROR, "Subscripting duckdb.row with an empty subscript is not supported");
-	}
+static SubscriptRoutines duckdb_map_subscript_routines = {
+    .transform = DuckdbMapSubscriptTransform,
+    .exec_setup = DuckdbMapSubscriptExecSetup,
+    .fetch_strict = false,
+    .fetch_leakproof = true,
+    .store_leakproof = true,
+};
 
-	// Transform each subscript expression
-	foreach_node(A_Indices, subscript, indirection) {
-		AddSubscriptExpressions(sbsref, pstate, subscript, isSlice);
-	}
+DECLARE_PG_FUNCTION(duckdb_map_subscript) {
+	PG_RETURN_POINTER(&duckdb_map_subscript_routines);
+}
 
-	// Set the result type of the subscripting operation
-	sbsref->refrestype = pgduckdb::DuckdbUnresolvedTypeOid();
-	sbsref->reftypmod = -1;
+void
+DuckdbUnresolvedTypeSubscriptTransform(SubscriptingRef *sbsref, List *indirection, struct ParseState *pstate,
+                                       bool isSlice, bool isAssignment) {
+	DuckdbSubscriptTransform(sbsref, indirection, pstate, isSlice, isAssignment, "duckdb.unresolved_type");
 }
 
-/*
- * DuckdbUnresolvedTypeSubscriptExecSetup is called by the executor when a
- * subscripting operation is performed on a duckdb.unresolved_type. This should
- * never happen, because any query that contains a duckdb.unresolved_type should
- * automatically be use DuckDB execution.
- */
 void
-DuckdbUnresolvedTypeSubscriptExecSetup(const SubscriptingRef * /*sbsref*/, SubscriptingRefState * /*sbsrefstate*/,
-                                       SubscriptExecSteps * /*exprstate*/) {
-	elog(ERROR, "Subscripting duckdb.unresolved_type is not supported in the Postgres Executor");
+DuckdbUnresolvedTypeSubscriptExecSetup(const SubscriptingRef *sbsref, SubscriptingRefState *sbsrefstate,
+                                       SubscriptExecSteps *methods) {
+	DuckdbSubscriptExecSetup(sbsref, sbsrefstate, methods, "duckdb.unresolved_type");
 }
 
 static SubscriptRoutines duckdb_unresolved_type_subscript_routines = {

test/pycheck/motherduck_test.py

@@ -6,7 +6,9 @@
 library for that purpose.
 """
 
-from .utils import Cursor, Postgres, PG_MAJOR_VERSION
+import datetime
+
+from .utils import Cursor, Postgres, Duckdb, PG_MAJOR_VERSION
 from .motherduck_token_helper import (
     can_run_md_multi_user_tests,
     can_run_md_tests,
@@ -196,3 +198,41 @@ def test_md_alter_table(md_cur: Cursor):
         # duckdb after the table is created
         with pytest.raises(psycopg.errors.FeatureNotSupported):
             md_cur.sql("ALTER TABLE t SET ACCESS METHOD duckdb")
+
+
+def test_md_duckdb_only_types(md_cur: Cursor, ddb: Duckdb):
+    ddb.sql("""
+            CREATE TABLE t1(
+                m MAP(INT, VARCHAR),
+                s STRUCT(v VARCHAR, i INTEGER),
+                u UNION(t time, d date),
+            )""")
+    ddb.sql("""
+        INSERT INTO t1 VALUES (
+            MAP{1: 'abc'},
+            {'v': 'struct abc', 'i': 123},
+            '12:00'::time,
+        ), (
+            MAP{2: 'def'},
+            {'v': 'struct def', 'i': 456},
+            '2023-10-01'::date,
+        )
+    """)
+    md_cur.wait_until_table_exists("t1")
+    assert md_cur.sql("""select * from t1""") == [
+        ("{1=abc}", "{'v': struct abc, 'i': 123}", "12:00:00"),
+        ("{2=def}", "{'v': struct def, 'i': 456}", "2023-10-01"),
+    ]
+
+    assert md_cur.sql("""select m[1] from t1""") == ["abc", None]
+    assert md_cur.sql("""select s['v'] from t1""") == ["struct abc", "struct def"]
+    assert md_cur.sql("""select s['i'] from t1""") == [123, 456]
+    assert md_cur.sql("""select union_extract(u,'t') from t1""") == [
+        datetime.time(12, 0),
+        None,
+    ]
+    assert md_cur.sql("""select union_extract(u, 'd') from t1""") == [
+        None,
+        datetime.date(2023, 10, 1),
+    ]
+    assert md_cur.sql("""select union_tag(u) from t1""") == ["t", "d"]