radixdb_executor/mutation/
pk_fast_path.rs

1// Copyright 2026 RadixDB Contributors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Fast-path execution for simple PK lookups
16//!
17//! This module provides an optimized execution path for simple queries like:
18//! - `SELECT * FROM table WHERE pk_col = $1`
19//! - `SELECT col1, col2 FROM table WHERE pk_col = 5`
20//!
21//! By detecting these patterns early, we bypass the full query planner and
22//! go directly to index lookup, reducing per-query overhead from ~2µs to ~200ns.
23//!
24//! # Performance Impact
25//!
26//! For Index Nested Loop joins that perform thousands of PK lookups,
27//! this fast-path can provide significant speedups by amortizing less overhead.
28
29use std::sync::RwLock;
30
31use radixdb_core::{CompactArc, SmartString};
32use radixdb_core::{DataType, Result, Row, RowVec, Schema, Value};
33use radixdb_sql::ast::{Expression, SelectStatement};
34use radixdb_storage::traits::{Engine, QueryResult};
35
36use crate::compiled_plan::{CompiledExecution, CompiledPkLookup, PkValueSource};
37use crate::context::ExecutionContext;
38use crate::lookup_key::{integer_pk_admission, IntegerPkAdmission};
39use crate::mutation::host::MutationHost;
40use crate::result::ExecutorResult;
41
42/// Information extracted from a simple PK lookup query
43#[doc(hidden)]
44pub struct PkLookupInfo {
45    /// Table name (already lowercased for storage lookups)
46    table_name: String,
47    /// PK value to look up
48    pk_value: IntegerPkAdmission,
49    /// How to extract the PK value (for caching)
50    pk_value_source: PkValueSource,
51    /// Cached schema to avoid second lookup
52    schema: CompactArc<Schema>,
53}
54
55#[doc(hidden)]
56pub trait PkFastPathExt: MutationHost {
57    /// Try to execute a SELECT as a fast PK lookup
58    ///
59    /// Returns Some(result) if the query is a simple PK lookup that was executed.
60    /// Returns None if the query doesn't qualify for fast-path.
61    fn try_fast_pk_lookup(
62        &self,
63        stmt: &SelectStatement,
64        ctx: &ExecutionContext,
65    ) -> Option<Result<Box<dyn QueryResult>>> {
66        // Quick reject: if we're in an explicit transaction, skip fast path
67        // The fast path uses fetch_rows_by_ids which only sees committed data,
68        // so it wouldn't see uncommitted changes from the current transaction.
69        // This could return stale data if Transaction A updates a row and then
70        // queries for it - the fast path would return the old committed value.
71        // Use try_lock for faster rejection under contention.
72        {
73            let active_tx = match self.mutation_active_transaction().try_lock() {
74                Ok(guard) => guard,
75                Err(_) => return None, // Lock contention - fall back to normal path
76            };
77            if active_tx.is_some() {
78                return None; // Let normal execution path handle transaction context
79            }
80        }
81
82        // Quick reject: must have WHERE clause and table_expr
83        let where_clause = stmt.where_clause.as_ref()?;
84        let table_expr = stmt.table_expr.as_ref()?;
85
86        // Quick reject: no GROUP BY, no HAVING, no CTEs, no set operations, no DISTINCT
87        if !stmt.group_by.columns.is_empty()
88            || stmt.having.is_some()
89            || !stmt.set_operations.is_empty()
90            || stmt.with.is_some()
91            || stmt.distinct
92        {
93            return None;
94        }
95
96        // Quick reject: no ORDER BY (PK lookup returns single row anyway, but skip for simplicity)
97        if !stmt.order_by.is_empty() {
98            return None;
99        }
100
101        // Must be SELECT * (for now - column projection adds complexity)
102        if stmt.columns.len() != 1 || !matches!(&stmt.columns[0], Expression::Star(_)) {
103            return None;
104        }
105
106        // Extract table name (must be a simple table reference, not a join or subquery)
107        // Use pre-computed lowercase from Identifier (avoids allocation and case conversion)
108        let table_name: &str = match table_expr.as_ref() {
109            Expression::TableSource(ts) if ts.as_of.is_none() => ts.name.value_lower.as_str(),
110            _ => return None, // Join, subquery, or other complex source
111        };
112
113        // Try to extract PK lookup info from WHERE clause
114        let lookup_info = self.extract_pk_lookup_info(table_name, where_clause, ctx)?;
115
116        // Execute the fast-path lookup
117        Some(self.execute_pk_lookup(lookup_info))
118    }
119
120    /// Extract PK lookup information from a WHERE clause
121    fn extract_pk_lookup_info(
122        &self,
123        table_name: &str,
124        where_clause: &Expression,
125        ctx: &ExecutionContext,
126    ) -> Option<PkLookupInfo> {
127        // Get table schema to find PK column
128        let schema = self.mutation_engine().get_table_schema(table_name).ok()?;
129        let pk_indices = schema.primary_key_indices();
130
131        // Only support single-column PK for now
132        if pk_indices.len() != 1 {
133            return None;
134        }
135        let pk_idx = pk_indices[0];
136        if schema.columns[pk_idx].data_type != DataType::Integer {
137            return None;
138        }
139        let pk_column = &schema.columns[pk_idx].name;
140
141        // Extract comparison info from WHERE clause
142        let (col_name, pk_value, pk_value_source) =
143            self.extract_pk_equality(where_clause, pk_column, ctx)?;
144
145        // Column must match PK (case-insensitive)
146        // Use schema's pre-computed lowercase for pk_column
147        let col_lower = col_name.to_lowercase();
148        let pk_lower = &schema.columns[pk_idx].name_lower;
149
150        // Handle qualified names like "users.id"
151        let matches_pk = col_lower == *pk_lower || col_lower.ends_with(&format!(".{}", pk_lower));
152
153        if !matches_pk {
154            return None;
155        }
156
157        Some(PkLookupInfo {
158            table_name: table_name.to_string(),
159            pk_value,
160            pk_value_source,
161            schema,
162        })
163    }
164
165    /// Extract PK equality from WHERE clause
166    /// Returns (column_name, pk_value, pk_value_source) if WHERE is `pk_col = literal` or `pk_col = $param`
167    fn extract_pk_equality(
168        &self,
169        expr: &Expression,
170        _pk_column: &str,
171        ctx: &ExecutionContext,
172    ) -> Option<(String, IntegerPkAdmission, PkValueSource)> {
173        match expr {
174            Expression::Infix(infix) => {
175                // Must be equality operator
176                if infix.operator != "=" {
177                    return None;
178                }
179
180                // Try column = value pattern
181                if let Some((col, val, source)) =
182                    self.extract_col_eq_val(&infix.left, &infix.right, ctx)
183                {
184                    return Some((col, val, source));
185                }
186
187                // Try value = column pattern
188                if let Some((col, val, source)) =
189                    self.extract_col_eq_val(&infix.right, &infix.left, ctx)
190                {
191                    return Some((col, val, source));
192                }
193
194                None
195            }
196            _ => None,
197        }
198    }
199
200    /// Extract column name, integer value, and value source from col = val pattern
201    fn extract_col_eq_val(
202        &self,
203        col_expr: &Expression,
204        val_expr: &Expression,
205        ctx: &ExecutionContext,
206    ) -> Option<(String, IntegerPkAdmission, PkValueSource)> {
207        // Get column name
208        let col_name = match col_expr {
209            Expression::Identifier(id) => id.value.to_string(),
210            Expression::QualifiedIdentifier(q) => format!("{}.{}", q.qualifier, q.name),
211            _ => return None,
212        };
213
214        // Get integer value and source
215        let (pk_value, pk_value_source) = match val_expr {
216            Expression::IntegerLiteral(lit) => (
217                IntegerPkAdmission::Exact(lit.value),
218                PkValueSource::Literal(lit.value),
219            ),
220            Expression::FloatLiteral(lit) => {
221                let admission = integer_pk_admission(&Value::Float(lit.value))?;
222                let cached_literal = match admission {
223                    IntegerPkAdmission::Exact(value) => value,
224                    IntegerPkAdmission::NoMatch => 0,
225                };
226                (admission, PkValueSource::Literal(cached_literal))
227            }
228            Expression::Parameter(param) => {
229                // Resolve parameter from context
230                // Named parameters (e.g., :name) use get_named_param()
231                // Positional parameters ($1, $2, ...) are 1-indexed, array is 0-indexed
232                if param.name.starts_with(':') {
233                    let name = &param.name[1..];
234                    let value = ctx.get_named_param(name)?;
235                    let pk_value = integer_pk_admission(value)?;
236                    (
237                        pk_value,
238                        PkValueSource::NamedParameter(SmartString::new(name)),
239                    )
240                } else {
241                    let params = ctx.params();
242                    let param_idx = if param.index > 0 {
243                        param.index - 1
244                    } else {
245                        return None;
246                    };
247                    if param_idx >= params.len() {
248                        return None;
249                    }
250                    let pk_value = integer_pk_admission(&params[param_idx])?;
251                    (pk_value, PkValueSource::Parameter(param_idx))
252                }
253            }
254            _ => return None,
255        };
256
257        Some((col_name, pk_value, pk_value_source))
258    }
259
260    /// Normalize a row to match the current schema
261    ///
262    /// This handles schema evolution (ALTER TABLE ADD/DROP COLUMN):
263    /// - If row has fewer columns than schema, append default values (or NULLs) for missing columns
264    /// - If row has more columns than schema, truncate the row
265    #[inline]
266    fn normalize_row_to_schema(mut row: Row, schema: &Schema) -> Row {
267        let schema_cols = schema.columns.len();
268        let row_cols = row.len();
269
270        if row_cols < schema_cols {
271            // Row has fewer columns - add default values (or NULLs) for new columns
272            for i in row_cols..schema_cols {
273                let col = &schema.columns[i];
274                // Use pre-computed default value if available, otherwise use NULL
275                if let Some(ref default_val) = col.default_value {
276                    row.push(default_val.clone());
277                } else {
278                    row.push(Value::null(col.data_type));
279                }
280            }
281        } else if row_cols > schema_cols {
282            // Row has more columns - truncate (columns were dropped)
283            row.truncate(schema_cols);
284        }
285
286        row
287    }
288
289    /// Execute the fast-path PK lookup using Engine::fetch_rows_by_ids
290    fn execute_pk_lookup(&self, info: PkLookupInfo) -> Result<Box<dyn QueryResult>> {
291        // Use cached schema for column names - Arc clone is O(1)
292        let columns = info.schema.column_names_arc();
293
294        let IntegerPkAdmission::Exact(pk_value) = info.pk_value else {
295            return Ok(Box::new(ExecutorResult::with_arc_columns(
296                columns,
297                RowVec::new(),
298            )));
299        };
300
301        // Use engine's fetch_rows_by_ids for direct MVCC lookup
302        // This bypasses the full query planner and goes straight to version store
303        // Note: table_name is already lowercased, so storage layer won't call to_lowercase again
304        let rows = self
305            .mutation_engine()
306            .fetch_rows_by_ids(&info.table_name, &[pk_value])?;
307
308        // Extract Row values and normalize to current schema (handles ADD/DROP COLUMN)
309        let result_rows: RowVec = rows
310            .into_iter()
311            .enumerate()
312            .map(|(i, (_, row))| (i as i64, Self::normalize_row_to_schema(row, &info.schema)))
313            .collect();
314
315        Ok(Box::new(ExecutorResult::with_arc_columns(
316            columns,
317            result_rows,
318        )))
319    }
320
321    // ============================================================================
322    // COMPILED EXECUTION METHODS - Use pre-compiled state for fast repeated queries
323    // ============================================================================
324
325    /// Try fast PK lookup using pre-compiled state (if available)
326    ///
327    /// This is the preferred entry point for queries that may be executed multiple times.
328    /// First execution compiles and caches the state, subsequent executions use the cache.
329    fn try_fast_pk_lookup_compiled(
330        &self,
331        stmt: &SelectStatement,
332        ctx: &ExecutionContext,
333        compiled: &RwLock<CompiledExecution>,
334    ) -> Option<Result<Box<dyn QueryResult>>> {
335        // Quick reject: explicit transaction (use try_lock for fast rejection)
336        {
337            let active_tx = match self.mutation_active_transaction().try_lock() {
338                Ok(guard) => guard,
339                Err(_) => return None, // Lock contention - fall back to normal path
340            };
341            if active_tx.is_some() {
342                return None;
343            }
344        }
345
346        // Try read lock first - check if already compiled
347        {
348            let compiled_guard = match compiled.read() {
349                Ok(guard) => guard,
350                Err(_) => return None,
351            };
352            match &*compiled_guard {
353                CompiledExecution::NotOptimizable(epoch)
354                    if self.mutation_engine().schema_epoch() == *epoch =>
355                {
356                    return None
357                }
358                CompiledExecution::PkLookup(lookup) => {
359                    // Fast validation using schema epoch (~1ns vs ~7ns for HashMap lookup)
360                    // If epoch matches, no DDL has occurred since compilation
361                    if self.mutation_engine().schema_epoch() == lookup.cached_epoch {
362                        // Fast path: extract value and execute
363                        let pk_value = self.extract_pk_value_fast(&lookup.pk_value_source, ctx)?;
364                        return Some(self.execute_compiled_pk_lookup(lookup, pk_value));
365                    }
366                    // Epoch changed - some DDL occurred, need to recompile
367                    // Fall through to recompile path
368                }
369                CompiledExecution::NotOptimizable(_) | CompiledExecution::Unknown => {} // Epoch changed or first run - fall through to recompile
370                // These variants are for UPDATE/DELETE/INSERT/COUNT DISTINCT/COUNT(*) - not PK lookups
371                CompiledExecution::PkUpdate(_)
372                | CompiledExecution::PkDelete(_)
373                | CompiledExecution::Insert(_)
374                | CompiledExecution::CountDistinct(_)
375                | CompiledExecution::CountStar(_) => return None,
376            }
377        }
378
379        // First execution or schema changed - compile and cache (write lock)
380        self.compile_and_execute_pk_lookup(stmt, ctx, compiled)
381    }
382
383    /// Extract PK value using pre-compiled source (very fast - just array access)
384    fn extract_pk_value_fast(
385        &self,
386        source: &PkValueSource,
387        ctx: &ExecutionContext,
388    ) -> Option<IntegerPkAdmission> {
389        match source {
390            PkValueSource::NamedParameter(name) => integer_pk_admission(ctx.get_named_param(name)?),
391            _ => Self::extract_pk_value_from_slice(source, ctx.params()),
392        }
393    }
394
395    /// Extract PK value from params slice directly (avoids ExecutionContext overhead)
396    #[inline]
397    fn extract_pk_value_from_slice(
398        source: &PkValueSource,
399        params: &[Value],
400    ) -> Option<IntegerPkAdmission> {
401        match source {
402            PkValueSource::Literal(v) => Some(IntegerPkAdmission::Exact(*v)),
403            PkValueSource::Parameter(idx) => {
404                if *idx >= params.len() {
405                    return None;
406                }
407                integer_pk_admission(&params[*idx])
408            }
409            PkValueSource::NamedParameter(_) => None, // No ctx available in slice path
410        }
411    }
412
413    /// Try fast PK lookup with borrowed params slice (avoids Arc allocation)
414    fn try_fast_pk_lookup_with_params(
415        &self,
416        _stmt: &SelectStatement,
417        params: &[Value],
418        compiled: &RwLock<CompiledExecution>,
419    ) -> Option<Result<Box<dyn QueryResult>>> {
420        // Try read lock first - check if already compiled
421        let compiled_guard = compiled.read().ok()?;
422        match &*compiled_guard {
423            CompiledExecution::NotOptimizable(_) => None,
424            CompiledExecution::PkLookup(lookup) => {
425                // Fast validation using schema epoch
426                if self.mutation_engine().schema_epoch() == lookup.cached_epoch {
427                    // Fast path: extract value from slice directly
428                    let pk_value =
429                        Self::extract_pk_value_from_slice(&lookup.pk_value_source, params)?;
430                    Some(self.execute_compiled_pk_lookup(lookup, pk_value))
431                } else {
432                    // Epoch changed - need recompile, use normal path
433                    None
434                }
435            }
436            CompiledExecution::Unknown => None, // Not compiled yet, use normal path
437            _ => None,
438        }
439    }
440
441    /// Execute using pre-compiled lookup (skip schema lookup, column name building)
442    fn execute_compiled_pk_lookup(
443        &self,
444        lookup: &CompiledPkLookup,
445        pk_value: IntegerPkAdmission,
446    ) -> Result<Box<dyn QueryResult>> {
447        let IntegerPkAdmission::Exact(pk_value) = pk_value else {
448            return Ok(Box::new(ExecutorResult::with_arc_columns(
449                lookup.column_names.clone(),
450                RowVec::new(),
451            )));
452        };
453        let rows = self
454            .mutation_engine()
455            .fetch_rows_by_ids(&lookup.table_name, &[pk_value])?;
456        // Normalize rows to current schema (handles ADD/DROP COLUMN)
457        // Pre-allocate with capacity 1 for single PK lookup (avoids realloc)
458        let mut result_rows = RowVec::with_capacity(1);
459        for (row_id, (_, row)) in rows.into_iter().enumerate() {
460            result_rows.push((
461                row_id as i64,
462                Self::normalize_row_to_schema(row, &lookup.schema),
463            ));
464        }
465        // Use Arc columns - O(1) clone since column_names is CompactArc<Vec<String>>
466        Ok(Box::new(ExecutorResult::with_arc_columns(
467            lookup.column_names.clone(),
468            result_rows,
469        )))
470    }
471
472    /// Compile and execute PK lookup, caching the compiled state
473    fn compile_and_execute_pk_lookup(
474        &self,
475        stmt: &SelectStatement,
476        ctx: &ExecutionContext,
477        compiled: &RwLock<CompiledExecution>,
478    ) -> Option<Result<Box<dyn QueryResult>>> {
479        // Acquire write lock
480        let mut compiled_guard = match compiled.write() {
481            Ok(guard) => guard,
482            Err(_) => return None,
483        };
484
485        // Double-check (another thread may have compiled while we waited)
486        // But also re-validate schema version to handle schema changes
487        match &*compiled_guard {
488            CompiledExecution::NotOptimizable(epoch)
489                if self.mutation_engine().schema_epoch() == *epoch =>
490            {
491                return None
492            }
493            CompiledExecution::PkLookup(lookup) => {
494                // Re-validate epoch: another thread may have compiled before DDL
495                if self.mutation_engine().schema_epoch() == lookup.cached_epoch {
496                    let pk_value = self.extract_pk_value_fast(&lookup.pk_value_source, ctx)?;
497                    return Some(self.execute_compiled_pk_lookup(lookup, pk_value));
498                }
499                // Epoch changed since last compilation - fall through to recompile
500            }
501            CompiledExecution::NotOptimizable(_) | CompiledExecution::Unknown => {} // Epoch changed or first run - recompile
502            // These variants are for UPDATE/DELETE/INSERT/COUNT DISTINCT/COUNT(*) - not PK lookups
503            CompiledExecution::PkUpdate(_)
504            | CompiledExecution::PkDelete(_)
505            | CompiledExecution::Insert(_)
506            | CompiledExecution::CountDistinct(_)
507            | CompiledExecution::CountStar(_) => return None,
508        }
509
510        // Do full pattern detection (same as try_fast_pk_lookup)
511        let where_clause = stmt.where_clause.as_ref()?;
512        let table_expr = stmt.table_expr.as_ref()?;
513
514        // Quick reject: no GROUP BY, no HAVING, no CTEs, no set operations, no DISTINCT
515        if !stmt.group_by.columns.is_empty()
516            || stmt.having.is_some()
517            || !stmt.set_operations.is_empty()
518            || stmt.with.is_some()
519            || stmt.distinct
520        {
521            *compiled_guard =
522                CompiledExecution::NotOptimizable(self.mutation_engine().schema_epoch());
523            return None;
524        }
525
526        // Quick reject: no ORDER BY
527        if !stmt.order_by.is_empty() {
528            *compiled_guard =
529                CompiledExecution::NotOptimizable(self.mutation_engine().schema_epoch());
530            return None;
531        }
532
533        // Must be SELECT *
534        // Don't set NotOptimizable here - other fast paths (like COUNT DISTINCT) may handle this
535        if stmt.columns.len() != 1 || !matches!(&stmt.columns[0], Expression::Star(_)) {
536            return None;
537        }
538
539        // Extract table name (use pre-computed lowercase)
540        let table_name: &str = match table_expr.as_ref() {
541            Expression::TableSource(ts) if ts.as_of.is_none() => ts.name.value_lower.as_str(),
542            _ => {
543                *compiled_guard =
544                    CompiledExecution::NotOptimizable(self.mutation_engine().schema_epoch());
545                return None;
546            }
547        };
548
549        // Try to extract PK lookup info
550        match self.extract_pk_lookup_info(table_name, where_clause, ctx) {
551            Some(info) => {
552                // `PkValueSource::Literal` cannot represent an impossible Float.
553                // Execute the proven empty result without caching a false key.
554                if info.pk_value == IntegerPkAdmission::NoMatch
555                    && matches!(&info.pk_value_source, PkValueSource::Literal(_))
556                {
557                    drop(compiled_guard);
558                    return Some(self.execute_pk_lookup(info));
559                }
560                // Build and cache compiled lookup
561                // Use schema's column_names_arc() directly - O(1) Arc clone on execution
562                let column_names = info.schema.column_names_arc();
563                let cached_epoch = self.mutation_engine().schema_epoch();
564                let compiled_lookup = CompiledPkLookup {
565                    table_name: SmartString::new(&info.table_name),
566                    schema: info.schema.clone(),
567                    column_names,
568                    pk_value_source: info.pk_value_source.clone(),
569                    cached_epoch,
570                };
571                *compiled_guard = CompiledExecution::PkLookup(compiled_lookup);
572                drop(compiled_guard);
573
574                // Execute
575                Some(self.execute_pk_lookup(info))
576            }
577            None => {
578                *compiled_guard =
579                    CompiledExecution::NotOptimizable(self.mutation_engine().schema_epoch());
580                None
581            }
582        }
583    }
584}
585
586impl<T: MutationHost + ?Sized> PkFastPathExt for T {}
587
588#[cfg(test)]
589mod tests {
590    use super::{integer_pk_admission, IntegerPkAdmission};
591    use crate::lookup_key::exact_integer_pk_value;
592    use radixdb_core::Value;
593
594    #[test]
595    fn exact_integer_pk_value_rejects_truncation_and_saturation() {
596        for value in [
597            Value::Float(1.5),
598            Value::Float(-1.5),
599            Value::Float(f64::MIN_POSITIVE),
600            Value::Float(f64::INFINITY),
601            Value::Float(f64::NEG_INFINITY),
602            Value::Float(f64::NAN),
603            Value::Float(i64::MAX as f64),
604        ] {
605            assert_eq!(exact_integer_pk_value(&value), None, "value={value:?}");
606            assert_eq!(
607                integer_pk_admission(&value),
608                Some(IntegerPkAdmission::NoMatch),
609                "value={value:?}"
610            );
611        }
612    }
613
614    #[test]
615    fn exact_integer_pk_value_accepts_only_canonical_integer_identity() {
616        let two_pow_53 = 1_i64 << 53;
617        for (value, expected) in [
618            (Value::Integer(i64::MIN), i64::MIN),
619            (Value::Integer(i64::MAX), i64::MAX),
620            (Value::Float(i64::MIN as f64), i64::MIN),
621            (Value::Float(-0.0), 0),
622            (Value::Float(0.0), 0),
623            (Value::Float(42.0), 42),
624            (Value::Float(two_pow_53 as f64), two_pow_53),
625            (Value::Float((two_pow_53 + 2) as f64), two_pow_53 + 2),
626        ] {
627            assert_eq!(
628                exact_integer_pk_value(&value),
629                Some(expected),
630                "value={value:?}"
631            );
632        }
633    }
634}
radixdb_executor/mutation/pk_fast_path.rs

radixdb_executor/mutation/
pk_fast_path.rs
