From 8ced6c036acb2da281700fe97436f5d47a9a66a8 Mon Sep 17 00:00:00 2001 From: Kostia R Date: Thu, 23 Jul 2026 21:44:01 +0000 Subject: [PATCH 1/2] Push down compatible UNION queries --- README.md | 2 - doc/pg_clickhouse.md | 32 + src/deparse.c | 127 ++- src/fdw.c | 1476 +++++++++++++++++++++++++++++- src/include/fdw.h | 26 + src/include/setop.h | 49 + src/option.c | 34 +- src/setop.c | 595 ++++++++++++ test/expected/result_map.txt | 11 + test/expected/union_pushdown.out | 574 ++++++++++++ test/sql/union_pushdown.sql | 781 ++++++++++++++++ 11 files changed, 3670 insertions(+), 37 deletions(-) create mode 100644 src/include/setop.h create mode 100644 src/setop.c create mode 100644 test/expected/union_pushdown.out create mode 100644 test/sql/union_pushdown.sql diff --git a/README.md b/README.md index b81b513e..12ae24c1 100644 --- a/README.md +++ b/README.md @@ -248,8 +248,6 @@ adding DML features. Our road map: * Support batch insertion via COPY * Add a function to execute an arbitrary ClickHouse query and return its results as a tables -* Add support for pushdown of UNION queries when they all query the remote - database ## Authors diff --git a/doc/pg_clickhouse.md b/doc/pg_clickhouse.md index 8d2e560b..f58a86f5 100644 --- a/doc/pg_clickhouse.md +++ b/doc/pg_clickhouse.md @@ -577,6 +577,38 @@ try=# EXPLAIN (ANALYZE, VERBOSE) the number of rows that must be pulled back into Postgres from 1000 (all of them) to just 8, one for each node. +### UNION and CTEs + +Plain `SELECT DISTINCT` is pushed down for the same conservative set of +equality-compatible output types listed below. `DISTINCT ON`, ordered +`DISTINCT`, and `DISTINCT` over grouping or window results stay in PostgreSQL. + +pg_clickhouse can combine `UNION ALL` and `UNION` (`UNION DISTINCT`) into one +remote query when every arm is already fully pushable through the same foreign +server and user mapping. If an arm requires local evaluation or uses another +server, PostgreSQL executes the set operation locally. + +`UNION DISTINCT` is pushed down only for output types whose equality semantics +match on both servers: `boolean`, `smallint`, `integer`, `bigint`, `date`, +`text` with a deterministic collation, unlimited `varchar`, `bytea`, and +`uuid`. Other output types keep duplicate elimination in PostgreSQL. +Direct `UNION DISTINCT` queries with `ORDER BY` also stay local so PostgreSQL +can apply the requested ordering. + +PostgreSQL inlines eligible CTEs before pg_clickhouse plans the query, so +single-use CTEs and CTEs declared `NOT MATERIALIZED` can participate in +pushdown. Multiply referenced CTEs using PostgreSQL's default materialization, +and CTEs declared `MATERIALIZED`, retain evaluate-once semantics and stay +local. + +Grouping and aggregation over a pushed `UNION ALL` can also be executed in +the same remote query. This requires at least two fully remote arms using the +same server and user mapping with the default table engine, and no local +conditions, `HAVING`, outer ordering or limiting, grouping sets, window +functions, row locking, or set-returning functions. Group keys and aggregate +`DISTINCT` inputs use the same conservative equality checks as `UNION +DISTINCT`; unsupported shapes stay local. + ### Partitioned Tables A PostgreSQL [partitioned table] can mix local partitions with foreign diff --git a/src/deparse.c b/src/deparse.c index 66d50d0b..2428fa09 100644 --- a/src/deparse.c +++ b/src/deparse.c @@ -131,6 +131,14 @@ typedef struct deparse_expr_cxt { */ bool has_inlined_subplan; + /* + * A composable set-operation query is exposed as a derived table with + * positional cN columns. When set, Vars belonging to scanrel are mapped + * through derived_tlist instead of being resolved as range-table columns. + */ + List* derived_tlist; + const char* derived_alias; + /* * SubPlan scope. While deparsing the body of a pushed-down SubPlan, * subplan points at it and root points at the SubPlan's own PlannerInfo @@ -245,6 +253,8 @@ static void deparseExpr(Expr* expr, deparse_expr_cxt* context); static void deparseVar(Var* node, deparse_expr_cxt* context); +static int +derived_column_number(Var* node, deparse_expr_cxt* context); static void deparseConst(Const* node, deparse_expr_cxt* context, int showtype); static void @@ -2157,6 +2167,46 @@ chfdw_deparse_select_stmt_for_rel( } } +/* + * Deparse grouping or aggregation over a fully remote set operation. The + * set operation is already a complete query, so expose it as a derived table + * and let the normal expression and GROUP BY deparsers handle the upper + * target. + */ +void +chfdw_deparse_setop_grouping_stmt( + StringInfo buf, + PlannerInfo* root, + RelOptInfo* rel, + List* tlist, + List* setop_tlist, + const char* setop_sql, + List** retrieved_attrs, + List** params_list +) { + deparse_expr_cxt context; + CHFdwRelationInfo* fpinfo = (CHFdwRelationInfo*)rel->fdw_private; + + Assert(IS_UPPER_REL(rel)); + Assert(fpinfo != NULL && fpinfo->outerrel != NULL); + Assert(setop_tlist != NIL && setop_sql != NULL); + + memset(&context, 0, sizeof(context)); + context.buf = buf; + context.root = root; + context.foreignrel = rel; + context.scanrel = fpinfo->outerrel; + context.params_list = params_list; + context.fpinfo = fpinfo; + context.derived_tlist = setop_tlist; + context.derived_alias = "setop_input"; + context.no_sort_parens = false; + + deparseSelectSql(tlist, false, retrieved_attrs, &context); + appendStringInfo(buf, " FROM (%s) AS %s", setop_sql, context.derived_alias); + appendGroupByClause(tlist, &context); +} + /* * Construct a simple SELECT statement that retrieves desired columns * of the specified foreign table, and append it to "buf". The output @@ -2190,6 +2240,9 @@ deparseSelectSql( * Construct SELECT list */ appendStringInfoString(buf, "SELECT "); + if (IS_UPPER_REL(foreignrel) && fpinfo->stage == UPPERREL_DISTINCT) { + appendStringInfoString(buf, "DISTINCT "); + } if (is_subquery) { /* @@ -2985,6 +3038,46 @@ deparseExpr(Expr* node, deparse_expr_cxt* context) { * Otherwise, it's effectively a Param (and will in fact be a Param at * run time). Handle it the same way we handle plain Params. */ +static int +derived_column_number(Var* node, deparse_expr_cxt* context) { + ListCell* lc; + int column = 0; + + if (context->derived_tlist == NIL || node->varlevelsup != 0) { + return 0; + } + + foreach (lc, context->derived_tlist) { + TargetEntry* tle = lfirst_node(TargetEntry, lc); + + if (tle->resjunk) { + continue; + } + column++; + if (equal(node, tle->expr)) { + return column; + } + } + + /* + * A flattened UNION ALL is represented by a subquery RTE, so Vars above + * it normally refer directly to its positional output attributes. + */ + if ((node->varno == 0 || (int64)node->varno == (int64)context->scanrel->relid) && + node->varattno > 0) { + column = 0; + foreach (lc, context->derived_tlist) { + TargetEntry* tle = lfirst_node(TargetEntry, lc); + + if (!tle->resjunk && ++column == node->varattno) { + return column; + } + } + return 0; + } + return 0; +} + static void deparseVar(Var* node, deparse_expr_cxt* context) { CustomObjectDef* cdef; @@ -2992,6 +3085,17 @@ deparseVar(Var* node, deparse_expr_cxt* context) { int relno; int colno; + colno = derived_column_number(node, context); + if (colno > 0) { + appendStringInfo(context->buf, "%s.c%d", context->derived_alias, colno); + return; + } + if (context->derived_tlist != NIL && + (node->varno == 0 || bms_is_member(node->varno, context->scanrel->relids)) && + node->varlevelsup == 0) { + elog(ERROR, "could not map derived set-operation column"); + } + /* * Qualify columns when multiple relations are involved, or when a SubPlan * is inlined anywhere in the statement (unqualified outer columns would @@ -5679,11 +5783,16 @@ appendAggOrderBy(List* orderList, List* targetList, deparse_expr_cxt* context) { */ static bool aggref_on_aggregate_function(Aggref* node, deparse_expr_cxt* context) { - List* vars = pull_var_clause((Node*)node->args, 0); + List* vars; ListCell* lc; Relids relids = context->scanrel->relids; bool found = false; + if (context->derived_tlist != NIL) { + return false; + } + + vars = pull_var_clause((Node*)node->args, 0); foreach (lc, vars) { Var* var = (Var*)lfirst(lc); @@ -5801,12 +5910,13 @@ static void deparseAggref(Aggref* node, deparse_expr_cxt* context) { StringInfo buf = context->buf; CustomObjectDef* cdef; - CHFdwRelationInfo* fpinfo = context->scanrel->fdw_private; - bool aggfilter = false; - bool sign_count_filter = false; - uint8 brcount = 1; - char* name = get_func_name(node->aggfnoid); - bool omit_star = false; /* Explained below. */ + CHFdwRelationInfo* fpinfo = + context->derived_tlist != NIL ? context->fpinfo : context->scanrel->fdw_private; + bool aggfilter = false; + bool sign_count_filter = false; + uint8 brcount = 1; + char* name = get_func_name(node->aggfnoid); + bool omit_star = false; /* Explained below. */ bool use_variadic; /* Simple aggregates push down directly */ @@ -6515,7 +6625,8 @@ appendFunctionName(Oid funcid, deparse_expr_cxt* context) { Form_pg_proc procform; const char* proname; CustomObjectDef* cdef; - CHFdwRelationInfo* fpinfo = context->scanrel->fdw_private; + CHFdwRelationInfo* fpinfo = + context->derived_tlist != NIL ? context->fpinfo : context->scanrel->fdw_private; cdef = chfdw_check_for_custom_function(funcid); if (cdef && cdef->custom_name[0] != '\0') { diff --git a/src/fdw.c b/src/fdw.c index 9de9db09..f28733ee 100644 --- a/src/fdw.c +++ b/src/fdw.c @@ -42,6 +42,7 @@ /* extension includes. */ #include "fdw.h" +#include "setop.h" #include "utils/builtins.h" #include "version.h" @@ -83,7 +84,9 @@ enum FdwScanPrivateIndex { * String describing join i.e. names of relations being joined and types * of join, added when the scan is join */ - FdwScanPrivateRelations + FdwScanPrivateRelations, + /* Effective user OID for scans without a representative foreign-table RTE */ + FdwScanPrivateUserId }; /* @@ -178,6 +181,10 @@ typedef struct { int64 offset_est; } ChFdwPathExtraData; +typedef struct ChFdwSetopValidationState { + int32 fetch_size; +} ChFdwSetopValidationState; + /* * SQL functions */ @@ -331,6 +338,12 @@ add_foreign_grouping_paths( GroupPathExtraData* extra ); static void +add_foreign_distinct_paths( + PlannerInfo* root, + RelOptInfo* input_rel, + RelOptInfo* distinct_rel +); +static void add_foreign_window_paths( PlannerInfo* root, RelOptInfo* input_rel, @@ -361,6 +374,70 @@ static int get_fetch_size_option(DefElem* def); static DefElem* ch_get_table_or_server_option(CHFdwRelationInfo* fpinfo, char* name); +static bool +setop_exec_param_walker(Node* node, void* context); +static bool +setop_restrictinfos_require_local_gate(List* conditions); +static bool +setop_relation_requires_local_gate(RelOptInfo* rel); +static bool +distinct_expr_is_safe(Node* expr); +static bool +distinct_target_is_safe(PathTarget* target); +static bool +setop_grouping_target_is_safe( + PlannerInfo* root, + RelOptInfo* input_rel, + RelOptInfo* grouped_rel, + List* setop_tlist +); +static bool +setop_relation_uses_default_engine(RelOptInfo* rel); +static bool +setop_foreign_path_ok(const ForeignPath* path, void* arg); +static bool +is_flattened_union_all(PlannerInfo* root, RelOptInfo* input_rel); +static bool +is_flattened_union_all_grouping( + PlannerInfo* root, + RelOptInfo* input_rel, + GroupPathExtraData* extra +); +static AppendPath* +make_setop_append_path(PlannerInfo* root, AppendPath* append_path); +static bool +find_setop_path( + PlannerInfo* root, + RelOptInfo* rel, + AppendPath** append_path, + bool* is_distinct, + ChFdwSetopPathInfo* path_info, + int32* fetch_size, + Cost* startup_cost, + Cost* total_cost +); +static void +add_foreign_setop_path( + PlannerInfo* root, + RelOptInfo* output_rel, + AppendPath* append_path, + bool is_distinct, + const ChFdwSetopPathInfo* path_info, + int32 fetch_size, + Cost startup_cost, + Cost total_cost +); +static void +add_foreign_setop_grouping_path( + PlannerInfo* root, + RelOptInfo* input_rel, + RelOptInfo* grouped_rel, + AppendPath* append_path, + const ChFdwSetopPathInfo* path_info, + int32 fetch_size, + Cost startup_cost, + Cost total_cost +); /* Make one query and close the connection */ Datum @@ -849,6 +926,42 @@ typedef struct WindowFuncSubstState { static Node* replace_windowfuncs_mutator(Node* node, WindowFuncSubstState* state); +static ForeignScan* +clickhouseGetForeignSetopPlan(RelOptInfo* foreignrel, List* tlist, Plan* outer_plan); +static ForeignScan* +clickhouseGetForeignSetopGroupingPlan( + PlannerInfo* root, + RelOptInfo* foreignrel, + List* tlist, + Plan* outer_plan +); +static void +build_foreign_setop_query( + RelOptInfo* foreignrel, + CHFdwRelationInfo* fpinfo, + Plan* outer_plan, + List* setop_tlist, + StringInfo sql, + List** params_list +); +static void +append_setop_query( + StringInfo buf, + ForeignScan* scan, + int arm_number, + int param_offset, + int expected_columns +); +static void +append_renumbered_query(StringInfo buf, const char* query, int param_offset); +static List* +setop_projection_positions( + ForeignScan* scan, + int expected_columns, + int* source_columns +); +static List* +make_setop_scan_tlist(List* source_tlist); static Node* replace_windowfuncs_mutator_callback(Node* node, void* state) { @@ -889,6 +1002,388 @@ replace_windowfuncs_mutator(Node* node, WindowFuncSubstState* state) { return expression_tree_mutator(node, replace_windowfuncs_mutator_callback, state); } +/* + * Append query while shifting ClickHouse parameter indexes. Each arm was + * planned independently and therefore starts numbering at p1. + */ +static void +append_renumbered_query(StringInfo buf, const char* query, int param_offset) { + const char* cursor = query; + char quote = '\0'; + + while (*cursor) { + if (quote != '\0') { + appendStringInfoChar(buf, *cursor); + + if (*cursor == '\\' && cursor[1] != '\0') { + appendStringInfoChar(buf, cursor[1]); + cursor += 2; + continue; + } + if (*cursor == quote) { + if (cursor[1] == quote) { + appendStringInfoChar(buf, cursor[1]); + cursor += 2; + continue; + } + quote = '\0'; + } + cursor++; + continue; + } + + if (*cursor == '\'' || *cursor == '"' || *cursor == '`') { + quote = *cursor; + appendStringInfoChar(buf, *cursor++); + continue; + } + + if (cursor[0] == '{' && cursor[1] == 'p' && cursor[2] >= '0' && + cursor[2] <= '9') { + const char* number = cursor + 2; + const char* end = number; + unsigned long index; + + while (*end >= '0' && *end <= '9') { + end++; + } + if (*end == ':') { + index = strtoul(number, NULL, 10); + appendStringInfo(buf, "{p%lu", index + param_offset); + cursor = end; + continue; + } + } + + appendStringInfoChar(buf, *cursor++); + } +} + +/* + * Map a ForeignScan's projected output to positions in its remote result. + * A base scan reconstructs a table row according to retrieved_attrs before + * applying its plan target list. Join and upper scans expose fdw_scan_tlist + * directly. Set-operation arms may only use those already-remote columns; + * any local expression keeps the UNION local. + */ +static List* +setop_projection_positions( + ForeignScan* scan, + int expected_columns, + int* source_columns +) { + List* positions = NIL; + ListCell* lc; + + if (scan->scan.plan.qual != NIL || scan->scan.plan.initPlan != NIL) { + return NIL; + } + + if (scan->scan.scanrelid > 0) { + List* retrieved_attrs = + (List*)list_nth(scan->fdw_private, FdwScanPrivateRetrievedAttrs); + + *source_columns = list_length(retrieved_attrs); + foreach (lc, scan->scan.plan.targetlist) { + TargetEntry* tle = lfirst_node(TargetEntry, lc); + Var* var; + ListCell* attr_lc; + int position = 0; + bool found = false; + + if (tle->resjunk || !IsA(tle->expr, Var)) { + return NIL; + } + var = castNode(Var, tle->expr); + if ((int64)var->varno != (int64)scan->scan.scanrelid || + var->varlevelsup != 0 || var->varattno <= 0) { + return NIL; + } + + foreach (attr_lc, retrieved_attrs) { + position++; + if (lfirst_int(attr_lc) == var->varattno) { + found = true; + break; + } + } + if (!found) { + return NIL; + } + positions = lappend_int(positions, position); + } + } else { + List* remote_tlist = scan->fdw_scan_tlist; + + *source_columns = list_length(remote_tlist); + foreach (lc, scan->scan.plan.targetlist) { + TargetEntry* tle = lfirst_node(TargetEntry, lc); + TargetEntry* remote_tle; + + if (tle->resjunk) { + return NIL; + } + remote_tle = tlist_member(tle->expr, remote_tlist); + if (remote_tle == NULL || remote_tle->resjunk) { + return NIL; + } + positions = lappend_int(positions, remote_tle->resno); + } + } + + if (*source_columns <= 0 || list_length(positions) != expected_columns) { + return NIL; + } + return positions; +} + +/* + * Normalize one independently-planned remote arm to positional cN columns. + * tuple(*) avoids depending on source column names, while tupleElement keeps + * this compatible with ClickHouse versions that predate derived-table column + * alias lists. The explicit output aliases make names deterministic across + * reordered base columns, joins, nested UNIONs, and expressions. + */ +static void +append_setop_query( + StringInfo buf, + ForeignScan* scan, + int arm_number, + int param_offset, + int expected_columns +) { + List* positions; + ListCell* lc; + int source_columns; + int output_column = 0; + const char* query = strVal(list_nth(scan->fdw_private, FdwScanPrivateSelectSql)); + + positions = setop_projection_positions(scan, expected_columns, &source_columns); + if (positions == NIL) { + ereport( + ERROR, + errcode(ERRCODE_FDW_ERROR), + errmsg("pg_clickhouse: UNION arm requires local projection") + ); + } + + appendStringInfoChar(buf, '('); + appendStringInfoString(buf, "SELECT "); + foreach (lc, positions) { + if (output_column > 0) { + appendStringInfoString(buf, ", "); + } + output_column++; + appendStringInfo( + buf, "tupleElement(tuple(*), %d) AS c%d", lfirst_int(lc), output_column + ); + } + + appendStringInfoString(buf, " FROM ("); + append_renumbered_query(buf, query, param_offset); + appendStringInfo(buf, ") AS u%d)", arm_number); +} + +/* + * Set-operation targets use planner-internal Vars with varno 0, which cannot + * appear in a finished ForeignScan. Replace them with unique, typed + * placeholders in both plan and scan target lists; setrefs then rewrites the + * plan target to INDEX_VAR references to the remote scan tuple. + */ +static List* +make_setop_scan_tlist(List* source_tlist) { + List* result = NIL; + ListCell* lc; + int column = 0; + + foreach (lc, source_tlist) { + TargetEntry* source_tle = lfirst_node(TargetEntry, lc); + CoerceViaIO* placeholder; + RelabelType* relabel; + TargetEntry* tle; + Oid type; + int32 typmod; + Oid collation; + + if (source_tle->resjunk) { + continue; + } + + column++; + type = exprType((Node*)source_tle->expr); + typmod = exprTypmod((Node*)source_tle->expr); + collation = exprCollation((Node*)source_tle->expr); + + placeholder = makeNode(CoerceViaIO); + placeholder->arg = (Expr*)makeConst( + INT4OID, -1, InvalidOid, sizeof(int32), Int32GetDatum(column), false, true + ); + placeholder->resulttype = type; + placeholder->resultcollid = collation; + placeholder->coerceformat = COERCE_IMPLICIT_CAST; + placeholder->location = -1; + + relabel = makeRelabelType( + (Expr*)placeholder, type, typmod, collation, COERCE_IMPLICIT_CAST + ); + tle = makeTargetEntry( + (Expr*)relabel, + column, + source_tle->resname ? pstrdup(source_tle->resname) : NULL, + false + ); + result = lappend(result, tle); + } + + return result; +} + +/* + * Build a UNION query from the independently planned pg_clickhouse scans + * under the planner's Append path. + */ +static void +build_foreign_setop_query( + RelOptInfo* foreignrel, + CHFdwRelationInfo* fpinfo, + Plan* outer_plan, + List* setop_tlist, + StringInfo sql, + List** params_list +) { + List* foreign_scans = NIL; + ListCell* lc; + int arm_number = 0; + int expected_columns; + int param_offset = 0; + + if (outer_plan == NULL || + !chfdw_setop_extract_foreign_scans( + outer_plan, foreignrel->serverid, &foreign_scans + ) || + list_length(foreign_scans) < 2) { + ereport( + ERROR, + errcode(ERRCODE_FDW_ERROR), + errmsg("pg_clickhouse: invalid UNION pushdown plan") + ); + } + + expected_columns = list_length(setop_tlist); + if (expected_columns <= 0) { + ereport( + ERROR, + errcode(ERRCODE_FDW_ERROR), + errmsg("pg_clickhouse: UNION has no remote output columns") + ); + } + + initStringInfo(sql); + foreach (lc, foreign_scans) { + ForeignScan* scan = lfirst_node(ForeignScan, lc); + ListCell* param_lc; + + if (arm_number > 0) { + appendStringInfoString( + sql, fpinfo->setop_all ? " UNION ALL " : " UNION DISTINCT " + ); + } + arm_number++; + append_setop_query(sql, scan, arm_number, param_offset, expected_columns); + + foreach (param_lc, scan->fdw_exprs) { + *params_list = lappend(*params_list, copyObject(lfirst(param_lc))); + } + param_offset += list_length(scan->fdw_exprs); + } +} + +/* + * Build a single ForeignScan for a UNION from the independently planned + * pg_clickhouse scans under the planner's Append path. + */ +static ForeignScan* +clickhouseGetForeignSetopPlan(RelOptInfo* foreignrel, List* tlist, Plan* outer_plan) { + CHFdwRelationInfo* fpinfo = (CHFdwRelationInfo*)foreignrel->fdw_private; + List* params_list = NIL; + List* retrieved_attrs = NIL; + List* fdw_private; + StringInfoData sql; + List* scan_tlist; + + (void)tlist; + + build_foreign_setop_query( + foreignrel, fpinfo, outer_plan, fpinfo->grouped_tlist, &sql, ¶ms_list + ); + + for (int column = 1; column <= list_length(fpinfo->grouped_tlist); column++) { + retrieved_attrs = lappend_int(retrieved_attrs, column); + } + + fdw_private = list_make3( + makeString(sql.data), retrieved_attrs, makeInteger(fpinfo->fetch_size) + ); + fdw_private = lappend(fdw_private, makeString(fpinfo->relation_name->data)); + fdw_private = + lappend(fdw_private, makeString(psprintf("%u", fpinfo->setop_userid))); + + scan_tlist = make_setop_scan_tlist(fpinfo->grouped_tlist); + + return make_foreignscan( + copyObject(scan_tlist), NIL, 0, params_list, fdw_private, scan_tlist, NIL, NULL + ); +} + +/* + * Build grouping or aggregation as a wrapper around a remote UNION ALL. The + * Append input plan is used only to obtain the independently deparsed arm + * queries and parameters; the returned scan replaces that input completely. + */ +static ForeignScan* +clickhouseGetForeignSetopGroupingPlan( + PlannerInfo* root, + RelOptInfo* foreignrel, + List* tlist, + Plan* outer_plan +) { + CHFdwRelationInfo* fpinfo = (CHFdwRelationInfo*)foreignrel->fdw_private; + List* params_list = NIL; + List* retrieved_attrs = NIL; + List* fdw_private; + List* scan_tlist; + StringInfoData setop_sql; + StringInfoData sql; + + build_foreign_setop_query( + foreignrel, fpinfo, outer_plan, fpinfo->setop_tlist, &setop_sql, ¶ms_list + ); + + initStringInfo(&sql); + chfdw_deparse_setop_grouping_stmt( + &sql, + root, + foreignrel, + fpinfo->grouped_tlist, + fpinfo->setop_tlist, + setop_sql.data, + &retrieved_attrs, + ¶ms_list + ); + + fdw_private = list_make3( + makeString(sql.data), retrieved_attrs, makeInteger(fpinfo->fetch_size) + ); + fdw_private = lappend(fdw_private, makeString(fpinfo->relation_name->data)); + fdw_private = + lappend(fdw_private, makeString(psprintf("%u", fpinfo->setop_userid))); + scan_tlist = copyObject(fpinfo->grouped_tlist); + + return make_foreignscan( + tlist, NIL, 0, params_list, fdw_private, scan_tlist, NIL, NULL + ); +} + /* * clickhouseGetForeignPlan * Create ForeignScan plan node which implements selected best path @@ -920,6 +1415,24 @@ clickhouseGetForeignPlan( gettimeofday(&time1, NULL); + if (fpinfo->is_setop_grouping) { + ForeignScan* plan = + clickhouseGetForeignSetopGroupingPlan(root, foreignrel, tlist, outer_plan); + + gettimeofday(&time2, NULL); + time_used += time_diff(&time1, &time2); + return plan; + } + + if (fpinfo->is_setop) { + ForeignScan* plan = + clickhouseGetForeignSetopPlan(foreignrel, tlist, outer_plan); + + gettimeofday(&time2, NULL); + time_used += time_diff(&time1, &time2); + return plan; + } + /* * Get FDW private data created by clickhouseGetForeignUpperPaths(), if * any. @@ -1138,7 +1651,7 @@ clickhouseBeginForeignScan(ForeignScanState* node, int eflags) { RangeTblEntry* rte; int rtindex; #endif - Oid userid; + Oid userid = InvalidOid; UserMapping* user; int numParams; @@ -1165,7 +1678,14 @@ clickhouseBeginForeignScan(ForeignScanState* node, int eflags) { * fs_relids can also contain synthetic join RTEs, which do not have a * relation OID. */ - if (fsplan->scan.scanrelid > 0) { + if (list_length(fsplan->fdw_private) > FdwScanPrivateUserId) { + userid = (Oid)strtoul( + strVal(list_nth(fsplan->fdw_private, FdwScanPrivateUserId)), NULL, 10 + ); + if (!OidIsValid(userid)) { + userid = GetUserId(); + } + } else if (fsplan->scan.scanrelid > 0) { rtindex = fsplan->scan.scanrelid; } else { rtindex = -1; @@ -1179,10 +1699,13 @@ clickhouseBeginForeignScan(ForeignScanState* node, int eflags) { elog(ERROR, "could not find foreign table for pg_clickhouse scan"); } } - rte = rt_fetch(rtindex, estate->es_range_table); - userid = rte->checkAsUser ? rte->checkAsUser : GetUserId(); + if (list_length(fsplan->fdw_private) <= FdwScanPrivateUserId) { + rte = rt_fetch(rtindex, estate->es_range_table); + userid = rte->checkAsUser ? rte->checkAsUser : GetUserId(); + } #endif + Assert(OidIsValid(userid)); user = GetUserMapping(userid, fsplan->fs_server); /* @@ -2893,27 +3416,822 @@ foreign_grouping_ok(PlannerInfo* root, RelOptInfo* grouped_rel, Node* havingQual return true; } -/* - * clickhouseGetForeignUpperPaths - * Add paths for post-join operations like aggregation, grouping etc. if - * corresponding operations are safe to push down. - * - * Right now, we only support aggregate, grouping and having clause pushdown. - */ -static void -clickhouseGetForeignUpperPaths( - PlannerInfo* root, - UpperRelationKind stage, - RelOptInfo* input_rel, - RelOptInfo* output_rel, - void* extra -) { - CHFdwRelationInfo* fpinfo; - struct timeval time1, time2; +static bool +setop_exec_param_walker(Node* node, void* context) { + (void)context; - gettimeofday(&time1, NULL); + if (node == NULL) { + return false; + } + if (IsA(node, Param) && ((Param*)node)->paramkind == PARAM_EXEC) { + return true; + } + return expression_tree_walker(node, setop_exec_param_walker, context); +} - /* +static bool +setop_restrictinfos_require_local_gate(List* conditions) { + ListCell* lc; + + foreach (lc, conditions) { + RestrictInfo* rinfo = lfirst_node(RestrictInfo, lc); + + if (rinfo->pseudoconstant || + setop_exec_param_walker((Node*)rinfo->clause, NULL)) { + return true; + } + } + return false; +} + +/* + * Result gating quals and InitPlan parameters are added above a scan only + * after path selection. A combined remote UNION drops its input plan, so + * reject such arms before offering that path. + */ +static bool +setop_relation_requires_local_gate(RelOptInfo* rel) { + CHFdwRelationInfo* fpinfo; + + if (rel == NULL || rel->fdw_private == NULL) { + return true; + } + + fpinfo = (CHFdwRelationInfo*)rel->fdw_private; + if (setop_restrictinfos_require_local_gate(rel->baserestrictinfo) || + setop_restrictinfos_require_local_gate(fpinfo->remote_conds) || + setop_restrictinfos_require_local_gate(fpinfo->local_conds) || + setop_restrictinfos_require_local_gate(fpinfo->joinclauses) || + setop_exec_param_walker((Node*)rel->reltarget->exprs, NULL) || + setop_exec_param_walker((Node*)fpinfo->grouped_tlist, NULL) || + setop_exec_param_walker((Node*)fpinfo->final_remote_exprs, NULL)) { + return true; + } + + if (fpinfo->outerrel != NULL && fpinfo->outerrel != rel && + setop_relation_requires_local_gate(fpinfo->outerrel)) { + return true; + } + if (fpinfo->innerrel != NULL && fpinfo->innerrel != rel && + setop_relation_requires_local_gate(fpinfo->innerrel)) { + return true; + } + return false; +} + +/* + * DISTINCT must use the same equality semantics on both servers. Keep this + * list intentionally narrow; notably, bpchar ignores trailing spaces in + * PostgreSQL while ClickHouse String does not. + */ +static bool +distinct_expr_is_safe(Node* expr) { + Oid type; + + if (expr == NULL) { + return false; + } + type = exprType(expr); + switch (type) { + case BOOLOID: + case INT2OID: + case INT4OID: + case INT8OID: + case DATEOID: + case BYTEAOID: + case UUIDOID: + return true; + case VARCHAROID: + if (exprTypmod(expr) >= 0) { + return false; + } + /* fall through */ + case TEXTOID: { + Oid collation = exprCollation(expr); + + return !OidIsValid(collation) || get_collation_isdeterministic(collation); + } + default: + return false; + } +} + +static bool +distinct_target_is_safe(PathTarget* target) { + ListCell* lc; + + if (target == NULL || target->exprs == NIL) { + return false; + } + + foreach (lc, target->exprs) { + if (!distinct_expr_is_safe((Node*)lfirst(lc))) { + return false; + } + } + return true; +} + +typedef struct ChFdwSetopGroupingValidationState { + List* setop_tlist; + int setop_columns; +} ChFdwSetopGroupingValidationState; + +static TargetEntry* +setop_grouping_var_target( + const Var* var, + const ChFdwSetopGroupingValidationState* state +) { + ListCell* lc; + + foreach (lc, state->setop_tlist) { + TargetEntry* tle = lfirst_node(TargetEntry, lc); + + if (!tle->resjunk && equal(var, tle->expr)) { + return tle; + } + } + if (var->varno == 0 && var->varattno > 0 && var->varattno <= state->setop_columns) { + return get_tle_by_resno(state->setop_tlist, var->varattno); + } + return NULL; +} + +static bool +setop_grouping_expr_walker(Node* node, void* context) { + ChFdwSetopGroupingValidationState* state = + (ChFdwSetopGroupingValidationState*)context; + + if (node == NULL) { + return false; + } + if (IsA(node, SubPlan) || + (IsA(node, Param) && ((Param*)node)->paramkind == PARAM_EXEC)) { + return true; + } + if (IsA(node, Var)) { + Var* var = (Var*)node; + TargetEntry* tle = setop_grouping_var_target(var, state); + + return var->varlevelsup != 0 || tle == NULL || + exprType((Node*)var) != exprType((Node*)tle->expr) || + exprTypmod((Node*)var) != exprTypmod((Node*)tle->expr) || + exprCollation((Node*)var) != exprCollation((Node*)tle->expr); + } + return expression_tree_walker(node, setop_grouping_expr_walker, context); +} + +static bool +setop_grouping_distinct_walker(Node* node, void* context) { + (void)context; + + if (node == NULL) { + return false; + } + if (IsA(node, Aggref)) { + Aggref* agg = (Aggref*)node; + ListCell* lc; + + foreach (lc, agg->aggdistinct) { + SortGroupClause* sort_clause = lfirst_node(SortGroupClause, lc); + TargetEntry* tle = + get_sortgroupref_tle(sort_clause->tleSortGroupRef, agg->args); + + if (!distinct_expr_is_safe((Node*)tle->expr)) { + return true; + } + } + } + return expression_tree_walker(node, setop_grouping_distinct_walker, context); +} + +static bool +setop_grouping_target_is_safe( + PlannerInfo* root, + RelOptInfo* input_rel, + RelOptInfo* grouped_rel, + List* setop_tlist +) { + ChFdwSetopGroupingValidationState state = { + .setop_tlist = setop_tlist, + .setop_columns = list_length(setop_tlist), + }; + PathTarget* target = grouped_rel->reltarget; + ListCell* lc; + int index = 0; + + (void)input_rel; + + if (setop_tlist == NIL || target == NULL || target->exprs == NIL || + setop_grouping_expr_walker((Node*)target->exprs, &state) || + setop_grouping_distinct_walker((Node*)target->exprs, NULL)) { + return false; + } + + foreach (lc, target->exprs) { + Node* expr = (Node*)lfirst(lc); + Index sortref = get_pathtarget_sortgroupref(target, index++); + + if (sortref != 0 && + get_sortgroupref_clause_noerr(sortref, root->parse->groupClause) != NULL && + !distinct_expr_is_safe(expr)) { + return false; + } + } + return true; +} + +static bool +setop_relation_uses_default_engine(RelOptInfo* rel) { + CHFdwRelationInfo* fpinfo; + + if (rel == NULL || rel->fdw_private == NULL) { + return false; + } + fpinfo = (CHFdwRelationInfo*)rel->fdw_private; + if (fpinfo->ch_table_engine != CH_DEFAULT) { + return false; + } + if (fpinfo->outerrel != NULL && fpinfo->outerrel != rel && + !setop_relation_uses_default_engine(fpinfo->outerrel)) { + return false; + } + if (fpinfo->innerrel != NULL && fpinfo->innerrel != rel && + !setop_relation_uses_default_engine(fpinfo->innerrel)) { + return false; + } + return true; +} + +static bool +setop_foreign_path_ok(const ForeignPath* path, void* arg) { + ChFdwSetopValidationState* state = (ChFdwSetopValidationState*)arg; + RelOptInfo* rel = path->path.parent; + CHFdwRelationInfo* fpinfo; + ListCell* lc; + + if (rel == NULL || rel->fdwroutine == NULL || + rel->fdwroutine->GetForeignPlan != clickhouseGetForeignPlan || + rel->fdw_private == NULL) { + return false; + } + + fpinfo = (CHFdwRelationInfo*)rel->fdw_private; + if (!fpinfo->pushdown_safe || fpinfo->local_conds != NIL || + (path->fdw_outerpath != NULL && !fpinfo->is_setop) || + setop_relation_requires_local_gate(rel)) { + return false; + } + + /* + * Base and join scans only project Vars remotely. Upper scans have + * already built and validated an explicit remote target list. + */ + if (!IS_UPPER_REL(rel)) { + foreach (lc, rel->reltarget->exprs) { + Var* var; + + if (!IsA(lfirst(lc), Var)) { + return false; + } + var = (Var*)lfirst(lc); + if (var->varlevelsup != 0 || var->varattno <= 0) { + return false; + } + } + } + + state->fetch_size = Max(state->fetch_size, fpinfo->fetch_size); + return true; +} + +static bool +is_flattened_union_all_relation(PlannerInfo* root, RelOptInfo* input_rel) { + Query* parse = root->parse; + RangeTblEntry* rte; + ListCell* lc; + int children = 0; + + if (input_rel == NULL || input_rel->reloptkind != RELOPT_BASEREL || + input_rel->rtekind != RTE_SUBQUERY || input_rel->relid <= 0 || + input_rel->baserestrictinfo != NIL || + !bms_is_empty(input_rel->lateral_relids) || parse->jointree == NULL || + list_length(parse->jointree->fromlist) != 1 || + !IsA(linitial(parse->jointree->fromlist), RangeTblRef) || + (int64)((RangeTblRef*)linitial(parse->jointree->fromlist))->rtindex != + (int64)input_rel->relid) { + return false; + } + + rte = planner_rt_fetch(input_rel->relid, root); + if (rte->rtekind != RTE_SUBQUERY || !rte->inh || /* codespell:ignore inh */ + rte->securityQuals != NIL) { + return false; + } + + foreach (lc, root->append_rel_list) { + AppendRelInfo* appinfo = lfirst_node(AppendRelInfo, lc); + + if (appinfo->parent_relid != input_rel->relid) { + continue; + } + if (OidIsValid(appinfo->parent_reltype) || OidIsValid(appinfo->child_reltype) || + OidIsValid(appinfo->parent_reloid)) { + return false; + } + children++; + } + + return children >= 2; +} + +static bool +is_flattened_union_all(PlannerInfo* root, RelOptInfo* input_rel) { + Query* parse = root->parse; + + return is_flattened_union_all_relation(root, input_rel) && + parse->commandType == CMD_SELECT && parse->setOperations == NULL && + parse->groupClause == NIL && parse->groupingSets == NIL && + parse->havingQual == NULL && parse->windowClause == NIL && + parse->distinctClause == NIL && parse->sortClause == NIL && + parse->limitOffset == NULL && parse->limitCount == NULL && + parse->rowMarks == NIL && !parse->hasAggs && !parse->hasTargetSRFs; +} + +static bool +is_flattened_union_all_grouping( + PlannerInfo* root, + RelOptInfo* input_rel, + GroupPathExtraData* extra +) { + Query* parse = root->parse; + + return extra != NULL && extra->havingQual == NULL && + is_flattened_union_all_relation(root, input_rel) && + parse->commandType == CMD_SELECT && parse->setOperations == NULL && + (parse->groupClause != NIL || parse->hasAggs) && + parse->groupingSets == NIL && parse->havingQual == NULL && + parse->windowClause == NIL && parse->distinctClause == NIL && + parse->sortClause == NIL && parse->limitOffset == NULL && + parse->limitCount == NULL && parse->rowMarks == NIL && !parse->hasTargetSRFs; +} + +static AppendPath* +make_setop_append_path(PlannerInfo* root, AppendPath* append_path) { +#if PG_VERSION_NUM >= 190000 + AppendPathInput input = { + .subpaths = list_copy(append_path->subpaths), + .partial_subpaths = NIL, + .child_append_relid_sets = list_copy(append_path->child_append_relid_sets), + }; + + return create_append_path( + root, + append_path->path.parent, + input, + append_path->path.pathkeys, + NULL, + 0, + false, + append_path->path.rows + ); +#else + return create_append_path( + root, + append_path->path.parent, + list_copy(append_path->subpaths), + NIL, + append_path->path.pathkeys, + NULL, + 0, + false, +#if PG_VERSION_NUM < 140000 + NIL, +#endif + append_path->path.rows + ); +#endif +} + +static bool +find_setop_path( + PlannerInfo* root, + RelOptInfo* rel, + AppendPath** append_path, + bool* is_distinct, + ChFdwSetopPathInfo* path_info, + int32* fetch_size, + Cost* startup_cost, + Cost* total_cost +) { + ListCell* lc; + bool found = false; + + foreach (lc, rel->pathlist) { + Path* path = (Path*)lfirst(lc); + AppendPath* candidate; + ChFdwSetopPathInfo candidate_info; + ChFdwSetopValidationState state = { 0 }; + bool candidate_distinct; + + candidate = chfdw_setop_find_append_path(path, &candidate_distinct); + if (candidate == NULL) { + Path* subpath = NULL; + +#if PG_VERSION_NUM >= 190000 + if (IsA(path, UniquePath)) { + subpath = ((UniquePath*)path)->subpath; + } +#else + if (IsA(path, UpperUniquePath)) { + subpath = ((UpperUniquePath*)path)->subpath; + } +#endif + if (subpath != NULL && IsA(subpath, MergeAppendPath)) { + MergeAppendPath* merge_path = (MergeAppendPath*)subpath; + AppendPath merge_append; + + memset(&merge_append, 0, sizeof(merge_append)); + merge_append.path.parent = rel; + merge_append.path.pathkeys = NIL; + merge_append.path.rows = path->rows; + merge_append.subpaths = merge_path->subpaths; +#if PG_VERSION_NUM >= 190000 + merge_append.child_append_relid_sets = + merge_path->child_append_relid_sets; +#endif + candidate = make_setop_append_path(root, &merge_append); + candidate_distinct = true; + } + } + if (candidate == NULL || + !chfdw_setop_extract_foreign_paths( + candidate, setop_foreign_path_ok, &state, &candidate_info + ) || + (candidate_distinct && + !distinct_target_is_safe(candidate->path.pathtarget))) { + continue; + } + + if (!found || path->total_cost < *total_cost) { + /* + * add_path() can discard the original candidate while retaining + * this ForeignPath. Keep a private AppendPath wrapper around its + * child paths for create_foreignscan_plan(). + */ + *append_path = make_setop_append_path(root, candidate); + *is_distinct = candidate_distinct; + *path_info = candidate_info; + *fetch_size = state.fetch_size; + *startup_cost = path->startup_cost; + *total_cost = path->total_cost; + found = true; + } + } + + return found; +} + +static void +add_foreign_setop_path( + PlannerInfo* root, + RelOptInfo* output_rel, + AppendPath* append_path, + bool is_distinct, + const ChFdwSetopPathInfo* path_info, + int32 fetch_size, + Cost startup_cost, + Cost total_cost +) { + CHFdwRelationInfo* fpinfo; + CHFdwRelationInfo* input_fpinfo; + ForeignPath* setop_path; + int arm_count = list_length(path_info->foreign_paths); + + if (output_rel->fdw_private != NULL || arm_count < 2) { + return; + } + + fpinfo = (CHFdwRelationInfo*)palloc0(sizeof(CHFdwRelationInfo)); + input_fpinfo = + (CHFdwRelationInfo*)((ForeignPath*)linitial(path_info->foreign_paths)) + ->path.parent->fdw_private; + + fpinfo->pushdown_safe = true; + fpinfo->stage = UPPERREL_SETOP; + fpinfo->is_setop = true; + fpinfo->setop_all = !is_distinct; + fpinfo->setop_userid = path_info->userid; + fpinfo->grouped_tlist = make_tlist_from_pathtarget(append_path->path.pathtarget); + fpinfo->server = GetForeignServer(path_info->serverid); + fpinfo->relation_name = makeStringInfo(); + appendStringInfo( + fpinfo->relation_name, + "%s on (%d remote queries)", + is_distinct ? "Union distinct" : "Union all", + arm_count + ); + merge_fdw_options(fpinfo, input_fpinfo, NULL); + fpinfo->fetch_size = Max(fpinfo->fetch_size, fetch_size); + + output_rel->fdw_private = fpinfo; + output_rel->fdwroutine = path_info->fdwroutine; + output_rel->serverid = path_info->serverid; + output_rel->userid = path_info->userid; + output_rel->useridiscurrent = path_info->useridiscurrent; + + total_cost = + Max(startup_cost, total_cost - DEFAULT_FDW_STARTUP_COST * (arm_count - 1)); + setop_path = create_foreign_upper_path( + root, + output_rel, + append_path->path.pathtarget, + append_path->path.rows, +#if PG_VERSION_NUM >= 180000 + 0, +#endif + startup_cost, + total_cost, + NIL, + (Path*)append_path, +#if PG_VERSION_NUM >= 170000 + NIL, +#endif + NIL + ); + add_path(output_rel, (Path*)setop_path); +} + +/* + * Add a grouping path whose input is a flattened, fully remote UNION ALL. + * The input relation is a PostgreSQL subquery rather than an FDW upper + * relation, so give it enough synthetic FDW metadata for the existing + * grouping safety checks while retaining the original AppendPath for plan + * construction. + */ +static void +add_foreign_setop_grouping_path( + PlannerInfo* root, + RelOptInfo* input_rel, + RelOptInfo* grouped_rel, + AppendPath* append_path, + const ChFdwSetopPathInfo* path_info, + int32 fetch_size, + Cost startup_cost, + Cost total_cost +) { + CHFdwRelationInfo* setop_fpinfo; + CHFdwRelationInfo* grouped_fpinfo; + CHFdwRelationInfo* arm_fpinfo; + ForeignPath* grouped_path; + List* setop_tlist; + double rows; + int width; + int arm_count = list_length(path_info->foreign_paths); + void* old_input_private; + FdwRoutine* old_input_routine; + Oid old_input_serverid; + Oid old_input_userid; + bool old_input_useridiscurrent; + ListCell* lc; + + if (arm_count < 2 || input_rel->fdw_private != NULL || + grouped_rel->fdw_private != NULL || append_path->path.pathkeys != NIL || + append_path->path.pathtarget == NULL || + append_path->path.pathtarget->exprs == NIL) { + return; + } + + arm_fpinfo = (CHFdwRelationInfo*)((ForeignPath*)linitial(path_info->foreign_paths)) + ->path.parent->fdw_private; + setop_tlist = make_tlist_from_pathtarget(append_path->path.pathtarget); + if (!setop_grouping_target_is_safe(root, input_rel, grouped_rel, setop_tlist)) { + return; + } + foreach (lc, path_info->foreign_paths) { + ForeignPath* arm_path = lfirst_node(ForeignPath, lc); + + if (!setop_relation_uses_default_engine(arm_path->path.parent)) { + return; + } + } + + setop_fpinfo = palloc0(sizeof(CHFdwRelationInfo)); + setop_fpinfo->pushdown_safe = true; + setop_fpinfo->stage = UPPERREL_SETOP; + setop_fpinfo->is_setop = true; + setop_fpinfo->setop_all = true; + setop_fpinfo->setop_userid = path_info->userid; + setop_fpinfo->grouped_tlist = setop_tlist; + setop_fpinfo->server = GetForeignServer(path_info->serverid); + setop_fpinfo->relation_name = makeStringInfo(); + setop_fpinfo->fetch_size = fetch_size; + appendStringInfo( + setop_fpinfo->relation_name, "Union all on (%d remote queries)", arm_count + ); + merge_fdw_options(setop_fpinfo, arm_fpinfo, NULL); + setop_fpinfo->fetch_size = Max(setop_fpinfo->fetch_size, fetch_size); + + old_input_private = input_rel->fdw_private; + old_input_routine = input_rel->fdwroutine; + old_input_serverid = input_rel->serverid; + old_input_userid = input_rel->userid; + old_input_useridiscurrent = input_rel->useridiscurrent; + + input_rel->fdw_private = setop_fpinfo; + input_rel->fdwroutine = path_info->fdwroutine; + input_rel->serverid = path_info->serverid; + input_rel->userid = path_info->userid; + input_rel->useridiscurrent = path_info->useridiscurrent; + + grouped_fpinfo = palloc0(sizeof(CHFdwRelationInfo)); + grouped_fpinfo->stage = UPPERREL_GROUP_AGG; + grouped_fpinfo->outerrel = input_rel; + grouped_fpinfo->table = setop_fpinfo->table; + grouped_fpinfo->server = setop_fpinfo->server; + grouped_fpinfo->user = setop_fpinfo->user; + grouped_fpinfo->is_setop_grouping = true; + grouped_fpinfo->setop_all = true; + grouped_fpinfo->setop_userid = path_info->userid; + grouped_fpinfo->setop_tlist = copyObject(setop_tlist); + merge_fdw_options(grouped_fpinfo, setop_fpinfo, NULL); + grouped_fpinfo->fetch_size = Max(grouped_fpinfo->fetch_size, fetch_size); + grouped_rel->fdw_private = grouped_fpinfo; + + if (!foreign_grouping_ok(root, grouped_rel, NULL) || + setop_grouping_expr_walker( + (Node*)grouped_fpinfo->grouped_tlist, + &(ChFdwSetopGroupingValidationState){ + .setop_tlist = setop_tlist, + .setop_columns = list_length(setop_tlist), + } + )) { + grouped_rel->fdw_private = NULL; + input_rel->fdw_private = old_input_private; + input_rel->fdwroutine = old_input_routine; + input_rel->serverid = old_input_serverid; + input_rel->userid = old_input_userid; + input_rel->useridiscurrent = old_input_useridiscurrent; + return; + } + + input_rel->fdw_private = old_input_private; + input_rel->fdwroutine = old_input_routine; + input_rel->serverid = old_input_serverid; + input_rel->userid = old_input_userid; + input_rel->useridiscurrent = old_input_useridiscurrent; + + grouped_rel->fdwroutine = path_info->fdwroutine; + grouped_rel->serverid = path_info->serverid; + grouped_rel->userid = path_info->userid; + grouped_rel->useridiscurrent = path_info->useridiscurrent; + + estimate_path_cost_size(&rows, &width, &startup_cost, &total_cost, 0.1); + grouped_fpinfo->rows = rows; + grouped_fpinfo->width = width; + grouped_fpinfo->startup_cost = startup_cost; + grouped_fpinfo->total_cost = total_cost; + + grouped_path = create_foreign_upper_path( + root, + grouped_rel, + grouped_rel->reltarget, + rows, +#if PG_VERSION_NUM >= 180000 + 0, +#endif + startup_cost, + total_cost, + NIL, + (Path*)append_path, +#if PG_VERSION_NUM >= 170000 + NIL, +#endif + NIL + ); + add_path(grouped_rel, (Path*)grouped_path); +} + +/* + * PostgreSQL does not call an FDW's GetForeignUpperPaths callback for set + * operations, and it rewrites simple UNION ALL into an append relation. + * The global upper-path hook covers both planner representations. + */ +void +chfdw_create_upper_paths_hook( + PlannerInfo* root, + UpperRelationKind stage, + RelOptInfo* input_rel, + RelOptInfo* output_rel, + void* extra +) { + AppendPath* append_path; + ChFdwSetopPathInfo path_info; + bool is_distinct; + int32 fetch_size; + Cost startup_cost; + Cost total_cost; + + if (stage == UPPERREL_SETOP) { + if (!find_setop_path( + root, + output_rel, + &append_path, + &is_distinct, + &path_info, + &fetch_size, + &startup_cost, + &total_cost + )) { + return; + } + /* + * An ordered DISTINCT set operation needs the original setop target + * for its local pathkeys. Keep the whole operation local until the + * remote setop path can preserve that target exactly. + */ + if (is_distinct && root->parse->sortClause != NIL) { + return; + } + } else if (stage == UPPERREL_FINAL && is_flattened_union_all(root, input_rel)) { + if (!find_setop_path( + root, + input_rel, + &append_path, + &is_distinct, + &path_info, + &fetch_size, + &startup_cost, + &total_cost + ) || + is_distinct) { + return; + } + } else if ( + stage == UPPERREL_GROUP_AGG && + is_flattened_union_all_grouping(root, input_rel, (GroupPathExtraData*)extra) + ) { + if (!find_setop_path( + root, + input_rel, + &append_path, + &is_distinct, + &path_info, + &fetch_size, + &startup_cost, + &total_cost + ) || + is_distinct) { + return; + } + add_foreign_setop_grouping_path( + root, + input_rel, + output_rel, + append_path, + &path_info, + fetch_size, + startup_cost, + total_cost + ); + return; + } else { + return; + } + + add_foreign_setop_path( + root, + output_rel, + append_path, + is_distinct, + &path_info, + fetch_size, + startup_cost, + total_cost + ); +} + +/* + * clickhouseGetForeignUpperPaths + * Add paths for post-join operations like aggregation, grouping etc. if + * corresponding operations are safe to push down. + * + * Right now, we only support aggregate, grouping and having clause pushdown. + */ +static void +clickhouseGetForeignUpperPaths( + PlannerInfo* root, + UpperRelationKind stage, + RelOptInfo* input_rel, + RelOptInfo* output_rel, + void* extra +) { + CHFdwRelationInfo* fpinfo; + struct timeval time1, time2; + + gettimeofday(&time1, NULL); + + /* * If input rel is not safe to pushdown, then simply return as we cannot * perform any post-join operations on the foreign server. */ @@ -2922,10 +4240,22 @@ clickhouseGetForeignUpperPaths( return; } + /* + * Set operations are represented by a complete remote query. Keep later + * ORDER BY/LIMIT stages local until they have dedicated wrapper deparse. + */ + if (((CHFdwRelationInfo*)input_rel->fdw_private)->is_setop) { + return; + } + if (IS_UPPER_REL(input_rel) && + ((CHFdwRelationInfo*)input_rel->fdw_private)->stage == UPPERREL_DISTINCT) { + return; + } + /* Ignore stages we don't support; and skip any duplicate calls. */ if ((stage != UPPERREL_GROUP_AGG && stage != UPPERREL_PARTIAL_GROUP_AGG && - stage != UPPERREL_WINDOW && stage != UPPERREL_ORDERED && - stage != UPPERREL_FINAL) || + stage != UPPERREL_WINDOW && stage != UPPERREL_DISTINCT && + stage != UPPERREL_ORDERED && stage != UPPERREL_FINAL) || output_rel->fdw_private) { return; } @@ -2951,6 +4281,9 @@ clickhouseGetForeignUpperPaths( case UPPERREL_WINDOW: add_foreign_window_paths(root, input_rel, output_rel); break; + case UPPERREL_DISTINCT: + add_foreign_distinct_paths(root, input_rel, output_rel); + break; case UPPERREL_ORDERED: add_foreign_ordered_paths(root, input_rel, output_rel); break; @@ -3085,6 +4418,93 @@ add_foreign_grouping_paths( add_path(grouped_rel, (Path*)grouppath); } +/* + * Add a path for plain SELECT DISTINCT when every output expression has + * matching equality semantics on PostgreSQL and ClickHouse. + */ +static void +add_foreign_distinct_paths( + PlannerInfo* root, + RelOptInfo* input_rel, + RelOptInfo* distinct_rel +) { + Query* parse = root->parse; + CHFdwRelationInfo* ifpinfo = input_rel->fdw_private; + CHFdwRelationInfo* fpinfo = distinct_rel->fdw_private; + PathTarget* target = root->upper_targets[UPPERREL_DISTINCT]; + ForeignPath* distinct_path; + List* tlist = NIL; + ListCell* lc; + double rows; + int width; + Cost startup_cost; + Cost total_cost; + int index = 0; + + if ((input_rel->reloptkind != RELOPT_BASEREL && + input_rel->reloptkind != RELOPT_JOINREL) || + parse->distinctClause == NIL || parse->hasDistinctOn || parse->hasTargetSRFs || + ifpinfo->local_conds || !distinct_target_is_safe(target) || + !setop_relation_uses_default_engine(input_rel)) { + return; + } + + fpinfo->outerrel = input_rel; + fpinfo->table = ifpinfo->table; + fpinfo->server = ifpinfo->server; + fpinfo->user = ifpinfo->user; + merge_fdw_options(fpinfo, ifpinfo, NULL); + + foreach (lc, target->exprs) { + Expr* expr = lfirst_node(Expr, lc); + Index sgref = get_pathtarget_sortgroupref(target, index++); + TargetEntry* tle; + + if (!chfdw_is_foreign_expr(root, distinct_rel, expr, true) || + is_foreign_param(root, distinct_rel, expr)) { + return; + } + + tle = makeTargetEntry(expr, list_length(tlist) + 1, NULL, false); + tle->ressortgroupref = sgref; + tlist = lappend(tlist, tle); + } + + fpinfo->grouped_tlist = tlist; + fpinfo->pushdown_safe = true; + fpinfo->rel_startup_cost = -1; + fpinfo->rel_total_cost = -1; + fpinfo->relation_name = makeStringInfo(); + appendStringInfo( + fpinfo->relation_name, "Distinct on (%s)", ifpinfo->relation_name->data + ); + + estimate_path_cost_size(&rows, &width, &startup_cost, &total_cost, 0.1); + fpinfo->rows = rows; + fpinfo->width = width; + fpinfo->startup_cost = startup_cost; + fpinfo->total_cost = total_cost; + + distinct_path = create_foreign_upper_path( + root, + distinct_rel, + target, + rows, +#if PG_VERSION_NUM >= 180000 + 0, +#endif + startup_cost, + total_cost, + NIL, + NULL, +#if PG_VERSION_NUM >= 170000 + NIL, +#endif + NIL + ); + add_path(distinct_rel, (Path*)distinct_path); +} + /* * foreign_window_ok * Assess whether window functions in the query are safe to push down. @@ -3189,6 +4609,10 @@ add_foreign_window_paths( Cost startup_cost; Cost total_cost; + if (root->parse->distinctClause != NIL) { + return; + } + /* Save the input_rel as outerrel in fpinfo */ fpinfo->outerrel = input_rel; diff --git a/src/include/fdw.h b/src/include/fdw.h index 221ccc5f..c516a85a 100644 --- a/src/include/fdw.h +++ b/src/include/fdw.h @@ -223,6 +223,13 @@ typedef struct CHFdwRelationInfo { /* Grouping information */ List* grouped_tlist; + /* Set-operation information */ + bool is_setop; + bool setop_all; + bool is_setop_grouping; + Oid setop_userid; + List* setop_tlist; + /* Subquery information */ bool make_outerrel_subquery; /* do we deparse outerrel as a * subquery? */ @@ -248,6 +255,14 @@ typedef struct CHFdwRelationInfo { } CHFdwRelationInfo; /* in fdw.c */ +extern void +chfdw_create_upper_paths_hook( + PlannerInfo* root, + UpperRelationKind stage, + RelOptInfo* input_rel, + RelOptInfo* output_rel, + void* extra +); extern ForeignServer* chfdw_get_foreign_server(Relation rel); extern Expr* @@ -353,6 +368,17 @@ chfdw_deparse_select_stmt_for_rel( List** retrieved_attrs, List** params_list ); +extern void +chfdw_deparse_setop_grouping_stmt( + StringInfo buf, + PlannerInfo* root, + RelOptInfo* rel, + List* tlist, + List* setop_tlist, + const char* setop_sql, + List** retrieved_attrs, + List** params_list +); extern const char* chfdw_get_jointype_name(JoinType jointype); char* diff --git a/src/include/setop.h b/src/include/setop.h new file mode 100644 index 00000000..9f3164af --- /dev/null +++ b/src/include/setop.h @@ -0,0 +1,49 @@ +/*------------------------------------------------------------------------- + * + * setop.h + * Planner helpers for set-operation pushdown + * + * Copyright (c) 2025-2026, ClickHouse, Inc. + * + * IDENTIFICATION + * github.com/clickhouse/pg_clickhouse/src/include/setop.h + * + *------------------------------------------------------------------------- + */ + +#ifndef CLICKHOUSE_SETOP_H +#define CLICKHOUSE_SETOP_H + +#include "foreign/fdwapi.h" +#include "nodes/pathnodes.h" +#include "nodes/plannodes.h" + +typedef bool (*ChFdwSetopPathValidator)(const ForeignPath* path, void* arg); + +typedef struct ChFdwSetopPathInfo { + List* foreign_paths; + FdwRoutine* fdwroutine; + Oid serverid; + Oid userid; + bool useridiscurrent; +} ChFdwSetopPathInfo; + +extern AppendPath* +chfdw_setop_find_append_path(Path* path, bool* is_distinct); + +extern bool +chfdw_setop_extract_foreign_paths( + AppendPath* append_path, + ChFdwSetopPathValidator validator, + void* validator_arg, + ChFdwSetopPathInfo* path_info +); + +extern bool +chfdw_setop_extract_foreign_scans( + Plan* plan, + Oid expected_serverid, + List** foreign_scans +); + +#endif /* CLICKHOUSE_SETOP_H */ diff --git a/src/option.c b/src/option.c index 130eb8cd..03ea3c3b 100644 --- a/src/option.c +++ b/src/option.c @@ -24,16 +24,20 @@ #include "commands/defrem.h" #include "commands/extension.h" #include "nodes/makefuncs.h" +#include "optimizer/planner.h" #include "utils/builtins.h" #include "utils/guc.h" #include "utils/varlena.h" -static char* DEFAULT_DBNAME = "default"; +static char* DEFAULT_DBNAME = "default"; +static create_upper_paths_hook_type prev_create_upper_paths_hook = NULL; #if PG_VERSION_NUM < 160000 extern PGDLLEXPORT void _PG_init(void); #endif +extern PGDLLEXPORT void +_PG_fini(void); /* * Describes the valid options for objects that this wrapper uses. @@ -667,6 +671,21 @@ chfdw_settings_assign_hook(const char* newval, void* extra) { ch_session_settings_list = (kv_list*)extra; } +static void +chfdw_upper_paths_hook( + PlannerInfo* root, + UpperRelationKind stage, + RelOptInfo* input_rel, + RelOptInfo* output_rel, + void* extra +) { + if (prev_create_upper_paths_hook) { + prev_create_upper_paths_hook(root, stage, input_rel, output_rel, extra); + } + + chfdw_create_upper_paths_hook(root, stage, input_rel, output_rel, extra); +} + /* * Module load callback */ @@ -716,4 +735,17 @@ _PG_init(void) { #if PG_VERSION_NUM >= 150000 MarkGUCPrefixReserved("pg_clickhouse"); #endif + + prev_create_upper_paths_hook = create_upper_paths_hook; + create_upper_paths_hook = chfdw_upper_paths_hook; +} + +/* + * Module unload callback + */ +void +_PG_fini(void) { + if (create_upper_paths_hook == chfdw_upper_paths_hook) { + create_upper_paths_hook = prev_create_upper_paths_hook; + } } diff --git a/src/setop.c b/src/setop.c new file mode 100644 index 00000000..707993af --- /dev/null +++ b/src/setop.c @@ -0,0 +1,595 @@ +/*------------------------------------------------------------------------- + * + * setop.c + * Planner helpers for set-operation pushdown + * + * Copyright (c) 2025-2026, ClickHouse, Inc. + * + * IDENTIFICATION + * github.com/clickhouse/pg_clickhouse/src/setop.c + * + *------------------------------------------------------------------------- + */ + +#include "postgres.h" + +#include "setop.h" + +#include "nodes/nodeFuncs.h" +#include "optimizer/tlist.h" +#include "parser/parsetree.h" + +typedef struct ChFdwSetopPathState { + ChFdwSetopPathValidator validator; + void* validator_arg; + ChFdwSetopPathInfo path_info; +} ChFdwSetopPathState; + +static bool +chfdw_setop_extract_path_member(Path* path, ChFdwSetopPathState* state); +static bool +chfdw_setop_extract_plan_member( + Plan* plan, + List* expected_tlist, + Oid expected_serverid, + List** foreign_scans +); + +static bool +chfdw_setop_path_has_no_parameters(const Path* path) { + return path->param_info == NULL && !path->parallel_aware; +} + +static bool +chfdw_setop_plan_has_no_parameters(const Plan* plan) { + return plan->initPlan == NIL && bms_is_empty(plan->extParam) && + bms_is_empty(plan->allParam) && !plan->parallel_aware; +} + +static bool +chfdw_setop_path_targets_equal(const PathTarget* left, const PathTarget* right) { + return left != NULL && right != NULL && equal(left->exprs, right->exprs); +} + +static bool +chfdw_setop_expr_signatures_equal(const Node* left, const Node* right) { + return exprType(left) == exprType(right) && exprTypmod(left) == exprTypmod(right) && + exprCollation(left) == exprCollation(right); +} + +static bool +chfdw_setop_path_signatures_equal(const Path* left, const Path* right) { + ListCell* left_lc; + ListCell* right_lc; + + if (left->pathtarget == NULL || right->pathtarget == NULL || + list_length(left->pathtarget->exprs) != list_length(right->pathtarget->exprs)) { + return false; + } + + forboth(left_lc, left->pathtarget->exprs, right_lc, right->pathtarget->exprs) { + if (!chfdw_setop_expr_signatures_equal( + (Node*)lfirst(left_lc), (Node*)lfirst(right_lc) + )) { + return false; + } + } + + return true; +} + +static bool +chfdw_setop_tlist_signatures_equal(const List* left, const List* right) { + ListCell* left_lc; + ListCell* right_lc; + List* left_exprs = NIL; + List* right_exprs = NIL; + + foreach (left_lc, left) { + TargetEntry* tle = lfirst_node(TargetEntry, left_lc); + + if (!tle->resjunk) { + left_exprs = lappend(left_exprs, tle->expr); + } + } + + foreach (right_lc, right) { + TargetEntry* tle = lfirst_node(TargetEntry, right_lc); + + if (!tle->resjunk) { + right_exprs = lappend(right_exprs, tle->expr); + } + } + + if (list_length(left_exprs) != list_length(right_exprs)) { + return false; + } + + forboth(left_lc, left_exprs, right_lc, right_exprs) { + if (!chfdw_setop_expr_signatures_equal( + (Node*)lfirst(left_lc), (Node*)lfirst(right_lc) + )) { + return false; + } + } + + return true; +} + +static bool +chfdw_setop_tlists_equal(const List* left, const List* right) { + ListCell* left_lc; + ListCell* right_lc; + + if (list_length(left) != list_length(right)) { + return false; + } + + forboth(left_lc, left, right_lc, right) { + TargetEntry* left_tle = lfirst_node(TargetEntry, left_lc); + TargetEntry* right_tle = lfirst_node(TargetEntry, right_lc); + + if (left_tle->resjunk != right_tle->resjunk || + !equal(left_tle->expr, right_tle->expr)) { + return false; + } + } + + return true; +} + +static bool +chfdw_setop_subquery_path_is_trivial(const SubqueryScanPath* path) { + RelOptInfo* rel = path->path.parent; + ListCell* outer_lc; + bool identity = true; + int output_column = 0; + + if (rel == NULL || + (rel->reloptkind != RELOPT_BASEREL && + rel->reloptkind != RELOPT_OTHER_MEMBER_REL) || + rel->rtekind != RTE_SUBQUERY || rel->baserestrictinfo != NIL || + !bms_is_empty(rel->lateral_relids) || rel->subplan_params != NIL || + path->subpath == NULL || + list_length(path->path.pathtarget->exprs) != + list_length(path->subpath->pathtarget->exprs)) { + return false; + } + + foreach (outer_lc, path->path.pathtarget->exprs) { + Node* outer_expr = (Node*)lfirst(outer_lc); + Node* inner_expr; + Var* var; + + output_column++; + if (!IsA(outer_expr, Var)) { + return false; + } + + var = (Var*)outer_expr; + if ((int64)var->varno != (int64)rel->relid || var->varattno <= 0 || + var->varattno > list_length(path->subpath->pathtarget->exprs) || + var->varlevelsup != 0) { + return false; + } + inner_expr = + (Node*)list_nth(path->subpath->pathtarget->exprs, var->varattno - 1); + if (!chfdw_setop_expr_signatures_equal(outer_expr, inner_expr)) { + return false; + } + + if (var->varattno != output_column) { + identity = false; + } + } + + return identity || IsA(path->subpath, ForeignPath); +} + +static bool +chfdw_setop_subquery_plan_is_trivial(const SubqueryScan* scan) { + ListCell* outer_lc; + List* outer_tlist; + List* inner_tlist; + + if (scan->subplan == NULL || scan->scan.plan.qual != NIL || + !chfdw_setop_plan_has_no_parameters(&scan->scan.plan)) { + return false; + } + + outer_tlist = scan->scan.plan.targetlist; + inner_tlist = scan->subplan->targetlist; + + if (list_length(outer_tlist) != list_length(inner_tlist)) { + return false; + } + + foreach (outer_lc, outer_tlist) { + TargetEntry* outer_tle = lfirst_node(TargetEntry, outer_lc); + TargetEntry* inner_tle; + Var* var; + + if (!IsA(outer_tle->expr, Var)) { + return false; + } + + var = (Var*)outer_tle->expr; + if ((int64)var->varno != (int64)scan->scan.scanrelid || var->varattno <= 0 || + var->varlevelsup != 0) { + return false; + } + inner_tle = get_tle_by_resno(inner_tlist, var->varattno); + if (inner_tle == NULL || outer_tle->resjunk != inner_tle->resjunk || + var->varlevelsup != 0 || + !chfdw_setop_expr_signatures_equal( + (Node*)outer_tle->expr, (Node*)inner_tle->expr + )) { + return false; + } + } + + return true; +} + +static bool +chfdw_setop_fdw_routines_equal(const FdwRoutine* left, const FdwRoutine* right) { + return left == right || (left != NULL && right != NULL && + left->GetForeignPlan == right->GetForeignPlan && + left->BeginForeignScan == right->BeginForeignScan && + left->IterateForeignScan == right->IterateForeignScan); +} + +AppendPath* +chfdw_setop_find_append_path(Path* path, bool* is_distinct) { + Path* subpath; + + if (is_distinct == NULL || path == NULL || + !chfdw_setop_path_has_no_parameters(path)) { + return NULL; + } + + *is_distinct = false; + + if (IsA(path, AppendPath)) { + return (AppendPath*)path; + } + + if (IsA(path, AggPath)) { + AggPath* agg_path = (AggPath*)path; + + if (agg_path->aggstrategy != AGG_HASHED || + agg_path->aggsplit != AGGSPLIT_SIMPLE || agg_path->qual != NIL || + list_length(agg_path->groupClause) != + list_length(path->pathtarget->exprs)) { + return NULL; + } + + subpath = agg_path->subpath; + } +#if PG_VERSION_NUM >= 190000 + else if (IsA(path, UniquePath)) { + UniquePath* unique_path = (UniquePath*)path; + + if (unique_path->numkeys != list_length(path->pathkeys)) { + return NULL; + } + + subpath = unique_path->subpath; +#else + else if (IsA(path, UpperUniquePath)) { + UpperUniquePath* unique_path = (UpperUniquePath*)path; + + if (unique_path->numkeys != list_length(path->pathkeys)) { + return NULL; + } + + subpath = unique_path->subpath; +#endif + } else { + return NULL; + } + + if (IsA(subpath, SortPath)) { + SortPath* sort_path = (SortPath*)subpath; + + if (!chfdw_setop_path_has_no_parameters(subpath) || + !chfdw_setop_path_targets_equal( + subpath->pathtarget, sort_path->subpath->pathtarget + )) { + return NULL; + } + + subpath = sort_path->subpath; + } + + if (!IsA(subpath, AppendPath)) { + return NULL; + } + + *is_distinct = true; + return (AppendPath*)subpath; +} + +static bool +chfdw_setop_extract_append_path(AppendPath* append_path, ChFdwSetopPathState* state) { + ListCell* lc; + int path_count; + + path_count = list_length(append_path->subpaths); + if (path_count < 2 || IS_PARTITIONED_REL(append_path->path.parent) || + append_path->first_partial_path != path_count || + append_path->limit_tuples >= 0 || + !chfdw_setop_path_has_no_parameters(&append_path->path)) { + return false; + } + + foreach (lc, append_path->subpaths) { + Path* subpath = (Path*)lfirst(lc); + + if (!chfdw_setop_path_signatures_equal(&append_path->path, subpath) || + !chfdw_setop_extract_path_member(subpath, state)) { + return false; + } + } + + return true; +} + +static bool +chfdw_setop_extract_foreign_path( + ForeignPath* foreign_path, + ChFdwSetopPathState* state +) { + RelOptInfo* rel = foreign_path->path.parent; + ChFdwSetopPathInfo* path_info = &state->path_info; + + if (rel == NULL || rel->fdwroutine == NULL || !OidIsValid(rel->serverid) || + !chfdw_setop_path_has_no_parameters(&foreign_path->path) +#if PG_VERSION_NUM >= 170000 + || foreign_path->fdw_restrictinfo != NIL +#endif + || !state->validator(foreign_path, state->validator_arg)) { + return false; + } + + if (path_info->foreign_paths == NIL) { + path_info->fdwroutine = rel->fdwroutine; + path_info->serverid = rel->serverid; + path_info->userid = rel->userid; + path_info->useridiscurrent = rel->useridiscurrent; + } else if ( + path_info->serverid != rel->serverid || path_info->userid != rel->userid || + path_info->useridiscurrent != rel->useridiscurrent || + !chfdw_setop_fdw_routines_equal(path_info->fdwroutine, rel->fdwroutine) + ) { + return false; + } + + path_info->foreign_paths = lappend(path_info->foreign_paths, foreign_path); + return true; +} + +static bool +chfdw_setop_extract_path_member(Path* path, ChFdwSetopPathState* state) { + if (path == NULL || !chfdw_setop_path_has_no_parameters(path)) { + return false; + } + + if (IsA(path, AppendPath)) { + return chfdw_setop_extract_append_path((AppendPath*)path, state); + } + + if (IsA(path, MaterialPath)) { + MaterialPath* material_path = (MaterialPath*)path; + + return material_path->subpath != NULL && + chfdw_setop_path_targets_equal( + path->pathtarget, material_path->subpath->pathtarget + ) && + chfdw_setop_extract_path_member(material_path->subpath, state); + } + + if (IsA(path, ProjectionPath)) { + ProjectionPath* projection_path = (ProjectionPath*)path; + + return projection_path->subpath != NULL && + chfdw_setop_path_targets_equal( + path->pathtarget, projection_path->subpath->pathtarget + ) && + chfdw_setop_extract_path_member(projection_path->subpath, state); + } + + if (IsA(path, SubqueryScanPath)) { + SubqueryScanPath* subquery_path = (SubqueryScanPath*)path; + + return chfdw_setop_subquery_path_is_trivial(subquery_path) && + chfdw_setop_extract_path_member(subquery_path->subpath, state); + } + + if (IsA(path, ForeignPath)) { + return chfdw_setop_extract_foreign_path((ForeignPath*)path, state); + } + + return false; +} + +bool +chfdw_setop_extract_foreign_paths( + AppendPath* append_path, + ChFdwSetopPathValidator validator, + void* validator_arg, + ChFdwSetopPathInfo* path_info +) { + ChFdwSetopPathState state; + + if (append_path == NULL || validator == NULL || path_info == NULL) { + return false; + } + + memset(&state, 0, sizeof(state)); + state.validator = validator; + state.validator_arg = validator_arg; + + if (!chfdw_setop_extract_append_path(append_path, &state) || + list_length(state.path_info.foreign_paths) < 2) { + return false; + } + + *path_info = state.path_info; + return true; +} + +static bool +chfdw_setop_extract_append_plan( + Append* append, + Oid expected_serverid, + List** foreign_scans +) { + ListCell* lc; + + if ( + list_length(append->appendplans) < 2 || + append->first_partial_plan != list_length(append->appendplans) || + append->plan.qual != NIL || append->plan.lefttree != NULL || + append->plan.righttree != NULL || + !chfdw_setop_plan_has_no_parameters(&append->plan) +#if PG_VERSION_NUM >= 180000 + || append->part_prune_index >= 0 +#else + || append->part_prune_info != NULL +#endif + ) { + return false; + } + + foreach (lc, append->appendplans) { + Plan* subplan = lfirst_node(Plan, lc); + + if (!chfdw_setop_tlist_signatures_equal( + append->plan.targetlist, subplan->targetlist + ) || + !chfdw_setop_extract_plan_member( + subplan, append->plan.targetlist, expected_serverid, foreign_scans + )) { + return false; + } + } + + return true; +} + +static bool +chfdw_setop_extract_foreign_scan( + ForeignScan* foreign_scan, + Oid expected_serverid, + List** foreign_scans +) { + Plan* plan = &foreign_scan->scan.plan; + + if (foreign_scan->operation != CMD_SELECT || + foreign_scan->fs_server != expected_serverid || plan->qual != NIL || + plan->lefttree != NULL || plan->righttree != NULL || + !chfdw_setop_plan_has_no_parameters(plan) || foreign_scan->fdw_private == NIL) { + return false; + } + + *foreign_scans = lappend(*foreign_scans, foreign_scan); + return true; +} + +static bool +chfdw_setop_extract_plan_member( + Plan* plan, + List* expected_tlist, + Oid expected_serverid, + List** foreign_scans +) { + if (plan == NULL || !chfdw_setop_plan_has_no_parameters(plan) || + !chfdw_setop_tlist_signatures_equal(expected_tlist, plan->targetlist)) { + return false; + } + + if (IsA(plan, Append)) { + return chfdw_setop_extract_append_plan( + (Append*)plan, expected_serverid, foreign_scans + ); + } + + if (IsA(plan, Material)) { + return plan->qual == NIL && plan->lefttree != NULL && plan->righttree == NULL && + chfdw_setop_tlists_equal(plan->targetlist, plan->lefttree->targetlist) && + chfdw_setop_extract_plan_member( + plan->lefttree, plan->targetlist, expected_serverid, foreign_scans + ); + } + + if (IsA(plan, Result)) { + Result* result = (Result*)plan; + + return result->resconstantqual == NULL && plan->qual == NIL && + plan->lefttree != NULL && plan->righttree == NULL && + chfdw_setop_tlists_equal(plan->targetlist, plan->lefttree->targetlist) && + chfdw_setop_extract_plan_member( + plan->lefttree, plan->targetlist, expected_serverid, foreign_scans + ); + } + + if (IsA(plan, SubqueryScan)) { + SubqueryScan* subquery_scan = (SubqueryScan*)plan; + List* projected_tlist = NIL; + ListCell* lc; + + if (!chfdw_setop_subquery_plan_is_trivial(subquery_scan)) { + return false; + } + foreach (lc, subquery_scan->scan.plan.targetlist) { + TargetEntry* outer_tle = lfirst_node(TargetEntry, lc); + Var* var = castNode(Var, outer_tle->expr); + TargetEntry* inner_tle = + get_tle_by_resno(subquery_scan->subplan->targetlist, var->varattno); + TargetEntry* projected_tle = copyObject(inner_tle); + + projected_tle->resno = list_length(projected_tlist) + 1; + projected_tle->resname = + outer_tle->resname ? pstrdup(outer_tle->resname) : NULL; + projected_tle->resjunk = outer_tle->resjunk; + projected_tlist = lappend(projected_tlist, projected_tle); + } + /* + * This is the private outer plan built only for this ForeignPath. + * The combined ForeignScan replaces it, so normalize the throwaway + * child target list in place to retain the SubqueryScan projection. + */ + subquery_scan->subplan->targetlist = projected_tlist; + return chfdw_setop_extract_plan_member( + subquery_scan->subplan, projected_tlist, expected_serverid, foreign_scans + ); + } + + if (IsA(plan, ForeignScan)) { + return chfdw_setop_extract_foreign_scan( + (ForeignScan*)plan, expected_serverid, foreign_scans + ); + } + + return false; +} + +bool +chfdw_setop_extract_foreign_scans( + Plan* plan, + Oid expected_serverid, + List** foreign_scans +) { + List* scans = NIL; + + if (plan == NULL || !OidIsValid(expected_serverid) || foreign_scans == NULL || + !IsA(plan, Append) || + !chfdw_setop_extract_append_plan((Append*)plan, expected_serverid, &scans) || + list_length(scans) < 2) { + return false; + } + + *foreign_scans = scans; + return true; +} diff --git a/test/expected/result_map.txt b/test/expected/result_map.txt index 95b43f20..6bddfe8e 100644 --- a/test/expected/result_map.txt +++ b/test/expected/result_map.txt @@ -297,6 +297,17 @@ timezone.sql 25.8+ | timezone.out 23-25.7 | timezone_1.out +union_pushdown.sql +------------------ + + Postgres | File +----------|------------------- + 13-19 | union_pushdown.out + + ClickHouse | File +------------|------------------- + 23+ | union_pushdown.out + where_sub.sql ------------- diff --git a/test/expected/union_pushdown.out b/test/expected/union_pushdown.out new file mode 100644 index 00000000..f054321f --- /dev/null +++ b/test/expected/union_pushdown.out @@ -0,0 +1,574 @@ +-- Tests for UNION and safely inlined CTE pushdown. +-- +-- Plan checks use JSON instead of printing the full EXPLAIN tree so the +-- expected output remains stable across PostgreSQL 13-19. +CREATE SERVER union_svr FOREIGN DATA WRAPPER clickhouse_fdw + OPTIONS (dbname 'union_pushdown_test', driver 'binary'); +CREATE USER MAPPING FOR CURRENT_USER SERVER union_svr; +CREATE SERVER union_http_svr FOREIGN DATA WRAPPER clickhouse_fdw + OPTIONS (dbname 'union_pushdown_test', driver 'http'); +CREATE USER MAPPING FOR CURRENT_USER SERVER union_http_svr; +-- A distinct server object pointing at the same ClickHouse database verifies +-- that server identity, rather than matching connection options, gates UNION +-- pushdown. +CREATE SERVER union_other_svr FOREIGN DATA WRAPPER clickhouse_fdw + OPTIONS (dbname 'union_pushdown_test', driver 'binary'); +CREATE USER MAPPING FOR CURRENT_USER SERVER union_other_svr; +SELECT clickhouse_raw_query('DROP DATABASE IF EXISTS union_pushdown_test'); + clickhouse_raw_query +---------------------- + +(1 row) + +SELECT clickhouse_raw_query('CREATE DATABASE union_pushdown_test'); + clickhouse_raw_query +---------------------- + +(1 row) + +SELECT clickhouse_raw_query('CREATE TABLE union_pushdown_test.messages + (message_id Int32, is_selected Int32) + ENGINE = MergeTree ORDER BY message_id'); + clickhouse_raw_query +---------------------- + +(1 row) + +SELECT clickhouse_raw_query('CREATE TABLE union_pushdown_test.source_a + (message_id Int32, category_id Int32) + ENGINE = MergeTree ORDER BY message_id'); + clickhouse_raw_query +---------------------- + +(1 row) + +SELECT clickhouse_raw_query('CREATE TABLE union_pushdown_test.collapsing_source + (message_id Int32, category_id Int32, sign Int8) + ENGINE = CollapsingMergeTree(sign) ORDER BY message_id'); + clickhouse_raw_query +---------------------- + +(1 row) + +SELECT clickhouse_raw_query('CREATE TABLE union_pushdown_test.source_b + (message_id Int32, category_id Int32) + ENGINE = MergeTree ORDER BY message_id'); + clickhouse_raw_query +---------------------- + +(1 row) + +SELECT clickhouse_raw_query('CREATE TABLE union_pushdown_test.bpchar_values + (value String, arm Int32) + ENGINE = MergeTree ORDER BY arm'); + clickhouse_raw_query +---------------------- + +(1 row) + +SELECT clickhouse_raw_query($$ + INSERT INTO union_pushdown_test.messages VALUES + (1, 1), (2, 1), (3, 0), (4, 1) +$$); + clickhouse_raw_query +---------------------- + +(1 row) + +SELECT clickhouse_raw_query($$ + INSERT INTO union_pushdown_test.source_a VALUES + (1, 10), (2, 20), (3, 90), (4, 40) +$$); + clickhouse_raw_query +---------------------- + +(1 row) + +SELECT clickhouse_raw_query($$ + INSERT INTO union_pushdown_test.source_b VALUES + (1, 20), (2, 30), (3, 90), (4, 40) +$$); + clickhouse_raw_query +---------------------- + +(1 row) + +SELECT clickhouse_raw_query($$ + INSERT INTO union_pushdown_test.bpchar_values VALUES + ('x', 1), ('x ', 2) +$$); + clickhouse_raw_query +---------------------- + +(1 row) + +CREATE SCHEMA union_pushdown_test; +CREATE FOREIGN TABLE union_pushdown_test.messages ( + message_id integer NOT NULL, + is_selected integer NOT NULL +) SERVER union_svr OPTIONS (table_name 'messages'); +CREATE FOREIGN TABLE union_pushdown_test.source_a ( + message_id integer NOT NULL, + category_id integer NOT NULL +) SERVER union_svr OPTIONS (table_name 'source_a'); +CREATE FOREIGN TABLE union_pushdown_test.source_a_collapsing ( + message_id integer NOT NULL, + category_id integer NOT NULL, + sign smallint NOT NULL +) SERVER union_svr OPTIONS ( + table_name 'collapsing_source', + engine 'CollapsingMergeTree(sign)' +); +CREATE FOREIGN TABLE union_pushdown_test.source_b ( + message_id integer NOT NULL, + category_id integer NOT NULL +) SERVER union_svr OPTIONS (table_name 'source_b'); +CREATE FOREIGN TABLE union_pushdown_test.source_b_other ( + message_id integer NOT NULL, + category_id integer NOT NULL +) SERVER union_other_svr OPTIONS (table_name 'source_b'); +CREATE FOREIGN TABLE union_pushdown_test.bpchar_values ( + value character(4) NOT NULL, + arm integer NOT NULL +) SERVER union_svr OPTIONS (table_name 'bpchar_values'); +CREATE FOREIGN TABLE union_pushdown_test.source_a_http ( + message_id integer NOT NULL, + category_id integer NOT NULL +) SERVER union_http_svr OPTIONS (table_name 'source_a'); +CREATE FOREIGN TABLE union_pushdown_test.source_b_http ( + message_id integer NOT NULL, + category_id integer NOT NULL +) SERVER union_http_svr OPTIONS (table_name 'source_b'); +SET SESSION search_path = union_pushdown_test,public; +SET SESSION enable_hashjoin = false; +SET SESSION enable_mergejoin = false; +\unset ECHO +NOTICE: same-server UNION ALL fully pushed down: t +NOTICE: plain integer DISTINCT fully pushed down: t +NOTICE: UNION ALL of DISTINCT arms fully pushed down: t +NOTICE: ordered plain DISTINCT not fully pushed down: t +NOTICE: DISTINCT ON not fully pushed down: t +NOTICE: bpchar plain DISTINCT not fully pushed down: t +NOTICE: non-default engine plain DISTINCT not fully pushed down: t +NOTICE: DISTINCT after window not fully pushed down: t +NOTICE: same-server UNION DISTINCT fully pushed down: t +NOTICE: HTTP same-server UNION ALL fully pushed down: t +NOTICE: ordered UNION DISTINCT not fully pushed down: t +NOTICE: default single-reference CTE fully pushed down: t +NOTICE: NOT MATERIALIZED multi-reference CTE fully pushed down: t +NOTICE: GROUP BY over same-server UNION ALL fully pushed down: t +NOTICE: GROUP BY over filtered join UNION ALL fully pushed down: t +NOTICE: GROUP BY pruned second UNION column fully pushed down: t +NOTICE: three-arm reordered UNION ALL fully pushed down: t +NOTICE: default multi-reference CTE not fully pushed down: t +NOTICE: MATERIALIZED multi-reference CTE not fully pushed down: t +NOTICE: different-server UNION ALL not fully pushed down: t +NOTICE: uncorrelated InitPlans not fully pushed down: t +NOTICE: zero-width UNION aggregate not fully pushed down: t +NOTICE: bpchar GROUP BY over UNION ALL not fully pushed down: t +NOTICE: bpchar aggregate DISTINCT over UNION ALL not fully pushed down: t +NOTICE: non-default engine GROUP BY over UNION ALL not fully pushed down: t +NOTICE: HAVING over UNION ALL not fully pushed down: t +NOTICE: bpchar UNION DISTINCT not fully pushed down: t +-- Materialize each positive set operation before the local ORDER BY so these +-- result checks execute the remote UNION while keeping output deterministic. +WITH pushed AS MATERIALIZED ( + SELECT DISTINCT message_id + FROM messages + WHERE is_selected = 1 +) +SELECT message_id FROM pushed +ORDER BY message_id; + message_id +------------ + 1 + 2 + 4 +(3 rows) + +-- UNION ALL preserves duplicates. +WITH pushed AS MATERIALIZED ( + SELECT category_id FROM source_a WHERE message_id <= 2 + UNION ALL + SELECT category_id FROM source_b WHERE message_id <= 2 +) +SELECT category_id FROM pushed +ORDER BY category_id; + category_id +-------- + 10 + 20 + 20 + 30 +(4 rows) + +-- Distinct generic-plan parameters verify parameter renumbering across arms. +SET plan_cache_mode = force_generic_plan; +PREPARE union_param(integer, integer) AS +WITH pushed AS MATERIALIZED ( + SELECT category_id FROM source_a WHERE message_id <= $1 + UNION ALL + SELECT category_id FROM source_b WHERE message_id <= $2 +) +SELECT category_id FROM pushed +ORDER BY category_id; +EXECUTE union_param(1, 2); + category_id +-------- + 10 + 20 + 30 +(3 rows) + +DEALLOCATE union_param; +PREPARE grouped_union_param(integer, integer, integer) AS +WITH pushed AS MATERIALIZED ( + SELECT + category_id, + count(DISTINCT message_id) FILTER (WHERE category_id <= $3) AS message_count + FROM ( + SELECT message_id, category_id FROM source_a WHERE message_id <= $1 + UNION ALL + SELECT message_id, category_id FROM source_b WHERE message_id <= $2 + ) AS combined_rows + GROUP BY category_id +) +SELECT category_id, message_count FROM pushed +ORDER BY category_id; +EXECUTE grouped_union_param(1, 2, 20); + category_id | message_count +--------+--------------- + 10 | 1 + 20 | 1 + 30 | 0 +(3 rows) + +DEALLOCATE grouped_union_param; +PREPARE union_http_param(integer, integer) AS +WITH pushed AS MATERIALIZED ( + SELECT category_id FROM source_a_http WHERE message_id <= $1 + UNION ALL + SELECT category_id FROM source_b_http WHERE message_id <= $2 +) +SELECT category_id FROM pushed +ORDER BY category_id; +EXECUTE union_http_param(1, 2); + category_id +-------- + 10 + 20 + 30 +(3 rows) + +DEALLOCATE union_http_param; +RESET plan_cache_mode; +-- PostgreSQL's bare UNION means UNION DISTINCT. +WITH pushed AS MATERIALIZED ( + SELECT category_id FROM source_a WHERE message_id <= 2 + UNION + SELECT category_id FROM source_b WHERE message_id <= 2 +) +SELECT category_id FROM pushed +ORDER BY category_id; + category_id +-------- + 10 + 20 + 30 +(3 rows) + +-- Direct ORDER BY keeps DISTINCT local; NULLS FIRST exercises a pathkey that +-- the remote set-operation target cannot currently preserve. +SELECT category_id FROM source_a WHERE message_id <= 2 +UNION +SELECT category_id FROM source_b WHERE message_id <= 2 +ORDER BY category_id NULLS FIRST; + category_id +-------- + 10 + 20 + 30 +(3 rows) + +-- A side-effect-free default CTE referenced once is inlined normally. +WITH pushed AS MATERIALIZED ( + WITH source_values AS ( + SELECT category_id FROM source_a WHERE message_id <= 2 + ) + SELECT category_id FROM source_values + UNION ALL + SELECT category_id FROM source_b WHERE message_id <= 2 +) +SELECT category_id FROM pushed +ORDER BY category_id; + category_id +-------- + 10 + 20 + 20 + 30 +(4 rows) + +-- NOT MATERIALIZED permits PostgreSQL to inline both references before the +-- same-server joins and UNION are considered for pushdown. +WITH pushed AS MATERIALIZED ( + WITH selected_messages AS NOT MATERIALIZED ( + SELECT message_id FROM messages WHERE is_selected = 1 + ) + SELECT selected_messages.message_id, source_a.category_id + FROM selected_messages + JOIN source_a USING (message_id) + UNION ALL + SELECT selected_messages.message_id, source_b.category_id + FROM selected_messages + JOIN source_b USING (message_id) +) +SELECT message_id, category_id FROM pushed +ORDER BY message_id, category_id; + message_id | category_id +------------+-------- + 1 | 10 + 1 | 20 + 2 | 20 + 2 | 30 + 4 | 40 + 4 | 40 +(6 rows) + +-- Default multi-reference and explicit MATERIALIZED CTEs retain PostgreSQL's +-- evaluate-once semantics and therefore stay local. +SELECT message_id, category_id +FROM ( + WITH selected_messages AS ( + SELECT message_id FROM messages WHERE is_selected = 1 + ) + SELECT selected_messages.message_id, source_a.category_id + FROM selected_messages + JOIN source_a USING (message_id) + UNION ALL + SELECT selected_messages.message_id, source_b.category_id + FROM selected_messages + JOIN source_b USING (message_id) +) AS pushed +ORDER BY message_id, category_id; + message_id | category_id +------------+-------- + 1 | 10 + 1 | 20 + 2 | 20 + 2 | 30 + 4 | 40 + 4 | 40 +(6 rows) + +SELECT message_id, category_id +FROM ( + WITH selected_messages AS MATERIALIZED ( + SELECT message_id FROM messages WHERE is_selected = 1 + ) + SELECT selected_messages.message_id, source_a.category_id + FROM selected_messages + JOIN source_a USING (message_id) + UNION ALL + SELECT selected_messages.message_id, source_b.category_id + FROM selected_messages + JOIN source_b USING (message_id) +) AS pushed +ORDER BY message_id, category_id; + message_id | category_id +------------+-------- + 1 | 10 + 1 | 20 + 2 | 20 + 2 | 30 + 4 | 40 + 4 | 40 +(6 rows) + +-- Matching connection options do not make distinct foreign servers safe to +-- combine into one remote query. +SELECT category_id +FROM ( + SELECT category_id FROM source_a WHERE message_id <= 2 + UNION ALL + SELECT category_id FROM source_b_other WHERE message_id <= 2 +) AS pushed +ORDER BY category_id; + category_id +-------- + 10 + 20 + 20 + 30 +(4 rows) + +-- InitPlans in UNION arms stay local and continue to execute normally. +WITH pushed AS MATERIALIZED ( + SELECT category_id FROM source_a + WHERE category_id > (SELECT 1) + UNION ALL + SELECT category_id FROM source_b + WHERE category_id > (SELECT 1) +) +SELECT category_id FROM pushed +ORDER BY category_id; + category_id +-------- + 10 + 20 + 20 + 30 + 40 + 40 + 90 + 90 +(8 rows) + +-- PostgreSQL bpchar equality ignores trailing spaces; keep DISTINCT local. +SELECT + count(*) AS distinct_rows, + bool_and(value = 'x'::character(4)) AS trailing_space_equal +FROM ( + SELECT value FROM bpchar_values WHERE arm = 1 + UNION + SELECT value FROM bpchar_values WHERE arm = 2 +) AS bpchar_union; + distinct_rows | trailing_space_equal +---------------+---------------------- + 1 | t +(1 row) + +-- GROUP BY and COUNT(DISTINCT) execute after the remote UNION ALL. Duplicate +-- rows within and across arms must not inflate the per-category message count. +SELECT clickhouse_raw_query($$ + INSERT INTO union_pushdown_test.source_a VALUES + (1, 10), (4, 40) +$$); + clickhouse_raw_query +---------------------- + +(1 row) + +WITH grouped AS MATERIALIZED ( + WITH selected_messages AS NOT MATERIALIZED ( + SELECT message_id FROM messages WHERE is_selected = 1 + ) + SELECT category_id, count(DISTINCT message_id) AS message_count + FROM ( + SELECT selected_messages.message_id, source_a.category_id + FROM selected_messages + JOIN source_a USING (message_id) + UNION ALL + SELECT selected_messages.message_id, source_b.category_id + FROM selected_messages + JOIN source_b USING (message_id) + ) AS combined_rows + GROUP BY category_id +) +SELECT category_id, message_count +FROM grouped +ORDER BY category_id; + category_id | message_count +--------+--------------- + 10 | 1 + 20 | 2 + 30 | 1 + 40 | 1 +(4 rows) + +-- Pruning message_id leaves a compact one-column remote UNION target whose +-- remaining value originated as the subquery's second output column. +WITH grouped AS MATERIALIZED ( + SELECT category_id, count(*) AS row_count + FROM ( + SELECT message_id, category_id FROM source_a + UNION ALL + SELECT message_id, category_id FROM source_b + ) AS combined_rows + GROUP BY category_id +) +SELECT category_id, row_count +FROM grouped +ORDER BY category_id; + category_id | row_count +--------+----------- + 10 | 2 + 20 | 2 + 30 | 1 + 40 | 3 + 90 | 2 +(5 rows) + +-- Direct SET ROLE plans use that role's mapping. A view instead plans its +-- foreign scans with the view owner's mapping, even for another current user. +CREATE ROLE union_view_owner; +GRANT USAGE, CREATE ON SCHEMA union_pushdown_test TO union_view_owner; +GRANT USAGE ON FOREIGN SERVER union_svr TO union_view_owner; +GRANT SELECT ON source_a, source_b TO union_view_owner; +CREATE USER MAPPING FOR union_view_owner SERVER union_svr; +CREATE VIEW owner_grouped_union AS +SELECT category_id, count(DISTINCT message_id) AS message_count +FROM ( + SELECT message_id, category_id FROM source_a + UNION ALL + SELECT message_id, category_id FROM source_b +) AS combined_rows +GROUP BY category_id; +ALTER VIEW owner_grouped_union OWNER TO union_view_owner; +DROP USER MAPPING FOR CURRENT_USER SERVER union_svr; +SET ROLE union_view_owner; +WITH grouped AS MATERIALIZED ( + SELECT category_id, count(DISTINCT message_id) AS message_count + FROM ( + SELECT message_id, category_id FROM source_a + UNION ALL + SELECT message_id, category_id FROM source_b + ) AS combined_rows + GROUP BY category_id +) +SELECT count(*) AS categories, sum(message_count) AS total FROM grouped; + categories | total +------+------- + 5 | 6 +(1 row) + +RESET ROLE; +SELECT count(*) AS categories, sum(message_count) AS total +FROM owner_grouped_union; + categories | total +------+------- + 5 | 6 +(1 row) + +CREATE USER MAPPING FOR CURRENT_USER SERVER union_svr; +DROP VIEW owner_grouped_union; +DROP USER MAPPING FOR union_view_owner SERVER union_svr; +REVOKE SELECT ON source_a, source_b FROM union_view_owner; +REVOKE USAGE ON FOREIGN SERVER union_svr FROM union_view_owner; +REVOKE USAGE, CREATE ON SCHEMA union_pushdown_test FROM union_view_owner; +DROP ROLE union_view_owner; +RESET enable_hashjoin; +RESET enable_mergejoin; +SET SESSION search_path = public; +DROP FOREIGN TABLE union_pushdown_test.source_b_http; +DROP FOREIGN TABLE union_pushdown_test.source_a_http; +DROP FOREIGN TABLE union_pushdown_test.bpchar_values; +DROP FOREIGN TABLE union_pushdown_test.source_b_other; +DROP FOREIGN TABLE union_pushdown_test.source_b; +DROP FOREIGN TABLE union_pushdown_test.source_a_collapsing; +DROP FOREIGN TABLE union_pushdown_test.source_a; +DROP FOREIGN TABLE union_pushdown_test.messages; +DROP SCHEMA union_pushdown_test; +DROP USER MAPPING FOR CURRENT_USER SERVER union_other_svr; +DROP SERVER union_other_svr; +DROP USER MAPPING FOR CURRENT_USER SERVER union_http_svr; +DROP SERVER union_http_svr; +DROP USER MAPPING FOR CURRENT_USER SERVER union_svr; +DROP SERVER union_svr; +SELECT clickhouse_raw_query('DROP DATABASE union_pushdown_test'); + clickhouse_raw_query +---------------------- + +(1 row) + +-- End of UNION pushdown tests. diff --git a/test/sql/union_pushdown.sql b/test/sql/union_pushdown.sql new file mode 100644 index 00000000..042fa0c7 --- /dev/null +++ b/test/sql/union_pushdown.sql @@ -0,0 +1,781 @@ +-- Tests for UNION and safely inlined CTE pushdown. +-- +-- Plan checks use JSON instead of printing the full EXPLAIN tree so the +-- expected output remains stable across PostgreSQL 13-19. +CREATE SERVER union_svr FOREIGN DATA WRAPPER clickhouse_fdw + OPTIONS (dbname 'union_pushdown_test', driver 'binary'); +CREATE USER MAPPING FOR CURRENT_USER SERVER union_svr; + +CREATE SERVER union_http_svr FOREIGN DATA WRAPPER clickhouse_fdw + OPTIONS (dbname 'union_pushdown_test', driver 'http'); +CREATE USER MAPPING FOR CURRENT_USER SERVER union_http_svr; + +-- A distinct server object pointing at the same ClickHouse database verifies +-- that server identity, rather than matching connection options, gates UNION +-- pushdown. +CREATE SERVER union_other_svr FOREIGN DATA WRAPPER clickhouse_fdw + OPTIONS (dbname 'union_pushdown_test', driver 'binary'); +CREATE USER MAPPING FOR CURRENT_USER SERVER union_other_svr; + +SELECT clickhouse_raw_query('DROP DATABASE IF EXISTS union_pushdown_test'); +SELECT clickhouse_raw_query('CREATE DATABASE union_pushdown_test'); +SELECT clickhouse_raw_query('CREATE TABLE union_pushdown_test.messages + (message_id Int32, is_selected Int32) + ENGINE = MergeTree ORDER BY message_id'); +SELECT clickhouse_raw_query('CREATE TABLE union_pushdown_test.source_a + (message_id Int32, category_id Int32) + ENGINE = MergeTree ORDER BY message_id'); +SELECT clickhouse_raw_query('CREATE TABLE union_pushdown_test.collapsing_source + (message_id Int32, category_id Int32, sign Int8) + ENGINE = CollapsingMergeTree(sign) ORDER BY message_id'); +SELECT clickhouse_raw_query('CREATE TABLE union_pushdown_test.source_b + (message_id Int32, category_id Int32) + ENGINE = MergeTree ORDER BY message_id'); +SELECT clickhouse_raw_query('CREATE TABLE union_pushdown_test.bpchar_values + (value String, arm Int32) + ENGINE = MergeTree ORDER BY arm'); + +SELECT clickhouse_raw_query($$ + INSERT INTO union_pushdown_test.messages VALUES + (1, 1), (2, 1), (3, 0), (4, 1) +$$); +SELECT clickhouse_raw_query($$ + INSERT INTO union_pushdown_test.source_a VALUES + (1, 10), (2, 20), (3, 90), (4, 40) +$$); +SELECT clickhouse_raw_query($$ + INSERT INTO union_pushdown_test.source_b VALUES + (1, 20), (2, 30), (3, 90), (4, 40) +$$); +SELECT clickhouse_raw_query($$ + INSERT INTO union_pushdown_test.bpchar_values VALUES + ('x', 1), ('x ', 2) +$$); + +CREATE SCHEMA union_pushdown_test; +CREATE FOREIGN TABLE union_pushdown_test.messages ( + message_id integer NOT NULL, + is_selected integer NOT NULL +) SERVER union_svr OPTIONS (table_name 'messages'); +CREATE FOREIGN TABLE union_pushdown_test.source_a ( + message_id integer NOT NULL, + category_id integer NOT NULL +) SERVER union_svr OPTIONS (table_name 'source_a'); +CREATE FOREIGN TABLE union_pushdown_test.source_a_collapsing ( + message_id integer NOT NULL, + category_id integer NOT NULL, + sign smallint NOT NULL +) SERVER union_svr OPTIONS ( + table_name 'collapsing_source', + engine 'CollapsingMergeTree(sign)' +); +CREATE FOREIGN TABLE union_pushdown_test.source_b ( + message_id integer NOT NULL, + category_id integer NOT NULL +) SERVER union_svr OPTIONS (table_name 'source_b'); +CREATE FOREIGN TABLE union_pushdown_test.source_b_other ( + message_id integer NOT NULL, + category_id integer NOT NULL +) SERVER union_other_svr OPTIONS (table_name 'source_b'); +CREATE FOREIGN TABLE union_pushdown_test.bpchar_values ( + value character(4) NOT NULL, + arm integer NOT NULL +) SERVER union_svr OPTIONS (table_name 'bpchar_values'); +CREATE FOREIGN TABLE union_pushdown_test.source_a_http ( + message_id integer NOT NULL, + category_id integer NOT NULL +) SERVER union_http_svr OPTIONS (table_name 'source_a'); +CREATE FOREIGN TABLE union_pushdown_test.source_b_http ( + message_id integer NOT NULL, + category_id integer NOT NULL +) SERVER union_http_svr OPTIONS (table_name 'source_b'); + +SET SESSION search_path = union_pushdown_test,public; +SET SESSION enable_hashjoin = false; +SET SESSION enable_mergejoin = false; + +\unset ECHO +CREATE FUNCTION pg_temp.assert_union_plan( + test_name text, + query_sql text, + expected_full_pushdown boolean, + expected_remote_fragment text +) RETURNS void +LANGUAGE plpgsql +AS $function$ +DECLARE + plan jsonb; + remote_sqls jsonb; + fully_pushed_down boolean; +BEGIN + EXECUTE 'EXPLAIN (VERBOSE, COSTS OFF, FORMAT JSON) ' || query_sql + INTO plan; + remote_sqls := + jsonb_path_query_array(plan, 'strict $.**."Remote SQL"'); + fully_pushed_down := + COALESCE(plan #>> '{0,Plan,Node Type}' = 'Foreign Scan', false) + AND jsonb_array_length(remote_sqls) = 1; + + IF fully_pushed_down AND expected_remote_fragment IS NOT NULL THEN + fully_pushed_down := + position(expected_remote_fragment IN (remote_sqls ->> 0)) > 0; + END IF; + + RAISE NOTICE '%: %', test_name, + fully_pushed_down IS NOT DISTINCT FROM expected_full_pushdown; +END; +$function$; + +DO $do$ +BEGIN + PERFORM pg_temp.assert_union_plan( + 'same-server UNION ALL fully pushed down', + $query$ + SELECT category_id FROM source_a WHERE message_id <= 2 + UNION ALL + SELECT category_id FROM source_b WHERE message_id <= 2 + $query$, + true, + 'UNION ALL' + ); + + PERFORM pg_temp.assert_union_plan( + 'plain integer DISTINCT fully pushed down', + $query$ + SELECT DISTINCT message_id + FROM messages + WHERE is_selected = 1 + $query$, + true, + 'SELECT DISTINCT' + ); + + PERFORM pg_temp.assert_union_plan( + 'UNION ALL of DISTINCT arms fully pushed down', + $query$ + SELECT DISTINCT message_id + FROM source_a + WHERE category_id <= 20 + UNION ALL + SELECT DISTINCT message_id + FROM source_b + WHERE category_id <= 20 + $query$, + true, + 'UNION ALL' + ); + + PERFORM pg_temp.assert_union_plan( + 'ordered plain DISTINCT not fully pushed down', + $query$ + SELECT DISTINCT message_id + FROM messages + WHERE is_selected = 1 + ORDER BY message_id + $query$, + false, + NULL + ); + + PERFORM pg_temp.assert_union_plan( + 'DISTINCT ON not fully pushed down', + $query$ + SELECT DISTINCT ON (is_selected) is_selected, message_id + FROM messages + $query$, + false, + NULL + ); + + PERFORM pg_temp.assert_union_plan( + 'bpchar plain DISTINCT not fully pushed down', + $query$ + SELECT DISTINCT value FROM bpchar_values + $query$, + false, + NULL + ); + + PERFORM pg_temp.assert_union_plan( + 'non-default engine plain DISTINCT not fully pushed down', + $query$ + SELECT DISTINCT message_id FROM source_a_collapsing + $query$, + false, + NULL + ); + + PERFORM pg_temp.assert_union_plan( + 'DISTINCT after window not fully pushed down', + $query$ + SELECT DISTINCT row_number() OVER (ORDER BY message_id) + FROM source_a + $query$, + false, + NULL + ); + + PERFORM pg_temp.assert_union_plan( + 'same-server UNION DISTINCT fully pushed down', + $query$ + SELECT category_id FROM source_a WHERE message_id <= 2 + UNION + SELECT category_id FROM source_b WHERE message_id <= 2 + $query$, + true, + 'UNION DISTINCT' + ); + + PERFORM pg_temp.assert_union_plan( + 'HTTP same-server UNION ALL fully pushed down', + $query$ + SELECT category_id FROM source_a_http WHERE message_id <= 2 + UNION ALL + SELECT category_id FROM source_b_http WHERE message_id <= 2 + $query$, + true, + 'UNION ALL' + ); + + PERFORM pg_temp.assert_union_plan( + 'ordered UNION DISTINCT not fully pushed down', + $query$ + SELECT category_id FROM source_a WHERE message_id <= 2 + UNION + SELECT category_id FROM source_b WHERE message_id <= 2 + ORDER BY category_id NULLS FIRST + $query$, + false, + NULL + ); + + PERFORM pg_temp.assert_union_plan( + 'default single-reference CTE fully pushed down', + $query$ + WITH source_values AS ( + SELECT category_id FROM source_a WHERE message_id <= 2 + ) + SELECT category_id FROM source_values + UNION ALL + SELECT category_id FROM source_b WHERE message_id <= 2 + $query$, + true, + 'UNION ALL' + ); + + PERFORM pg_temp.assert_union_plan( + 'NOT MATERIALIZED multi-reference CTE fully pushed down', + $query$ + WITH selected_messages AS NOT MATERIALIZED ( + SELECT message_id FROM messages WHERE is_selected = 1 + ) + SELECT selected_messages.message_id, source_a.category_id + FROM selected_messages + JOIN source_a USING (message_id) + UNION ALL + SELECT selected_messages.message_id, source_b.category_id + FROM selected_messages + JOIN source_b USING (message_id) + $query$, + true, + 'UNION ALL' + ); + + PERFORM pg_temp.assert_union_plan( + 'GROUP BY over same-server UNION ALL fully pushed down', + $query$ + WITH selected_messages AS NOT MATERIALIZED ( + SELECT message_id FROM messages WHERE is_selected = 1 + ) + SELECT category_id, count(DISTINCT message_id) + FROM ( + SELECT selected_messages.message_id, source_a.category_id + FROM selected_messages + JOIN source_a USING (message_id) + UNION ALL + SELECT selected_messages.message_id, source_b.category_id + FROM selected_messages + JOIN source_b USING (message_id) + ) AS combined_rows + GROUP BY category_id + $query$, + true, + 'count(DISTINCT' + ); + + PERFORM pg_temp.assert_union_plan( + 'GROUP BY over filtered join UNION ALL fully pushed down', + $query$ + SELECT category_id, count(DISTINCT message_id) + FROM ( + SELECT messages.message_id, source_a.category_id + FROM messages + JOIN source_a USING (message_id) + WHERE messages.is_selected = 1 + AND source_a.category_id <= 40 + UNION ALL + SELECT messages.message_id, source_b.category_id + FROM messages + JOIN source_b USING (message_id) + WHERE messages.is_selected = 1 + AND source_b.category_id <= 40 + ) AS combined_rows + GROUP BY 1 + $query$, + true, + 'count(DISTINCT' + ); + + PERFORM pg_temp.assert_union_plan( + 'GROUP BY pruned second UNION column fully pushed down', + $query$ + SELECT category_id, count(*) + FROM ( + SELECT message_id, category_id FROM source_a + UNION ALL + SELECT message_id, category_id FROM source_b + ) AS combined_rows + GROUP BY category_id + $query$, + true, + 'GROUP BY' + ); + + PERFORM pg_temp.assert_union_plan( + 'three-arm reordered UNION ALL fully pushed down', + $query$ + SELECT category_id, message_id FROM source_a + UNION ALL + SELECT category_id, message_id FROM source_b + UNION ALL + SELECT category_id, message_id FROM source_a + $query$, + true, + 'UNION ALL' + ); + + PERFORM pg_temp.assert_union_plan( + 'default multi-reference CTE not fully pushed down', + $query$ + WITH selected_messages AS ( + SELECT message_id FROM messages WHERE is_selected = 1 + ) + SELECT selected_messages.message_id, source_a.category_id + FROM selected_messages + JOIN source_a USING (message_id) + UNION ALL + SELECT selected_messages.message_id, source_b.category_id + FROM selected_messages + JOIN source_b USING (message_id) + $query$, + false, + NULL + ); + + PERFORM pg_temp.assert_union_plan( + 'MATERIALIZED multi-reference CTE not fully pushed down', + $query$ + WITH selected_messages AS MATERIALIZED ( + SELECT message_id FROM messages WHERE is_selected = 1 + ) + SELECT selected_messages.message_id, source_a.category_id + FROM selected_messages + JOIN source_a USING (message_id) + UNION ALL + SELECT selected_messages.message_id, source_b.category_id + FROM selected_messages + JOIN source_b USING (message_id) + $query$, + false, + NULL + ); + + PERFORM pg_temp.assert_union_plan( + 'different-server UNION ALL not fully pushed down', + $query$ + SELECT category_id FROM source_a WHERE message_id <= 2 + UNION ALL + SELECT category_id FROM source_b_other WHERE message_id <= 2 + $query$, + false, + NULL + ); + + PERFORM pg_temp.assert_union_plan( + 'uncorrelated InitPlans not fully pushed down', + $query$ + SELECT category_id FROM source_a + WHERE category_id > (SELECT 1) + UNION ALL + SELECT category_id FROM source_b + WHERE category_id > (SELECT 1) + $query$, + false, + NULL + ); + + PERFORM pg_temp.assert_union_plan( + 'zero-width UNION aggregate not fully pushed down', + $query$ + SELECT count(*) + FROM ( + SELECT message_id FROM source_a + UNION ALL + SELECT message_id FROM source_b + ) AS message_ids + $query$, + false, + NULL + ); + + PERFORM pg_temp.assert_union_plan( + 'bpchar GROUP BY over UNION ALL not fully pushed down', + $query$ + SELECT value, count(*) + FROM ( + SELECT value FROM bpchar_values WHERE arm = 1 + UNION ALL + SELECT value FROM bpchar_values WHERE arm = 2 + ) AS values + GROUP BY value + $query$, + false, + NULL + ); + + PERFORM pg_temp.assert_union_plan( + 'bpchar aggregate DISTINCT over UNION ALL not fully pushed down', + $query$ + SELECT arm, count(DISTINCT value) + FROM ( + SELECT arm, value FROM bpchar_values WHERE arm = 1 + UNION ALL + SELECT arm, value FROM bpchar_values WHERE arm = 2 + ) AS values + GROUP BY arm + $query$, + false, + NULL + ); + + PERFORM pg_temp.assert_union_plan( + 'non-default engine GROUP BY over UNION ALL not fully pushed down', + $query$ + SELECT category_id, count(*) + FROM ( + SELECT category_id FROM source_a_collapsing + UNION ALL + SELECT category_id FROM source_b + ) AS combined_rows + GROUP BY category_id + $query$, + false, + NULL + ); + + PERFORM pg_temp.assert_union_plan( + 'HAVING over UNION ALL not fully pushed down', + $query$ + SELECT category_id, count(*) + FROM ( + SELECT category_id FROM source_a + UNION ALL + SELECT category_id FROM source_b + ) AS combined_rows + GROUP BY category_id + HAVING count(*) > 1 + $query$, + false, + NULL + ); + + PERFORM pg_temp.assert_union_plan( + 'bpchar UNION DISTINCT not fully pushed down', + $query$ + SELECT value FROM bpchar_values WHERE arm = 1 + UNION + SELECT value FROM bpchar_values WHERE arm = 2 + $query$, + false, + NULL + ); +END; +$do$; +\set ECHO all + +-- Materialize each positive set operation before the local ORDER BY so these +-- result checks execute the remote UNION while keeping output deterministic. +WITH pushed AS MATERIALIZED ( + SELECT DISTINCT message_id + FROM messages + WHERE is_selected = 1 +) +SELECT message_id FROM pushed +ORDER BY message_id; + +-- UNION ALL preserves duplicates. +WITH pushed AS MATERIALIZED ( + SELECT category_id FROM source_a WHERE message_id <= 2 + UNION ALL + SELECT category_id FROM source_b WHERE message_id <= 2 +) +SELECT category_id FROM pushed +ORDER BY category_id; + +-- Distinct generic-plan parameters verify parameter renumbering across arms. +SET plan_cache_mode = force_generic_plan; +PREPARE union_param(integer, integer) AS +WITH pushed AS MATERIALIZED ( + SELECT category_id FROM source_a WHERE message_id <= $1 + UNION ALL + SELECT category_id FROM source_b WHERE message_id <= $2 +) +SELECT category_id FROM pushed +ORDER BY category_id; +EXECUTE union_param(1, 2); +DEALLOCATE union_param; + +PREPARE grouped_union_param(integer, integer, integer) AS +WITH pushed AS MATERIALIZED ( + SELECT + category_id, + count(DISTINCT message_id) FILTER (WHERE category_id <= $3) AS message_count + FROM ( + SELECT message_id, category_id FROM source_a WHERE message_id <= $1 + UNION ALL + SELECT message_id, category_id FROM source_b WHERE message_id <= $2 + ) AS combined_rows + GROUP BY category_id +) +SELECT category_id, message_count FROM pushed +ORDER BY category_id; +EXECUTE grouped_union_param(1, 2, 20); +DEALLOCATE grouped_union_param; + +PREPARE union_http_param(integer, integer) AS +WITH pushed AS MATERIALIZED ( + SELECT category_id FROM source_a_http WHERE message_id <= $1 + UNION ALL + SELECT category_id FROM source_b_http WHERE message_id <= $2 +) +SELECT category_id FROM pushed +ORDER BY category_id; +EXECUTE union_http_param(1, 2); +DEALLOCATE union_http_param; +RESET plan_cache_mode; + +-- PostgreSQL's bare UNION means UNION DISTINCT. +WITH pushed AS MATERIALIZED ( + SELECT category_id FROM source_a WHERE message_id <= 2 + UNION + SELECT category_id FROM source_b WHERE message_id <= 2 +) +SELECT category_id FROM pushed +ORDER BY category_id; + +-- Direct ORDER BY keeps DISTINCT local; NULLS FIRST exercises a pathkey that +-- the remote set-operation target cannot currently preserve. +SELECT category_id FROM source_a WHERE message_id <= 2 +UNION +SELECT category_id FROM source_b WHERE message_id <= 2 +ORDER BY category_id NULLS FIRST; + +-- A side-effect-free default CTE referenced once is inlined normally. +WITH pushed AS MATERIALIZED ( + WITH source_values AS ( + SELECT category_id FROM source_a WHERE message_id <= 2 + ) + SELECT category_id FROM source_values + UNION ALL + SELECT category_id FROM source_b WHERE message_id <= 2 +) +SELECT category_id FROM pushed +ORDER BY category_id; + +-- NOT MATERIALIZED permits PostgreSQL to inline both references before the +-- same-server joins and UNION are considered for pushdown. +WITH pushed AS MATERIALIZED ( + WITH selected_messages AS NOT MATERIALIZED ( + SELECT message_id FROM messages WHERE is_selected = 1 + ) + SELECT selected_messages.message_id, source_a.category_id + FROM selected_messages + JOIN source_a USING (message_id) + UNION ALL + SELECT selected_messages.message_id, source_b.category_id + FROM selected_messages + JOIN source_b USING (message_id) +) +SELECT message_id, category_id FROM pushed +ORDER BY message_id, category_id; + +-- Default multi-reference and explicit MATERIALIZED CTEs retain PostgreSQL's +-- evaluate-once semantics and therefore stay local. +SELECT message_id, category_id +FROM ( + WITH selected_messages AS ( + SELECT message_id FROM messages WHERE is_selected = 1 + ) + SELECT selected_messages.message_id, source_a.category_id + FROM selected_messages + JOIN source_a USING (message_id) + UNION ALL + SELECT selected_messages.message_id, source_b.category_id + FROM selected_messages + JOIN source_b USING (message_id) +) AS pushed +ORDER BY message_id, category_id; + +SELECT message_id, category_id +FROM ( + WITH selected_messages AS MATERIALIZED ( + SELECT message_id FROM messages WHERE is_selected = 1 + ) + SELECT selected_messages.message_id, source_a.category_id + FROM selected_messages + JOIN source_a USING (message_id) + UNION ALL + SELECT selected_messages.message_id, source_b.category_id + FROM selected_messages + JOIN source_b USING (message_id) +) AS pushed +ORDER BY message_id, category_id; + +-- Matching connection options do not make distinct foreign servers safe to +-- combine into one remote query. +SELECT category_id +FROM ( + SELECT category_id FROM source_a WHERE message_id <= 2 + UNION ALL + SELECT category_id FROM source_b_other WHERE message_id <= 2 +) AS pushed +ORDER BY category_id; + +-- InitPlans in UNION arms stay local and continue to execute normally. +WITH pushed AS MATERIALIZED ( + SELECT category_id FROM source_a + WHERE category_id > (SELECT 1) + UNION ALL + SELECT category_id FROM source_b + WHERE category_id > (SELECT 1) +) +SELECT category_id FROM pushed +ORDER BY category_id; + +-- PostgreSQL bpchar equality ignores trailing spaces; keep DISTINCT local. +SELECT + count(*) AS distinct_rows, + bool_and(value = 'x'::character(4)) AS trailing_space_equal +FROM ( + SELECT value FROM bpchar_values WHERE arm = 1 + UNION + SELECT value FROM bpchar_values WHERE arm = 2 +) AS bpchar_union; + +-- GROUP BY and COUNT(DISTINCT) execute after the remote UNION ALL. Duplicate +-- rows within and across arms must not inflate the per-category message count. +SELECT clickhouse_raw_query($$ + INSERT INTO union_pushdown_test.source_a VALUES + (1, 10), (4, 40) +$$); + +WITH grouped AS MATERIALIZED ( + WITH selected_messages AS NOT MATERIALIZED ( + SELECT message_id FROM messages WHERE is_selected = 1 + ) + SELECT category_id, count(DISTINCT message_id) AS message_count + FROM ( + SELECT selected_messages.message_id, source_a.category_id + FROM selected_messages + JOIN source_a USING (message_id) + UNION ALL + SELECT selected_messages.message_id, source_b.category_id + FROM selected_messages + JOIN source_b USING (message_id) + ) AS combined_rows + GROUP BY category_id +) +SELECT category_id, message_count +FROM grouped +ORDER BY category_id; + +-- Pruning message_id leaves a compact one-column remote UNION target whose +-- remaining value originated as the subquery's second output column. +WITH grouped AS MATERIALIZED ( + SELECT category_id, count(*) AS row_count + FROM ( + SELECT message_id, category_id FROM source_a + UNION ALL + SELECT message_id, category_id FROM source_b + ) AS combined_rows + GROUP BY category_id +) +SELECT category_id, row_count +FROM grouped +ORDER BY category_id; + +-- Direct SET ROLE plans use that role's mapping. A view instead plans its +-- foreign scans with the view owner's mapping, even for another current user. +CREATE ROLE union_view_owner; +GRANT USAGE, CREATE ON SCHEMA union_pushdown_test TO union_view_owner; +GRANT USAGE ON FOREIGN SERVER union_svr TO union_view_owner; +GRANT SELECT ON source_a, source_b TO union_view_owner; +CREATE USER MAPPING FOR union_view_owner SERVER union_svr; + +CREATE VIEW owner_grouped_union AS +SELECT category_id, count(DISTINCT message_id) AS message_count +FROM ( + SELECT message_id, category_id FROM source_a + UNION ALL + SELECT message_id, category_id FROM source_b +) AS combined_rows +GROUP BY category_id; +ALTER VIEW owner_grouped_union OWNER TO union_view_owner; + +DROP USER MAPPING FOR CURRENT_USER SERVER union_svr; +SET ROLE union_view_owner; +WITH grouped AS MATERIALIZED ( + SELECT category_id, count(DISTINCT message_id) AS message_count + FROM ( + SELECT message_id, category_id FROM source_a + UNION ALL + SELECT message_id, category_id FROM source_b + ) AS combined_rows + GROUP BY category_id +) +SELECT count(*) AS categories, sum(message_count) AS total FROM grouped; +RESET ROLE; + +SELECT count(*) AS categories, sum(message_count) AS total +FROM owner_grouped_union; +CREATE USER MAPPING FOR CURRENT_USER SERVER union_svr; + +DROP VIEW owner_grouped_union; +DROP USER MAPPING FOR union_view_owner SERVER union_svr; +REVOKE SELECT ON source_a, source_b FROM union_view_owner; +REVOKE USAGE ON FOREIGN SERVER union_svr FROM union_view_owner; +REVOKE USAGE, CREATE ON SCHEMA union_pushdown_test FROM union_view_owner; +DROP ROLE union_view_owner; + +RESET enable_hashjoin; +RESET enable_mergejoin; +SET SESSION search_path = public; + +DROP FOREIGN TABLE union_pushdown_test.source_b_http; +DROP FOREIGN TABLE union_pushdown_test.source_a_http; +DROP FOREIGN TABLE union_pushdown_test.bpchar_values; +DROP FOREIGN TABLE union_pushdown_test.source_b_other; +DROP FOREIGN TABLE union_pushdown_test.source_b; +DROP FOREIGN TABLE union_pushdown_test.source_a_collapsing; +DROP FOREIGN TABLE union_pushdown_test.source_a; +DROP FOREIGN TABLE union_pushdown_test.messages; +DROP SCHEMA union_pushdown_test; +DROP USER MAPPING FOR CURRENT_USER SERVER union_other_svr; +DROP SERVER union_other_svr; +DROP USER MAPPING FOR CURRENT_USER SERVER union_http_svr; +DROP SERVER union_http_svr; +DROP USER MAPPING FOR CURRENT_USER SERVER union_svr; +DROP SERVER union_svr; +SELECT clickhouse_raw_query('DROP DATABASE union_pushdown_test'); + +-- End of UNION pushdown tests. From e27d04bfa29ab36b9e4d61dac7650034fc6872b7 Mon Sep 17 00:00:00 2001 From: Kostia R Date: Thu, 23 Jul 2026 21:47:31 +0000 Subject: [PATCH 2/2] Document UNION pushdown --- CHANGELOG.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 453c6112..99b1d0c1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -78,6 +78,9 @@ All notable changes to this project will be documented in this file. It uses the rather than understating foreign scan cost ([#310]). * Added mapping to push down the re2 v0.4 `@~` operator to ClickHouse as the `match()` function ([#318]). +* Added pushdown for compatible `UNION ALL`, `UNION DISTINCT`, plain + `DISTINCT`, and grouping or aggregation over flattened `UNION ALL` + queries ([#324]). ### 🐞 Bug Fixes @@ -160,6 +163,8 @@ All notable changes to this project will be documented in this file. It uses the "ClickHouse/pg_clickhouse#317 Push down the array IN family unconditionally" [#319]: https://github.com/ClickHouse/pg_clickhouse/pull/319 "ClickHouse/pg_clickhouse#319 Fix foreign scan RTE selection" + [#324]: https://github.com/ClickHouse/pg_clickhouse/pull/324 + "ClickHouse/pg_clickhouse#324 Push down compatible UNION queries and grouped set operations" ## [v0.3.2] — 2026-06-16