Skip to content

Commit a09392f

Browse files
codethewebcodetheweb
committed
[ENH]: add batch get version file paths method to Sysdb
1 parent b7cd2b0 commit a09392f

File tree

11 files changed

+215
-7
lines changed

11 files changed

+215
-7
lines changed

‎go/pkg/sysdb/coordinator/coordinator.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -263,6 +263,10 @@ func (s *Coordinator) DeleteCollectionVersion(ctx context.Context, req *coordina
263263
return s.catalog.DeleteCollectionVersion(ctx, req)
264264
}
265265

266+
func (s *Coordinator) BatchGetCollectionVersionFilePaths(ctx context.Context, req *coordinatorpb.BatchGetCollectionVersionFilePathsRequest) (*coordinatorpb.BatchGetCollectionVersionFilePathsResponse, error) {
267+
return s.catalog.BatchGetCollectionVersionFilePaths(ctx, req.CollectionIds)
268+
}
269+
266270
// SetDeleteMode sets the delete mode for testing
267271
func (c *Coordinator) SetDeleteMode(mode DeleteMode) {
268272
c.deleteMode = mode

‎go/pkg/sysdb/coordinator/coordinator_test.go

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1657,6 +1657,37 @@ func (suite *APIsTestSuite) TestForkCollection() {
16571657
suite.True(exists, "Lineage file should exist in S3")
16581658
}
16591659

1660+
func (suite *APIsTestSuite) TestBatchGetCollectionVersionFilePaths() {
1661+
ctx := context.Background()
1662+
1663+
// Create a new collection
1664+
newCollection := &model.CreateCollection{
1665+
ID: types.NewUniqueID(),
1666+
Name: "test_batch_get_collection_version_file_paths",
1667+
TenantID: suite.tenantName,
1668+
DatabaseName: suite.databaseName,
1669+
}
1670+
1671+
newSegments := []*model.Segment{}
1672+
1673+
// Create the collection
1674+
suite.coordinator.catalog.versionFileEnabled = true
1675+
_, _, err := suite.coordinator.CreateCollectionAndSegments(ctx, newCollection, newSegments)
1676+
suite.NoError(err)
1677+
1678+
// Get the version file paths for the collection
1679+
versionFilePaths, err := suite.coordinator.BatchGetCollectionVersionFilePaths(ctx, &coordinatorpb.BatchGetCollectionVersionFilePathsRequest{
1680+
CollectionIds: []string{newCollection.ID.String()},
1681+
})
1682+
suite.NoError(err)
1683+
suite.Len(versionFilePaths.CollectionIdToVersionFilePath, 1)
1684+
1685+
// Verify version file exists in S3
1686+
exists, err := suite.s3MetaStore.HasObjectWithPrefix(ctx, versionFilePaths.CollectionIdToVersionFilePath[newCollection.ID.String()])
1687+
suite.NoError(err)
1688+
suite.True(exists, "Version file should exist in S3")
1689+
}
1690+
16601691
func (suite *APIsTestSuite) TestCountForks() {
16611692
ctx := context.Background()
16621693

‎go/pkg/sysdb/coordinator/table_catalog.go

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2094,6 +2094,20 @@ func (tc *Catalog) DeleteCollectionVersion(ctx context.Context, req *coordinator
20942094
return &result, nil
20952095
}
20962096

2097+
func (tc *Catalog) BatchGetCollectionVersionFilePaths(ctx context.Context, collectionIds []string) (*coordinatorpb.BatchGetCollectionVersionFilePathsResponse, error) {
2098+
result := coordinatorpb.BatchGetCollectionVersionFilePathsResponse{
2099+
CollectionIdToVersionFilePath: make(map[string]string),
2100+
}
2101+
2102+
paths, err := tc.metaDomain.CollectionDb(ctx).BatchGetCollectionVersionFilePaths(collectionIds)
2103+
if err != nil {
2104+
return nil, err
2105+
}
2106+
result.CollectionIdToVersionFilePath = paths
2107+
2108+
return &result, nil
2109+
}
2110+
20972111
func (tc *Catalog) GetVersionFileNamesForCollection(ctx context.Context, tenantID string, collectionID string) (string, error) {
20982112
collectionIDPtr := &collectionID
20992113
isDeleted := false

‎go/pkg/sysdb/grpc/collection_service.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -548,3 +548,11 @@ func (s *Server) DeleteCollectionVersion(ctx context.Context, req *coordinatorpb
548548
}
549549
return res, nil
550550
}
551+
552+
func (s *Server) BatchGetCollectionVersionFilePaths(ctx context.Context, req *coordinatorpb.BatchGetCollectionVersionFilePathsRequest) (*coordinatorpb.BatchGetCollectionVersionFilePathsResponse, error) {
553+
res, err := s.coordinator.BatchGetCollectionVersionFilePaths(ctx, req)
554+
if err != nil {
555+
return nil, grpcutils.BuildInternalGrpcError(err.Error())
556+
}
557+
return res, nil
558+
}

‎go/pkg/sysdb/metastore/db/dao/collection.go

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -558,3 +558,20 @@ func (s *collectionDb) UpdateCollectionLineageFilePath(collectionID string, curr
558558
}).Error
559559

560560
}
561+
562+
func (s *collectionDb) BatchGetCollectionVersionFilePaths(collectionIDs []string) (map[string]string, error) {
563+
var collections []dbmodel.Collection
564+
err := s.read_db.Model(&dbmodel.Collection{}).
565+
Select("id, version_file_name").
566+
Where("id IN ?", collectionIDs).
567+
Find(&collections).Error
568+
if err != nil {
569+
return nil, err
570+
}
571+
572+
result := make(map[string]string)
573+
for _, collection := range collections {
574+
result[collection.ID] = collection.VersionFileName
575+
}
576+
return result, nil
577+
}

