Skip to content

perf(compaction): skip building row-address maps when index remapping is not needed #7778

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
everySympathy wants to merge 1 commit into
base: main
Choose a base branch
from
+177 −41
Open

perf(compaction): skip building row-address maps when index remapping is not needed #7778

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 75 additions & 12 deletions rust/lance/src/dataset/index.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use arrow_schema::DataType;
use async_trait::async_trait;
use lance_core::{Error, Result};
use lance_encoding::version::LanceFileVersion;
use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME;
use lance_index::is_system_index;
use lance_index::pb::VectorIndexDetails;
use lance_index::scalar::lance_format::LanceIndexStore;
use lance_table::format::IndexMetadata;
Expand All @@ -29,19 +29,39 @@ use super::optimize::{IndexRemapper, IndexRemapperOptions};
#[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct DatasetIndexRemapperOptions {}

/// Loads index metadata when compaction has at least one index to remap.
///
/// Returns all index metadata, including system indices, so the remapper uses a
/// consistent snapshot. Returns `None` when there are no non-system indices.
pub(crate) async fn load_indices_for_remapping(
dataset: &Dataset,
) -> Result<Option<Arc<Vec<IndexMetadata>>>> {
if dataset.manifest.index_section.is_none() {
return Ok(None);
}

let indices = dataset.load_indices().await?;
let has_remappable_index = indices.iter().any(|index| !is_system_index(index));
Ok(has_remappable_index.then_some(indices))
}

