[ENH]: Enable rebuilds for sharded collections by tanujnay112 · Pull Request #6916 · chroma-core/chroma

tanujnay112 · 2026-04-15T20:02:08Z

Description of changes

This change adds rebuild capabilities for sharded collections. A rebuild command can now specify a shard index it wants to rebuild. An unspecified one defaults to 0 (first shard). This shard's filepaths are set to empty during rebuild compaction, analogous to all of the segment's filepaths being set to empty before this change. A new struct packaging the rebuild command parameters called RebuildInfo is passed to CompactionContext now. This contains the desired shard index to be rebuilt.

Improvements & Bug fixes
- ...
New functionality
- ...

Test plan

Tests have been added in compact.rs. They test rebuilding various shards in the full and partial rebuild scenarios.

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

Migration plan

Are there any migrations, or any forwards/backwards compatibility changes needed in order to make sure this change deploys reliably?

Observability plan

What is the plan to instrument and monitor this change?

Documentation Changes

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

github-actions · 2026-04-15T20:02:17Z

propel-code-bot

Review found no issues; shard-aware rebuild enhancements appear well-implemented and sufficiently covered.

Status: No Issues Found | Risk: Low

Review Details

📁 12 files reviewed | 💬 0 comments

propel-code-bot · 2026-04-15T20:04:04Z

Enable shard-targeted rebuilds for sharded collections across compaction orchestration

This PR adds shard-aware rebuild functionality end-to-end in the Rust compaction pipeline. Rebuild requests can now optionally include a shard index via RebuildRequest.shard_index and CLI --shard, with default behavior targeting shard 0 when unspecified. Rebuild intent is refactored into a new RebuildInfo object that flows through CompactionContext, LogFetchOrchestrator, and CompactionManager instead of separate is_rebuild and apply_segment_scopes fields.

The rebuild execution path now operates on specific shards by selecting SegmentShard using the requested index, materializing only that shard while emitting empty placeholders for other shard slots, and clearing file paths only for the targeted shard using Segment.clear_shard_file_paths(). The PR also updates segment/shard helpers, prefetch behavior, and extensive integration-style tests in rust/worker/src/execution/orchestration/compact.rs to validate vector-only and all-segment rebuilds across multi-shard scenarios.

This summary was automatically generated by @propel-code-bot

tanujnay112 · 2026-04-15T20:07:17Z

[ENH]: Enable rebuilds for sharded collections #6916 👈 (View in Graphite)
main

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

sanketkedia · 2026-04-16T19:03:37Z

        segment_scopes: Vec<String>,
+        /// Specify which shard to rebuild (defaults to 0)
+        #[arg(long)]
+        shard: Option<u32>,


would be nice to have the flexibility of specifying different shards for different collections of the batch so Vec<Option<u32>>?

sanketkedia · 2026-04-16T19:04:22Z

  // Segment scopes to rebuild. If empty, rebuilds all segments (metadata + vector).
  repeated SegmentScope segment_scopes = 2;
+  // Optional shard index to rebuild. If not specified, defaults to shard 0.
+  optional uint32 shard_index = 3;


since this is a batch api, why not have this a repeated field?

sanketkedia · 2026-04-16T22:58:51Z

-                )
-            }
+        let log_task = if self.context.is_rebuild() {
+            let shard_count = output


nit: there should be a num_shards method in Segment type that you can use here

sanketkedia · 2026-04-16T23:01:54Z

            )
            .await;

