Push down JSONB existence predicates

doc/pg_clickhouse.md

@@ -1313,6 +1313,7 @@ equivalents as follows:
 *   `!~*` (case insensitive regexp not match): [match](https://clickhouse.com/docs/sql-reference/functions/string-search-functions#match)
 *   `->>` (JSON/JSONB extract element as text): [sub-column syntax](https://clickhouse.com/docs/sql-reference/data-types/newjson#reading-json-paths-as-sub-columns)
 *   `->` (JSON/JSONB extract): [toJSONString](https://clickhouse.com/docs/sql-reference/functions/json-functions#toJSONString) + [sub-column syntax](https://clickhouse.com/docs/sql-reference/data-types/newjson#reading-json-paths-as-sub-columns)
+*   `?` (JSONB top-level key or array string existence): [JSONHas](https://clickhouse.com/docs/sql-reference/functions/json-functions#jsonhas) for objects and [JSONExtractArrayRaw](https://clickhouse.com/docs/sql-reference/functions/json-functions#jsonextractarrayraw) for arrays. This also supports compatibility views that parse a JSON document stored in a ClickHouse `String`.
 
 ### IN and NULL Semantics
 

src/custom_types.c

@@ -946,6 +946,7 @@ chfdw_check_for_custom_type(Oid typeoid) {
 #define OID_TEXT_REGEX_NE_OP 642
 #define OID_TEXT_IREGEX_NE_OP 1229
 #define OID_JSONB_FETCHVAL_OP 3211
+#define OID_JSONB_EXISTS_OP 3247
 #define OID_JSONB_FETCHVAL_TEXT_OP 3477
 #define OID_JSON_FETCHVAL_OP 3962
 #define OID_JSON_FETCHVAL_TEXT_OP 3963
@@ -977,6 +978,8 @@ classify_builtin_operator(Oid opoid) {
     case OID_JSONB_FETCHVAL_TEXT_OP:
     case OID_JSON_FETCHVAL_TEXT_OP:
         return CF_JSON_FETCHVAL_TEXT;
+    case OID_JSONB_EXISTS_OP:
+        return CF_JSON_EXISTS;
     case OID_ARRAY_CONTAINS_OP:
         return CF_ARRAY_CONTAINS;
     case OID_ARRAY_CONTAINED_OP:

src/deparse.c

@@ -198,6 +198,12 @@ typedef enum ExprTruthCtx {
     EXPR_CTX_EXACT,
 } ExprTruthCtx;
 
+typedef enum JsonbDocumentKind {
+    JSONB_DOCUMENT_UNSUPPORTED,
+    JSONB_DOCUMENT_NATIVE,
+    JSONB_DOCUMENT_STRING,
+} JsonbDocumentKind;
+
 /*
  * Functions to determine whether an expression can be evaluated safely on
  * remote server.
@@ -258,6 +264,13 @@ deparseSQLValueFunction(SQLValueFunction* node, deparse_expr_cxt* context);
 static void
 deparseOpExpr(OpExpr* node, deparse_expr_cxt* context);
 static void
+deparseJsonbExists(
+    Expr* document,
+    Expr* key,
+    JsonbDocumentKind document_kind,
+    deparse_expr_cxt* context
+);
+static void
 deparseOperatorName(StringInfo buf, Form_pg_operator opform);
 static void
 deparseDistinctExpr(DistinctExpr* node, deparse_expr_cxt* context);
@@ -382,6 +395,8 @@ get_relation_column_alias_ids(
     int* relno,
     int* colno
 );
+static JsonbDocumentKind
+classifyJsonbDocument(Expr* expr, Expr** document);
 
 /*
  * Examine each qual clause in input_conds, and classify them into two groups,
@@ -894,6 +909,55 @@ classify_notin_subplan(SubPlan* subplan, PlannerInfo* root, Relids relids) {
     return NOTIN_SHIP_NONE;
 }
 
+/*
+ * Recognize either a native ClickHouse JSON column imported as jsonb or the
+ * compatibility-view shape used for a String containing a JSON document:
+ *
+ *     jsonb_in(foreign_text_column::cstring)
+ *
+ * Arbitrary jsonb expressions stay local because their remote representation
+ * is not known.
+ */
+static JsonbDocumentKind
+classifyJsonbDocument(Expr* expr, Expr** document) {
+    while (IsA(expr, RelabelType)) {
+        expr = ((RelabelType*)expr)->arg;
+    }
+
+    if (IsA(expr, Var) && exprType((Node*)expr) == JSONBOID) {
+        *document = expr;
+        return JSONB_DOCUMENT_NATIVE;
+    }
+
+    if (IsA(expr, FuncExpr)) {
+        FuncExpr* func = (FuncExpr*)expr;
+
+        if (func->funcid == F_JSONB_IN && list_length(func->args) == 1) {
+            Expr* input = (Expr*)linitial(func->args);
+
+            while (IsA(input, RelabelType)) {
+                input = ((RelabelType*)input)->arg;
+            }
+            if (IsA(input, CoerceViaIO)) {
+                CoerceViaIO* cast = (CoerceViaIO*)input;
+                Expr* raw         = cast->arg;
+
+                while (IsA(raw, RelabelType)) {
+                    raw = ((RelabelType*)raw)->arg;
+                }
+                if (cast->resulttype == CSTRINGOID &&
+                    exprType((Node*)raw) == TEXTOID) {
+                    *document = raw;
+                    return JSONB_DOCUMENT_STRING;
+                }
+            }
+        }
+    }
+
+    *document = NULL;
+    return JSONB_DOCUMENT_UNSUPPORTED;
+}
+
 /*
  * Check if expression is safe to execute remotely, and return true if so.
  *
@@ -1055,6 +1119,7 @@ foreign_expr_walker(Node* node, foreign_glob_cxt* glob_cxt, ExprTruthCtx ctx) {
     case T_DistinctExpr: /* struct-equivalent to OpExpr */
     {
         OpExpr* oe = (OpExpr*)node;
+        CustomObjectDef* cdef;
 
         /*
          * Similarly, only shippable operators can be sent to remote.
@@ -1065,6 +1130,24 @@ foreign_expr_walker(Node* node, foreign_glob_cxt* glob_cxt, ExprTruthCtx ctx) {
             return false;
         }
 
+        cdef = chfdw_check_for_custom_operator(oe->opno, NULL);
+        if (cdef && cdef->cf_type == CF_JSON_EXISTS) {
+            Expr* document;
+            JsonbDocumentKind document_kind =
+                classifyJsonbDocument((Expr*)linitial(oe->args), &document);
+
+            if (document_kind == JSONB_DOCUMENT_UNSUPPORTED ||
+                !foreign_expr_walker(
+                    (Node*)document, glob_cxt, EXPR_CTX_EXACT
+                ) ||
+                !foreign_expr_walker(
+                    (Node*)lsecond(oe->args), glob_cxt, EXPR_CTX_EXACT
+                )) {
+                return false;
+            }
+            break;
+        }
+
         /*
          * Recurse to input subexpressions.
          */
@@ -4869,6 +4952,81 @@ findFunction(Oid typoid, char* name) {
     return result;
 }
 
+static void
+deparseJsonbDocument(
+    Expr* document,
+    JsonbDocumentKind document_kind,
+    deparse_expr_cxt* context
+) {
+    if (document_kind == JSONB_DOCUMENT_NATIVE) {
+        appendStringInfoString(context->buf, "toJSONString(");
+    }
+    deparseExpr(document, context);
+    if (document_kind == JSONB_DOCUMENT_NATIVE) {
+        appendStringInfoChar(context->buf, ')');
+    }
+}
+
+static void
+deparseJsonbNonnullDocument(
+    Expr* document,
+    JsonbDocumentKind document_kind,
+    deparse_expr_cxt* context
+) {
+    appendStringInfoString(context->buf, "ifNull(");
+    deparseJsonbDocument(document, document_kind, context);
+    appendStringInfoString(context->buf, ", 'null')");
+}
+
+static void
+deparseJsonbExists(
+    Expr* document,
+    Expr* key,
+    JsonbDocumentKind document_kind,
+    deparse_expr_cxt* context
+) {
+    StringInfo buf = context->buf;
+
+    appendStringInfoString(buf, "(if(isNull(");
+    deparseJsonbDocument(document, document_kind, context);
+    appendStringInfoString(buf, ") OR isNull(");
+    deparseExpr(key, context);
+    appendStringInfoString(buf, "), NULL, ");
+
+    if (document_kind == JSONB_DOCUMENT_STRING) {
+        appendStringInfoString(buf, "if(NOT isValidJSON(");
+        deparseJsonbNonnullDocument(document, document_kind, context);
+        appendStringInfoString(
+            buf,
+            "), throwIf(1, 'invalid input syntax for type json'), "
+        );
+    }
+
+    appendStringInfoString(buf, "multiIf(JSONType(");
+    deparseJsonbNonnullDocument(document, document_kind, context);
+    appendStringInfoString(buf, ") = 'Object', JSONHas(");
+    deparseJsonbNonnullDocument(document, document_kind, context);
+    appendStringInfoString(buf, ", ");
+    deparseExpr(key, context);
+    appendStringInfoString(buf, "), JSONType(");
+    deparseJsonbNonnullDocument(document, document_kind, context);
+    appendStringInfoString(
+        buf,
+        ") = 'Array', arrayExists(jsonb_exists_element -> "
+        "JSONType(jsonb_exists_element) = 'String' AND "
+        "JSONExtractString(jsonb_exists_element) = "
+    );
+    deparseExpr(key, context);
+    appendStringInfoString(buf, ", JSONExtractArrayRaw(");
+    deparseJsonbNonnullDocument(document, document_kind, context);
+    appendStringInfoString(buf, ")), 0)");
+
+    if (document_kind == JSONB_DOCUMENT_STRING) {
+        appendStringInfoChar(buf, ')');
+    }
+    appendStringInfoString(buf, "))");
+}
+
 /*
  * Deparse given operator expression. To avoid problems around priority of
  * operations, we always parenthesize the arguments.
@@ -4999,6 +5157,17 @@ deparseOpExpr(OpExpr* node, deparse_expr_cxt* context) {
             appendStringInfoChar(buf, ')');
             goto cleanup;
         } break;
+        case CF_JSON_EXISTS: {
+            Expr* document;
+            JsonbDocumentKind document_kind =
+                classifyJsonbDocument(linitial(node->args), &document);
+
+            Assert(document_kind != JSONB_DOCUMENT_UNSUPPORTED);
+            deparseJsonbExists(
+                document, lsecond(node->args), document_kind, context
+            );
+            goto cleanup;
+        } break;
         case CF_JSON_FETCHVAL:
         case CF_JSON_FETCHVAL_TEXT: {
             Expr* arg1 = linitial(node->args);

src/include/fdw.h

@@ -417,6 +417,7 @@ typedef enum {
     CF_REGEX_ICASE_NO_MATCH,   /* !~* case-insensitive regex operator */
     CF_JSON_FETCHVAL,          /* -> operator on json & jsonb */
     CF_JSON_FETCHVAL_TEXT,     /* ->> operator on json & jsonb */
+    CF_JSON_EXISTS,            /* ? operator on jsonb */
     CF_JSON_EXTRACT_PATH_TEXT, /* jsonb?_extract_path_text() →
                                 * col."k1"."k2" */
     CF_JSON_EXTRACT_PATH,      /* json?_extract_path() →