From 9c979856c5a5e43d0684588a70995f72b936cb32 Mon Sep 17 00:00:00 2001
From: Clemens Winter <clemenswinter1@gmail.com>
Date: Sat, 27 Jul 2024 08:59:00 -0700
Subject: [PATCH] Add missing compact operator

---
 src/engine/operators/compact_with_nullable.rs | 43 +++++++++++++++++++
 src/engine/operators/mod.rs                   |  1 +
 src/engine/operators/vector_operator.rs       |  7 +++
 tests/query_tests.rs                          |  4 +-
 4 files changed, 53 insertions(+), 2 deletions(-)
 create mode 100644 src/engine/operators/compact_with_nullable.rs

diff --git a/src/engine/operators/compact_with_nullable.rs b/src/engine/operators/compact_with_nullable.rs
new file mode 100644
index 00000000..bbc8ca08
--- /dev/null
+++ b/src/engine/operators/compact_with_nullable.rs
@@ -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)
+    }
+}
+
diff --git a/src/engine/operators/mod.rs b/src/engine/operators/mod.rs
index 46c7e13d..73a36457 100644
--- a/src/engine/operators/mod.rs
+++ b/src/engine/operators/mod.rs
@@ -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;
diff --git a/src/engine/operators/vector_operator.rs b/src/engine/operators/vector_operator.rs
index 1b36e1fc..ecccd38e 100644
--- a/src/engine/operators/vector_operator.rs
+++ b/src/engine/operators/vector_operator.rs
@@ -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::*;
@@ -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";
diff --git a/tests/query_tests.rs b/tests/query_tests.rs
index 88543e22..a6e3f40b 100644
--- a/tests/query_tests.rs
+++ b/tests/query_tests.rs
@@ -1389,7 +1389,7 @@ 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,
@@ -1397,7 +1397,7 @@ fn test_long_nullable() {
             "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();