[ENH] Add support for delete in blockfile

rust/worker/src/blockstore/arrow_blockfile/block/delta.rs

@@ -2,6 +2,7 @@ use super::{Block, BlockBuilderOptions, BlockData, BlockDataBuilder};
 use crate::blockstore::{
     arrow_blockfile::{blockfile::MAX_BLOCK_SIZE, provider::ArrowBlockProvider},
     types::{BlockfileKey, KeyType, Value, ValueType},
+    BlockfileError,
 };
 use arrow::util::bit_util;
 use parking_lot::RwLock;
@@ -41,9 +42,9 @@ impl BlockDelta {
     }
 
     /// Deletes a key from the block delta.
-    pub fn delete(&self, key: BlockfileKey) {
+    pub fn delete(&self, key: BlockfileKey) -> Result<(), Box<BlockfileError>> {
         let mut inner = self.inner.write();
-        inner.delete(key);
+        inner.delete(key)
     }
 
     /// Gets the minimum key in the block delta.
@@ -131,10 +132,12 @@ impl BlockDeltaInner {
         self.new_data.insert(key, value);
     }
 
-    fn delete(&mut self, key: BlockfileKey) {
+    fn delete(&mut self, key: BlockfileKey) -> Result<(), Box<BlockfileError>> {
         if self.new_data.contains_key(&key) {
             self.new_data.remove(&key);
+            return Ok(());
         }
+        Err(Box::new(BlockfileError::NotFoundError))
     }
 
     fn get_block_size(

rust/worker/src/blockstore/arrow_blockfile/block/iterator.rs

@@ -8,15 +8,18 @@ use arrow::array::{Array, BooleanArray, Int32Array, ListArray, StringArray, UInt
 pub(super) struct BlockIterator {
     block: Block,
     index: usize,
+    length: usize,
     key_type: KeyType,
     value_type: ValueType,
 }
 
 impl BlockIterator {
     pub fn new(block: Block, key_type: KeyType, value_type: ValueType) -> Self {
+        let len = block.len();
         Self {
             block,
             index: 0,
+            length: len,
             key_type,
             value_type,
         }
@@ -31,6 +34,9 @@ impl Iterator for BlockIterator {
         if data.is_none() {
             return None;
         }
+        if self.index >= self.length {
+            return None;
+        }
 
         // Arrow requires us to downcast the array to the specific type we want to work with.
         // This is a bit awkward, but it's the way Arrow works to allow for dynamic typing.

rust/worker/src/blockstore/arrow_blockfile/blockfile.rs

@@ -5,6 +5,7 @@ use super::sparse_index::SparseIndex;
 use crate::blockstore::arrow_blockfile::block::delta::BlockDelta;
 use crate::blockstore::BlockfileError;
 use crate::errors::ChromaError;
+use crate::errors::IntoResult;
 use parking_lot::Mutex;
 use std::sync::Arc;
 use thiserror::Error;
@@ -94,6 +95,48 @@ impl Blockfile for ArrowBlockfile {
         }
     }
 
+    fn delete(&mut self, key: BlockfileKey) -> Result<(), Box<dyn ChromaError>> {
+        if !self.in_transaction() {
+            return Err(Box::new(BlockfileError::TransactionNotInProgress));
+        }
+
+        self.validate_key(&key)?;
+
+        let transaction_state = match &self.transaction_state {
+            None => return Err(Box::new(BlockfileError::TransactionNotInProgress)),
+            Some(transaction_state) => transaction_state,
+        };
+
+        // Note: The code to get the target block as well as get or create the delta is duplicated
+        // in the set and delete methods. This is due to the lock management making it clunky
+        // to abstract this into a separate method. This likely should be refactored
+        // for cleaner reuse, but for now its ok to repeat ourselves.
+
+        // Get the target block id for the key
+        let transaction_sparse_index = transaction_state.sparse_index.lock();
+        let target_block_id = match *transaction_sparse_index {
+            None => self.sparse_index.lock().get_target_block_id(&key),
+            Some(ref index) => index.lock().get_target_block_id(&key),
+        };
+
+        // See if a delta for the target block already exists, if not create a new one and add it to the transaction state
+        // Creating a delta loads the block entirely into memory
+        let delta = match transaction_state.get_delta_for_block(&target_block_id) {
+            None => {
+                let target_block = match self.block_provider.get_block(&target_block_id) {
+                    None => return Err(Box::new(ArrowBlockfileError::BlockNotFoundError)),
+                    Some(block) => block,
+                };
+                let delta = BlockDelta::from(target_block);
+                transaction_state.add_delta(delta.clone());
+                delta
+            }
+            Some(delta) => delta,
+        };
+
+        delta.delete(key).into_result()
+    }
+
     fn get_by_prefix(
         &self,
         prefix: String,
@@ -144,29 +187,7 @@ impl Blockfile for ArrowBlockfile {
             return Err(Box::new(BlockfileError::TransactionNotInProgress));
         }
 
-        // Validate key type
-        match key.key {
-            Key::String(_) => {
-                if self.key_type != KeyType::String {
-                    return Err(Box::new(BlockfileError::InvalidKeyType));
-                }
-            }
-            Key::Float(_) => {
-                if self.key_type != KeyType::Float {
-                    return Err(Box::new(BlockfileError::InvalidKeyType));
-                }
-            }
-            Key::Bool(_) => {
-                if self.key_type != KeyType::Bool {
-                    return Err(Box::new(BlockfileError::InvalidKeyType));
-                }
-            }
-            Key::Uint(_) => {
-                if self.key_type != KeyType::Uint {
-                    return Err(Box::new(BlockfileError::InvalidKeyType));
-                }
-            }
-        }
+        self.validate_key(&key)?;
 
         // Validate value type
         match value {
@@ -388,6 +409,32 @@ impl ArrowBlockfile {
     fn in_transaction(&self) -> bool {
         self.transaction_state.is_some()
     }
+
+    fn validate_key(&self, key: &BlockfileKey) -> Result<(), Box<dyn ChromaError>> {
+        match key.key {
+            Key::String(_) => {
+                if self.key_type != KeyType::String {
+                    return Err(Box::new(BlockfileError::InvalidKeyType));
+                }
+            }
+            Key::Float(_) => {
+                if self.key_type != KeyType::Float {
+                    return Err(Box::new(BlockfileError::InvalidKeyType));
+                }
+            }
+            Key::Bool(_) => {
+                if self.key_type != KeyType::Bool {
+                    return Err(Box::new(BlockfileError::InvalidKeyType));
+                }
+            }
+            Key::Uint(_) => {
+                if self.key_type != KeyType::Uint {
+                    return Err(Box::new(BlockfileError::InvalidKeyType));
+                }
+            }
+        }
+        Ok(())
+    }
 }
 
 #[cfg(test)]
@@ -605,4 +652,48 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn test_delete() {
+        let block_provider = ArrowBlockProvider::new();
+        let mut blockfile =
+            ArrowBlockfile::new(KeyType::String, ValueType::Int32Array, block_provider);
+
+        blockfile.begin_transaction().unwrap();
+        let key1 = BlockfileKey::new("key".to_string(), Key::String("zzzz".to_string()));
+        blockfile
+            .set(
+                key1.clone(),
+                Value::Int32ArrayValue(Int32Array::from(vec![1, 2, 3])),
+            )
+            .unwrap();
+        let key2 = BlockfileKey::new("key".to_string(), Key::String("aaaa".to_string()));
+        blockfile
+            .set(
+                key2.clone(),
+                Value::Int32ArrayValue(Int32Array::from(vec![4, 5, 6])),
+            )
+            .unwrap();
+        blockfile.commit_transaction().unwrap();
+
+        blockfile.begin_transaction().unwrap();
+        blockfile.delete(key1.clone()).unwrap();
+        blockfile.commit_transaction().unwrap();
+
+        let res = blockfile.get(key1);
+        match res {
+            Err(err) => {
+                assert_eq!(err.code(), crate::errors::ErrorCodes::NotFound);
+            }
+            _ => panic!("Expected not found error"),
+        }
+
+        let res = blockfile.get(key2);
+        match res {
+            Ok(_) => {}
+            Err(err) => {
+                panic!("Expected key2 to be found, got error: {:?}", err);
+            }
+        }
+    }
 }

rust/worker/src/blockstore/types.rs

@@ -26,9 +26,10 @@ pub(crate) enum BlockfileError {
 impl ChromaError for BlockfileError {
     fn code(&self) -> ErrorCodes {
         match self {
-            BlockfileError::NotFoundError
-            | BlockfileError::InvalidKeyType
-            | BlockfileError::InvalidValueType => ErrorCodes::InvalidArgument,
+            BlockfileError::NotFoundError => ErrorCodes::NotFound,
+            BlockfileError::InvalidKeyType | BlockfileError::InvalidValueType => {
+                ErrorCodes::InvalidArgument
+            }
             BlockfileError::TransactionInProgress | BlockfileError::TransactionNotInProgress => {
                 ErrorCodes::FailedPrecondition
             }
@@ -266,6 +267,7 @@ pub(crate) trait Blockfile: BlockfileClone {
     ) -> Result<Vec<(BlockfileKey, Value)>, Box<dyn ChromaError>>;
 
     fn set(&mut self, key: BlockfileKey, value: Value) -> Result<(), Box<dyn ChromaError>>;
+    fn delete(&mut self, key: BlockfileKey) -> Result<(), Box<dyn ChromaError>>;
 
     fn get_gt(
         &self,
@@ -350,6 +352,13 @@ impl Blockfile for HashMapBlockfile {
         Ok(())
     }
 
+    fn delete(&mut self, key: BlockfileKey) -> Result<(), Box<dyn ChromaError>> {
+        match self.map.write().remove(&key) {
+            Some(_) => Ok(()),
+            None => Err(Box::new(BlockfileError::NotFoundError)),
+        }
+    }
+
     fn get_gt(
         &self,
         prefix: String,

rust/worker/src/errors.rs

@@ -2,6 +2,7 @@
 // gRPC spec. https://grpc.github.io/grpc/core/md_doc_statuscodes.html
 // Custom errors can use these codes in order to allow for generic handling
 use std::error::Error;
+use std::result::Result;
 
 #[derive(PartialEq, Debug)]
 pub(crate) enum ErrorCodes {
@@ -44,3 +45,19 @@ pub(crate) enum ErrorCodes {
 pub(crate) trait ChromaError: Error + Send {
     fn code(&self) -> ErrorCodes;
 }
+
+pub(crate) trait IntoResult<T> {
+    fn into_result(self) -> Result<T, Box<dyn ChromaError>>;
+}
+
+/// Converts a Result<R, E> into a Result<R, Box<dyn ChromaError>>.
+impl<R, E: ChromaError + 'static> IntoResult<R> for Result<R, Box<E>> {
+    fn into_result(self) -> Result<R, Box<dyn ChromaError>> {
+        if self.is_err() {
+            let as_chroma_error: Box<dyn ChromaError> = self.err().unwrap();
+            return Err(as_chroma_error);
+        } else {
+            return Ok(self.unwrap());
+        }
+    }
+}