JingsongLi commented on code in PR #800:
URL: https://github.com/apache/paimon-rust/pull/800#discussion_r3986220832
##########
crates/paimon/src/table/vindex_index_build_builder/writer.rs:
##########
@@ -87,245 +93,571 @@ impl<'a> VindexIndexBuildBuilder<'a> {
let expected_bytes = checked_vector_bytes(row_count_usize,
dimension_usize)?;
let training_vector_count =
checked_training_vector_count(row_count_usize,
options.train_sample_ratio)?;
- let training_buffer_rows =
- (VECTOR_BUFFER_BYTES / checked_vector_bytes(1,
dimension_usize)?).max(1);
- let training_buffer_floats = training_buffer_rows
- .checked_mul(dimension_usize)
- .ok_or_else(|| Error::DataInvalid {
- message: "vindex training buffer length overflows
usize".to_string(),
- source: None,
- })?;
+ let training_rows_retained = if self.index_type == DISKANN_IDENTIFIER {
+ 0
+ } else {
+ default_training_vector_count(training_vector_count,
options.config.nlist()).map_err(
+ |e| Error::DataInvalid {
+ message: format!("Failed to calculate IVF training vector
count: {e}"),
+ source: Some(Box::new(e)),
+ },
+ )?
+ };
+ let mut sparse_ranges =
+ if training_rows_retained > 0 && training_rows_retained <
row_count_usize {
+ Some(plan_ivf_training_ranges(shard, training_rows_retained)?)
+ } else {
+ None
+ };
+ let capability_start = (timing_enabled &&
sparse_ranges.is_some()).then(Instant::now);
+ if let Some(ranges) = sparse_ranges.as_ref() {
+ let mut checks = Vec::new();
+ let mut usable = true;
+ for file in &shard.files {
+ let path = file.data_file_path(&shard.bucket_path);
+ if !path.to_ascii_lowercase().ends_with(".parquet") {
+ usable = false;
+ break;
+ }
+ let Some((file_start, file_end)) = file.row_id_range() else {
+ usable = false;
+ break;
+ };
+ let local_ranges = ranges
+ .iter()
+ .filter_map(|range| {
+ let from = range.from().max(file_start);
+ let to = range.to().min(file_end);
+ (from <= to).then(|| RowRange::new(from - file_start,
to - file_start))
+ })
+ .collect::<Vec<_>>();
+ if local_ranges.is_empty()
+ || file
+ .write_cols
+ .as_ref()
+ .is_some_and(|columns| !columns.iter().any(|column|
column == index_column))
+ {
+ continue;
+ }
+ let file_size = u64::try_from(file.file_size).map_err(|e|
Error::DataInvalid {
+ message: format!(
+ "Invalid data file size for '{}': {}",
+ file.file_name, file.file_size
+ ),
+ source: Some(Box::new(e)),
+ })?;
+ checks.push((path, file_size, local_ranges));
+ }
+ let found_vector_file = !checks.is_empty();
+ if usable && found_vector_file {
+ let concurrency = self
+ .table
+ .schema()
+ .core_options()
+ .parquet_row_group_parallelism()?;
+ let file_io = self.table.file_io();
+ let mut checks = futures::stream::iter(checks)
+ .map(|(path, file_size, local_ranges)| async move {
+ let input = file_io.new_input(&path)?;
+ has_usable_offset_index(
+ Box::new(input.reader().await?),
+ file_size,
+ index_column,
+ &local_ranges,
+ )
+ .await
Review Comment:
Rechecked at 28d6670: this remains reproducible with the current 512 short,
jittered ranges. For the same default 100,000-row shard with 128-dimensional
vectors, the sample and full scan still each fetch exactly 47,788,791 data
bytes, excluding metadata. The latest concurrency and selected-page budget
changes do not prevent the extra effectively full source pass. The sparse-path
gate at `writer.rs:173` still only checks offset-index availability; please
include actual page savings after range coalescing, or fall back when the
sample reads essentially the entire projection.
##########
crates/paimon/src/table/vindex_index_build_builder/planning.rs:
##########
@@ -45,3 +51,78 @@ pub(super) fn plan_vindex_shards(
"vindex",
)
}
+
+pub(super) fn plan_ivf_training_ranges(
+ shard: &VindexIndexShard,
+ training_rows: usize,
+) -> Result<Option<Vec<RowRange>>> {
+ let shard_rows = usize::try_from(checked_row_count(
+ shard.row_range_start,
+ shard.row_range_end,
+ )?)
+ .map_err(|error| Error::DataInvalid {
+ message: "vindex shard row count does not fit usize".to_string(),
+ source: Some(Box::new(error)),
+ })?;
+ if training_rows == 0 || training_rows > shard_rows {
+ return Err(Error::DataInvalid {
+ message: format!(
+ "Invalid IVF training row count: {training_rows}; shard
contains {shard_rows} rows"
+ ),
+ source: None,
+ });
+ }
+ let range_count = training_rows.div_ceil(MAX_IVF_TRAINING_RANGE_ROWS);
Review Comment:
[P2] Keep small training samples distributed across the shard
When `training_rows <= 128`, this calculation produces just one contiguous
range. A realistic incremental shard with 1,000 new rows and
`train.sample-ratio=0.1` therefore trains on 100 adjacent rows, whereas the
previous implementation sampled every tenth row across the whole shard.
Using the actual planner and vindex 0.4 trainer, I reproduced this with 450
vectors `[0]`, 450 `[1]`, and 100 `[100]`, IVF-SQ, L2, and `nlist=nprobe=1`.
For snapshot 1, bucket 0, and an empty partition, the new range is `[385,484]`,
which excludes the entire final distribution. Querying `[100]` returns 10/10
results from the final cluster with the baseline, but 0/10 with this plan: SQ
clamps those vectors to the older upper bound and returns rows 450–459. Please
retain multiple strata for small samples, or fall back to the original sampling
path, and cover this incremental-shard case in a recall regression test.
##########
crates/paimon/src/arrow/format/parquet.rs:
##########
@@ -642,13 +729,50 @@ impl FormatFileReader for ParquetFormatReader {
}
}
-fn projected_row_group_bytes(row_group: &RowGroupMetaData, projection:
&ProjectionMask) -> u64 {
+fn projected_row_group_bytes(
+ row_group: &RowGroupMetaData,
+ projection: &ProjectionMask,
+ selection: Option<&RowSelection>,
+ offset_index: Option<&[OffsetIndexMetaData]>,
+) -> u64 {
row_group
.columns()
.iter()
.enumerate()
.filter(|(leaf_index, _)| projection.leaf_included(*leaf_index))
- .filter_map(|(_, column)|
u64::try_from(column.uncompressed_size()).ok())
+ .filter_map(|(leaf_index, column)| {
+ let uncompressed_bytes =
u64::try_from(column.uncompressed_size()).ok()?;
+ let selected_bytes = selection
+ .zip(offset_index.and_then(|index| index.get(leaf_index)))
+ .and_then(|(selection, page_index)| {
+ let page_locations = page_index.page_locations();
+ let column_start = u64::try_from(
+ column
+ .dictionary_page_offset()
+ .unwrap_or_else(|| column.data_page_offset()),
+ )
+ .ok()?;
+ let dictionary_bytes =
u64::try_from(page_locations.first()?.offset)
+ .ok()?
+ .checked_sub(column_start)?;
+ let selected_compressed_bytes = selection
+ .scan_ranges(page_locations)
+ .into_iter()
+ .try_fold(dictionary_bytes, |total, range| {
+
total.checked_add(range.end.checked_sub(range.start)?)
+ })?;
Review Comment:
[P2] Include retained coalesced buffers in sparse read admission
`scan_ranges` sums the selected pages before `ArrowFileReader` merges byte
ranges separated by at most 1 MiB. Its returned `Bytes` slices retain the
larger coalesced allocations, but the new concurrent row-group admission uses
only the smaller selected-page estimate. This leaves live source buffers out of
the memory estimate used to increase parallelism.
I verified this with four valid Parquet row groups of 16,384 rows × 128
floats, default Zstd/page settings, and one selected row every 4,096 rows. With
a 20,971,520-byte budget, all four groups were admitted and simultaneously
retained 25,390,984 bytes of owned read buffers. Allocation ownership and
release were tracked with `Bytes::from_owner` and `Drop`; these counts exclude
decoded Arrow arrays, and all tracked buffers were released on completion.
Please account for the coalesced allocations when calculating admission costs,
or retain conservative full-column admission when they cannot be estimated.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