‎go/pkg/sysdb/metastore/db/dbmodel/collection.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,4 +69,5 @@ type ICollectionDb interface {
6969
UpdateVersionRelatedFields(collectionID, existingVersionFileName, newVersionFileName string, oldestVersionTs *time.Time, numActiveVersions *int) (int64, error)
7070
LockCollection(collectionID string) (*bool, error)
7171
UpdateCollectionLineageFilePath(collectionID string, currentLineageFilePath *string, newLineageFilePath string) error
72+
BatchGetCollectionVersionFilePaths(collectionIDs []string) (map[string]string, error)
7273
}

‎go/pkg/sysdb/metastore/db/dbmodel/mocks/ICollectionDb.go

Lines changed: 30 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎idl/chromadb/proto/coordinator.proto

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -467,6 +467,14 @@ message DeleteCollectionVersionResponse {
467467
map<string, bool> collection_id_to_success = 1;
468468
}
469469

470+
message BatchGetCollectionVersionFilePathsRequest {
471+
repeated string collection_ids = 1;
472+
}
473+
474+
message BatchGetCollectionVersionFilePathsResponse {
475+
map<string, string> collection_id_to_version_file_path = 1;
476+
}
477+
470478
service SysDB {
471479
rpc CreateDatabase(CreateDatabaseRequest) returns (CreateDatabaseResponse) {}
472480
rpc GetDatabase(GetDatabaseRequest) returns (GetDatabaseResponse) {}
@@ -498,4 +506,5 @@ service SysDB {
498506
rpc ListCollectionsToGc(ListCollectionsToGcRequest) returns (ListCollectionsToGcResponse) {}
499507
rpc MarkVersionForDeletion(MarkVersionForDeletionRequest) returns (MarkVersionForDeletionResponse) {}
500508
rpc DeleteCollectionVersion(DeleteCollectionVersionRequest) returns (DeleteCollectionVersionResponse) {}
509+
rpc BatchGetCollectionVersionFilePaths(BatchGetCollectionVersionFilePathsRequest) returns (BatchGetCollectionVersionFilePathsResponse) {}
501510
}

‎rust/sysdb/src/sysdb.rs

Lines changed: 47 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,9 @@ use chroma_types::{
1919
UpdateCollectionConfiguration, UpdateCollectionError, VectorIndexConfiguration,
2020
};
2121
use chroma_types::{
22-
Collection, CollectionConversionError, CollectionUuid, CountForksError,
23-
FlushCompactionResponse, FlushCompactionResponseConversionError, ForkCollectionError, Segment,
24-
SegmentConversionError, SegmentScope, Tenant,
22+
BatchGetCollectionVersionFilePathsError, Collection, CollectionConversionError, CollectionUuid,
23+
CountForksError, FlushCompactionResponse, FlushCompactionResponseConversionError,
24+
ForkCollectionError, Segment, SegmentConversionError, SegmentScope, Tenant,
2525
};
2626
use std::collections::HashMap;
2727
use std::fmt::Debug;
@@ -400,6 +400,23 @@ impl SysDb {
400400
}
401401
}
402402

