Skip to content

Commit

Permalink
Add missing compact operator
Browse files Browse the repository at this point in the history
  • Loading branch information
cswinter committed Jul 27, 2024
1 parent e9045c7 commit 9c97985
Show file tree
Hide file tree
Showing 4 changed files with 53 additions and 2 deletions.
43 changes: 43 additions & 0 deletions src/engine/operators/compact_with_nullable.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
use crate::engine::*;
use crate::bitvec::*;

#[derive(Debug)]
pub struct CompactWithNullable<T, U> {
pub data: BufferRef<T>,
pub select: BufferRef<Nullable<U>>,
pub compacted: BufferRef<T>,
}

impl<'a, T: VecData<T> + 'a, U: GenericIntVec<U>> VecOperator<'a> for CompactWithNullable<T, U> {
fn execute(&mut self, _: bool, scratchpad: &mut Scratchpad<'a>) -> Result<(), QueryError> {
let mut data = scratchpad.get_mut(self.data);
let (select, select_present) = scratchpad.get_nullable(self.select);
// Remove all unmodified entries
let mut j = 0;
for (i, &s) in select.iter().take(data.len()).enumerate() {
if s > U::zero() && (&*select_present).is_set(i) {
data[j] = data[i];
j += 1;
}
}
data.truncate(j);
Ok(())
}

fn init(&mut self, _: usize, _: usize, scratchpad: &mut Scratchpad<'a>) {
scratchpad.alias(self.data, self.compacted);
}

fn inputs(&self) -> Vec<BufferRef<Any>> { vec![self.data.any(), self.select.any()] }
fn inputs_mut(&mut self) -> Vec<&mut usize> { vec![&mut self.data.i, &mut self.select.i] }
fn outputs(&self) -> Vec<BufferRef<Any>> { vec![self.compacted.any()] }
fn can_stream_input(&self, _: usize) -> bool { false }
fn can_stream_output(&self, _: usize) -> bool { false }
fn mutates(&self, i: usize) -> bool { i == self.data.i }
fn allocates(&self) -> bool { false }

fn display_op(&self, _: bool) -> String {
format!("{}[{} > 0]", self.data, self.select)
}
}

1 change: 1 addition & 0 deletions src/engine/operators/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ mod column_ops;
mod combine_null_maps;
mod compact_nullable_nullable;
mod compact_nullable;
mod compact_with_nullable;
mod compact;
mod comparison_operators;
mod constant;
Expand Down
7 changes: 7 additions & 0 deletions src/engine/operators/vector_operator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ use super::collect::Collect;
use super::column_ops::*;
use super::combine_null_maps::CombineNullMaps;
use super::compact::Compact;
use super::compact_with_nullable::CompactWithNullable;
use super::compact_nullable::CompactNullable;
use super::compact_nullable_nullable::CompactNullableNullable;
use super::comparison_operators::*;
Expand Down Expand Up @@ -1570,6 +1571,12 @@ pub mod operator {
data, compacted: NullablePrimitive, select: Integer;
Ok(Box::new(CompactNullable { data, select, compacted }))
}
} else if select.is_nullable() {
reify_types! {
"compact_with_nullable";
data, compacted: Primitive, select: NullableInteger;
Ok(Box::new(CompactWithNullable { data, select, compacted }))
}
} else {
reify_types! {
"compact";
Expand Down
4 changes: 2 additions & 2 deletions tests/query_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1389,15 +1389,15 @@ fn test_column_with_null_partitions() {
fn test_long_nullable() {
let _ = env_logger::try_init();
let locustdb = LocustDB::memory_only();
let _ = block_on(locustdb.gen_table(locustdb::colgen::GenTable {
block_on(locustdb.gen_table(locustdb::colgen::GenTable {
name: "test".to_string(),
partitions: 8,
partition_size: 2 << 14,
columns: vec![(
"nullable_int".to_string(),
locustdb::colgen::nullable_ints(vec![None, Some(1), Some(-10)], vec![0.9, 0.05, 0.05]),
)],
}));
})).unwrap();
let query = "SELECT nullable_int FROM test LIMIT 0;";
let expected_rows: Vec<[Value; 1]> = vec![];
let result = block_on(locustdb.run_query(query, true, true, vec![])).unwrap();
Expand Down

0 comments on commit 9c97985

Please sign in to comment.