From cbde1eda47809499b97556723c830c7f2bcb1b12 Mon Sep 17 00:00:00 2001
From: hammadb <hammad@trychroma.com>
Date: Tue, 5 Mar 2024 17:31:54 -0800
Subject: [PATCH] wip

---
 .../src/blockstore/arrow_blockfile/mod.rs     |  1 +
 .../blockstore/arrow_blockfile/provider.rs    | 84 +++++++++++++++++++
 rust/worker/src/blockstore/mod.rs             |  1 +
 3 files changed, 86 insertions(+)
 create mode 100644 rust/worker/src/blockstore/arrow_blockfile/mod.rs
 create mode 100644 rust/worker/src/blockstore/arrow_blockfile/provider.rs

diff --git a/rust/worker/src/blockstore/arrow_blockfile/mod.rs b/rust/worker/src/blockstore/arrow_blockfile/mod.rs
new file mode 100644
index 000000000000..be23ddf3d8d0
--- /dev/null
+++ b/rust/worker/src/blockstore/arrow_blockfile/mod.rs
@@ -0,0 +1 @@
+mod provider;
diff --git a/rust/worker/src/blockstore/arrow_blockfile/provider.rs b/rust/worker/src/blockstore/arrow_blockfile/provider.rs
new file mode 100644
index 000000000000..15d1b5eb1a2c
--- /dev/null
+++ b/rust/worker/src/blockstore/arrow_blockfile/provider.rs
@@ -0,0 +1,84 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use parking_lot::RwLock;
+use uuid::Uuid;
+
+use super::super::provider::BlockfileProvider;
+use crate::blockstore::arrow_blockfile::block::Block;
+// use crate::blockstore::arrow_blockfile::blockfile::ArrowBlockfile;
+use crate::blockstore::provider::{CreateError, OpenError};
+use crate::blockstore::types::{Blockfile, KeyType, ValueType};
+
+pub(super) struct ArrowBlockfileProvider {
+    block_provider: ArrowBlockProvider,
+    files: HashMap<String, Box<dyn Blockfile>>,
+}
+
+impl BlockfileProvider for ArrowBlockfileProvider {
+    fn new() -> Self {
+        Self {
+            block_provider: ArrowBlockProvider::new(),
+            files: HashMap::new(),
+        }
+    }
+
+    fn open(&self, path: &str) -> Result<Box<dyn Blockfile>, Box<OpenError>> {
+        match self.files.get(path) {
+            Some(file) => Ok(file.clone()),
+            None => Err(Box::new(OpenError::NotFound)),
+        }
+    }
+
+    fn create(
+        &mut self,
+        path: &str,
+        key_type: KeyType,
+        value_type: ValueType,
+    ) -> Result<Box<dyn Blockfile>, Box<CreateError>> {
+        match self.files.get(path) {
+            Some(_) => Err(Box::new(CreateError::AlreadyExists)),
+            None => {
+                let blockfile = Box::new(ArrowBlockfile::new(
+                    key_type,
+                    value_type,
+                    self.block_provider.clone(),
+                ));
+                self.files.insert(path.to_string(), blockfile);
+                Ok(self.files.get(path).unwrap().clone())
+            }
+        }
+    }
+}
+
+struct ArrowBlockProviderInner {
+    blocks: HashMap<Uuid, Arc<Block>>,
+}
+
+#[derive(Clone)]
+pub(super) struct ArrowBlockProvider {
+    inner: Arc<RwLock<ArrowBlockProviderInner>>,
+}
+
+impl ArrowBlockProvider {
+    pub(super) fn new() -> Self {
+        Self {
+            inner: Arc::new(RwLock::new(ArrowBlockProviderInner {
+                blocks: HashMap::new(),
+            })),
+        }
+    }
+
+    pub(super) fn create_block(&self, key_type: KeyType, value_type: ValueType) -> Arc<Block> {
+        let block = Arc::new(Block::new(Uuid::new_v4(), key_type, value_type));
+        self.inner
+            .write()
+            .blocks
+            .insert(block.get_id(), block.clone());
+        block
+    }
+
+    pub(super) fn get_block(&self, id: &Uuid) -> Option<Arc<Block>> {
+        self.inner.read().blocks.get(id).cloned()
+    }
+}
diff --git a/rust/worker/src/blockstore/mod.rs b/rust/worker/src/blockstore/mod.rs
index 7a9d5d38b7ee..4facc35f85d4 100644
--- a/rust/worker/src/blockstore/mod.rs
+++ b/rust/worker/src/blockstore/mod.rs
@@ -1,3 +1,4 @@
+mod arrow_blockfile;
 mod positional_posting_list_value;
 mod types;