// Copyright (c) 2023-2026 ParadeDB, Inc.
//
// This file is part of ParadeDB - Postgres for Search and Analytics
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see .
//! JoinScan: Custom scan operator for optimizing joins with BM25 full-text search.
//!
//! JoinScan intercepts PostgreSQL join operations and executes them using Tantivy's
//! search capabilities combined with a join algorithm, providing significant
//! performance improvements for queries that combine full-text search with joins.
//!
//! # Activation Conditions
//!
//! JoinScan is proposed by the planner when **all** of the following conditions are met.
//! These restrictions ensure that we can execute the join efficiently using Tantivy's
//! columnar storage (fast fields) and minimize expensive heap access.
//!
//! The core strategy is **late materialization**:
//! 1. Execute the search and join using ONLY the index (fast fields).
//! 2. Apply sorting and limits on the joined index data.
//! 3. Only access the PostgreSQL heap (materialize) for the final result rows (Top K).
//!
//! This strategy requires that all data needed for the join, filter, and sort phases
//! resides in fast fields, and that the result set size is small enough (via LIMIT)
//! that the random heap access cost doesn't outweigh the join benefit.
//!
//! 1. **GUC enabled**: `paradedb.enable_join_custom_scan = on` (default: off)
//!
//! 2. **Join type**: INNER, SEMI, and ANTI joins are supported
//! - LEFT, RIGHT, and FULL joins are planned for future work
//!
//! 3. **LIMIT clause**: Query must have a LIMIT clause
//! - This ensures we only pay the cost of "late materialization" (random heap access)
//! for a small number of rows. Without LIMIT, scanning the entire index and fetching
//! all rows from the heap is often slower than PostgreSQL's native execution.
//! - Future work will allow no-limit joins when both sides have search predicates.
//!
//! 4. **Search predicate**: At least one side must have:
//! - A BM25 index on the table
//! - A `@@@` search predicate in the WHERE clause
//!
//! 5. **Multi-level Joins**: JoinScan supports multi-level joins (e.g., `(A JOIN B) JOIN C`).
//! It achieves this by reconstructing the join tree from PostgreSQL's plan or by nesting
//! multiple JoinScan operators.
//!
//! 6. **Fast-field columns**: All columns used in the join must be fast fields in their
//! respective BM25 indexes. This allows the join to be executed entirely within the index:
//! - Equi-join keys (e.g., `a.id = b.id`) must be fast fields for join execution
//! - Multi-table predicates (e.g., `a.price > b.min_price`) must reference fast fields
//! - ORDER BY columns must be fast fields for efficient sorting
//! - If any required column is not a fast field, we would need to access the heap
//! during the join, breaking the late materialization strategy.
//!
//! 7. **Equi-join keys required**: At least one equi-join key (e.g., `a.id = b.id`) is
//! required. Cross joins (cartesian products) fall back to PostgreSQL
//!
//! # Example Queries
//!
//! ```sql
//! -- JoinScan IS proposed (has LIMIT, has @@@ predicate)
//! SELECT p.name, s.name
//! FROM products p
//! JOIN suppliers s ON p.supplier_id = s.id
//! WHERE p.description @@@ 'wireless'
//! LIMIT 10;
//!
//! -- JoinScan is NOT proposed (no LIMIT)
//! SELECT p.name, s.name
//! FROM products p
//! JOIN suppliers s ON p.supplier_id = s.id
//! WHERE p.description @@@ 'wireless';
//!
//! -- JoinScan is NOT proposed (LEFT JOIN not supported)
//! SELECT p.name, s.name
//! FROM products p
//! LEFT JOIN suppliers s ON p.supplier_id = s.id
//! WHERE p.description @@@ 'wireless'
//! LIMIT 10;
//!
//! -- JoinScan IS proposed if price/min_price are fast fields in BM25 indexes
//! SELECT p.name, s.name
//! FROM products p
//! JOIN suppliers s ON p.supplier_id = s.id
//! WHERE p.description @@@ 'wireless' AND p.price > s.min_price
//! LIMIT 10;
//!
//! -- JoinScan is NOT proposed if price is NOT a fast field
//! -- (falls back to PostgreSQL's native join)
//! ```
//!
//! # Architecture
//!
//! ```text
//! ┌─────────────────┐ ┌──────────────────┐ ┌─────────────────────┐
//! │ PostgreSQL │ │ JoinScan │ │ DataFusion │
//! │ Planner │────▶│ Custom Scan │────▶│ Execution Plan │
//! │ (hook) │ │ (planning + │ │ │
//! │ │ │ execution) │ │ │
//! └─────────────────┘ └──────────────────┘ └─────────────────────┘
//! │
//! ▼
//! ┌─────────────────────┐
//! │ Tantivy │
//! │ (Scan + Search) │
//! └─────────────────────┘
//! ```
//!
//! ## Execution Strategy
//!
//! 1. **Planning**: During PostgreSQL planning, `JoinScan` hooks into the join path list.
//! It identifies potential search joins (including reconstructing multi-level joins from
//! PostgreSQL's optimal paths), extracts predicates, and builds a `JoinCSClause`.
//! 2. **Execution**: A DataFusion logical plan is constructed from the `JoinCSClause`.
//! This plan defines the join, filters, sorts, and limits.
//! 3. **DataFusion**: The plan is executed by DataFusion, which chooses the best join algorithm.
//! - Scans results from Tantivy for all relations, filtering by search predicates where applicable.
//! 4. **Result**: Joined tuples are returned to PostgreSQL via the Custom Scan interface.
//!
//! # Submodules
//!
//! - [`build`]: Data structures for planning serialization.
//! - [`planning`]: Cost estimation, condition extraction, field collection, pathkey handling.
//! - [`predicate`]: Transform PostgreSQL expressions to evaluable expression trees.
//! - [`scan_state`]: Execution state and DataFusion plan building.
//! - `translator`: Maps PostgreSQL expressions/columns to DataFusion expressions.
//! - [`privdat`]: Private data serialization between planning and execution.
//! - `explain`: EXPLAIN output formatting.
pub mod build;
pub mod planning;
pub mod predicate;
pub mod privdat;
pub mod range_partitioning_rule;
pub mod scan_state;
pub mod visibility_filter;
pub use self::build::CtidColumn;
use self::build::{DistinctMode, JoinCSClause, RelNode, RelationAlias};
use self::planning::{
collect_join_sources, collect_join_sources_base_rel, collect_required_fields,
ensure_score_bubbling, expr_uses_scores_from_source, extract_join_conditions, extract_orderby,
extract_orderby_from_parse_sort_clause, get_score_func_rti, order_by_columns_are_fast_fields,
order_by_sort_clause_is_fast_fields, pathkey_uses_scores_from_source,
};
use self::predicate::{extract_join_level_conditions, resolve_join_conditions};
use self::privdat::PrivateData;
use crate::postgres::customscan::datafusion::explain::{
explain_physical_plan, format_join_level_expr, get_attname_safe, get_plan_with_merged_metrics,
};
use crate::postgres::customscan::pullup::resolve_fast_field;
use self::scan_state::{
build_joinscan_logical_plan, build_physical_plan, build_task_context,
create_datafusion_session_context, JoinScanState,
};
use crate::api::HashSet;
use crate::api::OrderByFeature;
use crate::index::mvcc::{MvccSatisfies, SegmentView};
use crate::index::reader::index::SearchIndexManifest;
use crate::postgres::customscan::builders::custom_path::{CustomPathBuilder, Flags};
use crate::postgres::customscan::builders::custom_scan::CustomScanBuilder;
use crate::postgres::customscan::builders::custom_state::{
CustomScanStateBuilder, CustomScanStateWrapper,
};
use crate::postgres::customscan::explainer::Explainer;
use crate::postgres::customscan::joinscan::planning::{
distinct_collations_are_deterministic, distinct_columns_are_fast_fields,
};
use crate::postgres::customscan::limit_offset::LimitOffset;
use crate::postgres::customscan::mpp::glue::query_allows_parallel_mode;
use crate::postgres::customscan::mpp::interrupt::block_on_next;
use crate::postgres::customscan::mpp::launch::mpp_eligible;
use crate::postgres::customscan::mpp::launch::MppLifecycle;
use crate::postgres::customscan::mpp::worker_fragments::mpp_plan_has_data_parallelism;
use arrow_array::Array;
use datafusion_distributed::shm::MppMesh;
use crate::postgres::customscan::parameterized_value::ParameterizedValue;
use crate::postgres::customscan::solve_expr::SolvePostgresExpressions;
use crate::postgres::customscan::{CustomScan, JoinPathlistHookArgs};
use crate::postgres::heap::VisibilityChecker;
use crate::postgres::rel::PgSearchRelation;
use crate::postgres::ParallelScanArgs;
use crate::scan::codec::{deserialize_logical_plan_with_runtime, serialize_logical_plan};
use crate::DEFAULT_PARAMETERIZED_LIMIT_ESTIMATE;
use datafusion::physical_plan::ExecutionPlan;
use datafusion_distributed::DistributedExt;
use pgrx::{pg_guard, pg_sys, PgList};
use std::ffi::CStr;
use std::sync::Arc;
#[derive(Default)]
pub struct JoinScan;
/// Output of [`JoinScan::try_build_join_custom_path`] when activation succeeds.
struct BuiltJoinPath {
path: pg_sys::CustomPath,
aliases: Vec,
multi_table_clauses: Vec<*mut pg_sys::Expr>,
}
/// Why the join scan declined to produce a custom path.
///
/// `Quiet` is for the early "this isn't even a candidate join" gates;
/// `Warn` is for validation failures past the "considered interesting"
/// boundary, where we owe the planner a `WARNING`.
enum JoinPathDecline {
Quiet,
Warn {
reason: JoinDeclineReason,
aliases: Vec,
},
}
/// Specific reason a `JoinPathDecline::Warn` was raised. Wraps a specific
/// warning message and any optional details (e.g., unsupported join types)
/// to be emitted as a planner warning.
#[derive(Clone, Debug)]
pub enum JoinDeclineReason {
ContainsAggregate,
Message {
message: String,
details: Option>,
},
}
impl JoinDeclineReason {
pub fn new(message: impl Into) -> Self {
Self::Message {
message: message.into(),
details: None,
}
}
pub fn with_details(self, new_details: Vec) -> Self {
match self {
Self::ContainsAggregate => {
debug_assert!(false, "Cannot add details to ContainsAggregate");
self
}
Self::Message {
message,
details: _,
} => Self::Message {
message,
details: Some(new_details),
},
}
}
fn emit(&self, aliases: &[String]) {
if crate::gucs::planner_warnings() == crate::gucs::PlannerWarnings::Off {
return;
}
match self {
Self::ContainsAggregate => {
if crate::gucs::enable_aggregate_custom_scan() {
// We currently suppress this warning if the aggregate scan is enabled, assuming
// it will handle the query. If the aggregate scan later declines it, the user
// won't see a JoinScan warning.
// TODO: https://github.com/paradedb/paradedb/issues/5285 will fix this by allowing
// the joinscan to propose itself even when aggregates are present.
return;
}
JoinScan::add_planner_warning(
"JoinScan not used: aggregates are not supported. Enable paradedb.enable_aggregate_custom_scan to optimize this query.",
aliases,
);
}
Self::Message { message, details } => {
if let Some(details) = details {
JoinScan::add_detailed_planner_warning(message, aliases, details.clone());
} else {
JoinScan::add_planner_warning(message, aliases);
}
}
}
}
}
/// Recursively walk an expression tree and collect the `plan_id` of every
/// `T_SubPlan` node found at any depth. Uses Postgres's
/// `expression_tree_walker` so it handles all node types automatically.
unsafe fn collect_all_subplan_ids_from_expr(node: *mut pg_sys::Node, ids: &mut HashSet) {
if node.is_null() {
return;
}
#[pg_guard]
unsafe extern "C-unwind" fn walker(
node: *mut pg_sys::Node,
context: *mut std::ffi::c_void,
) -> bool {
if node.is_null() {
return false;
}
if (*node).type_ == pg_sys::NodeTag::T_SubPlan {
let subplan = node as *mut pg_sys::SubPlan;
let ids = &mut *(context as *mut HashSet);
ids.insert((*subplan).plan_id);
}
pg_sys::expression_tree_walker(node, Some(walker), context)
}
walker(node, ids as *mut HashSet as *mut std::ffi::c_void);
}
/// Collect all SubPlan `plan_id`s present in `baserestrictinfo` of the
/// given base relations.
unsafe fn collect_all_subplan_ids_from_baserestrictinfo(
root: *mut pg_sys::PlannerInfo,
absorbed_rtis: &[pg_sys::Index],
) -> HashSet {
let mut all_ids = HashSet::default();
for rti in absorbed_rtis {
let rel = pg_sys::find_base_rel(root, *rti as i32);
let ri_list = PgList::::from_pg((*rel).baserestrictinfo);
for ri in ri_list.iter_ptr() {
let clause = (*ri).clause as *mut pg_sys::Node;
collect_all_subplan_ids_from_expr(clause, &mut all_ids);
}
}
all_ids
}
/// Collect the `plan_id`s of SubPlans that JoinScan absorbed into
/// Semi/Anti/LeftMark join nodes in the `RelNode` tree.
fn collect_absorbed_subplan_ids(plan: &RelNode) -> HashSet {
let mut ids = HashSet::default();
walk_relnode_for_subplan_ids(plan, &mut ids);
ids
}
fn walk_relnode_for_subplan_ids(node: &RelNode, ids: &mut HashSet) {
match node {
RelNode::Join(j) => {
if let Some(plan_id) = j.subplan_id {
ids.insert(plan_id);
}
walk_relnode_for_subplan_ids(&j.left, ids);
walk_relnode_for_subplan_ids(&j.right, ids);
}
RelNode::Filter(f) => walk_relnode_for_subplan_ids(&f.input, ids),
RelNode::Unnest(u) => walk_relnode_for_subplan_ids(&u.input, ids),
RelNode::Scan(_) => {}
}
}
/// Check whether it is safe to push LIMIT into the JoinScan plan.
///
/// Returns `true` when ALL of:
/// 1. No plan node above the join consumes the full row set: window
/// functions, set-returning functions in the target list, and
/// GROUP BY / GROUPING SETS / HAVING all need every joined row, so a
/// LIMIT applied inside the scan starves them (issue #5561: an
/// unpartitioned `count(*) OVER ()` returned the LIMIT instead of
/// the true match count). `grouping_planner` sets `limit_tuples = -1` for exactly
/// these queries; the parse flags are checked directly because
/// `limit_tuples == -1` also means "parameterized LIMIT", which is
/// safe to push.
/// 2. JoinScan absorbed every base relation in the query (no outer
/// relations that could add post-filters above JoinScan).
/// 3. Every SubPlan in `baserestrictinfo` of absorbed relations was also
/// absorbed into the `RelNode` tree (Semi/Anti/LeftMark joins).
/// Un-absorbed SubPlans would become Postgres post-filters above
/// the capped output.
/// 4. No volatile functions in `baserestrictinfo` of absorbed relations
/// (volatile functions can never be pushed into Tantivy).
///
/// DISTINCT is not declined here: JoinScan absorbs and implements it
/// (#4669). `hasAggs` queries were already declined before this point.
unsafe fn is_limit_pushdown_safe(
root: *mut pg_sys::PlannerInfo,
join_clause: &JoinCSClause,
) -> Result<(), JoinDeclineReason> {
// 1. Nothing above the join may need more rows than the LIMIT keeps.
let parse = (*root).parse;
if (*parse).hasWindowFuncs {
return Err(JoinDeclineReason::new(
"JoinScan not used: LIMIT pushdown is unsafe due to window functions",
));
}
if (*parse).hasTargetSRFs {
return Err(JoinDeclineReason::new(
"JoinScan not used: LIMIT pushdown is unsafe due to set-returning functions in the target list",
));
}
if !(*parse).groupClause.is_null() || !(*parse).groupingSets.is_null() || (*root).hasHavingQual
{
return Err(JoinDeclineReason::new(
"JoinScan not used: LIMIT pushdown is unsafe due to GROUP BY or HAVING",
));
}
let absorbed_rtis: Vec = join_clause.plan.absorbed_rtis();
// 2. Did JoinScan absorb ALL base relations? `all_baserels`, not
// `all_query_rels`: the latter also carries outer-join relids (PG16+),
// which are join identities rather than relations and can never appear
// in `absorbed_rtis`.
let all_rels = (*root).all_baserels;
let mut absorbed_bms: *mut pg_sys::Bitmapset = std::ptr::null_mut();
for rti in &absorbed_rtis {
absorbed_bms = pg_sys::bms_add_member(absorbed_bms, *rti as i32);
}
if !pg_sys::bms_is_subset(all_rels, absorbed_bms) {
return Err(JoinDeclineReason::new(
"JoinScan not used: LIMIT pushdown is unsafe due to un-absorbed relations",
));
}
// 3. Every SubPlan in baserestrictinfo must have been absorbed.
let all_subplan_ids = collect_all_subplan_ids_from_baserestrictinfo(root, &absorbed_rtis);
let absorbed_subplan_ids = collect_absorbed_subplan_ids(&join_clause.plan);
for id in &all_subplan_ids {
if !absorbed_subplan_ids.contains(id) {
return Err(JoinDeclineReason::new(
"JoinScan not used: LIMIT pushdown is unsafe due to un-absorbed SubPlans",
));
}
}
// 4. No volatile functions (these can never be absorbed).
for rti in &absorbed_rtis {
let rel = pg_sys::find_base_rel(root, *rti as i32);
let ri_list = PgList::::from_pg((*rel).baserestrictinfo);
for ri in ri_list.iter_ptr() {
let clause = (*ri).clause as *mut pg_sys::Node;
if pg_sys::contain_volatile_functions(clause) {
return Err(JoinDeclineReason::new(
"JoinScan not used: LIMIT pushdown is unsafe due to volatile functions",
));
}
}
}
Ok(())
}
/// Try to create JoinScan `CustomPath`s for a single base relation that contains
/// SubPlan-based join opportunities (e.g. `col IN (SELECT ...) OR col IS NULL`).
///
/// Called from `set_rel_pathlist_hook` after BaseScan has been considered.
/// When PostgreSQL keeps a subquery as a SubPlan instead of flattening it into
/// a join, `set_join_pathlist_hook` never fires. This function gives JoinScan
/// a chance to handle those patterns.
pub unsafe fn try_create_subplan_join_paths(
root: *mut pg_sys::PlannerInfo,
rel: *mut pg_sys::RelOptInfo,
rti: pg_sys::Index,
) -> Vec {
use crate::postgres::customscan::range_table::bms_iter;
// Do not generate subplan join paths for relations proven empty (dummy rels).
if pg_sys::is_dummy_rel(rel) {
return Vec::new();
}
if !crate::gucs::enable_join_custom_scan() {
return Vec::new();
}
// Only consider base relations (single RTI).
let relids = (*rel).relids;
if relids.is_null() || pg_sys::bms_num_members(relids) != 1 {
return Vec::new();
}
let base_rti = match bms_iter(relids).next() {
Some(r) => r,
None => return Vec::new(),
};
if base_rti != rti {
return Vec::new();
}
// Quick pre-check: only proceed if baserestrictinfo contains an OR expression
// with a SubPlan inside. This avoids interfering with normal queries where
// set_join_pathlist_hook already handles SubPlans.
{
use crate::postgres::customscan::qual_inspect::is_subplan;
let bri = PgList::::from_pg((*rel).baserestrictinfo);
let has_or_subplan = bri.iter_ptr().any(|ri| {
let clause = (*ri).clause as *mut pg_sys::Node;
if clause.is_null() {
return false;
}
// Check if the clause itself is an OR BoolExpr containing a SubPlan.
// Top-level SubPlans are handled by the normal join_pathlist hook.
if (*clause).type_ == pg_sys::NodeTag::T_BoolExpr {
let bexpr = clause as *mut pg_sys::BoolExpr;
(*bexpr).boolop == pg_sys::BoolExprType::OR_EXPR && is_subplan(clause, root)
} else {
false
}
});
if !has_or_subplan {
return Vec::new();
}
}
// Try to extract SubPlan-based joins from baserestrictinfo.
let collected = match collect_join_sources_base_rel(root, rel, rti) {
Some(res) => res,
None => return Vec::new(),
};
let plan = collected.plan;
let join_keys = collected.join_keys;
// Only proceed if the plan actually contains a join (from SubPlan extraction).
if !plan.has_semi_or_anti() {
return Vec::new();
}
// Phase 1: validate + build JoinCSClause.
let join_clause = match JoinScan::validate_and_build_clause(root, rel, &plan, &join_keys) {
Ok(res) => res,
Err(reason) => {
let aliases: Vec = plan
.sources()
.iter()
.map(|s| {
RelationAlias::new(s.scan_info.alias.as_deref())
.warning_context(s.scan_info.heaprelid)
})
.collect();
reason.emit(&aliases);
return Vec::new();
}
};
// No join-level predicate extraction needed for SubPlan-based paths.
// Phase 2: finalize into CustomPath.
match JoinScan::finalize_clause_into_path(root, rel, join_clause) {
Some(path) => vec![path],
None => Vec::new(),
}
}
impl JoinScan {
/// Phase 1: Validate a `RelNode` plan against JoinScan activation requirements
/// and build a `JoinCSClause` with score bubbling and partitioning applied.
///
/// Returns `Err` with a descriptive message if any activation check fails.
/// The caller can then optionally perform join-level predicate extraction on
/// the returned clause before passing it to `finalize_clause_into_path`.
unsafe fn validate_and_build_clause(
root: *mut pg_sys::PlannerInfo,
rel: *mut pg_sys::RelOptInfo,
plan: &RelNode,
join_keys: &[build::JoinKeyPair],
) -> Result {
let all_sources = plan.sources();
// --- Activation checks ---
// NOTE: We do NOT check has_search_predicate here. The caller is
// responsible for that check because the join-hook path also considers
// join_conditions.has_search_predicate, which is not available to us.
if all_sources
.iter()
.any(|s| s.scan_info.indexrelid == pg_sys::InvalidOid)
{
return Err(JoinDeclineReason::new(
"JoinScan not used: at least one relation lacks a BM25 index",
));
}
if (*(*root).parse).hasAggs {
return Err(JoinDeclineReason::ContainsAggregate);
}
// Require LIMIT for top-level queries (without it, JoinScan's TopK
// optimization has no bound). Subqueries are exempt because the parent
// plan provides the cardinality constraint. Either a static or
// parameterized LIMIT is sufficient — the latter is resolved at
// execution time from EState::es_param_list_info.
let limit_offset = LimitOffset::from_root(root);
let is_subquery = !(*root).parent_root.is_null();
if limit_offset.is_none() && !is_subquery {
return Err(JoinDeclineReason::new(
"JoinScan not used: LIMIT is required for top-level queries",
));
}
let parse = (*root).parse;
let query_has_distinct = !(*parse).distinctClause.is_null();
let distinct_list_len = if query_has_distinct {
PgList::::from_pg((*parse).distinctClause)
.iter_ptr()
.count()
} else {
0
};
let scan_tlist_len =
if rel.is_null() || (*rel).reltarget.is_null() || (*(*rel).reltarget).exprs.is_null() {
0
} else {
PgList::::from_pg((*(*rel).reltarget).exprs).len()
};
// Decide whether DISTINCT can be absorbed by JoinScan or must be deferred to PostgreSQL:
// JoinScan can execute expressions in DISTINCT (e.g. `SELECT DISTINCT upper(col)`) as
// long as each distinct expression is represented in `reltarget`.
// However, when a query selects both a column and expressions derived from it (e.g.
// `SELECT DISTINCT col, col IS NULL`), PostgreSQL strips the derived expressions from
// the scan's `reltarget`, requesting only the underlying base columns from CustomScan
// and planning an upper `Result` node to evaluate `col IS NULL`.
// When `distinct_list_len > scan_tlist_len`, DataFusion's output projection lacks
// slots for those extra expressions, so DataFusion cannot group by all DISTINCT keys.
// We defer DISTINCT to Postgres's upper Unique/HashAggregate node and sort only by `sortClause`.
let distinct_mode = if query_has_distinct {
if distinct_list_len > scan_tlist_len {
DistinctMode::Deferred
} else {
DistinctMode::Active
}
} else {
DistinctMode::None
};
let is_distinct_active = distinct_mode == DistinctMode::Active;
if is_distinct_active && !distinct_collations_are_deterministic(root) {
return Err(JoinDeclineReason::new(
"JoinScan not used: DISTINCT on a nondeterministic collation is not supported",
));
}
let output_sources = plan.output_sources();
if is_distinct_active && distinct_columns_are_fast_fields(root, &output_sources).is_none() {
return Err(JoinDeclineReason::new(
"JoinScan not used: DISTINCT columns must be fast fields",
));
}
if distinct_mode == DistinctMode::Deferred {
order_by_sort_clause_is_fast_fields(root, &output_sources)?;
} else {
order_by_columns_are_fast_fields(root, &output_sources, is_distinct_active)?;
}
for jk in join_keys {
let outer_source = all_sources.iter().find(|s| s.contains_rti(jk.outer_rti));
let inner_source = all_sources.iter().find(|s| s.contains_rti(jk.inner_rti));
match (outer_source, inner_source) {
(Some(outer), Some(inner)) => {
let outer_hr = PgSearchRelation::open(outer.scan_info.heaprelid);
let outer_ir = PgSearchRelation::open(outer.scan_info.indexrelid);
let inner_hr = PgSearchRelation::open(inner.scan_info.heaprelid);
let inner_ir = PgSearchRelation::open(inner.scan_info.indexrelid);
if resolve_fast_field(jk.outer_attno as i32, &outer_hr.tuple_desc(), &outer_ir)
.is_none()
|| resolve_fast_field(
jk.inner_attno as i32,
&inner_hr.tuple_desc(),
&inner_ir,
)
.is_none()
{
return Err(JoinDeclineReason::new(
"JoinScan not used: join conditions must reference columnar indexed fields",
));
}
}
_ => {
return Err(JoinDeclineReason::new(
"JoinScan not used: failed to resolve join keys to base relations",
))
}
}
}
// Why "derived" expressions in DISTINCT prevent LIMIT pushdown:
//
// JoinScan can evaluate expressions in DISTINCT (e.g. `SELECT DISTINCT upper(col)`)
// whenever every DISTINCT item maps 1-to-1 to a target list slot in the scan.
// However, when a query combines a base column and expressions computed from that
// same column (e.g. `SELECT DISTINCT col, col IS NULL`, or function calls alongside
// base columns), PostgreSQL optimizes `rel->reltarget` by stripping out the computed
// expression, asking the scan to emit only the underlying base column `col`.
// PostgreSQL then plans an upper `Result` node above the scan to evaluate `col IS NULL`,
// followed by an upper `Unique` or `HashAggregate` node to deduplicate across all keys.
//
// Because a CustomScan's output tuple must conform to `rel->reltarget`, JoinScan has
// no output slot to emit the upper expression, and therefore cannot perform the full
// DISTINCT deduplication inside DataFusion (`distinct_mode == DistinctMode::Deferred`).
//
// In turn, pushing down LIMIT below PostgreSQL's upper `Unique`/`Sort` node is unsound:
// upper deduplication collapses rows (risking fewer emitted rows than LIMIT), and upper
// Sort breaks ties using DISTINCT keys that an early LIMIT would arbitrarily truncate.
// Without LIMIT pushdown, a top-level JoinScan cannot perform Top-K early stopping and
// would force an unbounded scan followed by a full sort in PostgreSQL, losing its primary
// performance advantage. JoinScan therefore declines and falls back to PostgreSQL.
if distinct_mode == DistinctMode::Deferred && !is_subquery {
return Err(JoinDeclineReason::new(
"JoinScan not used: LIMIT cannot be pushed down when DISTINCT contains derived expressions (e.g. IS NULL or function calls)",
));
}
let pushdown_limit_offset = if distinct_mode == DistinctMode::Deferred {
None
} else {
limit_offset
};
// --- Build JoinCSClause ---
let mut join_clause = JoinCSClause::new(plan.clone())
.with_limit_offset(pushdown_limit_offset)
.with_distinct(distinct_mode);
for source in join_clause.plan.sources_mut() {
let score_in_tlist =
expr_uses_scores_from_source((*root).processed_tlist.cast(), source);
let score_in_pathkey = pathkey_uses_scores_from_source(root, source);
if score_in_tlist || score_in_pathkey {
ensure_score_bubbling(source);
}
}
// Safety check: bail out if an upper plan node needs the full row set,
// or if un-absorbed relations/SubPlans or volatile predicates could
// change the capped output.
if join_clause.limit_offset.is_some() {
is_limit_pushdown_safe(root, &join_clause)?;
}
Ok(join_clause)
}
/// Phase 2: Finalize a validated `JoinCSClause` into a `CustomPath` by
/// extracting ORDER BY, computing costs/parallel workers, and building the
/// `pg_sys::CustomPath` struct.
///
/// Row count estimates are NOT divided by the number of parallel workers here.
/// Since the custom scan natively manages its own parallel execution via MPP and
/// communicates directly with the leader without a Postgres `Gather` node, Postgres
/// sees a single path emitting the full `result_rows`. Similarly, DataFusion's
/// optimizer requires the undivided `estimated_rows_per_worker` to accurately assess
/// the full scale of the data and choose partitioned joins appropriately.
///
/// Returns `None` if ORDER BY extraction fails.
#[allow(clippy::needless_update)]
unsafe fn finalize_clause_into_path(
root: *mut pg_sys::PlannerInfo,
rel: *mut pg_sys::RelOptInfo,
mut join_clause: JoinCSClause,
) -> Option {
let output_rtis = join_clause.plan.output_rtis();
let current_sources = join_clause.plan.sources();
let order_by = if join_clause.is_distinct_deferred() {
// DISTINCT was deferred to a parent node; query_pathkeys still lists keys for the full
// DISTINCT row. Sort only by parse->sortClause.
extract_orderby_from_parse_sort_clause(root, ¤t_sources, &output_rtis)?
} else {
extract_orderby(
root,
¤t_sources,
&output_rtis,
join_clause.has_distinct(),
)?
};
join_clause = join_clause.with_order_by(order_by);
// --- Cost estimation ---
let startup_cost = crate::DEFAULT_STARTUP_COST;
let total_cost = startup_cost + 1.0;
let result_rows = join_clause
.limit_offset
.as_ref()
.map(|lo| lo.planning_estimate())
.unwrap_or_else(|| {
if !rel.is_null() && (*rel).rows > 0.0 {
(*rel).rows
} else {
DEFAULT_PARAMETERIZED_LIMIT_ESTIMATE
}
});
let has_order_by = !join_clause.order_by.is_empty();
let order_by_len = join_clause.order_by.len();
let is_distinct_deferred = join_clause.is_distinct_deferred();
let query_pathkeys_len = PgList::::from_pg((*root).query_pathkeys).len();
let private_data = PrivateData::new(join_clause, query_allows_parallel_mode(&*root));
let mut custom_path = pg_sys::CustomPath {
path: pg_sys::Path {
type_: pg_sys::NodeTag::T_CustomPath,
pathtype: pg_sys::NodeTag::T_CustomScan,
parent: rel,
pathtarget: (*rel).reltarget,
param_info: pg_sys::get_baserel_parampathinfo(
root,
rel,
pg_sys::bms_copy((*rel).lateral_relids),
),
rows: result_rows,
startup_cost,
total_cost,
..Default::default()
},
flags: Flags::Force as u32,
methods: JoinScan::custom_path_methods(),
custom_private: private_data.into(),
custom_paths: std::ptr::null_mut(),
..Default::default()
};
if !is_distinct_deferred && has_order_by && order_by_len == query_pathkeys_len {
custom_path.path.pathkeys = (*root).query_pathkeys;
}
Some(custom_path)
}
}
impl JoinScan {
/// Capture lightweight segment manifests for all join sources.
///
/// Uses `SearchIndexManifest::capture` instead of opening full `SearchIndexReader`s:
/// manifests are cheap, hold the Tantivy segment pins the launch and workers rely on,
/// and avoid calling `into_tantivy_query` on each source's `scan_info.query` — which
/// fails for parameterized predicates (prepared statements, initplan-backed subqueries)
/// that require a `PlanState` to evaluate.
///
/// Manifests also provide consistent segment counts for both DSM sizing and DSM
/// population, avoiding the divergence that can occur when planning-time counts
/// (from `scan_info.segment_count`) differ from execution-time counts due to
/// concurrent inserts.
fn ensure_source_manifests(state: &mut CustomScanStateWrapper) {
if !state.custom_state().source_manifests.is_empty() {
return;
}
let manifests = state
.custom_state()
.join_clause
.plan
.sources()
.iter()
.map(|source| {
let rel = PgSearchRelation::open(source.scan_info.indexrelid);
SearchIndexManifest::capture(&rel, MvccSatisfies::Snapshot).unwrap_or_else(|e| {
panic!(
"Failed to capture source manifest for indexrelid {}: {e}",
source.scan_info.indexrelid
)
})
})
.collect();
state.custom_state_mut().source_manifests = manifests;
}
/// Build the plan_position → segment view list every reader of a source replays.
///
/// This is keyed by plan_position rather than indexrelid because it is a
/// per-source contract, not just a per-index one. The same index can appear
/// more than once in one JoinScan plan; in parallel execution those source
/// copies can also carry different canonical segment sets (partitioned vs
/// replicated). If this were keyed only by indexrelid, one source could
/// inject another source's segment set and make packed DocAddresses resolve
/// against the wrong segment ordering.
///
/// The manifests are also what populates the DSM, so the leader's providers and every
/// worker reader open the same view.
fn build_source_manifests(
state: &mut CustomScanStateWrapper,
_join_clause: &JoinCSClause,
plan_sources: &[&build::JoinSource],
) -> Vec {
Self::ensure_source_manifests(state);
(0..plan_sources.len())
.map(|plan_position| {
state
.custom_state()
.source_manifests
.get(plan_position)
.unwrap_or_else(|| {
panic!(
"missing source manifest for join source at plan_position \
{plan_position}"
)
})
.clone()
})
.collect()
}
/// Whether any `SearchQueryInput` reachable from `join_clause` — scan-node queries AND
/// `join_level_predicates` — still carries a Param or PostgresExpression(SubPlan) the
/// executor needs to resolve. Clause-wide (not per-source): a cross-relation predicate
/// (e.g. an `OR` spanning two joined tables) lives in `join_level_predicates`, outside
/// every source's `scan_info.query`, so a per-source-only check misses it. Read-only
/// (`&JoinCSClause`) so EXPLAIN's `&CustomScanStateWrapper` call site doesn't need a
/// mutable borrow just to ask a yes/no question.
fn source_queries_need_executor_state(join_clause: &JoinCSClause) -> bool {
join_clause.has_postgres_expressions() || join_clause.has_parameters()
}
}
impl JoinScan {
/// Build the leader's distributed session context for this JoinScan query. Thin wrapper
/// over the shared [`crate::postgres::customscan::mpp::exec_worker::build_mpp_session_context`]
/// that seeds with `create_datafusion_session_context(SessionContextProfile::Join)`.
/// `mesh = None` is the EXPLAIN-time path. See the shared helper's doc.
fn build_mpp_session_context(
mesh: Option>,
) -> datafusion::prelude::SessionContext {
crate::postgres::customscan::mpp::exec_worker::build_mpp_session_context(
create_datafusion_session_context(),
mesh,
)
}
/// First-exec solve of Param/SubPlan-backed expressions on the leader, rebaking the
/// logical plan afterward so the bytes used for planning (both MPP dispatch and the
/// leader's own execution) are fully resolved. Called before `build_plan` in
/// `exec_custom_scan`, ahead of the plan-first MPP launch (#5667). Rebaking is required once
/// expressions have been solved: failures abort at the operation that violated the invariant
/// rather than being represented as a recoverable absence of plan bytes.
fn maybe_solve_and_rebake(state: &mut CustomScanStateWrapper) {
if !Self::source_queries_need_executor_state(&state.custom_state().join_clause) {
return;
}
let planstate = state.planstate();
let expr_context = state.runtime_context;
state
.custom_state_mut()
.prepare_query_for_execution(planstate, expr_context);
let bytes = unsafe { Self::rebake_for_mpp(state) };
// Keep the leader's own execution in sync with what's dispatched to workers:
// exec_custom_scan's build_plan closure deserializes this same field.
state.custom_state_mut().logical_plan = Some(bytes::Bytes::from(bytes));
}
/// Plan-first MPP launch (#5667). Called with the leader's already-built physical plan:
/// sizes the producer pool from the plan's widest stage, builds the DSM, stamps the shared
/// scan state into the plan, and spawns exactly the needed workers. `None` means run
/// serially — the plan had nothing to distribute (no workers were forked at all) or the
/// launch fell back.
fn launch_mpp(
state: &mut CustomScanStateWrapper,
physical: &Arc,
) -> Option {
Self::ensure_source_manifests(state);
let all_sources: Vec = state
.custom_state()
.source_manifests
.iter()
.map(|manifest| manifest.segment_view())
.collect();
let args = ParallelScanArgs {
all_sources,
query: vec![],
with_aggregates: false,
with_segment_info: false,
};
crate::postgres::customscan::mpp::launch::launch_mpp_join(physical, args)
}
/// Re-bake the DataFusion logical plan from the current (post-solve) `join_clause`, so the
/// plan bytes used for dispatch and for the leader's own execution have every
/// Param/SubPlan-backed `SearchQueryInput` fully resolved.
unsafe fn rebake_for_mpp(state: &mut CustomScanStateWrapper) -> Vec {
Self::rebake_from_custom_exprs_string(state, false)
}
/// Re-bake with `mpp_source_idx` forced to `None` on every source after a short-launch
/// decline. Physical replanning does not rewrite provider metadata already serialized in
/// the logical plan, so the serial fallback needs its own logical shape.
unsafe fn rebake_for_mpp_fallback(state: &mut CustomScanStateWrapper) -> Vec {
Self::rebake_from_custom_exprs_string(state, true)
}
unsafe fn rebake_from_custom_exprs_string(
state: &mut CustomScanStateWrapper,
force_serial: bool,
) -> Vec {
let custom_exprs: *mut pg_sys::List = match &state.custom_state().custom_exprs_string {
Some(s) => {
let cstr = std::ffi::CString::new(s.as_str())
.expect("nodeToString snapshot cannot contain an interior NUL");
pg_sys::stringToNode(cstr.as_ptr()).cast()
}
None => std::ptr::null_mut(),
};
let mut private_data = PrivateData::new(
state.custom_state().join_clause.clone(),
state.custom_state().parallel_mode_ok,
);
private_data.output_columns = state.custom_state().output_columns.clone();
bake_logical_plan(&mut private_data, custom_exprs, force_serial);
private_data
.logical_plan
.expect("rebaking must produce serialized logical plan bytes")
.to_vec()
}
/// Join the MPP producer workers and destroy the parallel context once nothing
/// references the ring mesh. Drains any remaining worker metrics first, takes the
/// leader out (its mesh handle drops with it), then drops the stream/plan/runtime
/// (all carry mesh references) before `wait_for_finish` destroys the DSM. Used at
/// EndCustomScan and before a correlated rescan relaunches a fresh worker set.
fn finish_mpp_execution(state: &mut CustomScanStateWrapper) {
state.custom_state_mut().datafusion_stream = None;
if let Some(leader) = state.custom_state().mpp.leader() {
if let Some(plan) = state.custom_state().physical_plan.as_ref() {
crate::postgres::customscan::mpp::glue::drain_worker_metrics(
plan,
&leader.session.mesh,
);
}
}
let finish = match state.custom_state_mut().mpp.take_leader() {
Some(mut leader) => leader.finish.take(),
None => None,
};
if finish.is_some() {
let cs = state.custom_state_mut();
cs.current_batch = None;
cs.physical_plan = None;
cs.runtime = None;
}
if let Some(finish) = finish {
finish.wait_for_finish();
}
}
}
impl CustomScan for JoinScan {
const NAME: &'static CStr = c"ParadeDB Join Scan";
type Args = JoinPathlistHookArgs;
type State = JoinScanState;
type PrivateData = PrivateData;
fn exec_methods() -> pg_sys::CustomExecMethods {
pg_sys::CustomExecMethods {
CustomName: Self::NAME.as_ptr(),
BeginCustomScan: Some(crate::postgres::customscan::exec::begin_custom_scan::),
ExecCustomScan: Some(crate::postgres::customscan::exec::exec_custom_scan::),
EndCustomScan: Some(crate::postgres::customscan::exec::end_custom_scan::),
ReScanCustomScan: Some(crate::postgres::customscan::exec::rescan_custom_scan::),
MarkPosCustomScan: None,
RestrPosCustomScan: None,
EstimateDSMCustomScan: None,
InitializeDSMCustomScan: None,
ReInitializeDSMCustomScan: None,
InitializeWorkerCustomScan: None,
ShutdownCustomScan: Some(
crate::postgres::customscan::exec::shutdown_custom_scan::,
),
ExplainCustomScan: Some(crate::postgres::customscan::exec::explain_custom_scan::),
}
}
fn create_custom_path(builder: CustomPathBuilder) -> Vec {
unsafe {
match Self::try_build_join_custom_path(&builder) {
Ok(BuiltJoinPath {
path,
aliases,
multi_table_clauses,
}) => {
let mut path = path;
if !multi_table_clauses.is_empty() {
let mut private_list = PgList::::from_pg(path.custom_private);
for clause in multi_table_clauses {
private_list.push(clause.cast());
}
path.custom_private = private_list.into_pg();
}
Self::mark_contexts_successful(&aliases);
vec![path]
}
Err(JoinPathDecline::Quiet) => Vec::new(),
Err(JoinPathDecline::Warn { reason, aliases }) => {
reason.emit(&aliases);
Vec::new()
}
}
}
}
fn plan_custom_path(mut builder: CustomScanBuilder) -> pg_sys::CustomScan {
// For joins, scanrelid must be 0 (it's not scanning a single relation)
builder.set_scanrelid(0);
// Get best_path before builder is consumed
let best_path = builder.args().best_path;
let root = builder.args().root;
let mut node = builder.build();
unsafe {
// For joins, we need to set both custom_scan_tlist and scan.plan.targetlist
// to describe the output columns. We create a fresh copy of the target list
// to avoid corrupting the original.
let mut tlist = PgList::::from_pg(
pg_sys::copyObjectImpl(node.scan.plan.targetlist.cast()).cast(),
);
// If the parent plan (`processed_tlist` or `pathkeys`) needs `pdb.score(...)` but the
// planner didn't push it down into the target list, we must add it. Otherwise,
// the parent node will attempt to evaluate `pdb.score(...)` natively, which fails.
crate::postgres::utils::add_missing_search_operators_to_tlist(
root,
best_path as *mut pg_sys::Path,
&mut tlist,
&crate::postgres::customscan::score_funcoids(),
);
// Update node.scan.plan.targetlist so parent nodes can reference the outputs.
// We also set custom_scan_tlist, which is what setrefs.c uses to translate
// Vars in custom_exprs into INDEX_VAR references. We must provide a separate
// copy to custom_scan_tlist because setrefs.c may modify it in-place.
let tlist_ptr = tlist.into_pg();
node.scan.plan.targetlist = tlist_ptr;
let mut custom_scan_tlist = PgList::::from_pg(
pg_sys::copyObjectImpl(tlist_ptr.cast()).cast(),
);
let path_private_full = PgList::::from_pg((*best_path).custom_private);
for i in 1..path_private_full.len() {
if let Some(node_ptr) = path_private_full.get_ptr(i) {
crate::postgres::utils::add_vars_to_tlist(node_ptr, &mut custom_scan_tlist);
}
}
let custom_scan_tlist_ptr = custom_scan_tlist.into_pg();
node.custom_scan_tlist = custom_scan_tlist_ptr;
// For join custom scans, PostgreSQL doesn't pass clauses via the usual parameter.
// We stored the restrictlist in custom_private during create_custom_path.
//
// Note: We do NOT add restrictlist clauses to custom_exprs because setrefs would try
// to resolve their Vars using the child plans' target lists, which may not have all
// the needed columns. Instead, we keep the restrictlist in custom_private and handle
// join condition evaluation manually during execution using the original Var
// references.
// Add heap condition clauses to custom_exprs so they get transformed by
// set_customscan_references. The Vars in these expressions will be converted to
// INDEX_VAR references into custom_scan_tlist.
node.custom_exprs = splice_path_private_into_list(node.custom_exprs, best_path);
// Ensure all Vars referenced in custom_exprs (heap condition clauses)
// are present in custom_scan_tlist so setrefs can create INDEX_VAR references.
let path_private_full = PgList::::from_pg((*best_path).custom_private);
let mut scan_tlist = PgList::::from_pg(node.custom_scan_tlist);
for i in 1..path_private_full.len() {
if let Some(node_ptr) = path_private_full.get_ptr(i) {
crate::postgres::utils::add_vars_to_tlist(node_ptr, &mut scan_tlist);
}
}
node.custom_scan_tlist = scan_tlist.into_pg();
// Extract the column mappings from custom_scan_tlist (including any added
// custom_exprs Vars) so INDEX_VAR references can be resolved.
let mut private_data = PrivateData::from(node.custom_private);
private_data.output_columns =
compute_output_columns(&private_data.join_clause, node.custom_scan_tlist, root);
let updated_entries = PgList::::from_pg(tlist_ptr);
build_output_projection(&mut private_data, &updated_entries, root);
// Snapshot custom_exprs before setrefs rewrites it, for MPP re-baking.
// Needed for any of: maybe_solve_and_rebake (resolves Param/PostgresExpression
// nodes) or rebake_for_mpp_fallback (serial replan on a short-launch decline,
// which can hit any MPP query regardless of params). Both funnel into
// translate_custom_exprs, which also resolves JoinLevelExpr::MultiTablePredicate
// nodes (cross-relation predicates like `users.age < products.age`) by indexing
// into custom_exprs — so this is required whenever the plan has params,
// postgres expressions, OR multi-table predicates, not just the first two.
let needs_rebake_snapshot = private_data.join_clause.has_parameters()
|| private_data.join_clause.has_postgres_expressions()
|| private_data.join_clause.has_multi_table_predicates();
private_data.custom_exprs_string =
if !needs_rebake_snapshot || node.custom_exprs.is_null() {
None
} else {
// Non-lossy round-trip: fail loudly on invalid UTF-8 (should never happen —
// nodeToString produces ASCII-safe output) rather than to_string_lossy()
// silently mangling bytes stringToNode would later fail to parse, or worse,
// parse into a different node than what was captured. Mirrors
// PostgresPointer::serialize's handling of the same nodeToString/stringToNode
// round-trip in query/mod.rs.
let s = pg_sys::nodeToString(node.custom_exprs.cast());
let owned = std::ffi::CStr::from_ptr(s)
.to_str()
.expect("nodeToString output must be valid UTF-8")
.to_owned();
pg_sys::pfree(s.cast());
Some(owned)
};
// Collect all required fields for execution
collect_required_fields(
&mut private_data.join_clause,
&private_data.output_columns,
node.custom_exprs,
);
bake_logical_plan(&mut private_data, node.custom_exprs, false);
// Convert PrivateData back to a list and preserve the restrictlist.
let private_list = PrivateData::into(private_data);
node.custom_private = splice_path_private_into_list(private_list, best_path);
}
node
}
fn create_custom_scan_state(
mut builder: CustomScanStateBuilder,
) -> *mut CustomScanStateWrapper {
let join_clause = builder.custom_private().join_clause.clone();
let has_params = join_clause.has_parameters() || join_clause.has_postgres_expressions();
builder.custom_state().base_join_clause = if has_params {
Some(join_clause.clone())
} else {
None
};
builder.custom_state().join_clause = join_clause;
builder.custom_state().output_columns = builder.custom_private().output_columns.clone();
builder.custom_state().logical_plan = builder.custom_private().logical_plan.clone();
builder.custom_state().custom_exprs_string =
builder.custom_private().custom_exprs_string.clone();
builder.custom_state().parallel_mode_ok = builder.custom_private().parallel_mode_ok;
builder.build()
}
fn explain_custom_scan(
state: &CustomScanStateWrapper,
_ancestors: *mut pg_sys::List,
explainer: &mut Explainer,
) {
let join_clause = &state.custom_state().join_clause;
explainer.add_text("Relation Tree", join_clause.plan.explain());
let mut base_relations = Vec::new();
join_clause.collect_base_relations(&mut base_relations);
fn collect_join_cond_strings(node: &RelNode, acc: &mut Vec) {
match node {
RelNode::Scan(_) => {}
RelNode::Filter(filter) => collect_join_cond_strings(&filter.input, acc),
RelNode::Unnest(unnest) => collect_join_cond_strings(&unnest.input, acc),
RelNode::Join(join) => {
for jk in &join.equi_keys {
let ((left_source, left_attno), (right_source, right_attno)) = jk
.resolve_against(&join.left, &join.right)
.unwrap_or_else(|| {
panic!(
"Failed to resolve join key to current join sides: outer_rti={}, inner_rti={}",
jk.outer_rti, jk.inner_rti
)
});
let (outer_source, outer_attno, inner_source, inner_attno) =
if join.left.contains_rti(jk.outer_rti)
&& join.right.contains_rti(jk.inner_rti)
{
(left_source, left_attno, right_source, right_attno)
} else {
(right_source, right_attno, left_source, left_attno)
};
let outer_alias =
RelationAlias::new(outer_source.scan_info.alias.as_deref())
.display(outer_source.plan_position);
let inner_alias =
RelationAlias::new(inner_source.scan_info.alias.as_deref())
.display(inner_source.plan_position);
acc.push(format!(
"{} = {}",
get_attname_safe(
Some(outer_source.scan_info.heaprelid),
outer_attno,
&outer_alias
),
get_attname_safe(
Some(inner_source.scan_info.heaprelid),
inner_attno,
&inner_alias
)
));
}
collect_join_cond_strings(&join.left, acc);
collect_join_cond_strings(&join.right, acc);
}
}
}
fn collect_join_filter_strings(
node: &RelNode,
join_clause: &JoinCSClause,
explainer: &Explainer,
acc: &mut Vec,
) {
match node {
RelNode::Scan(_) => {}
RelNode::Filter(filter) => {
collect_join_filter_strings(&filter.input, join_clause, explainer, acc);
}
RelNode::Unnest(unnest) => {
collect_join_filter_strings(&unnest.input, join_clause, explainer, acc);
}
RelNode::Join(join) => {
if let Some(filter) = &join.filter {
acc.push(format_join_level_expr(filter, join_clause, explainer));
}
collect_join_filter_strings(&join.left, join_clause, explainer, acc);
collect_join_filter_strings(&join.right, join_clause, explainer, acc);
}
}
}
let mut keys_str = Vec::new();
collect_join_cond_strings(&join_clause.plan, &mut keys_str);
if !keys_str.is_empty() {
explainer.add_text("Join Cond", keys_str.join(", "));
}
let mut filters_str = Vec::new();
collect_join_filter_strings(&join_clause.plan, join_clause, explainer, &mut filters_str);
if !filters_str.is_empty() {
explainer.add_text("Join Filter", filters_str.join(" AND "));
}
if let Some(expr) = join_clause.plan.join_level_expr() {
explainer.add_text(
"Join Predicate",
format_join_level_expr(expr, join_clause, explainer),
);
}
if let Some(lo) = &join_clause.limit_offset {
explainer.add_text("Limit", lo.limit.to_string());
if let Some(off) = &lo.offset {
// Suppress an explicit "0" OFFSET to match prior EXPLAIN output.
let suppress = matches!(off, ParameterizedValue::Static(0));
if !suppress {
explainer.add_text("Offset", off.to_string());
}
}
}
if join_clause.has_distinct() {
explainer.add_text("Distinct", "true");
}
if !join_clause.order_by.is_empty() {
explainer.add_text(
"Order By",
join_clause
.order_by
.iter()
.map(|oi| match &oi.feature {
OrderByFeature::Field { name: f, .. } => {
format!("{} {}", f, oi.direction.as_ref())
}
OrderByFeature::Var { rti, attno, name } => {
if let Some(info) = base_relations.iter().find(|i| i.heap_rti == *rti) {
let col_name = get_attname_safe(
Some(info.heaprelid),
*attno,
info.alias.as_deref().unwrap_or("?"),
);
format!("{} {}", col_name, oi.direction.as_ref())
} else {
format!(
"{} {}",
name.as_deref().unwrap_or("?"),
oi.direction.as_ref()
)
}
}
other => {
format!("{other} {}", oi.direction.as_ref())
}
})
.collect::>()
.join(", "),
);
}
if explainer.is_analyze() {
// For EXPLAIN ANALYZE, render the plan with metrics inline.
// VERBOSE includes timing; without VERBOSE, timing is stripped for stable output.
if let Some(ref physical_plan) = state.custom_state().physical_plan {
// Under MPP the worker fragments reported their metrics as mesh frames when
// they exited; fold them in so the stage boxes show more than the leader's
// own nodes.
let plan = get_plan_with_merged_metrics(
physical_plan,
state.custom_state().mpp.leader().is_some(),
state.custom_state().runtime.is_some(),
explainer,
);
explain_physical_plan(&plan, explainer);
}
// The MPP launch floor (worker spawn, ring attach, plan dispatch) lives outside the
// DataFusion plan, so surface its per-phase breakdown separately when the query ran
// distributed.
if explainer.is_verbose() {
if let Some(t) = state.custom_state().launch_timing {
explainer.add_text("MPP Launch", t.explain_text());
}
}
} else if let Some(ref logical_plan) = state.custom_state().logical_plan {
// Plain EXPLAIN reconstructs the physical plan by deserializing the logical
// plan and calling PgSearchTableProvider::scan(), but without executor state
// (planstate=None). If any source query contains a PostgresExpression
// (e.g., prepared-statement parameter), scan() would fail with "missing
// planstate". Skip physical plan display for those cases.
if Self::source_queries_need_executor_state(&state.custom_state().join_clause) {
explainer.add_text(
"DataFusion Physical Plan",
"omitted for EXPLAIN because source queries require executor-time expression resolution",
);
return;
}
// For plain EXPLAIN, reconstruct the plan using the same session configuration
// that execution uses so `VisibilityFilterExec` appears in the displayed plan,
// matching EXPLAIN ANALYZE. When MPP is active, first build against
// `producer_worker_cap()` (`mesh = None`); if task discovery says launch will
// not run (#5784 / `max_producer_task_count < 2`), rebuild serially so the
// printed shape matches execution.
let expr_context = crate::postgres::utils::ExprContextGuard::new();
let runtime = tokio::runtime::Builder::new_current_thread()
.build()
.expect("Failed to create tokio runtime");
let build_with = |ctx: &datafusion::prelude::SessionContext| {
let logical_plan = deserialize_logical_plan_with_runtime(
logical_plan,
&ctx.task_ctx(),
None,
Some(expr_context.as_ptr()),
None,
vec![],
)
.expect("Failed to deserialize logical plan");
runtime
.block_on(build_physical_plan(ctx, logical_plan))
.expect("Failed to create execution plan")
};
let physical_plan = if mpp_eligible(
state.custom_state().parallel_mode_ok,
&state.custom_state().join_clause.plan,
) {
let mpp_plan = build_with(&Self::build_mpp_session_context(None));
if mpp_plan_has_data_parallelism(&mpp_plan) {
mpp_plan
} else {
build_with(&create_datafusion_session_context())
}
} else {
build_with(&create_datafusion_session_context())
};
explain_physical_plan(&physical_plan, explainer);
}
}
fn begin_custom_scan(
state: &mut CustomScanStateWrapper,
estate: *mut pg_sys::EState,
eflags: i32,
) {
if eflags & (pg_sys::EXEC_FLAG_EXPLAIN_ONLY as i32) == 0 {
unsafe {
let planstate = state.planstate();
// Always assign an ExprContext — heap filters, runtime
// expressions, and pushed-down predicates may all need it.
pg_sys::ExecAssignExprContext(estate, planstate);
state.custom_state_mut().result_slot = Some(state.csstate.ss.ps.ps_ResultTupleSlot);
state.runtime_context = state.csstate.ss.ps.ps_ExprContext;
}
// MPP: mark one launch attempt for the first exec call. The existing logical plan is
// resolved and rebaked at execution time before it is deserialized to build the
// physical plan. The finished stages then provide the exact dispatch payload before
// DSM allocation. Only the leader runs this branch (`ParallelWorkerNumber == -1`).
// The size gate decides here, before any MPP work: a gated query takes the plain
// serial path and never builds the distributed plan or captures manifests.
if mpp_eligible(
state.custom_state().parallel_mode_ok,
&state.custom_state().join_clause.plan,
) && unsafe { pg_sys::ParallelWorkerNumber } == -1
&& state.custom_state().logical_plan.is_some()
{
state.custom_state_mut().mpp = MppLifecycle::Pending;
}
}
}
fn rescan_custom_scan(state: &mut CustomScanStateWrapper) {
let relaunch_mpp = matches!(
&state.custom_state().mpp,
MppLifecycle::Pending | MppLifecycle::Launched(_)
);
Self::finish_mpp_execution(state);
state.custom_state_mut().relations.clear();
state.custom_state_mut().reset();
if relaunch_mpp {
assert!(
state.custom_state().logical_plan.is_some(),
"MPP rescan requires logical plan bytes"
);
state.custom_state_mut().mpp = MppLifecycle::Pending;
}
}
fn exec_custom_scan(state: &mut CustomScanStateWrapper) -> *mut pg_sys::TupleTableSlot {
let mut launch_us = crate::postgres::customscan::mpp::glue::MppLaunchTiming::default();
unsafe {
if state.custom_state().datafusion_stream.is_none() {
// Solve any Param/SubPlan-backed SearchQueryInputs on the leader and rebake the
// logical plan before it's cloned/deserialized below, so both `join_clause` and
// `plan_bytes` reflect the resolved query. Must run before the pending MPP
// marker is consumed below and before plan_bytes is cloned, or a stale unresolved
// plan gets planned/dispatched. Rebake failures are invariant violations and
// abort at their source rather than falling back to unresolved bytes.
Self::maybe_solve_and_rebake(state);
let runtime = tokio::runtime::Builder::new_current_thread()
.build()
.unwrap();
let mut join_clause = state.custom_state().join_clause.clone();
let snapshot = state.csstate.ss.ps.state.as_ref().unwrap().es_snapshot;
let plan_sources = join_clause.plan.sources();
for (plan_position, source) in plan_sources.iter().enumerate() {
let heaprelid = source.scan_info.heaprelid;
let heaprel = PgSearchRelation::open(heaprelid);
let visibility_checker =
VisibilityChecker::with_rel_and_snap(&heaprel, snapshot);
let fetch_slot =
pg_sys::MakeTupleTableSlot(heaprel.rd_att, &pg_sys::TTSOpsBufferHeapTuple);
state.custom_state_mut().relations.insert(
plan_position,
scan_state::RelationState {
_heaprel: heaprel,
visibility_checker,
fetch_slot,
ctid_col_idx: None,
},
);
}
// Deserialize the logical plan and convert to execution plan
let planstate = state.planstate();
// Clone plan_bytes to release the immutable borrow on `state`
// before the mutable borrow in ensure_source_manifests below.
let plan_bytes = state
.custom_state()
.logical_plan
.clone()
.expect("Logical plan is required");
let source_manifests =
Self::build_source_manifests(state, &join_clause, &plan_sources);
// For parameterized LIMIT/OFFSET, the planning-time logical plan has no Limit
// node. Resolve the fetch once; each planning pass below injects it (before
// physical planning) so SegmentedTopKRule can detect SortExec(fetch=K) and
// apply its TopK optimization. Static cases were already pushed in
// `build_clause_df`.
let runtime_fetch = {
let estate = state.csstate.ss.ps.state;
join_clause
.limit_offset
.as_mut()
.filter(|lo| lo.has_any_param())
.map(|lo| {
lo.resolve_mut(estate)
.expect("LIMIT must be resolvable from EState")
.static_fetch()
.expect("static_fetch must succeed after resolve_mut")
})
};
// Raw pointers precomputed so the planning closure below never borrows `state`.
let runtime_context = state.runtime_context;
let build_plan =
|ctx: &datafusion::prelude::SessionContext| -> Arc {
let logical_plan = deserialize_logical_plan_with_runtime(
&plan_bytes,
&ctx.task_ctx(),
// Plans build address-free (#5667): on the MPP path `launch_mpp`
// stamps the shared scan state into the built plan once the DSM
// exists; serial plans never have one.
None,
Some(runtime_context),
Some(planstate),
source_manifests.clone(),
)
.expect("Failed to deserialize logical plan");
let logical_plan = match runtime_fetch {
Some(fetch) => {
use datafusion::logical_expr::LogicalPlanBuilder;
LogicalPlanBuilder::from(logical_plan)
.limit(0, Some(fetch))
.expect("failed to add Limit to logical plan")
.build()
.expect("failed to build logical plan with Limit")
}
None => logical_plan,
};
runtime
.block_on(build_physical_plan(ctx, logical_plan))
.expect("Failed to create execution plan")
};
let mpp_pending = state.custom_state_mut().mpp.take_pending();
// Leader session context: on an MPP attempt, layer the DF-D fork's
// distributed-planner knobs over the Join profile so the resulting physical
// plan is a `DistributedExec`, with `producer_worker_cap()` acting as the
// planner's ceiling. The mesh and the dispatch source are execute-time
// concerns; the exec session below carries them once the workers are committed.
let plan_ctx = if mpp_pending {
Self::build_mpp_session_context(None)
} else {
create_datafusion_session_context()
};
let t_plan = std::time::Instant::now();
let plan = build_plan(&plan_ctx);
launch_us.plan_us = t_plan.elapsed().as_micros() as u64;
// On a launch fallback (nothing to distribute, or too few attached workers) no
// workers remain and the `DistributedExec` shape has no mesh to read from, so
// replan serially.
let (ctx, plan) = if mpp_pending {
match Self::launch_mpp(state, &plan) {
Some(leader) => {
let source = crate::postgres::customscan::mpp::glue::StagePlanDispatchSource::default();
let exec_ctx = Self::build_mpp_session_context(Some(Arc::clone(
&leader.session.mesh,
)))
.with_distributed_dispatch_plan_source(source);
launch_us.prepare_us = leader.timing.prepare_us;
launch_us.payload_us = leader.timing.payload_us;
launch_us.attach_us = leader.timing.attach_us;
launch_us.leader_setup_us = leader.timing.leader_setup_us;
launch_us.workers = leader.timing.workers;
state.custom_state_mut().mpp = MppLifecycle::Launched(leader);
(exec_ctx, plan)
}
None => {
// Short-launch decline: rebuild the logical plan with serial provider
// metadata. Merely changing SessionContext would replan the existing
// MPP-shaped logical providers without rewriting their source fields.
let serial_ctx = create_datafusion_session_context();
let fallback_bytes = Self::rebake_for_mpp_fallback(state);
let logical_plan = deserialize_logical_plan_with_runtime(
&fallback_bytes,
&serial_ctx.task_ctx(),
None,
Some(runtime_context),
Some(planstate),
source_manifests.clone(),
)
.expect("Failed to deserialize serial fallback logical plan");
let logical_plan = match runtime_fetch {
Some(fetch) => {
use datafusion::logical_expr::LogicalPlanBuilder;
LogicalPlanBuilder::from(logical_plan)
.limit(0, Some(fetch))
.expect("failed to add Limit to logical plan")
.build()
.expect("failed to build logical plan with Limit")
}
None => logical_plan,
};
let plan = runtime
.block_on(build_physical_plan(&serial_ctx, logical_plan))
.expect("Failed to create execution plan from serial fallback");
// Keep custom_state's logical_plan in sync: a later rescan reads it
// (see maybe_solve_and_rebake / JoinScanState::reset), and it should
// reflect the serial shape actually executed here, not the MPP-shaped
// bytes this fallback declined to use.
state.custom_state_mut().logical_plan =
Some(bytes::Bytes::from(fallback_bytes));
(serial_ctx, plan)
}
}
} else {
(plan_ctx, plan)
};
let task_ctx = build_task_context(
&ctx,
&plan,
pg_sys::work_mem as usize * 1024,
pg_sys::hash_mem_multiplier,
);
let t_exec = std::time::Instant::now();
let stream = {
let _guard = runtime.enter();
plan.execute(0, task_ctx)
.expect("Failed to execute DataFusion plan")
};
launch_us.exec_us = t_exec.elapsed().as_micros() as u64;
// Retain the executed plan so EXPLAIN ANALYZE can extract metrics. Record the
// launch timing only when the query actually ran distributed (workers attached);
// a serial fallback never reaches `Launched` and leaves `workers` at zero.
if state.custom_state().mpp.is_launched() {
state.custom_state_mut().launch_timing = Some(launch_us);
state.custom_state_mut().stream_built_at = Some(std::time::Instant::now());
}
state.custom_state_mut().physical_plan = Some(plan.clone());
let schema = plan.schema();
for (i, field) in schema.fields().iter().enumerate() {
if let Ok(ctid_col) = CtidColumn::try_from(field.name().as_str()) {
let plan_position = ctid_col.plan_position();
if let Some(rel_state) =
state.custom_state_mut().relations.get_mut(&plan_position)
{
rel_state.ctid_col_idx = Some(i);
}
}
}
let plan_sources = state.custom_state().join_clause.plan.sources();
let output_batch_col_indices: Vec