1use hawdb_core::{HawDBError, LogicalType, Result, Value, ValueRef};
18use hawdb_plan_cypher::ComparisonOp;
19use hawdb_storage::relational::{RelationalKey, RelationalValue};
20use std::sync::Arc;
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
23pub struct SlotId(pub u32);
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum ColumnType {
27 Bool,
28 Int64,
29 Float64,
30 Utf8,
31 NodeId,
32 RelationalRowLocator,
33 Dynamic,
34}
35
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum SlotType {
43 Logical(LogicalType),
44 NodeId,
45 RelationalRowLocator,
46}
47
48impl SlotType {
49 pub const fn logical(logical_type: LogicalType) -> Self {
50 Self::Logical(logical_type)
51 }
52
53 pub const fn logical_type(self) -> Option<LogicalType> {
54 match self {
55 Self::Logical(logical_type) => Some(logical_type),
56 Self::NodeId => Some(LogicalType::Int64),
57 Self::RelationalRowLocator => None,
58 }
59 }
60
61 pub const fn accepts(self, column_type: ColumnType) -> bool {
62 match (self, column_type) {
63 (
64 Self::Logical(LogicalType::Any),
65 ColumnType::Bool
66 | ColumnType::Int64
67 | ColumnType::Float64
68 | ColumnType::Utf8
69 | ColumnType::Dynamic,
70 )
71 | (Self::Logical(LogicalType::Boolean), ColumnType::Bool)
72 | (Self::Logical(LogicalType::Int64), ColumnType::Int64)
73 | (Self::Logical(LogicalType::Float64), ColumnType::Float64)
74 | (Self::Logical(LogicalType::String | LogicalType::Text), ColumnType::Utf8)
75 | (Self::NodeId, ColumnType::NodeId)
76 | (Self::RelationalRowLocator, ColumnType::RelationalRowLocator) => true,
77 (Self::Logical(logical_type), ColumnType::Dynamic) => {
78 !matches!(logical_type, LogicalType::Binary)
79 }
80 _ => false,
81 }
82 }
83}
84
85#[derive(Debug, Clone, PartialEq, Eq)]
86pub struct RelationalRowLocator {
87 table_id: u32,
88 primary_key: RelationalKey,
89}
90
91impl RelationalRowLocator {
92 pub fn new(table_id: u32, primary_key: RelationalKey) -> Self {
93 Self {
94 table_id,
95 primary_key,
96 }
97 }
98
99 pub fn table_id(&self) -> u32 {
100 self.table_id
101 }
102
103 pub fn primary_key(&self) -> &RelationalKey {
104 &self.primary_key
105 }
106
107 pub fn allocated_bytes(&self) -> usize {
108 self.primary_key
109 .0
110 .capacity()
111 .saturating_mul(std::mem::size_of::<RelationalValue>())
112 .saturating_add(
113 self.primary_key
114 .0
115 .iter()
116 .map(|value| match value {
117 RelationalValue::Text(value) => value.capacity(),
118 RelationalValue::Bytea(value) => value.capacity(),
119 RelationalValue::Vector(value) => value.retained_capacity(),
120 RelationalValue::Null
121 | RelationalValue::Boolean(_)
122 | RelationalValue::BigInt(_)
123 | RelationalValue::DoublePrecision(_)
124 | RelationalValue::Decimal(_)
125 | RelationalValue::Uuid(_)
126 | RelationalValue::Overflow(_) => 0,
127 })
128 .sum::<usize>(),
129 )
130 }
131}
132
133#[derive(Debug, Clone, PartialEq, Eq)]
134pub struct SlotDescriptor {
135 pub id: SlotId,
136 pub name: String,
137 pub slot_type: SlotType,
138}
139
140impl SlotDescriptor {
141 pub const fn logical_type(&self) -> Option<LogicalType> {
142 self.slot_type.logical_type()
143 }
144}
145
146#[derive(Debug, Clone, PartialEq, Eq)]
147pub struct BindingSchema {
148 slots: Arc<[SlotDescriptor]>,
149}
150
151impl BindingSchema {
152 pub fn try_new(slots: Vec<SlotDescriptor>) -> Result<Self> {
153 for (index, slot) in slots.iter().enumerate() {
154 if slot.id.0 as usize != index {
155 return Err(HawDBError::Execution(format!(
156 "columnar schema slot ids must be dense: expected {index}, got {}",
157 slot.id.0
158 )));
159 }
160 }
161 Ok(Self {
162 slots: slots.into(),
163 })
164 }
165
166 pub fn slots(&self) -> &[SlotDescriptor] {
167 &self.slots
168 }
169
170 pub fn slot(&self, id: SlotId) -> Option<&SlotDescriptor> {
171 self.slots.get(id.0 as usize)
172 }
173
174 pub fn slot_by_name(&self, name: &str) -> Option<&SlotDescriptor> {
175 self.slots.iter().find(|slot| slot.name == name)
176 }
177
178 pub fn len(&self) -> usize {
179 self.slots.len()
180 }
181
182 pub fn is_empty(&self) -> bool {
183 self.slots.is_empty()
184 }
185}
186
187#[derive(Debug, Clone, PartialEq, Eq)]
188pub enum Validity {
189 All { len: usize },
190 Bitmap { len: usize, words: Arc<[u64]> },
191}
192
193impl Validity {
194 pub fn all(len: usize) -> Self {
195 Self::All { len }
196 }
197
198 pub fn len(&self) -> usize {
199 match self {
200 Self::All { len } | Self::Bitmap { len, .. } => *len,
201 }
202 }
203
204 pub fn is_empty(&self) -> bool {
205 self.len() == 0
206 }
207
208 pub fn is_valid(&self, row: usize) -> bool {
209 if row >= self.len() {
210 return false;
211 }
212 match self {
213 Self::All { .. } => true,
214 Self::Bitmap { words, .. } => words
215 .get(row / u64::BITS as usize)
216 .is_some_and(|word| word & (1u64 << (row % u64::BITS as usize)) != 0),
217 }
218 }
219
220 pub fn valid_count(&self) -> usize {
221 match self {
222 Self::All { len } => *len,
223 Self::Bitmap { len, words } => count_bitmap_rows(words, *len),
224 }
225 }
226
227 pub fn view(&self) -> ValidityView<'_> {
228 match self {
229 Self::All { len } => ValidityView::All { len: *len },
230 Self::Bitmap { len, words } => ValidityView::Bitmap { len: *len, words },
231 }
232 }
233}
234
235#[derive(Debug, Clone, Copy)]
236pub enum ValidityView<'a> {
237 All { len: usize },
238 Bitmap { len: usize, words: &'a [u64] },
239}
240
241impl ValidityView<'_> {
242 pub fn len(self) -> usize {
243 match self {
244 Self::All { len } | Self::Bitmap { len, .. } => len,
245 }
246 }
247
248 pub fn is_empty(self) -> bool {
249 self.len() == 0
250 }
251
252 pub fn is_valid(self, row: usize) -> bool {
253 if row >= self.len() {
254 return false;
255 }
256 match self {
257 Self::All { .. } => true,
258 Self::Bitmap { words, .. } => words
259 .get(row / u64::BITS as usize)
260 .is_some_and(|word| word & (1u64 << (row % u64::BITS as usize)) != 0),
261 }
262 }
263
264 pub fn valid_count(self) -> usize {
265 match self {
266 Self::All { len } => len,
267 Self::Bitmap { len, words } => count_bitmap_rows(words, len),
268 }
269 }
270
271 pub fn to_owned(self) -> Validity {
272 match self {
273 Self::All { len } => Validity::All { len },
274 Self::Bitmap { len, words } => Validity::Bitmap {
275 len,
276 words: words.to_vec().into(),
277 },
278 }
279 }
280}
281
282#[derive(Debug, Default)]
283pub struct ValidityBuilder {
284 len: usize,
285 words: Vec<u64>,
286 word_capacity: usize,
287}
288
289impl ValidityBuilder {
290 pub fn with_capacity(rows: usize) -> Self {
291 Self {
292 len: 0,
293 words: Vec::new(),
294 word_capacity: rows.div_ceil(u64::BITS as usize),
295 }
296 }
297
298 pub fn push(&mut self, valid: bool) {
299 let word_index = self.len / u64::BITS as usize;
300 let bit_index = self.len % u64::BITS as usize;
301 if self.words.is_empty() {
302 if !valid {
303 if self.words.capacity() < self.word_capacity {
304 self.words.reserve(self.word_capacity);
305 }
306 self.words.resize(word_index + 1, u64::MAX);
307 self.words[word_index] = if bit_index == 0 {
308 0
309 } else {
310 (1u64 << bit_index) - 1
311 };
312 }
313 } else {
314 if word_index == self.words.len() {
315 self.words.push(0);
316 }
317 if valid {
318 self.words[word_index] |= 1u64 << bit_index;
319 }
320 }
321 self.len = self.len.saturating_add(1);
322 }
323
324 pub fn clear(&mut self) {
325 self.len = 0;
326 self.words.clear();
327 }
328
329 pub fn view(&self) -> ValidityView<'_> {
330 if self.words.is_empty() {
331 ValidityView::All { len: self.len }
332 } else {
333 ValidityView::Bitmap {
334 len: self.len,
335 words: &self.words,
336 }
337 }
338 }
339
340 pub fn finish(self) -> Validity {
341 if self.words.is_empty() {
342 Validity::All { len: self.len }
343 } else {
344 Validity::Bitmap {
345 len: self.len,
346 words: self.words.into(),
347 }
348 }
349 }
350}
351
352#[derive(Debug, Clone, PartialEq)]
353pub enum ColumnVector {
354 Bool {
355 values: Arc<[u8]>,
356 validity: Validity,
357 },
358 Int64 {
359 values: Arc<[i64]>,
360 validity: Validity,
361 },
362 Float64 {
363 values: Arc<[f64]>,
364 validity: Validity,
365 },
366 Utf8 {
367 values: Arc<[String]>,
368 validity: Validity,
369 },
370 NodeId(Arc<[u64]>),
371 RelationalRowLocator(Arc<[RelationalRowLocator]>),
372 Dynamic(Arc<[Value]>),
373}
374
375impl ColumnVector {
376 pub fn boolean(values: Vec<bool>, validity: Validity) -> Result<Self> {
377 Self::boolean_bytes(values.into_iter().map(u8::from).collect(), validity)
378 }
379
380 pub fn boolean_bytes(values: Vec<u8>, validity: Validity) -> Result<Self> {
381 ensure_column_len("Bool", values.len(), validity.len())?;
382 if values.iter().any(|value| *value > 1) {
383 return Err(HawDBError::Execution(
384 "Bool column contains a value other than 0 or 1".to_string(),
385 ));
386 }
387 Ok(Self::Bool {
388 values: values.into(),
389 validity,
390 })
391 }
392
393 pub fn int64(values: Vec<i64>, validity: Validity) -> Result<Self> {
394 ensure_column_len("Int64", values.len(), validity.len())?;
395 Ok(Self::Int64 {
396 values: values.into(),
397 validity,
398 })
399 }
400
401 pub fn float64(values: Vec<f64>, validity: Validity) -> Result<Self> {
402 ensure_column_len("Float64", values.len(), validity.len())?;
403 Ok(Self::Float64 {
404 values: values.into(),
405 validity,
406 })
407 }
408
409 pub fn node_ids(values: Vec<u64>) -> Self {
410 Self::NodeId(values.into())
411 }
412
413 pub fn relational_row_locators(values: Vec<RelationalRowLocator>) -> Self {
414 Self::RelationalRowLocator(values.into())
415 }
416
417 pub fn len(&self) -> usize {
418 match self {
419 Self::Bool { values, .. } => values.len(),
420 Self::Int64 { values, .. } => values.len(),
421 Self::Float64 { values, .. } => values.len(),
422 Self::Utf8 { values, .. } => values.len(),
423 Self::NodeId(values) => values.len(),
424 Self::RelationalRowLocator(values) => values.len(),
425 Self::Dynamic(values) => values.len(),
426 }
427 }
428
429 pub fn is_empty(&self) -> bool {
430 self.len() == 0
431 }
432
433 pub fn column_type(&self) -> ColumnType {
434 match self {
435 Self::Bool { .. } => ColumnType::Bool,
436 Self::Int64 { .. } => ColumnType::Int64,
437 Self::Float64 { .. } => ColumnType::Float64,
438 Self::Utf8 { .. } => ColumnType::Utf8,
439 Self::NodeId(_) => ColumnType::NodeId,
440 Self::RelationalRowLocator(_) => ColumnType::RelationalRowLocator,
441 Self::Dynamic(_) => ColumnType::Dynamic,
442 }
443 }
444
445 pub fn is_valid(&self, row: usize) -> bool {
446 match self {
447 Self::Bool { validity, .. }
448 | Self::Int64 { validity, .. }
449 | Self::Float64 { validity, .. }
450 | Self::Utf8 { validity, .. } => validity.is_valid(row),
451 Self::NodeId(values) => row < values.len(),
452 Self::RelationalRowLocator(values) => row < values.len(),
453 Self::Dynamic(values) => values.get(row).is_some_and(|value| *value != Value::Null),
454 }
455 }
456
457 pub fn value_ref(&self, row: usize) -> Option<ValueRef<'_>> {
458 match self {
459 Self::Bool { values, validity } if validity.is_valid(row) => {
460 values.get(row).map(|value| ValueRef::Bool(*value != 0))
461 }
462 Self::Int64 { values, validity } if validity.is_valid(row) => {
463 values.get(row).copied().map(ValueRef::Int)
464 }
465 Self::Float64 { values, validity } if validity.is_valid(row) => {
466 values.get(row).copied().map(ValueRef::Float)
467 }
468 Self::Utf8 { values, validity } if validity.is_valid(row) => {
469 values.get(row).map(|value| ValueRef::String(value))
470 }
471 Self::Bool { .. } | Self::Int64 { .. } | Self::Float64 { .. } | Self::Utf8 { .. }
472 if row < self.len() =>
473 {
474 Some(ValueRef::Null)
475 }
476 Self::Bool { .. } | Self::Int64 { .. } | Self::Float64 { .. } | Self::Utf8 { .. } => {
477 None
478 }
479 Self::NodeId(values) => values.get(row).map(|value| ValueRef::Int(*value as i64)),
480 Self::RelationalRowLocator(_) => None,
481 Self::Dynamic(values) => values.get(row).map(Value::as_ref),
482 }
483 }
484
485 pub fn value(&self, row: usize) -> Option<Value> {
486 self.value_ref(row).map(ValueRef::to_owned_value)
487 }
488
489 pub fn relational_row_locator(&self, row: usize) -> Option<&RelationalRowLocator> {
490 let Self::RelationalRowLocator(values) = self else {
491 return None;
492 };
493 values.get(row)
494 }
495
496 pub fn estimated_memory_bytes(&self) -> usize {
497 let validity_bytes = match self {
498 Self::Bool { validity, .. }
499 | Self::Int64 { validity, .. }
500 | Self::Float64 { validity, .. }
501 | Self::Utf8 { validity, .. } => match validity {
502 Validity::All { .. } => 0,
503 Validity::Bitmap { words, .. } => words.len() * std::mem::size_of::<u64>(),
504 },
505 Self::NodeId(_) | Self::RelationalRowLocator(_) | Self::Dynamic(_) => 0,
506 };
507 validity_bytes.saturating_add(match self {
508 Self::Bool { values, .. } => values.len(),
509 Self::Int64 { values, .. } => values.len() * std::mem::size_of::<i64>(),
510 Self::Float64 { values, .. } => values.len() * std::mem::size_of::<f64>(),
511 Self::Utf8 { values, .. } => values.iter().fold(
512 values.len() * std::mem::size_of::<String>(),
513 |total, value| total.saturating_add(value.len()),
514 ),
515 Self::NodeId(values) => values.len() * std::mem::size_of::<u64>(),
516 Self::RelationalRowLocator(values) => values.iter().fold(
517 values
518 .len()
519 .saturating_mul(std::mem::size_of::<RelationalRowLocator>()),
520 |total, locator| total.saturating_add(locator.allocated_bytes()),
521 ),
522 Self::Dynamic(values) => values.len() * std::mem::size_of::<Value>(),
523 })
524 }
525}
526
527fn ensure_column_len(kind: &str, values: usize, validity: usize) -> Result<()> {
528 if values != validity {
529 return Err(HawDBError::Execution(format!(
530 "{kind} column has {values} values but {validity} validity entries"
531 )));
532 }
533 Ok(())
534}
535
536#[derive(Debug, Clone, PartialEq, Eq)]
537pub enum Selection {
538 All {
539 len: usize,
540 },
541 Bitmap {
542 len: usize,
543 words: Arc<[u64]>,
544 selected: usize,
545 },
546 Indices {
547 len: usize,
548 rows: Arc<[u32]>,
549 },
550}
551
552impl Selection {
553 pub fn all(len: usize) -> Self {
554 Self::All { len }
555 }
556
557 pub fn none(len: usize) -> Self {
558 Self::Indices {
559 len,
560 rows: Arc::from([]),
561 }
562 }
563
564 pub fn len(&self) -> usize {
565 match self {
566 Self::All { len } | Self::Bitmap { len, .. } | Self::Indices { len, .. } => *len,
567 }
568 }
569
570 pub fn selected_count(&self) -> usize {
571 match self {
572 Self::All { len } => *len,
573 Self::Bitmap { selected, .. } => *selected,
574 Self::Indices { rows, .. } => rows.len(),
575 }
576 }
577
578 pub fn is_empty(&self) -> bool {
579 self.selected_count() == 0
580 }
581
582 pub fn iter(&self) -> SelectionIter<'_> {
583 SelectionIter {
584 selection: self,
585 cursor: 0,
586 }
587 }
588
589 pub fn limit(&self, offset: usize, max_rows: usize) -> Self {
590 let mut output = SelectionBuilder::new(self.len());
591 for row in self.iter().skip(offset).take(max_rows) {
592 output.select(row);
593 }
594 output.finish()
595 }
596
597 pub fn estimated_memory_bytes(&self) -> usize {
598 match self {
599 Self::All { .. } => 0,
600 Self::Bitmap { words, .. } => words.len() * std::mem::size_of::<u64>(),
601 Self::Indices { rows, .. } => rows.len() * std::mem::size_of::<u32>(),
602 }
603 }
604}
605
606struct SelectionBuilder {
607 len: usize,
608 selected: usize,
609 sparse_rows: Vec<u32>,
610 bitmap_words: Option<Vec<u64>>,
611}
612
613impl SelectionBuilder {
614 fn new(len: usize) -> Self {
615 Self {
616 len,
617 selected: 0,
618 sparse_rows: Vec::new(),
619 bitmap_words: (len > u32::MAX as usize)
620 .then(|| vec![0; len.div_ceil(u64::BITS as usize)]),
621 }
622 }
623
624 fn select(&mut self, row: usize) {
625 self.selected = self.selected.saturating_add(1);
626 if let Some(words) = &mut self.bitmap_words {
627 words[row / u64::BITS as usize] |= 1u64 << (row % u64::BITS as usize);
628 return;
629 }
630
631 self.sparse_rows.push(row as u32);
632 if self.selected.saturating_mul(8) > self.len {
633 let rows = std::mem::take(&mut self.sparse_rows);
634 let mut words = vec![0; self.len.div_ceil(u64::BITS as usize)];
635 for row in rows {
636 let row = row as usize;
637 words[row / u64::BITS as usize] |= 1u64 << (row % u64::BITS as usize);
638 }
639 self.bitmap_words = Some(words);
640 }
641 }
642
643 fn finish(self) -> Selection {
644 if self.selected == self.len {
645 return Selection::all(self.len);
646 }
647 if self.selected == 0 {
648 return Selection::none(self.len);
649 }
650 match self.bitmap_words {
651 Some(words) => Selection::Bitmap {
652 len: self.len,
653 words: words.into(),
654 selected: self.selected,
655 },
656 None => Selection::Indices {
657 len: self.len,
658 rows: self.sparse_rows.into(),
659 },
660 }
661 }
662}
663
664pub struct SelectionIter<'a> {
665 selection: &'a Selection,
666 cursor: usize,
667}
668
669impl Iterator for SelectionIter<'_> {
670 type Item = usize;
671
672 fn next(&mut self) -> Option<Self::Item> {
673 match self.selection {
674 Selection::All { len } => {
675 let row = self.cursor;
676 if row == *len {
677 return None;
678 }
679 self.cursor = self.cursor.saturating_add(1);
680 Some(row)
681 }
682 Selection::Indices { rows, .. } => {
683 let row = rows.get(self.cursor).copied()? as usize;
684 self.cursor = self.cursor.saturating_add(1);
685 Some(row)
686 }
687 Selection::Bitmap { len, words, .. } => {
688 while self.cursor < *len {
689 let row = self.cursor;
690 self.cursor = self.cursor.saturating_add(1);
691 if words[row / u64::BITS as usize] & (1u64 << (row % u64::BITS as usize)) != 0 {
692 return Some(row);
693 }
694 }
695 None
696 }
697 }
698 }
699}
700
701#[derive(Debug, Clone, PartialEq)]
702pub struct ColumnarBatch {
703 schema: Arc<BindingSchema>,
704 columns: Vec<Arc<ColumnVector>>,
705 selection: Selection,
706 row_count: usize,
707}
708
709#[derive(Debug, Clone, Copy)]
710pub struct ColumnarRowRef<'a> {
711 batch: &'a ColumnarBatch,
712 row: usize,
713}
714
715impl<'a> ColumnarRowRef<'a> {
716 pub fn row_index(self) -> usize {
717 self.row
718 }
719
720 pub fn schema(self) -> &'a BindingSchema {
721 self.batch.schema()
722 }
723
724 pub fn value(self, slot: SlotId) -> Option<ValueRef<'a>> {
725 self.batch.value_ref(slot, self.row)
726 }
727
728 pub fn get(self, name: &str) -> Option<ValueRef<'a>> {
729 let slot = self.schema().slot_by_name(name)?;
730 self.value(slot.id)
731 }
732
733 pub fn column(self, index: usize) -> Option<(&'a str, ValueRef<'a>)> {
734 let slot = self.schema().slots().get(index)?;
735 self.value(slot.id).map(|value| (slot.name.as_str(), value))
736 }
737}
738
739impl ColumnarBatch {
740 pub fn try_new(schema: Arc<BindingSchema>, columns: Vec<Arc<ColumnVector>>) -> Result<Self> {
741 if schema.len() != columns.len() {
742 return Err(HawDBError::Execution(format!(
743 "columnar batch has {} slots but {} columns",
744 schema.len(),
745 columns.len()
746 )));
747 }
748 let row_count = columns.first().map_or(0, |column| column.len());
749 for (slot, column) in schema.slots().iter().zip(&columns) {
750 if column.len() != row_count {
751 return Err(HawDBError::Execution(format!(
752 "columnar slot '{}' has {} rows, expected {row_count}",
753 slot.name,
754 column.len()
755 )));
756 }
757 if !slot.slot_type.accepts(column.column_type()) {
758 return Err(HawDBError::Execution(format!(
759 "columnar slot '{}' expects {:?}, got {:?}",
760 slot.name,
761 slot.slot_type,
762 column.column_type()
763 )));
764 }
765 }
766 Ok(Self {
767 schema,
768 columns,
769 selection: Selection::all(row_count),
770 row_count,
771 })
772 }
773
774 pub fn row_count(&self) -> usize {
775 self.row_count
776 }
777
778 pub fn selected_count(&self) -> usize {
779 self.selection.selected_count()
780 }
781
782 pub fn schema(&self) -> &BindingSchema {
783 &self.schema
784 }
785
786 pub fn column(&self, slot: SlotId) -> Option<&Arc<ColumnVector>> {
787 self.columns.get(slot.0 as usize)
788 }
789
790 pub fn value_ref(&self, slot: SlotId, row: usize) -> Option<ValueRef<'_>> {
791 self.column(slot)?.value_ref(row)
792 }
793
794 pub fn selection(&self) -> &Selection {
795 &self.selection
796 }
797
798 pub fn rows(&self) -> impl Iterator<Item = ColumnarRowRef<'_>> {
799 self.selection
800 .iter()
801 .map(|row| ColumnarRowRef { batch: self, row })
802 }
803
804 pub fn with_selection(mut self, selection: Selection) -> Result<Self> {
805 if selection.len() != self.row_count {
806 return Err(HawDBError::Execution(format!(
807 "columnar selection has {} rows, expected {}",
808 selection.len(),
809 self.row_count
810 )));
811 }
812 self.selection = selection;
813 Ok(self)
814 }
815
816 pub fn project(&self, slots: &[SlotId]) -> Result<Self> {
817 let mut descriptors = Vec::with_capacity(slots.len());
818 let mut columns = Vec::with_capacity(slots.len());
819 for (output_index, slot) in slots.iter().copied().enumerate() {
820 let descriptor = self.schema.slot(slot).ok_or_else(|| {
821 HawDBError::Execution(format!("unknown columnar slot {}", slot.0))
822 })?;
823 descriptors.push(SlotDescriptor {
824 id: SlotId(output_index as u32),
825 name: descriptor.name.clone(),
826 slot_type: descriptor.slot_type,
827 });
828 columns.push(Arc::clone(
829 self.column(slot).expect("schema and columns align"),
830 ));
831 }
832 Ok(Self {
833 schema: Arc::new(BindingSchema::try_new(descriptors)?),
834 columns,
835 selection: self.selection.clone(),
836 row_count: self.row_count,
837 })
838 }
839
840 pub fn filter_numeric(
841 &self,
842 slot: SlotId,
843 predicate: NumericPredicate,
844 expected: NumericLiteral,
845 ) -> Result<Self> {
846 let column = self.column(slot).ok_or_else(|| {
847 HawDBError::Execution(format!("unknown columnar filter slot {}", slot.0))
848 })?;
849 let selection = filter_numeric_column(column, &self.selection, predicate, expected)?;
850 self.clone().with_selection(selection)
851 }
852
853 pub fn filter_boolean(&self, slot: SlotId, expected: bool) -> Result<Self> {
854 let column = self.column(slot).ok_or_else(|| {
855 HawDBError::Execution(format!("unknown columnar filter slot {}", slot.0))
856 })?;
857 let selection = filter_boolean_column(column, &self.selection, expected)?;
858 self.clone().with_selection(selection)
859 }
860
861 pub fn limit(&self, offset: usize, max_rows: usize) -> Self {
862 Self {
863 schema: Arc::clone(&self.schema),
864 columns: self.columns.clone(),
865 selection: self.selection.limit(offset, max_rows),
866 row_count: self.row_count,
867 }
868 }
869
870 pub fn count_selected(&self) -> usize {
871 self.selection.selected_count()
872 }
873
874 pub fn count_valid(&self, slot: SlotId) -> Result<usize> {
875 let column = self.column(slot).ok_or_else(|| {
876 HawDBError::Execution(format!("unknown columnar aggregate slot {}", slot.0))
877 })?;
878 Ok(self
879 .selection
880 .iter()
881 .filter(|row| column.is_valid(*row))
882 .count())
883 }
884
885 pub fn sum_int64(&self, slot: SlotId) -> Result<Option<i64>> {
886 let column = self.column(slot).ok_or_else(|| {
887 HawDBError::Execution(format!("unknown columnar aggregate slot {}", slot.0))
888 })?;
889 let ColumnVector::Int64 { values, validity } = column.as_ref() else {
890 return Err(HawDBError::Execution(format!(
891 "columnar SUM requires Int64, got {:?}",
892 column.column_type()
893 )));
894 };
895 let mut sum = None::<i64>;
896 for row in self.selection.iter() {
897 if validity.is_valid(row) {
898 sum =
899 Some(sum.unwrap_or(0).checked_add(values[row]).ok_or_else(|| {
900 HawDBError::Execution("columnar Int64 SUM overflow".into())
901 })?);
902 }
903 }
904 Ok(sum)
905 }
906
907 pub fn estimated_memory_bytes(&self) -> usize {
908 let schema_bytes = self.schema.slots().iter().fold(
909 self.schema.len() * std::mem::size_of::<SlotDescriptor>(),
910 |total, slot| total.saturating_add(slot.name.len()),
911 );
912 self.columns
913 .iter()
914 .fold(schema_bytes, |total, column| {
915 total.saturating_add(column.estimated_memory_bytes())
916 })
917 .saturating_add(self.selection.estimated_memory_bytes())
918 }
919}
920
921#[derive(Debug, Clone, Copy, PartialEq)]
922pub enum NumericLiteral {
923 Int(i64),
924 Float(f64),
925}
926
927#[derive(Debug, Clone, Copy, PartialEq, Eq)]
928pub enum NumericPredicate {
929 Eq,
930 Compare(ComparisonOp),
931}
932
933impl NumericLiteral {
934 pub fn from_value(value: &Value) -> Option<Self> {
935 Self::from_value_ref(value.as_ref())
936 }
937
938 pub fn from_value_ref(value: ValueRef<'_>) -> Option<Self> {
939 match value {
940 ValueRef::Int(value) => Some(Self::Int(value)),
941 ValueRef::Float(value) => Some(Self::Float(value)),
942 _ => None,
943 }
944 }
945}
946
947pub fn filter_numeric_column(
948 column: &ColumnVector,
949 input: &Selection,
950 predicate: NumericPredicate,
951 expected: NumericLiteral,
952) -> Result<Selection> {
953 if column.len() != input.len() {
954 return Err(HawDBError::Execution(format!(
955 "numeric filter column has {} rows but input selection has {}",
956 column.len(),
957 input.len()
958 )));
959 }
960 match column {
961 ColumnVector::Int64 { values, validity } => {
962 filter_int64_values(values, validity, input, predicate, expected)
963 }
964 ColumnVector::Float64 { values, validity } => {
965 filter_float64_values(values, validity, input, predicate, expected)
966 }
967 other => Err(HawDBError::Execution(format!(
968 "numeric filter requires Int64 or Float64, got {:?}",
969 other.column_type()
970 ))),
971 }
972}
973
974pub fn filter_boolean_column(
975 column: &ColumnVector,
976 input: &Selection,
977 expected: bool,
978) -> Result<Selection> {
979 if column.len() != input.len() {
980 return Err(HawDBError::Execution(format!(
981 "boolean filter column has {} rows but input selection has {}",
982 column.len(),
983 input.len()
984 )));
985 }
986 let ColumnVector::Bool { values, validity } = column else {
987 return Err(HawDBError::Execution(format!(
988 "boolean filter requires Bool, got {:?}",
989 column.column_type()
990 )));
991 };
992 let expected = u8::from(expected);
993 let mut output = SelectionBuilder::new(input.len());
994 for row in input.iter() {
995 if validity.is_valid(row) && values[row] == expected {
996 output.select(row);
997 }
998 }
999 Ok(output.finish())
1000}
1001
1002pub fn filter_int64_values(
1003 values: &[i64],
1004 validity: &Validity,
1005 input: &Selection,
1006 predicate: NumericPredicate,
1007 expected: NumericLiteral,
1008) -> Result<Selection> {
1009 filter_int64_values_view(values, validity.view(), input, predicate, expected)
1010}
1011
1012pub fn filter_int64_values_view(
1013 values: &[i64],
1014 validity: ValidityView<'_>,
1015 input: &Selection,
1016 predicate: NumericPredicate,
1017 expected: NumericLiteral,
1018) -> Result<Selection> {
1019 match predicate {
1020 NumericPredicate::Eq => match expected {
1021 NumericLiteral::Int(expected) => {
1022 filter_numeric_values(values, validity, input, |actual| actual == expected)
1023 }
1024 NumericLiteral::Float(_) => filter_numeric_values(values, validity, input, |_| false),
1025 },
1026 NumericPredicate::Compare(op) => match expected {
1027 NumericLiteral::Int(expected) => {
1028 filter_ordered_values(values, validity, input, op, expected, |actual| actual)
1029 }
1030 NumericLiteral::Float(expected) => filter_ordered_values(
1031 values,
1032 validity,
1033 input,
1034 op,
1035 float_total_order_key(expected),
1036 |actual| float_total_order_key(actual as f64),
1037 ),
1038 },
1039 }
1040}
1041
1042pub fn filter_float64_values(
1043 values: &[f64],
1044 validity: &Validity,
1045 input: &Selection,
1046 predicate: NumericPredicate,
1047 expected: NumericLiteral,
1048) -> Result<Selection> {
1049 filter_float64_values_view(values, validity.view(), input, predicate, expected)
1050}
1051
1052pub fn filter_float64_values_view(
1053 values: &[f64],
1054 validity: ValidityView<'_>,
1055 input: &Selection,
1056 predicate: NumericPredicate,
1057 expected: NumericLiteral,
1058) -> Result<Selection> {
1059 match predicate {
1060 NumericPredicate::Eq => match expected {
1061 NumericLiteral::Int(_) => filter_numeric_values(values, validity, input, |_| false),
1062 NumericLiteral::Float(expected) => {
1063 filter_numeric_values(values, validity, input, |actual| {
1064 actual.total_cmp(&expected).is_eq()
1065 })
1066 }
1067 },
1068 NumericPredicate::Compare(op) => {
1069 let expected = match expected {
1070 NumericLiteral::Int(expected) => expected as f64,
1071 NumericLiteral::Float(expected) => expected,
1072 };
1073 filter_ordered_values(
1074 values,
1075 validity,
1076 input,
1077 op,
1078 float_total_order_key(expected),
1079 float_total_order_key,
1080 )
1081 }
1082 }
1083}
1084
1085pub fn select_int64_values_view(
1086 values: &[i64],
1087 validity: ValidityView<'_>,
1088 predicate: NumericPredicate,
1089 expected: NumericLiteral,
1090 selected_rows: &mut Vec<u32>,
1091) -> Result<()> {
1092 match predicate {
1093 NumericPredicate::Eq => match expected {
1094 NumericLiteral::Int(expected) => {
1095 select_numeric_values(values, validity, selected_rows, |actual| actual == expected)
1096 }
1097 NumericLiteral::Float(_) => {
1098 select_numeric_values(values, validity, selected_rows, |_| false)
1099 }
1100 },
1101 NumericPredicate::Compare(op) => match expected {
1102 NumericLiteral::Int(expected) => {
1103 select_ordered_values(values, validity, selected_rows, op, expected, |actual| {
1104 actual
1105 })
1106 }
1107 NumericLiteral::Float(expected) => select_ordered_values(
1108 values,
1109 validity,
1110 selected_rows,
1111 op,
1112 float_total_order_key(expected),
1113 |actual| float_total_order_key(actual as f64),
1114 ),
1115 },
1116 }
1117}
1118
1119pub fn select_float64_values_view(
1120 values: &[f64],
1121 validity: ValidityView<'_>,
1122 predicate: NumericPredicate,
1123 expected: NumericLiteral,
1124 selected_rows: &mut Vec<u32>,
1125) -> Result<()> {
1126 match predicate {
1127 NumericPredicate::Eq => match expected {
1128 NumericLiteral::Int(_) => {
1129 select_numeric_values(values, validity, selected_rows, |_| false)
1130 }
1131 NumericLiteral::Float(expected) => {
1132 select_numeric_values(values, validity, selected_rows, |actual| {
1133 actual.total_cmp(&expected).is_eq()
1134 })
1135 }
1136 },
1137 NumericPredicate::Compare(op) => {
1138 let expected = match expected {
1139 NumericLiteral::Int(expected) => expected as f64,
1140 NumericLiteral::Float(expected) => expected,
1141 };
1142 select_ordered_values(
1143 values,
1144 validity,
1145 selected_rows,
1146 op,
1147 float_total_order_key(expected),
1148 float_total_order_key,
1149 )
1150 }
1151 }
1152}
1153
1154fn filter_ordered_values<T, K>(
1155 values: &[T],
1156 validity: ValidityView<'_>,
1157 input: &Selection,
1158 op: ComparisonOp,
1159 expected: K,
1160 key: impl Fn(T) -> K + Copy,
1161) -> Result<Selection>
1162where
1163 T: Copy,
1164 K: Copy + Ord,
1165{
1166 match op {
1167 ComparisonOp::Lt => {
1168 filter_numeric_values(values, validity, input, |actual| key(actual) < expected)
1169 }
1170 ComparisonOp::Lte => {
1171 filter_numeric_values(values, validity, input, |actual| key(actual) <= expected)
1172 }
1173 ComparisonOp::Gt => {
1174 filter_numeric_values(values, validity, input, |actual| key(actual) > expected)
1175 }
1176 ComparisonOp::Gte => {
1177 filter_numeric_values(values, validity, input, |actual| key(actual) >= expected)
1178 }
1179 }
1180}
1181
1182fn select_ordered_values<T, K>(
1183 values: &[T],
1184 validity: ValidityView<'_>,
1185 selected_rows: &mut Vec<u32>,
1186 op: ComparisonOp,
1187 expected: K,
1188 key: impl Fn(T) -> K + Copy,
1189) -> Result<()>
1190where
1191 T: Copy,
1192 K: Copy + Ord,
1193{
1194 match op {
1195 ComparisonOp::Lt => select_numeric_values(values, validity, selected_rows, |actual| {
1196 key(actual) < expected
1197 }),
1198 ComparisonOp::Lte => select_numeric_values(values, validity, selected_rows, |actual| {
1199 key(actual) <= expected
1200 }),
1201 ComparisonOp::Gt => select_numeric_values(values, validity, selected_rows, |actual| {
1202 key(actual) > expected
1203 }),
1204 ComparisonOp::Gte => select_numeric_values(values, validity, selected_rows, |actual| {
1205 key(actual) >= expected
1206 }),
1207 }
1208}
1209
1210fn select_numeric_values<T: Copy>(
1211 values: &[T],
1212 validity: ValidityView<'_>,
1213 selected_rows: &mut Vec<u32>,
1214 mut matches: impl FnMut(T) -> bool,
1215) -> Result<()> {
1216 if values.len() != validity.len() {
1217 return Err(HawDBError::Execution(format!(
1218 "numeric selection has {} values and {} validity entries",
1219 values.len(),
1220 validity.len()
1221 )));
1222 }
1223 if values.len() > u32::MAX as usize {
1224 return Err(HawDBError::Execution(format!(
1225 "numeric selection batch has {} rows, exceeding the u32 row index limit",
1226 values.len()
1227 )));
1228 }
1229 selected_rows.clear();
1230 if selected_rows.capacity() < values.len() {
1231 selected_rows.reserve(values.len());
1232 }
1233 match validity {
1234 ValidityView::All { .. } => {
1235 for (row, value) in values.iter().copied().enumerate() {
1236 if matches(value) {
1237 selected_rows.push(row as u32);
1238 }
1239 }
1240 }
1241 ValidityView::Bitmap { .. } => {
1242 for (row, value) in values.iter().copied().enumerate() {
1243 if validity.is_valid(row) && matches(value) {
1244 selected_rows.push(row as u32);
1245 }
1246 }
1247 }
1248 }
1249 Ok(())
1250}
1251
1252fn filter_numeric_values<T: Copy>(
1253 values: &[T],
1254 validity: ValidityView<'_>,
1255 input: &Selection,
1256 mut matches: impl FnMut(T) -> bool,
1257) -> Result<Selection> {
1258 if values.len() != validity.len() || values.len() != input.len() {
1259 return Err(HawDBError::Execution(format!(
1260 "numeric filter has {} values, {} validity entries, and {} selected input rows",
1261 values.len(),
1262 validity.len(),
1263 input.len()
1264 )));
1265 }
1266 let mut output = SelectionBuilder::new(input.len());
1267 for row in input.iter() {
1268 if validity.is_valid(row) && matches(values[row]) {
1269 output.select(row);
1270 }
1271 }
1272 Ok(output.finish())
1273}
1274
1275#[inline]
1276pub fn int64_value_matches(
1277 actual: i64,
1278 predicate: NumericPredicate,
1279 expected: NumericLiteral,
1280) -> bool {
1281 match predicate {
1282 NumericPredicate::Eq => {
1283 matches!(expected, NumericLiteral::Int(expected) if actual == expected)
1284 }
1285 NumericPredicate::Compare(op) => match expected {
1286 NumericLiteral::Int(expected) => compare_ordering(actual.cmp(&expected), op),
1287 NumericLiteral::Float(expected) => compare_ordering(
1288 float_total_order_key(actual as f64).cmp(&float_total_order_key(expected)),
1289 op,
1290 ),
1291 },
1292 }
1293}
1294
1295#[inline]
1296pub fn float64_value_matches(
1297 actual: f64,
1298 predicate: NumericPredicate,
1299 expected: NumericLiteral,
1300) -> bool {
1301 match predicate {
1302 NumericPredicate::Eq => {
1303 matches!(expected, NumericLiteral::Float(expected) if actual.total_cmp(&expected).is_eq())
1304 }
1305 NumericPredicate::Compare(op) => {
1306 let expected = match expected {
1307 NumericLiteral::Int(expected) => expected as f64,
1308 NumericLiteral::Float(expected) => expected,
1309 };
1310 compare_ordering(
1311 float_total_order_key(actual).cmp(&float_total_order_key(expected)),
1312 op,
1313 )
1314 }
1315 }
1316}
1317
1318#[inline(always)]
1319fn float_total_order_key(value: f64) -> i64 {
1320 let bits = value.to_bits() as i64;
1321 bits ^ (((bits >> 63) as u64 >> 1) as i64)
1322}
1323
1324fn compare_ordering(ordering: std::cmp::Ordering, op: ComparisonOp) -> bool {
1325 match op {
1326 ComparisonOp::Lt => ordering == std::cmp::Ordering::Less,
1327 ComparisonOp::Lte => ordering != std::cmp::Ordering::Greater,
1328 ComparisonOp::Gt => ordering == std::cmp::Ordering::Greater,
1329 ComparisonOp::Gte => ordering != std::cmp::Ordering::Less,
1330 }
1331}
1332
1333fn count_bitmap_rows(words: &[u64], len: usize) -> usize {
1334 words
1335 .iter()
1336 .enumerate()
1337 .fold(0usize, |total, (index, word)| {
1338 let bits = if index + 1 == words.len() && !len.is_multiple_of(u64::BITS as usize) {
1339 word & ((1u64 << (len % u64::BITS as usize)) - 1)
1340 } else {
1341 *word
1342 };
1343 total.saturating_add(bits.count_ones() as usize)
1344 })
1345}
1346
1347#[cfg(test)]
1348mod tests {
1349 use super::*;
1350
1351 #[test]
1352 fn validity_uses_no_bitmap_for_dense_columns() {
1353 let mut builder = ValidityBuilder::with_capacity(3);
1354 builder.push(true);
1355 builder.push(true);
1356 builder.push(true);
1357
1358 assert!(builder.words.is_empty());
1359 assert_eq!(builder.finish(), Validity::All { len: 3 });
1360 }
1361
1362 #[test]
1363 fn validity_materializes_prior_rows_on_first_null() {
1364 let mut builder = ValidityBuilder::with_capacity(130);
1365 for _ in 0..65 {
1366 builder.push(true);
1367 }
1368 builder.push(false);
1369 builder.push(true);
1370
1371 let validity = builder.finish();
1372 assert_eq!(validity.valid_count(), 66);
1373 assert!(validity.is_valid(64));
1374 assert!(!validity.is_valid(65));
1375 assert!(validity.is_valid(66));
1376 }
1377
1378 #[test]
1379 fn value_ref_distinguishes_null_from_out_of_bounds() {
1380 let column = ColumnVector::int64(
1381 vec![1, 0],
1382 Validity::Bitmap {
1383 len: 2,
1384 words: Arc::from([0b01]),
1385 },
1386 )
1387 .unwrap();
1388
1389 assert_eq!(column.value_ref(0), Some(ValueRef::Int(1)));
1390 assert_eq!(column.value_ref(1), Some(ValueRef::Null));
1391 assert_eq!(column.value_ref(2), None);
1392 }
1393
1394 #[test]
1395 fn boolean_bytes_reject_non_canonical_values() {
1396 let error = ColumnVector::boolean_bytes(vec![0, 2], Validity::all(2)).unwrap_err();
1397
1398 assert!(error.to_string().contains("other than 0 or 1"));
1399 }
1400
1401 #[test]
1402 fn validity_builder_reuses_bitmap_storage_between_batches() {
1403 let mut builder = ValidityBuilder::with_capacity(130);
1404 builder.push(false);
1405 let words_ptr = builder.words.as_ptr();
1406 let words_capacity = builder.words.capacity();
1407
1408 builder.clear();
1409 builder.push(true);
1410 builder.push(false);
1411
1412 assert_eq!(builder.words.as_ptr(), words_ptr);
1413 assert_eq!(builder.words.capacity(), words_capacity);
1414 let validity = builder.view();
1415 assert!(!validity.is_empty());
1416 assert!(validity.is_valid(0));
1417 assert!(!validity.is_valid(1));
1418 }
1419
1420 #[test]
1421 fn numeric_filter_preserves_null_and_nan_semantics() {
1422 let mut validity = ValidityBuilder::with_capacity(4);
1423 validity.push(true);
1424 validity.push(false);
1425 validity.push(true);
1426 validity.push(true);
1427 let column =
1428 ColumnVector::float64(vec![1.0, 0.0, f64::NAN, 3.0], validity.finish()).unwrap();
1429
1430 let selection = filter_numeric_column(
1431 &column,
1432 &Selection::all(4),
1433 NumericPredicate::Compare(ComparisonOp::Gte),
1434 NumericLiteral::Float(2.0),
1435 )
1436 .unwrap();
1437
1438 assert_eq!(selection.iter().collect::<Vec<_>>(), vec![2, 3]);
1439 }
1440
1441 #[test]
1442 fn numeric_equality_filter_preserves_type_validity_and_float_ordering_semantics() {
1443 let mut validity = ValidityBuilder::with_capacity(5);
1444 validity.push(true);
1445 validity.push(false);
1446 validity.push(true);
1447 validity.push(true);
1448 validity.push(true);
1449 let nan = f64::from_bits(0x7ff8_0000_0000_0042);
1450 let column =
1451 ColumnVector::float64(vec![5.0, 5.0, nan, -0.0, 0.0], validity.finish()).unwrap();
1452
1453 let nan_selection = filter_numeric_column(
1454 &column,
1455 &Selection::all(5),
1456 NumericPredicate::Eq,
1457 NumericLiteral::Float(nan),
1458 )
1459 .unwrap();
1460 let negative_zero_selection = filter_numeric_column(
1461 &column,
1462 &Selection::all(5),
1463 NumericPredicate::Eq,
1464 NumericLiteral::Float(-0.0),
1465 )
1466 .unwrap();
1467 let cross_type_selection = filter_numeric_column(
1468 &column,
1469 &Selection::all(5),
1470 NumericPredicate::Eq,
1471 NumericLiteral::Int(5),
1472 )
1473 .unwrap();
1474
1475 assert_eq!(nan_selection.iter().collect::<Vec<_>>(), vec![2]);
1476 assert_eq!(negative_zero_selection.iter().collect::<Vec<_>>(), vec![3]);
1477 assert!(cross_type_selection.is_empty());
1478 }
1479
1480 #[test]
1481 fn numeric_equality_selection_reuses_row_index_storage() {
1482 let mut selected_rows = Vec::new();
1483 select_int64_values_view(
1484 &[1, 2, 1, 3],
1485 ValidityView::Bitmap {
1486 len: 4,
1487 words: &[0b1101],
1488 },
1489 NumericPredicate::Eq,
1490 NumericLiteral::Int(1),
1491 &mut selected_rows,
1492 )
1493 .unwrap();
1494 assert_eq!(selected_rows, [0, 2]);
1495 let rows_ptr = selected_rows.as_ptr();
1496
1497 select_int64_values_view(
1498 &[4, 5, 4, 6],
1499 ValidityView::All { len: 4 },
1500 NumericPredicate::Eq,
1501 NumericLiteral::Int(4),
1502 &mut selected_rows,
1503 )
1504 .unwrap();
1505
1506 assert_eq!(selected_rows, [0, 2]);
1507 assert_eq!(selected_rows.as_ptr(), rows_ptr);
1508 }
1509
1510 #[test]
1511 fn specialized_numeric_kernels_match_total_order_reference() {
1512 let int_values = [i64::MIN, -1, 0, 1, i64::MAX];
1513 let float_values = [
1514 f64::from_bits(0xfff8_0000_0000_0001),
1515 f64::NEG_INFINITY,
1516 -0.0,
1517 0.0,
1518 1.0,
1519 f64::INFINITY,
1520 f64::from_bits(0x7ff8_0000_0000_0001),
1521 ];
1522 let int_validity = Validity::Bitmap {
1523 len: int_values.len(),
1524 words: Arc::from([0b1_1101]),
1525 };
1526 let float_validity = Validity::Bitmap {
1527 len: float_values.len(),
1528 words: Arc::from([0b110_1111]),
1529 };
1530 let int_expected = [
1531 NumericLiteral::Int(0),
1532 NumericLiteral::Float(-0.0),
1533 NumericLiteral::Float(f64::INFINITY),
1534 NumericLiteral::Float(f64::from_bits(0x7ff8_0000_0000_0001)),
1535 ];
1536 let float_expected = [
1537 NumericLiteral::Int(0),
1538 NumericLiteral::Float(-0.0),
1539 NumericLiteral::Float(0.0),
1540 NumericLiteral::Float(f64::from_bits(0xfff8_0000_0000_0001)),
1541 NumericLiteral::Float(f64::from_bits(0x7ff8_0000_0000_0001)),
1542 ];
1543
1544 assert_specialized_numeric_kernel(
1545 &int_values,
1546 &int_validity,
1547 &int_expected,
1548 filter_int64_values,
1549 select_int64_values_view,
1550 int64_value_matches,
1551 reference_int64_matches,
1552 );
1553 assert_specialized_numeric_kernel(
1554 &float_values,
1555 &float_validity,
1556 &float_expected,
1557 filter_float64_values,
1558 select_float64_values_view,
1559 float64_value_matches,
1560 reference_float64_matches,
1561 );
1562 }
1563
1564 fn assert_specialized_numeric_kernel<T: Copy>(
1565 values: &[T],
1566 validity: &Validity,
1567 expected_literals: &[NumericLiteral],
1568 filter: impl Fn(
1569 &[T],
1570 &Validity,
1571 &Selection,
1572 NumericPredicate,
1573 NumericLiteral,
1574 ) -> Result<Selection>,
1575 select: impl for<'a> Fn(
1576 &[T],
1577 ValidityView<'a>,
1578 NumericPredicate,
1579 NumericLiteral,
1580 &mut Vec<u32>,
1581 ) -> Result<()>,
1582 point_matches: impl Fn(T, NumericPredicate, NumericLiteral) -> bool,
1583 reference_matches: impl Fn(T, ComparisonOp, NumericLiteral) -> bool,
1584 ) {
1585 for op in [
1586 ComparisonOp::Lt,
1587 ComparisonOp::Lte,
1588 ComparisonOp::Gt,
1589 ComparisonOp::Gte,
1590 ] {
1591 for &expected in expected_literals {
1592 let predicate = NumericPredicate::Compare(op);
1593 let reference = values
1594 .iter()
1595 .copied()
1596 .enumerate()
1597 .filter_map(|(row, actual)| {
1598 assert_eq!(
1599 point_matches(actual, predicate, expected),
1600 reference_matches(actual, op, expected)
1601 );
1602 (validity.is_valid(row) && reference_matches(actual, op, expected))
1603 .then_some(row)
1604 })
1605 .collect::<Vec<_>>();
1606 let filtered = filter(
1607 values,
1608 validity,
1609 &Selection::all(values.len()),
1610 predicate,
1611 expected,
1612 )
1613 .unwrap();
1614 let mut selected = Vec::new();
1615 select(values, validity.view(), predicate, expected, &mut selected).unwrap();
1616
1617 assert_eq!(filtered.iter().collect::<Vec<_>>(), reference);
1618 assert_eq!(
1619 selected,
1620 reference.iter().map(|row| *row as u32).collect::<Vec<_>>()
1621 );
1622 }
1623 }
1624 }
1625
1626 fn reference_int64_matches(actual: i64, op: ComparisonOp, expected: NumericLiteral) -> bool {
1627 let ordering = match expected {
1628 NumericLiteral::Int(expected) => actual.cmp(&expected),
1629 NumericLiteral::Float(expected) => (actual as f64).total_cmp(&expected),
1630 };
1631 compare_ordering(ordering, op)
1632 }
1633
1634 fn reference_float64_matches(actual: f64, op: ComparisonOp, expected: NumericLiteral) -> bool {
1635 let expected = match expected {
1636 NumericLiteral::Int(expected) => expected as f64,
1637 NumericLiteral::Float(expected) => expected,
1638 };
1639 compare_ordering(actual.total_cmp(&expected), op)
1640 }
1641
1642 #[test]
1643 fn numeric_selection_reuses_row_index_storage() {
1644 let mut selected_rows = Vec::new();
1645 select_int64_values_view(
1646 &[1, 2, 3, 4],
1647 ValidityView::All { len: 4 },
1648 NumericPredicate::Compare(ComparisonOp::Gte),
1649 NumericLiteral::Int(3),
1650 &mut selected_rows,
1651 )
1652 .unwrap();
1653 assert_eq!(selected_rows, [2, 3]);
1654 let rows_ptr = selected_rows.as_ptr();
1655
1656 select_int64_values_view(
1657 &[5, 6, 7, 8],
1658 ValidityView::Bitmap {
1659 len: 4,
1660 words: &[0b1101],
1661 },
1662 NumericPredicate::Compare(ComparisonOp::Lt),
1663 NumericLiteral::Int(8),
1664 &mut selected_rows,
1665 )
1666 .unwrap();
1667
1668 assert_eq!(selected_rows, [0, 2]);
1669 assert_eq!(selected_rows.as_ptr(), rows_ptr);
1670 }
1671
1672 #[test]
1673 fn sparse_filter_uses_index_selection() {
1674 let column = ColumnVector::int64((0..64).collect(), Validity::all(64)).unwrap();
1675 let selection = filter_numeric_column(
1676 &column,
1677 &Selection::all(64),
1678 NumericPredicate::Compare(ComparisonOp::Gte),
1679 NumericLiteral::Int(63),
1680 )
1681 .unwrap();
1682
1683 assert!(matches!(selection, Selection::Indices { .. }));
1684 assert_eq!(selection.iter().collect::<Vec<_>>(), vec![63]);
1685 }
1686
1687 #[test]
1688 fn dense_filter_uses_bitmap_selection() {
1689 let column = ColumnVector::int64((0..64).collect(), Validity::all(64)).unwrap();
1690 let selection = filter_numeric_column(
1691 &column,
1692 &Selection::all(64),
1693 NumericPredicate::Compare(ComparisonOp::Gte),
1694 NumericLiteral::Int(32),
1695 )
1696 .unwrap();
1697
1698 assert!(matches!(selection, Selection::Bitmap { .. }));
1699 assert_eq!(selection.selected_count(), 32);
1700 }
1701
1702 #[test]
1703 fn projection_reuses_column_storage() {
1704 let schema = Arc::new(
1705 BindingSchema::try_new(vec![SlotDescriptor {
1706 id: SlotId(0),
1707 name: "score".to_string(),
1708 slot_type: SlotType::logical(LogicalType::Int64),
1709 }])
1710 .unwrap(),
1711 );
1712 let column = Arc::new(ColumnVector::int64(vec![1, 2], Validity::all(2)).unwrap());
1713 let batch = ColumnarBatch::try_new(Arc::clone(&schema), vec![Arc::clone(&column)]).unwrap();
1714 let projected = batch.project(&[SlotId(0)]).unwrap();
1715
1716 assert!(Arc::ptr_eq(projected.column(SlotId(0)).unwrap(), &column));
1717 }
1718
1719 #[test]
1720 fn utf8_value_ref_borrows_column_storage_until_materialization() {
1721 let schema = Arc::new(
1722 BindingSchema::try_new(vec![SlotDescriptor {
1723 id: SlotId(0),
1724 name: "content".to_string(),
1725 slot_type: SlotType::logical(LogicalType::Text),
1726 }])
1727 .unwrap(),
1728 );
1729 let values: Arc<[String]> = vec!["borrowed content".to_string()].into();
1730 let source_ptr = values[0].as_ptr();
1731 let column = Arc::new(ColumnVector::Utf8 {
1732 values,
1733 validity: Validity::all(1),
1734 });
1735 let batch = ColumnarBatch::try_new(schema, vec![column]).unwrap();
1736
1737 let value = batch.value_ref(SlotId(0), 0).unwrap();
1738 assert_eq!(value.logical_type(), Some(LogicalType::String));
1739 assert_eq!(value.as_str().unwrap().as_ptr(), source_ptr);
1740 assert_eq!(
1741 value.to_owned_value(),
1742 Value::String("borrowed content".into())
1743 );
1744 }
1745
1746 #[test]
1747 fn logical_and_physical_slot_types_are_validated_separately() {
1748 assert!(SlotType::logical(LogicalType::Any).accepts(ColumnType::Bool));
1749 assert!(SlotType::logical(LogicalType::Any).accepts(ColumnType::Utf8));
1750 assert!(!SlotType::logical(LogicalType::Any).accepts(ColumnType::NodeId));
1751 assert!(SlotType::logical(LogicalType::String).accepts(ColumnType::Utf8));
1752 assert!(SlotType::logical(LogicalType::Text).accepts(ColumnType::Utf8));
1753 assert!(SlotType::NodeId.accepts(ColumnType::NodeId));
1754 assert!(!SlotType::NodeId.accepts(ColumnType::Int64));
1755 assert_eq!(SlotType::NodeId.logical_type(), Some(LogicalType::Int64));
1756 assert_eq!(SlotType::RelationalRowLocator.logical_type(), None);
1757 }
1758
1759 #[test]
1760 fn relational_locator_projection_preserves_compact_identity_storage() {
1761 let schema = Arc::new(
1762 BindingSchema::try_new(vec![SlotDescriptor {
1763 id: SlotId(0),
1764 name: "row_locator".to_string(),
1765 slot_type: SlotType::RelationalRowLocator,
1766 }])
1767 .unwrap(),
1768 );
1769 let locator = RelationalRowLocator::new(
1770 3,
1771 RelationalKey(vec![
1772 RelationalValue::BigInt(7),
1773 RelationalValue::Text("thread-1".to_string()),
1774 ]),
1775 );
1776 let nested_bytes = locator.allocated_bytes();
1777 let column = Arc::new(ColumnVector::relational_row_locators(vec![locator]));
1778 let batch = ColumnarBatch::try_new(schema, vec![Arc::clone(&column)]).unwrap();
1779 let projected = batch.project(&[SlotId(0)]).unwrap();
1780
1781 assert!(Arc::ptr_eq(projected.column(SlotId(0)).unwrap(), &column));
1782 let actual = projected
1783 .column(SlotId(0))
1784 .unwrap()
1785 .relational_row_locator(0)
1786 .unwrap();
1787 assert_eq!(actual.table_id(), 3);
1788 assert_eq!(actual.primary_key().0.len(), 2);
1789 assert_eq!(column.value(0), None);
1790 assert_eq!(
1791 column.estimated_memory_bytes(),
1792 std::mem::size_of::<RelationalRowLocator>() + nested_bytes
1793 );
1794 }
1795
1796 #[test]
1797 fn boolean_filter_limit_and_projection_share_storage() {
1798 let schema = Arc::new(
1799 BindingSchema::try_new(vec![
1800 SlotDescriptor {
1801 id: SlotId(0),
1802 name: "visible".to_string(),
1803 slot_type: SlotType::logical(LogicalType::Boolean),
1804 },
1805 SlotDescriptor {
1806 id: SlotId(1),
1807 name: "id".to_string(),
1808 slot_type: SlotType::NodeId,
1809 },
1810 ])
1811 .unwrap(),
1812 );
1813 let visible = Arc::new(
1814 ColumnVector::boolean(
1815 vec![true, false, true, true, false],
1816 Validity::Bitmap {
1817 len: 5,
1818 words: Arc::from([0b1_1011]),
1819 },
1820 )
1821 .unwrap(),
1822 );
1823 let ids = Arc::new(ColumnVector::node_ids(vec![10, 11, 12, 13, 14]));
1824 let batch = ColumnarBatch::try_new(schema, vec![Arc::clone(&visible), Arc::clone(&ids)])
1825 .unwrap()
1826 .filter_boolean(SlotId(0), true)
1827 .unwrap()
1828 .limit(1, 1)
1829 .project(&[SlotId(1)])
1830 .unwrap();
1831
1832 assert_eq!(batch.selection().iter().collect::<Vec<_>>(), vec![3]);
1833 assert!(Arc::ptr_eq(batch.column(SlotId(0)).unwrap(), &ids));
1834 }
1835
1836 #[test]
1837 fn count_and_sum_follow_selection_and_validity() {
1838 let schema = Arc::new(
1839 BindingSchema::try_new(vec![SlotDescriptor {
1840 id: SlotId(0),
1841 name: "token_count".to_string(),
1842 slot_type: SlotType::logical(LogicalType::Int64),
1843 }])
1844 .unwrap(),
1845 );
1846 let values = Arc::new(
1847 ColumnVector::int64(
1848 vec![4, 8, 16, 32],
1849 Validity::Bitmap {
1850 len: 4,
1851 words: Arc::from([0b1101]),
1852 },
1853 )
1854 .unwrap(),
1855 );
1856 let batch = ColumnarBatch::try_new(schema, vec![values])
1857 .unwrap()
1858 .filter_numeric(
1859 SlotId(0),
1860 NumericPredicate::Compare(ComparisonOp::Gte),
1861 NumericLiteral::Int(4),
1862 )
1863 .unwrap()
1864 .limit(1, 2);
1865
1866 assert_eq!(batch.count_selected(), 2);
1867 assert_eq!(batch.count_valid(SlotId(0)).unwrap(), 2);
1868 assert_eq!(batch.sum_int64(SlotId(0)).unwrap(), Some(48));
1869 }
1870
1871 #[test]
1872 fn int64_sum_fails_closed_on_overflow() {
1873 let schema = Arc::new(
1874 BindingSchema::try_new(vec![SlotDescriptor {
1875 id: SlotId(0),
1876 name: "value".to_string(),
1877 slot_type: SlotType::logical(LogicalType::Int64),
1878 }])
1879 .unwrap(),
1880 );
1881 let values = Arc::new(ColumnVector::int64(vec![i64::MAX, 1], Validity::all(2)).unwrap());
1882 let batch = ColumnarBatch::try_new(schema, vec![values]).unwrap();
1883
1884 let error = batch.sum_int64(SlotId(0)).unwrap_err();
1885 assert!(error.to_string().contains("SUM overflow"));
1886 }
1887}