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