403+
pub async fn batch_get_collection_version_file_paths(
404+
&mut self,
405+
collection_ids: Vec<CollectionUuid>,
406+
) -> Result<HashMap<CollectionUuid, String>, BatchGetCollectionVersionFilePathsError> {
407+
match self {
408+
SysDb::Grpc(grpc) => {
409+
grpc.batch_get_collection_version_file_paths(collection_ids)
410+
.await
411+
}
412+
SysDb::Sqlite(_) => todo!(),
413+
SysDb::Test(test) => {
414+
test.batch_get_collection_version_file_paths(collection_ids)
415+
.await
416+
}
417+
}
418+
}
419+
403420
pub async fn get_last_compaction_time(
404421
&mut self,
405422
tanant_ids: Vec<String>,
@@ -1114,6 +1131,33 @@ impl GrpcSysDb {
11141131
})
11151132
}
11161133

1134+
async fn batch_get_collection_version_file_paths(
1135+
&mut self,
1136+
collection_ids: Vec<CollectionUuid>,
1137+
) -> Result<HashMap<CollectionUuid, String>, BatchGetCollectionVersionFilePathsError> {
1138+
let res = self
1139+
.client
1140+
.batch_get_collection_version_file_paths(
1141+
chroma_proto::BatchGetCollectionVersionFilePathsRequest {
1142+
collection_ids: collection_ids
1143+
.into_iter()
1144+
.map(|id| id.0.to_string())
1145+
.collect(),
1146+
},
1147+
)
1148+
.await?;
1149+
let collection_id_to_path = res.into_inner().collection_id_to_version_file_path;
1150+
let mut result = HashMap::new();
1151+
for (key, value) in collection_id_to_path {
1152+
let collection_id = CollectionUuid(
1153+
Uuid::try_parse(&key)
1154+
.map_err(|err| BatchGetCollectionVersionFilePathsError::Uuid(err, key))?,
1155+
);
1156+
result.insert(collection_id, value);
1157+
}
1158+
Ok(result)
1159+
}
1160+
11171161
async fn get_last_compaction_time(
11181162
&mut self,
11191163
tenant_ids: Vec<String>,

‎rust/sysdb/src/test_sysdb.rs

Lines changed: 37 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
use chroma_types::{
2-
Collection, CollectionAndSegments, CollectionUuid, CountForksError, Database,
3-
FlushCompactionResponse, GetCollectionSizeError, GetCollectionWithSegmentsError,
4-
GetSegmentsError, ListDatabasesError, ListDatabasesResponse, Segment, SegmentFlushInfo,
5-
SegmentScope, SegmentType, Tenant,
2+
BatchGetCollectionVersionFilePathsError, Collection, CollectionAndSegments, CollectionUuid,
3+
CountForksError, Database, FlushCompactionResponse, GetCollectionSizeError,
4+
GetCollectionWithSegmentsError, GetSegmentsError, ListDatabasesError, ListDatabasesResponse,
5+
Segment, SegmentFlushInfo, SegmentScope, SegmentType, Tenant,
66
};
77
use chroma_types::{GetCollectionsError, SegmentUuid};
88
use parking_lot::Mutex;
@@ -93,6 +93,18 @@ impl TestSysDb {
9393
.insert(tenant, last_compaction_time);
9494
}
9595

96+
pub fn set_collection_version_file_path(
97+
&mut self,
98+
collection_id: CollectionUuid,
99+
version_file_path: String,
100+
) {
101+
let mut inner = self.inner.lock();
102+
103+
inner
104+
.collection_to_version_file_name
105+
.insert(collection_id, version_file_path);
106+
}
107+
96108
fn filter_collections(
97109
collection: &Collection,
98110
collection_id: Option<CollectionUuid>,
@@ -549,4 +561,25 @@ impl TestSysDb {
549561
) -> Result<usize, CountForksError> {
550562
Ok(10)
551563
}
564+
565+
pub(crate) async fn batch_get_collection_version_file_paths(
566+
&self,
567+
collection_ids: Vec<CollectionUuid>,
568+
) -> Result<HashMap<CollectionUuid, String>, BatchGetCollectionVersionFilePathsError> {
569+
let inner = self.inner.lock();
570+
let mut paths = HashMap::new();
571+
for collection_id in collection_ids {
572+
if let Some(path) = inner.collection_to_version_file_name.get(&collection_id) {
573+
paths.insert(collection_id, path.clone());
574+
} else {
575+
return Err(BatchGetCollectionVersionFilePathsError::Grpc(
576+
tonic::Status::not_found(format!(
577+
"Version file not found for collection: {}",
578+
collection_id
579+
)),
580+
));
581+
}
582+
}
583+
Ok(paths)
584+
}
552585
}

0 commit comments

Comments
 (0)