Skip to content

Commit 3eaa498

Browse files
committed
fix(pgwire): bypass cache for data-dependent point plans
1 parent 22c104d commit 3eaa498

6 files changed

Lines changed: 473 additions & 12 deletions

File tree

nodedb-sql/src/types/mod.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,8 @@ pub use collection::{CollectionInfo, ColumnInfo, IndexSpec, IndexState};
1414
pub use filter::{CompareOp, Filter, FilterExpr};
1515
pub use plan::{
1616
ArrayPrefilter, DistanceMetric, KvInsertIntent, MergeClauseKind, MergePlanAction,
17-
MergePlanClause, SqlPlan, VectorAnnOptions, VectorPrimaryRow, VectorQuantization,
17+
MergePlanClause, PlanCacheEligibility, SqlPlan, VectorAnnOptions, VectorPrimaryRow,
18+
VectorQuantization,
1819
};
1920
pub use query::{
2021
AggOutputSlot, AggregateExpr, EngineType, JoinType, Projection, SortKey, SpatialPredicate,
Lines changed: 218 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,218 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
3+
use crate::types::query::EngineType;
4+
5+
use super::SqlPlan;
6+
7+
/// Whether a logical plan may be lowered once and reused from the physical-plan cache.
8+
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9+
pub enum PlanCacheEligibility {
10+
/// Lowering depends only on schema/catalog descriptors tracked by the cache.
11+
Cacheable,
12+
/// Lowering consults mutable row identity and must run for every execution.
13+
DataDependent,
14+
}
15+
16+
impl PlanCacheEligibility {
17+
/// Whether the lowered physical tasks may be admitted to the plan cache.
18+
pub fn is_cacheable(self) -> bool {
19+
self == Self::Cacheable
20+
}
21+
22+
fn combine(self, other: Self) -> Self {
23+
if self == Self::DataDependent || other == Self::DataDependent {
24+
Self::DataDependent
25+
} else {
26+
Self::Cacheable
27+
}
28+
}
29+
}
30+
31+
impl SqlPlan {
32+
/// Classify dependencies that are not represented by descriptor versions.
33+
///
34+
/// Document point operations resolve primary-key bytes to a surrogate while
35+
/// lowering. That binding can appear after an earlier miss without any
36+
/// schema-version change, so those physical tasks cannot be cached.
37+
pub fn cache_eligibility(&self) -> PlanCacheEligibility {
38+
use PlanCacheEligibility::{Cacheable, DataDependent};
39+
40+
match self {
41+
Self::PointGet {
42+
engine: EngineType::DocumentSchemaless | EngineType::DocumentStrict,
43+
..
44+
} => DataDependent,
45+
Self::Update {
46+
engine,
47+
target_keys,
48+
..
49+
}
50+
| Self::Delete {
51+
engine,
52+
target_keys,
53+
..
54+
} if !target_keys.is_empty()
55+
&& matches!(
56+
engine,
57+
EngineType::DocumentSchemaless | EngineType::DocumentStrict
58+
) =>
59+
{
60+
DataDependent
61+
}
62+
Self::InsertSelect { source, .. }
63+
| Self::UpdateFrom { source, .. }
64+
| Self::Aggregate { input: source, .. }
65+
| Self::Merge { source, .. } => source.cache_eligibility(),
66+
Self::Join { left, right, .. }
67+
| Self::Intersect { left, right, .. }
68+
| Self::Except { left, right, .. } => {
69+
left.cache_eligibility().combine(right.cache_eligibility())
70+
}
71+
Self::Union { inputs, .. } => inputs.iter().fold(Cacheable, |eligibility, input| {
72+
eligibility.combine(input.cache_eligibility())
73+
}),
74+
Self::Cte { definitions, outer } => definitions
75+
.iter()
76+
.fold(outer.cache_eligibility(), |eligibility, (_, plan)| {
77+
eligibility.combine(plan.cache_eligibility())
78+
}),
79+
Self::LateralTopK { outer, .. } => outer.cache_eligibility(),
80+
Self::LateralLoop { outer, inner, .. } => {
81+
outer.cache_eligibility().combine(inner.cache_eligibility())
82+
}
83+
Self::ConstantResult { .. }
84+
| Self::Scan { .. }
85+
| Self::PointGet { .. }
86+
| Self::DocumentIndexLookup { .. }
87+
| Self::RangeScan { .. }
88+
| Self::Insert { .. }
89+
| Self::KvInsert { .. }
90+
| Self::Upsert { .. }
91+
| Self::Update { .. }
92+
| Self::Delete { .. }
93+
| Self::Truncate { .. }
94+
| Self::TimeseriesScan { .. }
95+
| Self::TimeseriesIngest { .. }
96+
| Self::VectorSearch { .. }
97+
| Self::MultiVectorSearch { .. }
98+
| Self::SparseSearch { .. }
99+
| Self::TextSearch { .. }
100+
| Self::HybridSearch { .. }
101+
| Self::HybridSearchTriple { .. }
102+
| Self::SpatialScan { .. }
103+
| Self::RecursiveScan { .. }
104+
| Self::RecursiveValue { .. }
105+
| Self::CreateArray { .. }
106+
| Self::DropArray { .. }
107+
| Self::AlterArray { .. }
108+
| Self::InsertArray { .. }
109+
| Self::DeleteArray { .. }
110+
| Self::ArraySlice { .. }
111+
| Self::ArrayProject { .. }
112+
| Self::ArrayAgg { .. }
113+
| Self::ArrayElementwise { .. }
114+
| Self::ArrayFlush { .. }
115+
| Self::ArrayCompact { .. }
116+
| Self::VectorPrimaryInsert { .. }
117+
| Self::CreateIndex { .. }
118+
| Self::DropIndex { .. } => Cacheable,
119+
}
120+
}
121+
}
122+
123+
#[cfg(test)]
124+
mod tests {
125+
use super::*;
126+
use crate::types_expr::SqlValue;
127+
128+
fn point_get(engine: EngineType) -> SqlPlan {
129+
SqlPlan::PointGet {
130+
collection: "docs".into(),
131+
alias: None,
132+
engine,
133+
key_column: "id".into(),
134+
key_value: SqlValue::String("k".into()),
135+
projection: Vec::new(),
136+
}
137+
}
138+
139+
fn update(engine: EngineType, target_keys: Vec<SqlValue>) -> SqlPlan {
140+
SqlPlan::Update {
141+
collection: "docs".into(),
142+
engine,
143+
assignments: Vec::new(),
144+
filters: Vec::new(),
145+
target_keys,
146+
returning: false,
147+
}
148+
}
149+
150+
fn delete(engine: EngineType, target_keys: Vec<SqlValue>) -> SqlPlan {
151+
SqlPlan::Delete {
152+
collection: "docs".into(),
153+
engine,
154+
filters: Vec::new(),
155+
target_keys,
156+
}
157+
}
158+
159+
#[test]
160+
fn document_point_get_is_data_dependent() {
161+
assert_eq!(
162+
point_get(EngineType::DocumentStrict).cache_eligibility(),
163+
PlanCacheEligibility::DataDependent
164+
);
165+
}
166+
167+
#[test]
168+
fn key_value_point_get_is_cacheable() {
169+
assert_eq!(
170+
point_get(EngineType::KeyValue).cache_eligibility(),
171+
PlanCacheEligibility::Cacheable
172+
);
173+
}
174+
175+
#[test]
176+
fn document_point_update_is_data_dependent() {
177+
assert_eq!(
178+
update(
179+
EngineType::DocumentSchemaless,
180+
vec![SqlValue::String("k".into())]
181+
)
182+
.cache_eligibility(),
183+
PlanCacheEligibility::DataDependent
184+
);
185+
}
186+
187+
#[test]
188+
fn document_predicate_update_is_cacheable() {
189+
assert_eq!(
190+
update(EngineType::DocumentStrict, Vec::new()).cache_eligibility(),
191+
PlanCacheEligibility::Cacheable
192+
);
193+
}
194+
195+
#[test]
196+
fn document_point_delete_is_data_dependent() {
197+
assert_eq!(
198+
delete(
199+
EngineType::DocumentStrict,
200+
vec![SqlValue::String("k".into())]
201+
)
202+
.cache_eligibility(),
203+
PlanCacheEligibility::DataDependent
204+
);
205+
}
206+
207+
#[test]
208+
fn nested_point_dependency_propagates() {
209+
let plan = SqlPlan::Cte {
210+
definitions: vec![("selected".into(), point_get(EngineType::DocumentStrict))],
211+
outer: Box::new(point_get(EngineType::KeyValue)),
212+
};
213+
assert_eq!(
214+
plan.cache_eligibility(),
215+
PlanCacheEligibility::DataDependent
216+
);
217+
}
218+
}

nodedb-sql/src/types/plan/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,13 @@
22

33
//! SqlPlan intermediate representation and supporting types.
44
5+
mod cacheability;
56
mod merge_types;
67
mod row_types;
78
mod variants;
89
mod vector_opts;
910

11+
pub use cacheability::PlanCacheEligibility;
1012
pub use merge_types::{MergeClauseKind, MergePlanAction, MergePlanClause};
1113
pub use row_types::{KvInsertIntent, VectorPrimaryRow};
1214
pub use variants::{DistanceMetric, SqlPlan};

nodedb/src/control/planner/context/query/planning.rs

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ impl QueryContext {
3838
OutputSchema,
3939
)> {
4040
self.plan_with_nodedb_sql(sql, tenant_id, database_id)
41-
.map(|(t, schema, _)| (t, schema))
41+
.map(|(t, schema, _, _)| (t, schema))
4242
}
4343

4444
/// Core planning via nodedb-sql: parse → plan → optimize → convert.
@@ -58,6 +58,7 @@ impl QueryContext {
5858
Vec<nodedb_physical::physical_task::PhysicalTask>,
5959
OutputSchema,
6060
crate::control::planner::descriptor_set::DescriptorVersionSet,
61+
nodedb_sql::types::PlanCacheEligibility,
6162
)> {
6263
let inputs = match &self.catalog_inputs {
6364
Some(i) => i,
@@ -146,8 +147,16 @@ impl QueryContext {
146147
&catalog,
147148
database_id,
148149
);
150+
let cache_eligibility = if plans
151+
.iter()
152+
.all(|plan| plan.cache_eligibility().is_cacheable())
153+
{
154+
nodedb_sql::types::PlanCacheEligibility::Cacheable
155+
} else {
156+
nodedb_sql::types::PlanCacheEligibility::DataDependent
157+
};
149158
let tasks = crate::control::planner::sql_plan_convert::convert(&plans, tenant_id, &ctx)?;
150-
Ok((tasks, output_schema, version_set))
159+
Ok((tasks, output_schema, version_set, cache_eligibility))
151160
}
152161

153162
/// Parse SQL, inject RLS predicates, convert to physical plan.
@@ -184,7 +193,7 @@ impl QueryContext {
184193
)> {
185194
self.plan_sql_with_rls_and_versions(sql, tenant_id, database_id, sec, returning)
186195
.await
187-
.map(|(tasks, schema, _)| (tasks, schema))
196+
.map(|(tasks, schema, _, _)| (tasks, schema))
188197
}
189198

190199
/// Variant of [`plan_sql_with_rls_returning`] that also
@@ -204,8 +213,9 @@ impl QueryContext {
204213
Vec<nodedb_physical::physical_task::PhysicalTask>,
205214
OutputSchema,
206215
crate::control::planner::descriptor_set::DescriptorVersionSet,
216+
nodedb_sql::types::PlanCacheEligibility,
207217
)> {
208-
let (mut tasks, output_schema, version_set) =
218+
let (mut tasks, output_schema, version_set, cache_eligibility) =
209219
self.plan_with_nodedb_sql(sql, tenant_id, database_id)?;
210220

211221
// Inject RLS predicates.
@@ -218,7 +228,7 @@ impl QueryContext {
218228
)?;
219229
}
220230

221-
Ok((tasks, output_schema, version_set))
231+
Ok((tasks, output_schema, version_set, cache_eligibility))
222232
}
223233

224234
/// Plan SQL with bound parameters and RLS injection.

nodedb/src/control/server/pgwire/handler/routing/planning.rs

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,7 @@ impl NodeDbPgHandler {
171171
let scope = self.state.acquire_plan_lease_scope(&versions);
172172
(tasks, output_schema, scope)
173173
} else {
174-
let (planned, output_schema, versions) =
174+
let (planned, output_schema, versions, cache_eligibility) =
175175
super::super::retry::retry_on_schema_change(|| async {
176176
let perm_cache = self.state.permission_cache.read().await;
177177
let sec = crate::control::planner::context::PlanSecurityContext {
@@ -203,11 +203,12 @@ impl NodeDbPgHandler {
203203
})?;
204204

205205
let scope = self.state.acquire_plan_lease_scope(&versions);
206-
// Do not cache a plan built under a strategy-knob override (force
207-
// shuffle, or a non-default broadcast threshold) — the cache key
208-
// does not encode the session knob, so caching it would leak a
209-
// strategy-specific plan into later default queries on this session.
210-
if !bypass_cache {
206+
// Strategy overrides and data-dependent identity lowering are not
207+
// represented by the cache key. Document point plans resolve a
208+
// mutable PK→surrogate binding while lowering, so caching either a
209+
// sentinel miss or a partially resolved target set would preserve
210+
// stale row identity across later writes.
211+
if !bypass_cache && cache_eligibility.is_cacheable() {
211212
self.sessions.put_cached_plan(
212213
addr,
213214
&clean_sql,

0 commit comments

Comments
 (0)