-        if record_segment_reader.is_none() {


how was this even correct? for e.g. if a collection is compacted for the first time then this will not fetch the logs and return early?

from_segment from above always returned a reader so we never hit this

what you linked is of RecordSegmentWriterShard and not of RecordSegmentReader

But I found the relevant that returns. It returns a Vec<Option<RecordSegmentReaderShard>> so it will return [None] for one shard and uninitialized reader

But you should still keep it for e.g. if ok_or_terminate returns None because the reader creation error'd out

sanketkedia · 2026-04-16T23:10:44Z

        // Prefetch segments
-        let prefetch_segments = match self.context.is_rebuild {
+        let prefetch_segments = match self.context.is_rebuild() {
            true => vec![output.record_segment],


should prefetch the shard that is getting rebuilt here instead of the active shard

sanketkedia · 2026-04-16T23:12:57Z

 }

 #[async_trait]
 impl Handler<TaskResult<SourceRecordSegmentOutput, SourceRecordSegmentError>>


add a todo to clean up V1

sanketkedia · 2026-04-16T23:18:28Z


+                // Create a vector with empty shards for all positions except shard_index
+                let mut shards = vec![];
+                for i in 0..self.shard_count {


this and the below block of code are exactly the same so can extract into a helper

sanketkedia · 2026-04-16T23:23:52Z

@@ -935,4 +925,4 @@
        collection_info.collection.total_records_post_compaction = output.total_records as u64;


this would be wrong since total_records is only of that shard

hmm, this might be a bit hard to compute now

should we just not compute it

sanketkedia · 2026-04-16T23:26:28Z

        collection_info.collection.total_records_post_compaction = output.total_records as u64;

        // If no records, terminate early
        if output.partitions.is_empty() {


this code needs to change now

sanketkedia · 2026-04-17T16:29:27Z

                            ids: id.iter().map(ToString::to_string).collect(),
                        }),
                        segment_scopes: proto_scopes,
+                        shard_index: *shard,


why deref here? seems unnecessary

sanketkedia · 2026-04-17T16:43:20Z

                            ids: id.iter().map(ToString::to_string).collect(),
                        }),
                        segment_scopes: proto_scopes,
+                        shard_index: *shard,


also would it be simpler for the server to just have an u32 instead of Option and in the client you set to 0 if None?

Effectively doing this at the server level. I think it makes sense to different on the client-server layer of whether or not a shard was specified. Say in the future you want a multi-shard rebuild API

sanketkedia · 2026-04-17T17:00:31Z

        // TODO(tanujnay112): This is awful, we need to find a better way to pass
        // the active collection info around.
        self.collection_info = OnceCell::new();
+        let rebuild_info = if is_getting_compacted_logs {


this is a bit weird. Clarifying my understanding - Is it to prevent abstraction violation that the attached functions call does not have to construct a RebuildInfo struct? Because you are using RebuildInfo as Some to implicitly mean "fetch logs from record segment" in LogFetchOrchestrator

Yeah, for abstraction boundaries. In the attached function backfill path I need to signal that backfilling is happening this way without trying to break the abstraction boundaries you mentioned.

sanketkedia · 2026-04-17T17:12:38Z

        blockfile_provider: &BlockfileProvider,
    ) {
-        // Get offset IDs from vector segment
+        // Get offset IDs from vector segment using distributed reader (handles all shards)


nit: can remove

sanketkedia · 2026-04-17T17:47:09Z

-        }
-
        // Also create RecordSegmentReaderShard for source operators
        let record_segment_shard = match self


the construction of this segment shard and its reader should be done in the conditional below in the if self.context.is_rebuild() path. No need to pay this cost in the general non rebuild case

sanketkedia · 2026-04-17T18:03:34Z

+            vec![output.record_segment]
+        } else {
+            let mut segments = vec![output.metadata_segment, output.record_segment];
+            if vector_segment.r#type == chroma_types::SegmentType::Spann {


also prefetch QuantizedSpann here

sanketkedia · 2026-04-17T18:13:13Z

-            let vec = vec![self.new_shard()];
-            return Ok(vec);
-        }
+        // If there are no file paths, return a single empty shard


Segment type also has num_shards() method. should use that here. It already has the validation logic for shard count

sanketkedia · 2026-04-17T18:26:45Z

