linkrs_query/executor/streaming/chunk/
columnar_common.rs

1//! Shared variant set and helpers for typed columnar layouts.
2//!
3//! [`TypedColumn`](super::typed::TypedColumn) (immutable per-chunk snapshot)
4//! and [`BatchColumn`](super::columnar_batch::BatchColumn) (accumulable
5//! blocking-operator state) intentionally share one variant set so a chunk's
6//! typed columns can move into a batch without re-inspecting values. The
7//! single source of truth lives here:
8//!
9//! - [`column_variants!`] declares the enum plus the mechanical `len`,
10//!   `is_empty`, and `estimated_size` methods for both types;
11//! - [`gather_column!`] builds either output from a `TypedColumn` source at
12//!   an index selection, backing both gather entry points;
13//! - the validity-bitmap helpers and [`gather_bitmap`] are shared instead of
14//!   being duplicated per module.
15//!
16//! `value_at`, comparison, and accumulation keep hand-written impls: their
17//! contracts differ per type (`Option<Value>` vs `Value`, `Empty`, row
18//! comparison), so generating them would obscure rather than unify.
19
20use std::sync::Arc;
21
22use linkrs_core::types::VertexId;
23use linkrs_core::{EdgeHeader, Value};
24
25/// Struct-of-arrays edge identity column: one `(src, dst, edge_type,
26/// ranking)` tuple per row with no property map and no storage-internal
27/// identifier. All four vectors always share one length; every constructor
28/// and mutator below maintains that invariant together, so element access
29/// stays index-aligned without per-access checks.
30#[derive(Debug, Clone)]
31pub struct EdgeHeaderColumn {
32    src: Vec<VertexId>,
33    dst: Vec<VertexId>,
34    edge_type: Vec<Arc<str>>,
35    ranking: Vec<i64>,
36}
37
38impl EdgeHeaderColumn {
39    /// Take over storage batch vectors. The storage batch contract guarantees
40    /// equal lengths across the four vectors.
41    pub(crate) fn from_parts(
42        src: Vec<VertexId>,
43        dst: Vec<VertexId>,
44        edge_type: Vec<String>,
45        ranking: Vec<i64>,
46    ) -> Self {
47        debug_assert!(
48            src.len() == dst.len() && src.len() == edge_type.len() && src.len() == ranking.len(),
49            "edge identity vectors must share one length"
50        );
51        let edge_type = edge_type
52            .into_iter()
53            .map(|s| Arc::from(s.as_str()))
54            .collect();
55        Self {
56            src,
57            dst,
58            edge_type,
59            ranking,
60        }
61    }
62
63    pub(crate) fn len(&self) -> usize {
64        self.src.len()
65    }
66
67    pub(crate) fn estimated_size(&self) -> usize {
68        self.src.capacity() * std::mem::size_of::<VertexId>()
69            + self.dst.capacity() * std::mem::size_of::<VertexId>()
70            + self.edge_type.iter().map(|s| s.len()).sum::<usize>()
71            + self.ranking.capacity() * std::mem::size_of::<i64>()
72    }
73
74    /// Materialize the header at `idx` (`None` when out of bounds).
75    pub(crate) fn header_at(&self, idx: usize) -> Option<EdgeHeader> {
76        match (
77            self.src.get(idx),
78            self.dst.get(idx),
79            self.edge_type.get(idx),
80            self.ranking.get(idx),
81        ) {
82            (Some(&src), Some(&dst), Some(edge_type), Some(&ranking)) => {
83                Some(EdgeHeader::new(src, dst, edge_type.to_string(), ranking))
84            }
85            _ => None,
86        }
87    }
88
89    /// Materialize every row (row-major fallback for row-path consumers).
90    pub(crate) fn to_values(&self) -> Vec<Value> {
91        self.src
92            .iter()
93            .zip(self.dst.iter())
94            .zip(self.edge_type.iter())
95            .zip(self.ranking.iter())
96            .map(|(((src, dst), edge_type), ranking)| {
97                Value::edge_header(EdgeHeader::new(*src, *dst, edge_type.to_string(), *ranking))
98            })
99            .collect()
100    }
101
102    /// Gather the entries at `indices`, preserving order.
103    pub(crate) fn gather(&self, indices: &[usize]) -> Self {
104        Self {
105            src: indices.iter().map(|&i| self.src[i]).collect(),
106            dst: indices.iter().map(|&i| self.dst[i]).collect(),
107            edge_type: indices.iter().map(|&i| self.edge_type[i].clone()).collect(),
108            ranking: indices.iter().map(|&i| self.ranking[i]).collect(),
109        }
110    }
111
112    pub(crate) fn truncate(&mut self, len: usize) {
113        self.src.truncate(len);
114        self.dst.truncate(len);
115        self.edge_type.truncate(len);
116        self.ranking.truncate(len);
117    }
118
119    /// Reorder rows by `perm` (`result[i] = old[perm[i]]`).
120    pub(crate) fn permute(&mut self, perm: &[usize]) {
121        let old = std::mem::replace(
122            self,
123            Self {
124                src: Vec::new(),
125                dst: Vec::new(),
126                edge_type: Vec::new(),
127                ranking: Vec::new(),
128            },
129        );
130        self.src = perm.iter().map(|&i| old.src[i]).collect();
131        self.dst = perm.iter().map(|&i| old.dst[i]).collect();
132        self.edge_type = perm.iter().map(|&i| old.edge_type[i].clone()).collect();
133        self.ranking = perm.iter().map(|&i| old.ranking[i]).collect();
134    }
135
136    /// Append the entries of `other` at `indices`.
137    pub(crate) fn extend_indexed(&mut self, other: &Self, indices: &[usize]) {
138        self.src.extend(indices.iter().map(|&i| other.src[i]));
139        self.dst.extend(indices.iter().map(|&i| other.dst[i]));
140        self.edge_type
141            .extend(indices.iter().map(|&i| other.edge_type[i].clone()));
142        self.ranking
143            .extend(indices.iter().map(|&i| other.ranking[i]));
144    }
145
146    /// Append one header.
147    pub(crate) fn push(&mut self, src: VertexId, dst: VertexId, edge_type: &str, ranking: i64) {
148        self.src.push(src);
149        self.dst.push(dst);
150        self.edge_type.push(Arc::from(edge_type));
151        self.ranking.push(ranking);
152    }
153}
154
155/// Whether row `idx` is valid in `bitmap` (bit set = valid, bit clear = NULL).
156#[inline]
157pub(crate) fn bitmap_is_valid(bitmap: &[u64], idx: usize) -> bool {
158    bitmap[idx / 64] & (1u64 << (idx % 64)) != 0
159}
160
161/// Set (`valid == true`) or clear the validity bit of row `idx`.
162#[inline]
163pub(crate) fn bitmap_set_bit(bitmap: &mut [u64], idx: usize, valid: bool) {
164    if valid {
165        bitmap[idx / 64] |= 1u64 << (idx % 64);
166    } else {
167        bitmap[idx / 64] &= !(1u64 << (idx % 64));
168    }
169}
170
171/// Gather the validity bits at `indices` into a new bitmap.
172pub(crate) fn gather_bitmap(bitmap: &[u64], indices: &[usize]) -> Vec<u64> {
173    let mut out = vec![0u64; indices.len().div_ceil(64)];
174    for (j, &i) in indices.iter().enumerate() {
175        bitmap_set_bit(&mut out, j, bitmap_is_valid(bitmap, i));
176    }
177    out
178}
179
180/// Declare a typed columnar enum with the shared variant set.
181///
182/// `$name` is the enum to declare; `$extra` lists additional leading
183/// variants (e.g. `BatchColumn::Empty`) that contribute `0` to `len` and
184/// `estimated_size`. The generated `len`, `is_empty`, and `estimated_size`
185/// are textually identical for both column types, so adding a kind touches
186/// this list once instead of drifting across two files.
187///
188/// The macro expands at the call site, so `Arc`, `Decimal128Value`,
189/// `EdgeHeaderColumn`, and `Value` must be in scope where it is invoked
190/// (`VertexId` is named by absolute path and needs no import).
191macro_rules! column_variants {
192    ($name:ident $(, $extra:ident)*) => {
193        #[derive(Debug, Clone)]
194        pub enum $name {
195            $($extra,)*
196            I64(Vec<i64>),
197            F64(Vec<f64>),
198            I32(Vec<i32>),
199            Bool(Vec<bool>),
200            /// Days since epoch per row (see [`DateValue::to_days`]).
201            Date(Vec<i64>),
202            /// Micros since epoch per row (see [`DateTimeValue::to_micros`]).
203            DateTime(Vec<i64>),
204            Utf8(Vec<Arc<str>>),
205            /// Decimal128 per row (decimal semantics, `Ord`).
206            Decimal(Vec<Decimal128Value>),
207            /// I64 column with a validity bitmap (bit set = valid).
208            NullableI64(Vec<i64>, Vec<u64>),
209            /// F64 column with a validity bitmap.
210            NullableF64(Vec<f64>, Vec<u64>),
211            /// I32 column with a validity bitmap.
212            NullableI32(Vec<i32>, Vec<u64>),
213            /// Bool column with a validity bitmap.
214            NullableBool(Vec<bool>, Vec<u64>),
215            /// Date column with a validity bitmap.
216            NullableDate(Vec<i64>, Vec<u64>),
217            /// DateTime column with a validity bitmap.
218            NullableDateTime(Vec<i64>, Vec<u64>),
219            /// Utf8 column with a validity bitmap.
220            NullableUtf8(Vec<Arc<str>>, Vec<u64>),
221            /// Decimal column with a validity bitmap.
222            NullableDecimal(Vec<Decimal128Value>, Vec<u64>),
223            /// Vertex identity column: one id per row, never NULL.
224            VertexIdentity(Vec<linkrs_core::types::VertexId>),
225            /// Edge identity column: struct-of-arrays headers, never NULL.
226            EdgeHeader(EdgeHeaderColumn),
227            Fallback(Vec<Value>),
228        }
229
230        impl $name {
231            pub fn len(&self) -> usize {
232                match self {
233                    $($name::$extra => 0,)*
234                    $name::I64(v) => v.len(),
235                    $name::F64(v) => v.len(),
236                    $name::I32(v) => v.len(),
237                    $name::Bool(v) => v.len(),
238                    $name::Date(v) => v.len(),
239                    $name::DateTime(v) => v.len(),
240                    $name::Utf8(v) => v.len(),
241                    $name::Decimal(v) => v.len(),
242                    $name::NullableI64(v, _) => v.len(),
243                    $name::NullableF64(v, _) => v.len(),
244                    $name::NullableI32(v, _) => v.len(),
245                    $name::NullableBool(v, _) => v.len(),
246                    $name::NullableDate(v, _) => v.len(),
247                    $name::NullableDateTime(v, _) => v.len(),
248                    $name::NullableUtf8(v, _) => v.len(),
249                    $name::NullableDecimal(v, _) => v.len(),
250                    $name::VertexIdentity(v) => v.len(),
251                    $name::EdgeHeader(v) => v.len(),
252                    $name::Fallback(v) => v.len(),
253                }
254            }
255
256            pub fn is_empty(&self) -> bool {
257                self.len() == 0
258            }
259
260            /// Estimated heap bytes of this column (for memory accounting).
261            pub fn estimated_size(&self) -> usize {
262                match self {
263                    $($name::$extra => 0,)*
264                    $name::I64(v) => v.capacity() * std::mem::size_of::<i64>(),
265                    $name::F64(v) => v.capacity() * std::mem::size_of::<f64>(),
266                    $name::I32(v) => v.capacity() * std::mem::size_of::<i32>(),
267                    $name::Bool(v) => v.capacity() * std::mem::size_of::<bool>(),
268                    $name::Date(v) => v.capacity() * std::mem::size_of::<i64>(),
269                    $name::DateTime(v) => v.capacity() * std::mem::size_of::<i64>(),
270                    $name::Utf8(v) => v.iter().map(|s| s.len()).sum(),
271                    $name::Decimal(v) => v.capacity() * std::mem::size_of::<Decimal128Value>(),
272                    $name::NullableI64(v, b) => {
273                        v.capacity() * std::mem::size_of::<i64>()
274                            + b.capacity() * std::mem::size_of::<u64>()
275                    }
276                    $name::NullableF64(v, b) => {
277                        v.capacity() * std::mem::size_of::<f64>()
278                            + b.capacity() * std::mem::size_of::<u64>()
279                    }
280                    $name::NullableI32(v, b) => {
281                        v.capacity() * std::mem::size_of::<i32>()
282                            + b.capacity() * std::mem::size_of::<u64>()
283                    }
284                    $name::NullableBool(v, b) => {
285                        v.capacity() * std::mem::size_of::<bool>()
286                            + b.capacity() * std::mem::size_of::<u64>()
287                    }
288                    $name::NullableDate(v, b) => {
289                        v.capacity() * std::mem::size_of::<i64>()
290                            + b.capacity() * std::mem::size_of::<u64>()
291                    }
292                    $name::NullableDateTime(v, b) => {
293                        v.capacity() * std::mem::size_of::<i64>()
294                            + b.capacity() * std::mem::size_of::<u64>()
295                    }
296                    $name::NullableUtf8(v, b) => {
297                        v.iter().map(|s| s.len()).sum::<usize>()
298                            + b.capacity() * std::mem::size_of::<u64>()
299                    }
300                    $name::NullableDecimal(v, b) => {
301                        v.capacity() * std::mem::size_of::<Decimal128Value>()
302                            + b.capacity() * std::mem::size_of::<u64>()
303                    }
304                    $name::VertexIdentity(v) => {
305                        v.capacity() * std::mem::size_of::<linkrs_core::types::VertexId>()
306                    }
307                    $name::EdgeHeader(v) => v.estimated_size(),
308                    $name::Fallback(v) => v.iter().map(Value::estimated_size).sum(),
309                }
310            }
311        }
312    };
313}
314
315pub(crate) use column_variants;
316
317/// Gather a `TypedColumn` source at `indices` into the `$out` column type.
318///
319/// Both column types share variant names, so one match serves the
320/// `TypedColumn` gather and the `BatchColumn` gather without duplication.
321/// The caller must have `TypedColumn` and `gather_bitmap` in scope.
322macro_rules! gather_column {
323    ($out:ident, $col:expr, $indices:expr) => {{
324        let indices: &[usize] = $indices;
325        match $col {
326            TypedColumn::I64(v) => $out::I64(indices.iter().map(|&i| v[i]).collect()),
327            TypedColumn::F64(v) => $out::F64(indices.iter().map(|&i| v[i]).collect()),
328            TypedColumn::I32(v) => $out::I32(indices.iter().map(|&i| v[i]).collect()),
329            TypedColumn::Bool(v) => $out::Bool(indices.iter().map(|&i| v[i]).collect()),
330            TypedColumn::Date(v) => $out::Date(indices.iter().map(|&i| v[i]).collect()),
331            TypedColumn::DateTime(v) => $out::DateTime(indices.iter().map(|&i| v[i]).collect()),
332            TypedColumn::Utf8(v) => $out::Utf8(indices.iter().map(|&i| v[i].clone()).collect()),
333            TypedColumn::Decimal(v) => {
334                $out::Decimal(indices.iter().map(|&i| v[i].clone()).collect())
335            }
336            TypedColumn::NullableI64(v, b) => $out::NullableI64(
337                indices.iter().map(|&i| v[i]).collect(),
338                gather_bitmap(b, indices),
339            ),
340            TypedColumn::NullableF64(v, b) => $out::NullableF64(
341                indices.iter().map(|&i| v[i]).collect(),
342                gather_bitmap(b, indices),
343            ),
344            TypedColumn::NullableI32(v, b) => $out::NullableI32(
345                indices.iter().map(|&i| v[i]).collect(),
346                gather_bitmap(b, indices),
347            ),
348            TypedColumn::NullableBool(v, b) => $out::NullableBool(
349                indices.iter().map(|&i| v[i]).collect(),
350                gather_bitmap(b, indices),
351            ),
352            TypedColumn::NullableDate(v, b) => $out::NullableDate(
353                indices.iter().map(|&i| v[i]).collect(),
354                gather_bitmap(b, indices),
355            ),
356            TypedColumn::NullableDateTime(v, b) => $out::NullableDateTime(
357                indices.iter().map(|&i| v[i]).collect(),
358                gather_bitmap(b, indices),
359            ),
360            TypedColumn::NullableUtf8(v, b) => $out::NullableUtf8(
361                indices.iter().map(|&i| v[i].clone()).collect(),
362                gather_bitmap(b, indices),
363            ),
364            TypedColumn::NullableDecimal(v, b) => $out::NullableDecimal(
365                indices.iter().map(|&i| v[i].clone()).collect(),
366                gather_bitmap(b, indices),
367            ),
368            TypedColumn::VertexIdentity(v) => {
369                $out::VertexIdentity(indices.iter().map(|&i| v[i]).collect())
370            }
371            TypedColumn::EdgeHeader(v) => $out::EdgeHeader(v.gather(indices)),
372            TypedColumn::Fallback(v) => {
373                $out::Fallback(indices.iter().map(|&i| v[i].clone()).collect())
374            }
375        }
376    }};
377}
378
379pub(crate) use gather_column;
linkrs_query/executor/streaming/chunk/columnar_common.rs

linkrs_query/executor/streaming/chunk/
columnar_common.rs
