[PERF] prefetch blocks for fulltext index writer #3370

codetheweb · 2024-12-30T19:36:42Z

Description of changes

Prefetchs blocks in parallel when writing the full text index rather than fetching the blocks sequentially on-demand. Fetching in parallel saves us quite a bit of time and having prefetch as a distinct operator also allows us to mark it as I/O.

Test plan

How are these changes tested?

Tests pass locally with pytest for python, yarn test for js, cargo test for rust

Tested by disabling the write-through cache and observing that the prefetch step pulled in the correct blocks.

Documentation Changes

Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the docs repository?

n/a

github-actions · 2024-12-30T19:36:52Z

rust/index/src/fulltext/types.rs

rust/blockstore/src/arrow/ordered_blockfile_writer.rs

HammadB

I think @sanketkedia should weigh in on this before merge, as we have operators for prefetch and established patterns for it

codetheweb · 2025-01-02T22:30:33Z

I think @sanketkedia should weigh in on this before merge, as we have operators for prefetch and established patterns for it

I will ask
I originally sent more context to Robert in a DM, should have included it here:

I didn't implement prefetching in an operator in this first pass because it doesn't fit well into the upstream SegmentWriter abstraction & associated operators

rust/blockstore/src/types/writer.rs

rust/blockstore/src/arrow/ordered_blockfile_writer.rs

rust/index/src/fulltext/types.rs

rust/blockstore/src/arrow/blockfile.rs

codetheweb · 2025-01-13T22:36:13Z

This stack of pull requests is managed by Graphite. Learn more about stacking.

rust/worker/src/execution/orchestration/compact.rs

rust/index/src/fulltext/types.rs

rust/worker/src/execution/operators/prefetch_keys.rs

rust/blockstore/src/arrow/ordered_blockfile_writer.rs

HammadB · 2025-01-21T21:27:42Z

rust/blockstore/src/arrow/ordered_blockfile_writer.rs

+        let num_block_cache_hits = self
+            .num_block_cache_hits
+            .load(std::sync::atomic::Ordering::Relaxed);
+        let num_block_cache_misses = self


Foyer already exposes the underlying hit rate and miss rate - i think this is redundant?

chroma/rust/cache/src/foyer.rs

Line 251 in 8f7a0ff

cache_hit: opentelemetry::metrics::Counter<u64>,

Also if we are doing it - i i think it should be a metric not a trace

ah yep missed that
I think it's actually kinda useful having it on the trace; if all you have is a metric point with a low cache hit rate it could be pretty hard to debug
will leave out for now

ah makes sense

rust/index/src/fulltext/types.rs

codetheweb · 2025-01-22T19:25:41Z

rust/worker/src/execution/orchestration/compact.rs