+                })
+                .collect();
+            MaterializeLogOutput {
+                result: PartitionedMaterializeLogsResult { shards },


as a cleanup we should rename this to Sharded* since Partition is an existing concept in the codebase which is different from Sharding

Will do in a followup change

sg, added to the burndown doc

sanketkedia · 2026-04-17T19:08:30Z

    pub fn get_shards(&self) -> Result<Vec<SegmentShard>, SegmentShardError> {
-        // If there are no file paths, return empty vector
+        let num_shards = self.num_shards()?;
        if self.file_path.is_empty() {


this case should already be handled in the SegmentShard::try_from. can you pls double check?

sanketkedia · 2026-04-17T19:30:28Z

+            vec![output.record_segment]
+        } else {
+            let mut segments = vec![output.metadata_segment, output.record_segment];
+            if vector_segment.r#type != chroma_types::SegmentType::HnswDistributed {


this won't work because it will try to prefetch the local hnsw or sqlite segment types

sanketkedia · 2026-04-17T19:33:49Z

These changes need so much care, great job at that

This PR cherry-picks the commit fdcd216 onto release/2026-04-03. If there are unresolved conflicts, please resolve them manually. Co-authored-by: tanujnay112 <tanujnay112@live.com>

tanujnay112 marked this pull request as ready for review April 15, 2026 20:02

propel-code-bot bot reviewed Apr 15, 2026

View reviewed changes

tanujnay112 requested a review from sanketkedia April 15, 2026 20:09

tanujnay112 force-pushed the rebuild_sharding branch 3 times, most recently from d0bdf23 to e98b237 Compare April 16, 2026 07:29

This comment has been minimized.

Sign in to view

tanujnay112 force-pushed the rebuild_sharding branch from e98b237 to 6506960 Compare April 16, 2026 08:56

This comment has been minimized.

Sign in to view

sanketkedia reviewed Apr 16, 2026

View reviewed changes

[ENH]: Enable rebuilds for sharded collections

398cb8e

tanujnay112 force-pushed the rebuild_sharding branch from 6506960 to 398cb8e Compare April 16, 2026 19:31

sanketkedia reviewed Apr 16, 2026

View reviewed changes

tanujnay112 added 2 commits April 16, 2026 17:11

sanket comments

7615f7b

removed hnsw changeS

da0abb0

tanujnay112 force-pushed the rebuild_sharding branch from 66ada8c to da0abb0 Compare April 17, 2026 00:54

sanketkedia reviewed Apr 17, 2026

View reviewed changes

sanket comments

32957fc

tanujnay112 requested a review from sanketkedia April 17, 2026 18:52

sanketkedia reviewed Apr 17, 2026

View reviewed changes

more sanket comments

8872e75

sanketkedia reviewed Apr 17, 2026

View reviewed changes

sanketkedia approved these changes Apr 17, 2026

View reviewed changes

tanujnay112 merged commit fdcd216 into main Apr 17, 2026
122 of 124 checks passed

		@@ -935,4 +925,4 @@
		collection_info.collection.total_records_post_compaction = output.total_records as u64;

Conversation

tanujnay112 commented Apr 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description of changes

Test plan

Migration plan

Observability plan

Documentation Changes

Uh oh!

github-actions bot commented Apr 15, 2026

Reviewer Checklist

Testing, Bugs, Errors, Logs, Documentation

System Compatibility

Quality

Uh oh!

propel-code-bot bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

propel-code-bot bot commented Apr 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tanujnay112 commented Apr 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

This comment has been minimized.

This comment has been minimized.

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sanketkedia Apr 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sanketkedia Apr 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

tanujnay112 commented Apr 15, 2026 •

edited

Loading

propel-code-bot bot left a comment •

edited

Loading

propel-code-bot bot commented Apr 15, 2026 •

edited

Loading

tanujnay112 commented Apr 15, 2026 •

edited

Loading

sanketkedia Apr 17, 2026 •

edited

Loading

sanketkedia Apr 17, 2026 •

edited

Loading