[CLN] Move max_block_size_bytes to block_manager_config (#2535)

## Description of changes *Summarize the changes made by this PR.* - Improvements & Bug fixes - Refactored the configuration key `max_block_size_bytes` from `ArrowBlockfileProviderConfig` to `BlockManagerConfig`. - New functionality - ... ## Test plan *How are these changes tested?* - [x] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust ## Documentation Changes *Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs repository](https://github.com/chroma-core/docs)?*
chroma-core · Jul 18, 2024 · 561e6f5 · 561e6f5
1 parent a4b8c3b
commit 561e6f5
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 18 deletions.
diff --git a/rust/worker/chroma_config.yaml b/rust/worker/chroma_config.yaml
@@ -40,8 +40,8 @@ query_service:
         worker_queue_size: 100
     blockfile_provider:
         Arrow:
-            max_block_size_bytes: 16384
             block_manager_config:
+                max_block_size_bytes: 16384
                 block_cache_config:
                     lru:
                         capacity: 1000
@@ -97,8 +97,8 @@ compaction_service:
         min_compaction_size: 10
     blockfile_provider:
         Arrow:
-            max_block_size_bytes: 16384
             block_manager_config:
+                max_block_size_bytes: 16384
                 block_cache_config:
                     lru:
                         capacity: 1000

diff --git a/rust/worker/src/blockstore/arrow/config.rs b/rust/worker/src/blockstore/arrow/config.rs
@@ -7,19 +7,14 @@ pub(crate) const TEST_MAX_BLOCK_SIZE_BYTES: usize = 16384;
 
 #[derive(Deserialize, Debug, Clone)]
 pub(crate) struct ArrowBlockfileProviderConfig {
-    // Note: This provider has two dependent components that
-    // are both internal to the arrow blockfile provider.
-    // The BlockManager and the SparseIndexManager.
-    // We could have a BlockManagerConfig and a SparseIndexManagerConfig
-    // but the only configuration that is needed is the max_block_size_bytes
-    // so for now we just hoid this configuration in the ArrowBlockfileProviderConfig.
-    pub(crate) max_block_size_bytes: usize,
+    // pub(crate) max_block_size_bytes: usize,
     pub(crate) block_manager_config: BlockManagerConfig,
     pub(crate) sparse_index_manager_config: SparseIndexManagerConfig,
 }
 
 #[derive(Deserialize, Debug, Clone)]
 pub(crate) struct BlockManagerConfig {
+    pub(crate) max_block_size_bytes: usize,
     pub(crate) block_cache_config: CacheConfig,
 }
 

diff --git a/rust/worker/src/blockstore/arrow/provider.rs b/rust/worker/src/blockstore/arrow/provider.rs
@@ -128,7 +128,7 @@ impl Configurable<(ArrowBlockfileProviderConfig, Storage)> for ArrowBlockfilePro
         };
         Ok(ArrowBlockfileProvider::new(
             storage.clone(),
-            blockfile_config.max_block_size_bytes,
+            blockfile_config.block_manager_config.max_block_size_bytes,
             block_cache,
             sparse_index_cache,
         ))

diff --git a/rust/worker/src/config.rs b/rust/worker/src/config.rs
@@ -196,8 +196,8 @@ mod tests {
                         worker_queue_size: 100
                     blockfile_provider:
                         Arrow:
-                            max_block_size_bytes: 16384
                             block_manager_config:
+                                max_block_size_bytes: 16384
                                 block_cache_config:
                                     lru:
                                         capacity: 1000
@@ -253,8 +253,8 @@ mod tests {
                         min_compaction_size: 10
                     blockfile_provider:
                         Arrow:
-                            max_block_size_bytes: 16384
                             block_manager_config:
+                                max_block_size_bytes: 16384
                                 block_cache_config:
                                     lru:
                                         capacity: 1000
@@ -326,8 +326,8 @@ mod tests {
                         worker_queue_size: 100
                     blockfile_provider:
                         Arrow:
-                            max_block_size_bytes: 16384
                             block_manager_config:
+                                max_block_size_bytes: 16384
                                 block_cache_config:
                                     lru:
                                         capacity: 1000
@@ -383,8 +383,8 @@ mod tests {
                         min_compaction_size: 10
                     blockfile_provider:
                         Arrow:
-                            max_block_size_bytes: 16384
                             block_manager_config:
+                                max_block_size_bytes: 16384
                                 block_cache_config:
                                     lru:
                                         capacity: 1000
@@ -474,8 +474,8 @@ mod tests {
                         worker_queue_size: 100
                     blockfile_provider:
                         Arrow:
-                            max_block_size_bytes: 16384
                             block_manager_config:
+                                max_block_size_bytes: 16384
                                 block_cache_config:
                                     lru:
                                         capacity: 1000
@@ -531,8 +531,8 @@ mod tests {
                         min_compaction_size: 10
                     blockfile_provider:
                         Arrow:
-                            max_block_size_bytes: 16384
                             block_manager_config:
+                                max_block_size_bytes: 16384
                                 block_cache_config:
                                     lru:
                                         capacity: 1000
@@ -616,8 +616,8 @@ mod tests {
                         worker_queue_size: 100
                     blockfile_provider:
                         Arrow:
-                            max_block_size_bytes: 16384
                             block_manager_config:
+                                max_block_size_bytes: 16384
                                 block_cache_config:
                                     lru:
                                         capacity: 1000
@@ -665,8 +665,8 @@ mod tests {
                         min_compaction_size: 10
                     blockfile_provider:
                         Arrow:
-                            max_block_size_bytes: 16384
                             block_manager_config:
+                                max_block_size_bytes: 16384
                                 block_cache_config:
                                     lru:
                                         capacity: 1000