reifydb_engine/vm/volcano/scan/
dictionary.rs1use std::sync::Arc;
5
6use arrow_array::{ArrayRef, RecordBatch};
7use arrow_schema::FieldRef;
8use postcard::from_bytes;
9use reifydb_core::{
10 interface::{catalog::dictionary::Dictionary, resolved::ResolvedDictionary, store::SingleVersionRange},
11 internal_error,
12 key::{any::TaggedKey, bound::TaggedKeyBoundRange, catalog::DictionaryEntryIndexKey},
13 value::{
14 batch::batch,
15 column::{
16 builder::ColumnBuilder,
17 factory::{none_typed, uint1, uint2, uint4, uint8, uint16},
18 headers::ColumnHeaders,
19 },
20 },
21};
22use reifydb_transaction::transaction::Transaction;
23use reifydb_value::{
24 fragment::Fragment,
25 reifydb_assertions,
26 value::{Value, dictionary::DictionaryEntryId, value_type::ValueType},
27};
28use tracing::instrument;
29
30use crate::{
31 Result,
32 vm::volcano::query::{QueryContext, QueryNode},
33};
34
35pub struct DictionaryScanNode {
36 dictionary: ResolvedDictionary,
37 context: Option<Arc<QueryContext>>,
38 headers: ColumnHeaders,
39 last_key: Option<TaggedKey>,
40 exhausted: bool,
41}
42
43impl DictionaryScanNode {
44 pub fn new(dictionary: ResolvedDictionary, context: Arc<QueryContext>) -> Result<Self> {
45 let headers = ColumnHeaders {
46 columns: vec![Fragment::internal("id"), Fragment::internal("value")],
47 };
48
49 Ok(Self {
50 dictionary,
51 context: Some(context),
52 headers,
53 last_key: None,
54 exhausted: false,
55 })
56 }
57
58 #[instrument(level = "trace", skip_all, name = "volcano::scan::dictionary::drain")]
59 fn drain_batch<'a>(
60 rx: &mut Transaction<'a>,
61 range: TaggedKeyBoundRange,
62 batch_size: u64,
63 dict_def: &Dictionary,
64 ) -> Result<(Vec<DictionaryEntryId>, Vec<Value>, Option<TaggedKey>)> {
65 let mut ids: Vec<DictionaryEntryId> = Vec::new();
66 let mut values: Vec<Value> = Vec::new();
67 let mut new_last_key = None;
68
69 let single = rx
70 .single()
71 .ok_or_else(|| internal_error!("single-version store is not available for dictionary scans"))?;
72 let store = single.read_store();
73 let batch = SingleVersionRange::range_batch(&store, range.encode(), batch_size)?;
74
75 for entry in batch.items {
76 let Some(key) = DictionaryEntryIndexKey::decode(&entry.key) else {
77 panic!(
78 "dictionary {} holds an entry index key that does not decode: {:?}",
79 dict_def.id, entry.key
80 );
81 };
82
83 let entry_id = DictionaryEntryId::from_u128(key.id, dict_def.id_type.clone())?;
84 new_last_key = Some(TaggedKey::from(key));
85
86 let value: Value = from_bytes(&entry.bytes)
87 .map_err(|e| internal_error!("Failed to deserialize dictionary value: {}", e))?;
88
89 ids.push(entry_id);
90 values.push(value);
91 }
92
93 Ok((ids, values, new_last_key))
94 }
95
96 #[instrument(level = "trace", skip_all, name = "volcano::scan::dictionary::empty_columns")]
97 fn empty_columns(dict_def: &Dictionary) -> Vec<(FieldRef, ArrayRef)> {
98 vec![none_typed("id", dict_def.id_type.clone(), 0), none_typed("value", dict_def.value_type.clone(), 0)]
99 }
100
101 #[instrument(level = "trace", skip_all, name = "volcano::scan::dictionary::assemble")]
102 fn assemble(ids: &[DictionaryEntryId], values: &[Value], dict_def: &Dictionary) -> Result<Option<RecordBatch>> {
103 let id_column = build_id_column(ids, dict_def.id_type.clone())?;
104 let value_column = build_value_column(values, dict_def.value_type.clone());
105
106 Ok(Some(batch(vec![id_column, value_column])?))
107 }
108}
109
110impl QueryNode for DictionaryScanNode {
111 #[instrument(name = "volcano::scan::dictionary::initialize", level = "trace", skip_all)]
112 fn initialize<'a>(&mut self, _rx: &mut Transaction<'a>, _ctx: &QueryContext) -> Result<()> {
113 Ok(())
114 }
115
116 #[instrument(name = "volcano::scan::dictionary::next", level = "trace", skip_all)]
117 fn next<'a>(&mut self, rx: &mut Transaction<'a>, _ctx: &mut QueryContext) -> Result<Option<RecordBatch>> {
118 reifydb_assertions! {
119 assert!(self.context.is_some(), "DictionaryScan::next() called before initialize()");
120 }
121 let stored_ctx = self.context.as_ref().unwrap();
122
123 if self.exhausted {
124 return Ok(None);
125 }
126
127 let batch_size = stored_ctx.batch_size;
128 let dict_def = self.dictionary.def();
129
130 let range = DictionaryEntryIndexKey::full_scan(dict_def.id).resume_after(self.last_key.as_ref());
131
132 let (ids, values, new_last_key) = Self::drain_batch(rx, range, batch_size, dict_def)?;
133
134 if ids.is_empty() {
135 self.exhausted = true;
136 if self.last_key.is_none() {
137 return Ok(Some(batch(Self::empty_columns(dict_def))?));
138 }
139 return Ok(None);
140 }
141
142 self.last_key = new_last_key;
143
144 Self::assemble(&ids, &values, dict_def)
145 }
146
147 fn headers(&self) -> Option<ColumnHeaders> {
148 Some(self.headers.clone())
149 }
150}
151
152fn build_id_column(ids: &[DictionaryEntryId], id_type: ValueType) -> Result<(FieldRef, ArrayRef)> {
153 let data = match id_type {
154 ValueType::Uint1 => {
155 let vals: Vec<u8> = ids.iter().map(|id| id.to_u128() as u8).collect();
156 uint1("id", vals)
157 }
158 ValueType::Uint2 => {
159 let vals: Vec<u16> = ids.iter().map(|id| id.to_u128() as u16).collect();
160 uint2("id", vals)
161 }
162 ValueType::Uint4 => {
163 let vals: Vec<u32> = ids.iter().map(|id| id.to_u128() as u32).collect();
164 uint4("id", vals)
165 }
166 ValueType::Uint8 => {
167 let vals: Vec<u64> = ids.iter().map(|id| id.to_u128() as u64).collect();
168 uint8("id", vals)
169 }
170 ValueType::Uint16 => {
171 let vals: Vec<u128> = ids.iter().map(|id| id.to_u128()).collect();
172 uint16("id", vals)
173 }
174 _ => return Err(internal_error!("Invalid dictionary id_type: {:?}", id_type)),
175 };
176
177 Ok(data)
178}
179
180fn build_value_column(values: &[Value], value_type: ValueType) -> (FieldRef, ArrayRef) {
181 let mut data = ColumnBuilder::with_capacity(value_type, values.len());
182 for value in values {
183 data.push_value(value.clone());
184 }
185 data.finish("value")
186}