radixdb_executor/operators/
count_pk_semijoin.rs

1//! Count-only unique PK semi-join.
2//!
3//! This physical operator is intentionally narrower than the general join
4//! operators: it consumes a filtered child key stream, checks the referenced
5//! integer primary keys in bounded batches, and produces one scalar count. It
6//! implements `INNER JOIN + COUNT(*)` and `INNER/LEFT JOIN +
7//! COUNT(right.not_null_column)`. It never constructs a joined row or fetches a
8//! parent payload row.
9
10use crate::context::ExecutionContext;
11use crate::operator::{ColumnInfo, Operator, RowRef};
12use radixdb_core::{Result, Row, Value};
13use radixdb_storage::traits::Table;
14
15/// Default bound for parent membership probes. It is deliberately independent
16/// of table size: memory is O(batch size), even when a child scan returns many
17/// millions of rows.
18pub const DEFAULT_PK_SEMIJOIN_BATCH_SIZE: usize = 1_024;
19
20pub struct CountPkSemiJoinOperator {
21    child: Box<dyn Operator>,
22    parent: Box<dyn Table>,
23    child_key_index: usize,
24    batch_size: usize,
25    schema: Vec<ColumnInfo>,
26    result: Option<Row>,
27    context: Option<ExecutionContext>,
28    opened: bool,
29}
30
31impl CountPkSemiJoinOperator {
32    pub fn new(child: Box<dyn Operator>, parent: Box<dyn Table>, child_key_index: usize) -> Self {
33        Self::with_batch_size(
34            child,
35            parent,
36            child_key_index,
37            DEFAULT_PK_SEMIJOIN_BATCH_SIZE,
38        )
39    }
40
41    pub fn with_batch_size(
42        child: Box<dyn Operator>,
43        parent: Box<dyn Table>,
44        child_key_index: usize,
45        batch_size: usize,
46    ) -> Self {
47        Self {
48            child,
49            parent,
50            child_key_index,
51            batch_size: batch_size.max(1),
52            schema: vec![ColumnInfo::new("count")],
53            result: None,
54            context: None,
55            opened: false,
56        }
57    }
58
59    /// Preserve cancellation/timeout checks while a large child stream is
60    /// consumed eagerly to produce the scalar aggregate.
61    pub fn with_context(mut self, context: &ExecutionContext) -> Self {
62        self.context = Some(context.clone());
63        self
64    }
65
66    fn flush_batch(&self, keys: &mut Vec<i64>, matches: &mut Vec<bool>) -> Result<u64> {
67        if keys.is_empty() {
68            return Ok(0);
69        }
70        matches.clear();
71        matches.resize(keys.len(), false);
72        let hits = self.parent.probe_visible_row_ids(keys, matches)?;
73        radixdb_storage::instrumentation::record_join_pk_probe(keys.len() as u64, hits as u64, 0);
74        keys.clear();
75        Ok(hits as u64)
76    }
77}
78
79impl Operator for CountPkSemiJoinOperator {
80    fn open(&mut self) -> Result<()> {
81        if let Err(error) = self.child.open() {
82            let _ = self.child.close();
83            return Err(error);
84        }
85        let mut keys = Vec::with_capacity(self.batch_size);
86        let mut matches = Vec::with_capacity(self.batch_size);
87        let mut child_rows = 0_u64;
88        let mut child_key_rows = 0_u64;
89        let computation = (|| -> Result<u64> {
90            let mut count = 0_u64;
91            if let Some(context) = &self.context {
92                context.check_cancelled()?;
93            }
94
95            while let Some(row) = self.child.next()? {
96                child_rows = child_rows.saturating_add(1);
97                if child_rows.is_multiple_of(256) {
98                    if let Some(context) = &self.context {
99                        context.check_cancelled()?;
100                    }
101                }
102                if let Some(Value::Integer(key)) = row.get(self.child_key_index) {
103                    child_key_rows = child_key_rows.saturating_add(1);
104                    keys.push(*key);
105                    if keys.len() == self.batch_size {
106                        count = count.saturating_add(self.flush_batch(&mut keys, &mut matches)?);
107                    }
108                }
109            }
110            Ok(count.saturating_add(self.flush_batch(&mut keys, &mut matches)?))
111        })();
112        let count = match computation {
113            Ok(count) => count,
114            Err(error) => {
115                // This operator consumes its input eagerly. Make the resource
116                // lifecycle explicit even if a scan, cancellation or parent
117                // membership probe fails halfway through the stream.
118                let _ = self.child.close();
119                return Err(error);
120            }
121        };
122        radixdb_storage::instrumentation::record_join_outer_rows(child_rows, child_key_rows);
123
124        self.result = Some(Row::from_values(vec![Value::Integer(
125            i64::try_from(count).unwrap_or(i64::MAX),
126        )]));
127        self.opened = true;
128        Ok(())
129    }
130
131    fn next(&mut self) -> Result<Option<RowRef>> {
132        if !self.opened {
133            return Err(radixdb_core::Error::internal(
134                "CountPkSemiJoinOperator::next called before open",
135            ));
136        }
137        Ok(self.result.take().map(RowRef::Owned))
138    }
139
140    fn close(&mut self) -> Result<()> {
141        self.child.close()
142    }
143
144    fn schema(&self) -> &[ColumnInfo] {
145        &self.schema
146    }
147
148    fn estimated_rows(&self) -> Option<usize> {
149        Some(1)
150    }
151
152    fn name(&self) -> &str {
153        "CountPkSemiJoin"
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160    use crate::operator::MaterializedOperator;
161    use std::sync::Arc;
162
163    use radixdb_storage::mvcc::engine::MVCCEngine;
164    use radixdb_storage::traits::Engine;
165
166    #[test]
167    fn counts_matching_keys_preserves_child_duplicates_and_skips_nulls() {
168        let engine = Arc::new(MVCCEngine::in_memory());
169        engine.open_engine().unwrap();
170        let executor = crate::Executor::new(Arc::clone(&engine));
171        executor
172            .execute("CREATE TABLE parent (id INTEGER PRIMARY KEY, payload TEXT)")
173            .unwrap();
174        drop(executor);
175
176        let mut write = engine.begin_transaction().unwrap();
177        let mut parent = write.get_table("parent").unwrap();
178        parent
179            .insert(Row::from_values(vec![
180                Value::Integer(10),
181                Value::from("must not be read"),
182            ]))
183            .unwrap();
184        parent
185            .insert(Row::from_values(vec![
186                Value::Integer(20),
187                Value::from("must not be read"),
188            ]))
189            .unwrap();
190        drop(parent);
191        write.commit().unwrap();
192
193        let mut read = engine.begin_transaction().unwrap();
194        let parent = read.get_table("parent").unwrap();
195        let child = MaterializedOperator::new(
196            vec![
197                Row::from_values(vec![Value::Integer(10)]),
198                Row::from_values(vec![Value::Integer(10)]),
199                Row::from_values(vec![Value::Integer(99)]),
200                Row::from_values(vec![Value::null_unknown()]),
201                Row::from_values(vec![Value::Integer(20)]),
202            ],
203            vec![ColumnInfo::new("parent_id")],
204        );
205        let mut op = CountPkSemiJoinOperator::with_batch_size(Box::new(child), parent, 0, 2);
206        op.open().unwrap();
207        let row = op.next().unwrap().unwrap().into_owned();
208        assert_eq!(row.get(0), Some(&Value::Integer(3)));
209        assert!(op.next().unwrap().is_none());
210        op.close().unwrap();
211        read.rollback().unwrap();
212        engine.close_engine().unwrap();
213    }
214
215    #[test]
216    fn checks_cancellation_before_consuming_child_stream() {
217        let engine = Arc::new(MVCCEngine::in_memory());
218        engine.open_engine().unwrap();
219        let executor = crate::Executor::new(Arc::clone(&engine));
220        executor
221            .execute("CREATE TABLE parent (id INTEGER PRIMARY KEY)")
222            .unwrap();
223        drop(executor);
224        let mut tx = engine.begin_transaction().unwrap();
225        let parent = tx.get_table("parent").unwrap();
226        let child = MaterializedOperator::new(
227            vec![Row::from_values(vec![Value::Integer(1)])],
228            vec![ColumnInfo::new("parent_id")],
229        );
230        let context = ExecutionContext::new();
231        context.cancel();
232        let mut op =
233            CountPkSemiJoinOperator::new(Box::new(child), parent, 0).with_context(&context);
234        assert!(op.open().is_err());
235        tx.rollback().unwrap();
236        engine.close_engine().unwrap();
237    }
238}
radixdb_executor/operators/count_pk_semijoin.rs

radixdb_executor/operators/
count_pk_semijoin.rs
