Converge PQ pivot loading #1276
+352
−161
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -29,6 +29,19 @@ type FullPivotDataType = Vec<f32>; | |
| type CentroidType = Vec<f32>; | ||
| type ChunkOffsetsType = Vec<usize>; | ||
|
|
||
| #[derive(Debug)] | ||
| struct PivotFileParts { | ||
| pivots: Matrix<f32>, | ||
| centroid: Matrix<f32>, | ||
| chunk_offsets: Matrix<usize>, | ||
| } | ||
|
|
||
| impl PivotFileParts { | ||
| fn dim(&self) -> usize { | ||
| self.pivots.ncols() | ||
| } | ||
| } | ||
|
|
||
| #[derive(Debug, Clone)] | ||
| pub struct PQStorage { | ||
| /// Pivot table path | ||
|
|
@@ -179,63 +192,18 @@ impl PQStorage { | |
| where | ||
| Storage: StorageReadProvider, | ||
| { | ||
| let parts = self.load_pivot_file_parts( | ||
| &self.pivot_data_path, | ||
| Some(*num_pq_chunks), | ||
| Some(*num_centers), | ||
| Some(*dim), | ||
| storage_provider, | ||
| )?; | ||
|
|
||
| Ok(( | ||
| parts.pivots.into_inner().into_vec(), | ||
| parts.centroid.into_inner().into_vec(), | ||
| parts.chunk_offsets.into_inner().into_vec(), | ||
| )) | ||
| } | ||
|
|
||
|
|
@@ -289,7 +257,43 @@ impl PQStorage { | |
|
|
||
| info!("Loading PQ pivots from {}...", pq_pivots); | ||
|
|
||
| let PivotFileParts { | ||
| mut pivots, | ||
| centroid, | ||
| chunk_offsets, | ||
| } = self.load_pivot_file_parts( | ||
| pq_pivots, | ||
| (num_pq_chunks != 0).then_some(num_pq_chunks), | ||
| None, | ||
| None, | ||
| storage_provider, | ||
| )?; | ||
|
|
||
| // If the centroid is non-zero, we need to add it to the pivots to restore the | ||
| // numeric behavior. | ||
| if centroid.as_slice().iter().any(|c| *c != 0.0) { | ||
| accum_row_inplace(pivots.as_mut_view(), centroid.as_slice()) | ||
| } | ||
|
|
||
| FixedChunkPQTable::new( | ||
| pivots.ncols(), | ||
| pivots.into_inner(), | ||
| chunk_offsets.into_inner(), | ||
| ) | ||
| } | ||
|
Contributor
There was a problem hiding this comment. Thank you for looking at this! I have a higher-level comment that I think we may want to think about instead. The current divergent loading paths are a symptom of larger architectural issues and while unifying the load paths is a step in the right direction, it is treating a symptom rather than the cause. Much of the complexity here comes from the (optional) inline invariant checks. However, these checks are not exhaustive. For example, the offsets also need to be strictly monotonically increasing, which is not covered by the checks here. However, duplicating these checks is not strictly necessary as This loader has two basic return paths:
To recap, if this loading function only has a single pivot loading method returning a |
||
|
|
||
| fn load_pivot_file_parts<Storage: StorageReadProvider>( | ||
| &self, | ||
| pq_pivots: &str, | ||
| expected_num_pq_chunks: Option<usize>, | ||
| expected_num_centers: Option<usize>, | ||
| expected_dim: Option<usize>, | ||
| storage_provider: &Storage, | ||
| ) -> ANNResult<PivotFileParts> { | ||
|
Comment on lines
+297
to
+305
Contributor
There was a problem hiding this comment. Can we add some documentation to this function to make sure we understand how the arguments here are used? nit: It feels a bit hacky to me to case on these values when performing the checks in this function. It looks like the last two are used to validate pivots and the first to validate the chunks_offset. Maybe you can make this check more explicit? |
||
| let mut reader = storage_provider.open_reader(pq_pivots)?; | ||
|
|
||
| // File layout: offset table(4*1) -> pivot data(num_centers*dim) -> | ||
| // centroid(dim*1) -> chunk offsets(num_chunks+1*1). | ||
| let offsets = read_bin_from::<u64>(&mut reader, 0)?; | ||
| if offsets.nrows() != 4 { | ||
| return Err(ANNError::log_pq_error(format_args!( | ||
|
|
@@ -301,31 +305,53 @@ impl PQStorage { | |
| } | ||
| let file_offset_data = offsets.map(|x| x.into_usize()); | ||
|
|
||
| info!(" Offset data: {:?}", file_offset_data.as_slice()); | ||
|
|
||
| let pivots = read_bin_from::<f32>(&mut reader, file_offset_data[(0, 0)])?; | ||
| if let Some(num_centers) = expected_num_centers { | ||
| if pivots.nrows() != num_centers { | ||
| return Err(ANNError::log_pq_error(format_args!( | ||
| "Error reading pq_pivots file {}. file_num_centers = {}, but expecting {} centers.", | ||
| pq_pivots, | ||
| pivots.nrows(), | ||
| num_centers | ||
| ))); | ||
|
Comment on lines
+326
to
+331
Contributor
There was a problem hiding this comment. We're moving away from this pattern - #1282. How can we get ahead of this here? Or maybe we don't need to if 1282 is checked in first :) |
||
| } | ||
| } else if pivots.nrows() > NUM_PQ_CENTROIDS { | ||
| return Err(ANNError::log_pq_error(format_args!( | ||
| "Error reading pq_pivots file {}. file_num_centers = {}, but expecting {} centers.", | ||
| pq_pivots, | ||
| pivots.nrows(), | ||
| NUM_PQ_CENTROIDS | ||
| ))); | ||
| } | ||
|
|
||
| if let Some(dim) = expected_dim { | ||
| if pivots.ncols() != dim { | ||
| return Err(ANNError::log_pq_error(format_args!( | ||
| "Error reading pq_pivots file {}. file_dim = {} but expecting {} dimensions.", | ||
| pq_pivots, | ||
| pivots.ncols(), | ||
| dim | ||
| ))); | ||
| } | ||
| } | ||
|
|
||
| let centroid = read_bin_from::<f32>(&mut reader, file_offset_data[(1, 0)])?; | ||
| if centroid.nrows() != pivots.ncols() || centroid.ncols() != 1 { | ||
| return Err(ANNError::log_pq_error(format_args!( | ||
| "Error reading pq_pivots file {}. file_dim = {}, file_cols = {} \ | ||
| but expecting {} entries in 1 dimension.", | ||
| pq_pivots, | ||
| centroid.nrows(), | ||
| centroid.ncols(), | ||
| pivots.ncols() | ||
| ))); | ||
| } | ||
|
|
||
| let chunk_offsets_m = read_bin_from::<u32>(&mut reader, file_offset_data[(2, 0)])?; | ||
| if expected_num_pq_chunks | ||
| .is_some_and(|num_pq_chunks| chunk_offsets_m.nrows() != num_pq_chunks + 1) | ||
| || chunk_offsets_m.ncols() != 1 | ||
| { | ||
|
Comment on lines
365
to
368
|
||
| return Err(ANNError::log_pq_error(format_args!( | ||
|
|
@@ -334,18 +360,27 @@ impl PQStorage { | |
| passed as 0 if we want to infer.", | ||
| chunk_offsets_m.nrows(), | ||
| chunk_offsets_m.ncols(), | ||
| expected_num_pq_chunks.map_or(0, |num_pq_chunks| num_pq_chunks + 1) | ||
| ))); | ||
| } | ||
| let chunk_offsets = chunk_offsets_m.map(|x| x.into_usize()); | ||
|
|
||
| let parts = PivotFileParts { | ||
| pivots, | ||
| centroid, | ||
| chunk_offsets, | ||
| }; | ||
| if parts.chunk_offsets.nrows() < 2 | ||
| || parts.chunk_offsets[(0, 0)] != 0 | ||
| || parts.chunk_offsets[(parts.chunk_offsets.nrows() - 1, 0)] != parts.dim() | ||
| { | ||
| return Err(ANNError::log_pq_error(format_args!( | ||
| "Error reading pq_pivots file at chunk offsets; chunk offsets must start at 0, end at dim {}, and contain at least two entries.", | ||
| parts.dim() | ||
| ))); | ||
| } | ||
|
|
||
| Ok(parts) | ||
| } | ||
|
|
||
| /// streams data from the file, and samples each vector with probability p_val | ||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.