#[async_trait]
impl IndexRemapperOptions for DatasetIndexRemapperOptions {
fn create_remapper(
&self,
dataset: &Dataset,
) -> crate::Result<Box<dyn super::optimize::IndexRemapper>> {
Ok(Box::new(DatasetIndexRemapper {
async fn create_remapper(&self, dataset: &Dataset) -> Result<Option<Box<dyn IndexRemapper>>> {
let Some(indices) = load_indices_for_remapping(dataset).await? else {
return Ok(None);
};
Comment thread
coderabbitai[bot] marked this conversation as resolved.

Ok(Some(Box::new(DatasetIndexRemapper {
dataset: Arc::new(dataset.clone()),
}))
indices,
})))
}
}

struct DatasetIndexRemapper {
dataset: Arc<Dataset>,
indices: Arc<Vec<IndexMetadata>>,
}

impl DatasetIndexRemapper {
Expand All @@ -62,10 +82,9 @@ impl IndexRemapper for DatasetIndexRemapper {
affected_fragment_ids: &[u64],
) -> Result<Vec<RemappedIndex>> {
let affected_frag_ids = HashSet::<u64>::from_iter(affected_fragment_ids.iter().copied());
let indices = self.dataset.load_indices().await?;
let mut remapped = Vec::with_capacity(indices.len());
for index in indices.iter() {
let needs_remapped = index.name != FRAG_REUSE_INDEX_NAME
let mut remapped = Vec::with_capacity(self.indices.len());
for index in self.indices.iter() {
let needs_remapped = !is_system_index(index)
&& match &index.fragment_bitmap {
None => true,
Some(fragment_bitmap) => fragment_bitmap
Expand Down Expand Up @@ -178,13 +197,55 @@ impl LanceIndexStoreExt for LanceIndexStore {
mod tests {
use super::*;
use crate::dataset::WriteParams;
use crate::dataset::transaction::{Operation, Transaction};
use crate::index::DatasetIndexExt;
use crate::index::frag_reuse::build_frag_reuse_index_metadata;
use crate::index::vector::VectorIndexParams;
use lance_datagen::{BatchCount, RowCount, array};
use lance_index::IndexType;
use lance_index::frag_reuse::{FRAG_REUSE_INDEX_NAME, FragReuseIndexDetails};
use lance_linalg::distance::MetricType;
use uuid::Uuid;

#[tokio::test]
async fn test_remapper_not_created_without_remappable_indices() {
let reader = lance_datagen::gen_batch()
.col("id", array::step::<arrow_array::types::Int32Type>())
.into_reader_rows(RowCount::from(1), BatchCount::from(1));
let mut dataset = Dataset::write(reader, "memory://", None).await.unwrap();
let options = DatasetIndexRemapperOptions::default();

assert!(options.create_remapper(&dataset).await.unwrap().is_none());

let frag_reuse_index = build_frag_reuse_index_metadata(
&dataset,
None,
FragReuseIndexDetails {
versions: Vec::new(),
},
Default::default(),
)
.await
.unwrap();
let transaction = Transaction::new(
dataset.manifest.version,
Operation::CreateIndex {
new_indices: vec![frag_reuse_index],
removed_indices: Vec::new(),
},
None,
);
dataset
.apply_commit(transaction, &Default::default(), &Default::default())
.await
.unwrap();

let indices = dataset.load_indices().await.unwrap();
assert_eq!(indices.len(), 1);
assert_eq!(indices[0].name, FRAG_REUSE_INDEX_NAME);
assert!(options.create_remapper(&dataset).await.unwrap().is_none());
}

#[tokio::test]
async fn test_remapper_only_touches_segments_with_affected_fragments() {
let test_dir = tempfile::tempdir().unwrap();
Expand Down Expand Up @@ -296,7 +357,9 @@ mod tests {

let remapper = DatasetIndexRemapperOptions::default()
.create_remapper(&dataset)
.unwrap();
.await
.unwrap()
.expect("vector index should require a remapper");
let remapped = remapper
.remap_indices(RowAddrRemap::empty(), &[target_fragments[0].id() as u64])
.await
Expand Down
119 changes: 93 additions & 26 deletions rust/lance/src/dataset/optimize.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ use std::ops::{AddAssign, Range};
use std::sync::Arc;

use super::fragment::FileFragment;
use super::index::DatasetIndexRemapperOptions;
use super::index::{DatasetIndexRemapperOptions, load_indices_for_remapping};
use super::rowids::load_row_id_sequences;
use super::transaction::{
Operation, RewriteGroup, RewrittenIndex, Transaction, TransactionBuilder,
Expand Down Expand Up @@ -1608,8 +1608,13 @@ async fn rewrite_files(
.iter()
.map(|f| f.physical_rows.unwrap() as u64)
.sum::<u64>();
// If we aren't using stable row ids, then we need to remap indices.
let needs_remapping = !dataset.manifest.uses_stable_row_ids();
// Capturing row addresses is only useful if something will consume them:
// an index to remap now, or a deferred remap through the FRI.
let capture_row_addrs = !dataset.manifest.uses_stable_row_ids()
&& (options.defer_index_remap
|| load_indices_for_remapping(dataset.as_ref())
.await?
.is_some());
let mut new_fragments: Vec<Fragment>;
let task_id = uuid::Uuid::new_v4();
log::info!(
Expand All @@ -1635,7 +1640,7 @@ async fn rewrite_files(
options.batch_size,
options.io_buffer_size,
true,
needs_remapping,
capture_row_addrs,
)
.await?;
row_ids_rx = rx_initial;
Expand Down Expand Up @@ -1739,7 +1744,7 @@ async fn rewrite_files(
));
}

if needs_remapping {
if capture_row_addrs {
let (tx, rx) = std::sync::mpsc::channel();
let mut addrs = RoaringTreemap::new();
for frag in &fragments {
Expand Down Expand Up @@ -2009,8 +2014,17 @@ pub async fn commit_compaction(
return Ok(CompactionMetrics::default());
}

// If we aren't using stable row ids, then we need to remap indices.
let needs_remapping = !dataset.manifest.uses_stable_row_ids() && !options.defer_index_remap;
let has_address_style = completed_tasks.iter().any(|t| t.row_addrs.is_some());
// Address-style results require immediate index remapping unless it is deferred.
let needs_remapping =
!dataset.manifest.uses_stable_row_ids() && !options.defer_index_remap && has_address_style;

// Confirm there is a remapper before materializing the potentially very large row address map.
let index_remapper = if needs_remapping {
remap_options.create_remapper(dataset).await?
} else {
None
};

// Determine the earliest version at which compaction tasks were planned/executed.
//
Expand All @@ -2034,7 +2048,6 @@ pub async fn commit_compaction(
let mut completed_tasks = completed_tasks;

// Single reserve_fragment_ids for all address-style tasks
let has_address_style = completed_tasks.iter().any(|t| t.row_addrs.is_some());
if has_address_style {
let frags: Vec<&mut Fragment> = completed_tasks
.iter_mut()
Expand Down Expand Up @@ -2084,7 +2097,7 @@ pub async fn commit_compaction(
new_fragments: task.new_fragments.clone(),
};

if needs_remapping {
if index_remapper.is_some() {
if let Some(row_addrs_bytes) = task.row_addrs {
let row_addrs =
RoaringTreemap::deserialize_from(&mut Cursor::new(&row_addrs_bytes))?;
Expand Down Expand Up @@ -2151,8 +2164,7 @@ pub async fn commit_compaction(
rewrite_groups.push(rewrite_group);
}

let rewritten_indices = if needs_remapping {
let index_remapper = remap_options.create_remapper(dataset)?;
let rewritten_indices = if let Some(index_remapper) = index_remapper {
let affected_ids = rewrite_groups
.iter()
.flat_map(|group| group.old_fragments.iter().map(|frag| frag.id))
Expand Down Expand Up @@ -2448,9 +2460,10 @@ mod tests {
}
}

#[async_trait]
impl IndexRemapperOptions for MockIndexRemapper {
fn create_remapper(&self, _: &Dataset) -> Result<Box<dyn IndexRemapper>> {
Ok(Box::new(self.clone()))
async fn create_remapper(&self, _: &Dataset) -> Result<Option<Box<dyn IndexRemapper>>> {
Ok(Some(Box::new(self.clone())))
}
}

Expand Down Expand Up @@ -2979,9 +2992,70 @@ mod tests {
}
}

#[async_trait]
impl IndexRemapperOptions for IgnoreRemap {
fn create_remapper(&self, _: &Dataset) -> Result<Box<dyn IndexRemapper>> {
Ok(Box::new(Self {}))
async fn create_remapper(&self, _: &Dataset) -> Result<Option<Box<dyn IndexRemapper>>> {
Ok(None)
}
}

#[rstest]
#[case::without_index(false)]
#[case::with_index(true)]
#[tokio::test]
async fn test_row_addrs_only_used_with_remappable_index(#[case] has_index: bool) {
let data = sample_data();
let reader = RecordBatchIterator::new(vec![Ok(data.slice(0, 9_000))], data.schema());
let mut dataset = Dataset::write(
reader,
"memory://",
Some(WriteParams {
max_rows_per_file: 3_000,
data_storage_version: Some(LanceFileVersion::Legacy),
..Default::default()
}),
)
.await
.unwrap();

if has_index {
create_scalar_index(&mut dataset, "a", false).await;
}

let options = CompactionOptions {
target_rows_per_fragment: 9_000,
..Default::default()
};
let plan = plan_compaction(&dataset, &options).await.unwrap();
assert_eq!(plan.tasks().len(), 1);

let mut result = rewrite_files(Cow::Borrowed(&dataset), plan.tasks()[0].clone(), &options)
.await
.unwrap();
assert_eq!(result.row_addrs.is_some(), has_index);

if has_index {
let row_addrs_bytes = result
.row_addrs
.as_ref()
.expect("indexed compaction should capture row addresses");
let row_addrs =
RoaringTreemap::deserialize_from(&mut Cursor::new(row_addrs_bytes)).unwrap();
assert_eq!(row_addrs.len(), 9_000);
} else {
// Simulate a stale worker result that captured row addresses before the
// dataset no longer needed a remapper. Invalid bytes ensure the commit
// does not attempt to deserialize or materialize the unused map.
result.row_addrs = Some(b"not a roaring treemap".to_vec());
commit_compaction(
&mut dataset,
vec![result],
Arc::new(DatasetIndexRemapperOptions::default()),
&options,
)
.await
.unwrap();
assert_eq!(dataset.get_fragments().len(), 1);
}
}

Expand Down Expand Up @@ -3322,17 +3396,10 @@ mod tests {
dataset.delete("i < 500").await.unwrap();
dataset2.delete("i < 500").await.unwrap();

// Create a scalar index to check this is not touched
dataset
.create_index(
&["i"],
IndexType::Scalar,
Some("scalar".into()),
&ScalarIndexParams::default(),
false,
)
.await
.unwrap();
// Create the same scalar index on both datasets so deferred and immediate
// remapping are compared under the same conditions.
create_scalar_index(&mut dataset, "i", false).await;
create_scalar_index(&mut dataset2, "i", false).await;

// Verify the initial state - no fragment reuse index should exist
let initial_indices = dataset.load_indices().await.unwrap();
Expand Down
12 changes: 9 additions & 3 deletions rust/lance/src/dataset/optimize/remapping.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,8 +61,13 @@ pub trait IndexRemapper: Send + Sync {
///
/// Currently we don't have any options but we may need options in the future and so we
/// want to keep a placeholder
#[async_trait]
pub trait IndexRemapperOptions: Send + Sync {
fn create_remapper(&self, dataset: &Dataset) -> Result<Box<dyn IndexRemapper>>;
/// Creates a remapper when the dataset has indices that need row address remapping.
///
/// Returns `None` when no remappable indices exist, allowing compaction to avoid
/// materializing an unused row address map.
async fn create_remapper(&self, dataset: &Dataset) -> Result<Option<Box<dyn IndexRemapper>>>;
Comment on lines +64 to +70

@coderabbitai coderabbitai Bot Jul 16, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
ast-grep outline rust/lance/src/dataset/optimize/remapping.rs --items all
rg -n -C3 'IndexRemapperOptions|create_remapper' rust

Repository: lance-format/lance

Length of output: 18111

🏁 Script executed:

#!/bin/bash
set -euo pipefail

sed -n '1,120p' rust/lance/src/dataset/optimize/remapping.rs
printf '\n--- optimize.rs export ---\n'
sed -n '120,150p' rust/lance/src/dataset/optimize.rs
printf '\n--- index.rs impl ---\n'
sed -n '1,70p' rust/lance/src/dataset/index.rs
printf '\n--- search deprecated/compat hook ---\n'
rg -n -C2 'deprecated|compat|create_remapper|IndexRemapperOptions' rust/lance/src/dataset

Repository: lance-format/lance

Length of output: 50376

Preserve the existing remapper trait method for downstream implementations.
This public trait is re-exported and implemented outside this module; changing create_remapper to async and Option<Box<_>> is a source break for external implementors and callers. Keep the current method as a deprecated compatibility hook and add a new async optional method with a default wrapper, or introduce a new trait.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@rust/lance/src/dataset/optimize/remapping.rs` around lines 64 - 70, Preserve
the existing public IndexRemapperOptions::create_remapper signature for
downstream implementors and callers, marking it deprecated if needed. Add a
separate async optional remapper method with a default implementation that
delegates to the compatibility hook, and update internal compaction code to use
the new method without breaking existing trait implementations.

Source: Coding guidelines

}

#[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)]
Expand All @@ -75,9 +80,10 @@ impl IndexRemapper for IgnoreRemap {
}
}

#[async_trait]
impl IndexRemapperOptions for IgnoreRemap {
fn create_remapper(&self, _: &Dataset) -> Result<Box<dyn IndexRemapper>> {
Ok(Box::new(Self {}))
async fn create_remapper(&self, _: &Dataset) -> Result<Option<Box<dyn IndexRemapper>>> {
Ok(None)
}
}

Expand Down
Loading