@@ -271,6 +277,21 @@ impl CompactOrchestrator {
        let input = PartitionInput::new(records, self.max_partition_size);
        let task = wrap(operator, input, ctx.receiver());
        self.send(task, ctx).await;
+
+        let segments = self.get_all_segments().await.unwrap();


(Not putting this in the initial_tasks() method from the Orchestrator trait as that would require making initial_tasks() async and my understanding is that it's intended to be a static list of tasks. A cleaner solution would perhaps be adding a GetSegment operator and dispatching that from initial_tasks().)

codetheweb · 2025-01-22T20:52:18Z

rust/blockstore/src/provider.rs

@@ -100,6 +100,15 @@ impl BlockfileProvider {
        };
        Ok(())
    }
+
+    pub async fn prefetch(&self, id: &uuid::Uuid) -> Result<usize, Box<dyn ChromaError>> {


this is a new top-level API, not attached to a reader or writer
I think this makes more sense semantically since we're not fetching the blocks for any particular reader or writer instance and it lets us avoid K & V generics

HammadB · 2025-01-23T04:32:54Z

rust/blockstore/src/arrow/provider.rs

@@ -62,6 +80,29 @@ impl ArrowBlockfileProvider {
        }
    }

+    pub async fn prefetch(&self, id: &Uuid) -> Result<usize, ArrowBlockfileProviderPrefetchError> {
+        let block_ids = self


Why not just fetch the root and then do this? The abstractions feel a bit off to me here.

Fetching the root requires K to be known, which would break a general-purpose prefetch operator that has no knowledge of the concrete K/V types for each file path.

ah that sucks, might be nice to comment that

HammadB · 2025-01-23T06:11:45Z

Barring any fallout from my question above this makes sense to me. Thanks for the attention to quality.

rescrv reviewed Jan 2, 2025

View reviewed changes

rust/index/src/fulltext/types.rs Outdated Show resolved Hide resolved

rust/blockstore/src/arrow/ordered_blockfile_writer.rs Outdated Show resolved Hide resolved

HammadB reviewed Jan 2, 2025

View reviewed changes

codetheweb commented Jan 10, 2025

View reviewed changes

rust/blockstore/src/types/writer.rs Outdated Show resolved Hide resolved

codetheweb commented Jan 10, 2025

View reviewed changes

rust/blockstore/src/arrow/ordered_blockfile_writer.rs Outdated Show resolved Hide resolved

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from 8b144d3 to 154c080 Compare January 10, 2025 23:36

codetheweb commented Jan 10, 2025

View reviewed changes

rust/index/src/fulltext/types.rs Outdated Show resolved Hide resolved

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch 2 times, most recently from 8ed8d3b to 931a5f2 Compare January 10, 2025 23:54

codetheweb commented Jan 10, 2025

View reviewed changes

rust/blockstore/src/arrow/blockfile.rs Outdated Show resolved Hide resolved

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from 931a5f2 to 323b15e Compare January 13, 2025 22:34

codetheweb changed the base branch from main to fix-arrow-cached-check-bug January 13, 2025 22:35

codetheweb mentioned this pull request Jan 13, 2025

[BUG]: BlockManager::cached() should not always return true #3480

Merged

1 task

codetheweb force-pushed the fix-arrow-cached-check-bug branch from ce5521e to 2c3884c Compare January 14, 2025 18:30

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from 323b15e to 31ee004 Compare January 14, 2025 18:31

codetheweb marked this pull request as ready for review January 14, 2025 18:35

codetheweb requested a review from sanketkedia January 14, 2025 18:35

codetheweb changed the base branch from fix-arrow-cached-check-bug to graphite-base/3370 January 15, 2025 02:09

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from 31ee004 to f7ecfd9 Compare January 15, 2025 02:09

codetheweb force-pushed the graphite-base/3370 branch from 2c3884c to 6ae644d Compare January 15, 2025 02:09

codetheweb changed the base branch from graphite-base/3370 to main January 15, 2025 02:10

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from f7ecfd9 to a13d405 Compare January 15, 2025 02:10

sanketkedia reviewed Jan 15, 2025

View reviewed changes

rust/worker/src/execution/orchestration/compact.rs Outdated Show resolved Hide resolved

sanketkedia reviewed Jan 15, 2025

View reviewed changes

rust/index/src/fulltext/types.rs Outdated Show resolved Hide resolved

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from a13d405 to 8c7288e Compare January 16, 2025 00:53

codetheweb commented Jan 16, 2025

View reviewed changes

rust/worker/src/execution/operators/prefetch_keys.rs Outdated Show resolved Hide resolved

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from e0f7533 to 577da0e Compare January 16, 2025 18:35

codetheweb requested a review from sanketkedia January 16, 2025 18:37

sanketkedia reviewed Jan 16, 2025

View reviewed changes

rust/blockstore/src/arrow/ordered_blockfile_writer.rs Outdated Show resolved Hide resolved

sanketkedia approved these changes Jan 16, 2025

View reviewed changes

codetheweb requested a review from HammadB January 17, 2025 00:51

HammadB reviewed Jan 21, 2025

View reviewed changes

rust/index/src/fulltext/types.rs Outdated Show resolved Hide resolved

codetheweb added 2 commits January 22, 2025 10:41

[PERF] prefetch blocks for fulltext index writer

f311949

Pass keys in prefetch input

afb02fa

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch 3 times, most recently from 77c6a76 to f7c73c4 Compare January 22, 2025 19:21

codetheweb mentioned this pull request Jan 22, 2025

[ENH]: improve interface for load_blocks_for_keys() #3534

Open

1 task

codetheweb commented Jan 22, 2025

View reviewed changes

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch 3 times, most recently from c547c2f to a264c8f Compare January 22, 2025 20:51

codetheweb commented Jan 22, 2025

View reviewed changes

Prefetch entire segment

8241a65

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from a264c8f to 8241a65 Compare January 22, 2025 21:01

codetheweb requested a review from HammadB January 22, 2025 21:02

HammadB reviewed Jan 23, 2025

View reviewed changes

HammadB approved these changes Jan 23, 2025

View reviewed changes

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch 2 times, most recently from 3440b66 to 8241a65 Compare January 23, 2025 21:46

Add comment

61d49f0

codetheweb enabled auto-merge (squash) January 23, 2025 21:49

codetheweb added 2 commits January 23, 2025 14:07

Merge branch 'main' into perf-prefetch-fulltext-index-writer-blocks

daf9a7d

Fix imports

c9f7fa0

codetheweb merged commit b848326 into main Jan 23, 2025
79 checks passed

codetheweb deleted the perf-prefetch-fulltext-index-writer-blocks branch January 23, 2025 23:07

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[PERF] prefetch blocks for fulltext index writer #3370

[PERF] prefetch blocks for fulltext index writer #3370

codetheweb commented Dec 30, 2024 •

edited

Loading

github-actions bot commented Dec 30, 2024

HammadB left a comment

codetheweb commented Jan 2, 2025

codetheweb commented Jan 13, 2025 •

edited

Loading

HammadB Jan 21, 2025

HammadB Jan 21, 2025

codetheweb Jan 22, 2025

HammadB Jan 23, 2025

codetheweb Jan 22, 2025

codetheweb Jan 22, 2025 •

edited

Loading

HammadB Jan 23, 2025

codetheweb Jan 23, 2025

HammadB Jan 23, 2025

HammadB commented Jan 23, 2025

[PERF] prefetch blocks for fulltext index writer #3370

[PERF] prefetch blocks for fulltext index writer #3370

Conversation

codetheweb commented Dec 30, 2024 • edited Loading

Description of changes

Test plan

Documentation Changes

github-actions bot commented Dec 30, 2024

Reviewer Checklist

Testing, Bugs, Errors, Logs, Documentation

System Compatibility

Quality

HammadB left a comment

Choose a reason for hiding this comment

codetheweb commented Jan 2, 2025

codetheweb commented Jan 13, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codetheweb Jan 22, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

HammadB commented Jan 23, 2025

codetheweb commented Dec 30, 2024 •

edited

Loading

codetheweb commented Jan 13, 2025 •

edited

Loading

codetheweb Jan 22, 2025 •

edited

Loading