fix: respect fragment filters during distributed worker training (#6358)

This fixes a bug in distributed vector index builds where any worker
that fell back to builder-local training sampled training data from the
entire dataset instead of the worker's selected fragments. The change
threads `fragment_filter` through IVF and quantizer training, and adds
an `IvfSq` regression test that verifies worker-local SQ bounds are
derived only from the shard fragments.

Author: Xuanwo

rust/lance/src/index/vector/builder.rs

@@ -385,7 +385,7 @@ impl<S: IvfSubIndex + 'static, Q: Quantization + 'static> IvfIndexBuilder<S, Q>
                     dim,
                     self.distance_type,
                     ivf_params,
-                    None,
+                    self.fragment_filter.as_deref(),
                     self.progress.clone(),
                 )
                 .await
@@ -414,9 +414,13 @@ impl<S: IvfSubIndex + 'static, Q: Quantization + 'static> IvfIndexBuilder<S, Q>
             "loading training data for quantizer. sample size: {}",
             sample_size_hint
         );
-        let training_data =
-            utils::maybe_sample_training_data(dataset, &self.column, sample_size_hint, None)
-                .await?;
+        let training_data = utils::maybe_sample_training_data(
+            dataset,
+            &self.column,
+            sample_size_hint,
+            self.fragment_filter.as_deref(),
+        )
+        .await?;
         info!(
             "Finished loading training data in {:02} seconds",
             start.elapsed().as_secs_f32()

rust/lance/src/index/vector/ivf/v2.rs

@@ -660,7 +660,9 @@ mod tests {
     use lance_index::vector::quantizer::QuantizerMetadata;
     use lance_index::vector::sq::builder::SQBuildParams;
     use lance_index::vector::{
-        pq::storage::ProductQuantizationMetadata, storage::STORAGE_METADATA_KEY,
+        pq::storage::ProductQuantizationMetadata,
+        sq::storage::{SQ_METADATA_KEY, ScalarQuantizationMetadata},
+        storage::STORAGE_METADATA_KEY,
     };
     use lance_index::{INDEX_AUXILIARY_FILE_NAME, metrics::NoOpMetricsCollector};
     use lance_index::{optimize::OptimizeOptions, scalar::IndexReader};
@@ -767,6 +769,37 @@ mod tests {
         serde_json::from_str(&metadata_entries[0]).unwrap()
     }
 
+    async fn get_sq_metadata(
+        dataset: &Dataset,
+        scheduler: Arc<ScanScheduler>,
+        index_uuid: &str,
+    ) -> ScalarQuantizationMetadata {
+        let index_path = dataset
+            .indices_dir()
+            .child(index_uuid)
+            .child(INDEX_AUXILIARY_FILE_NAME);
+        let file_scheduler = scheduler
+            .open_file(&index_path, &CachedFileSize::unknown())
+            .await
+            .unwrap();
+        let reader = FileReader::try_open(
+            file_scheduler,
+            None,
+            Arc::<DecoderPlugins>::default(),
+            &LanceCache::no_cache(),
+            FileReaderOptions::default(),
+        )
+        .await
+        .unwrap();
+        if let Some(metadata) = reader.schema().metadata.get(SQ_METADATA_KEY) {
+            serde_json::from_str(metadata).unwrap()
+        } else {
+            let metadata = reader.schema().metadata.get(STORAGE_METADATA_KEY).unwrap();
+            let metadata_entries: Vec<String> = serde_json::from_str(metadata).unwrap();
+            serde_json::from_str(&metadata_entries[0]).unwrap()
+        }
+    }
+
     async fn assert_rq_rotation_type(dataset: &Dataset, expected: RQRotationType) {
         let obj_store = Arc::new(ObjectStore::local());
         let scheduler = ScanScheduler::new(obj_store, SchedulerConfig::default_for_testing());
@@ -948,6 +981,49 @@ mod tests {
         )
     }
 
+    fn make_fragment_offset_batches(
+        rows_per_fragment: usize,
+        offsets: &[f32],
+    ) -> (Arc<Schema>, Vec<RecordBatch>) {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("id", DataType::UInt64, false),
+            Field::new(
+                "vector",
+                DataType::FixedSizeList(
+                    Arc::new(Field::new("item", DataType::Float32, true)),
+                    DIM as i32,
+                ),
+                false,
+            ),
+        ]));
+
+        let mut next_id = 0_u64;
+        let batches = offsets
+            .iter()
+            .map(|offset| {
+                let ids = Arc::new(UInt64Array::from_iter_values(
+                    next_id..next_id + rows_per_fragment as u64,
+                ));
+                next_id += rows_per_fragment as u64;
+
+                let mut values = Vec::with_capacity(rows_per_fragment * DIM);
+                for _ in 0..rows_per_fragment {
+                    for dim in 0..DIM {
+                        values.push(*offset + dim as f32);
+                    }
+                }
+
+                let vectors = Arc::new(
+                    FixedSizeListArray::try_new_from_values(Float32Array::from(values), DIM as i32)
+                        .unwrap(),
+                );
+                RecordBatch::try_new(schema.clone(), vec![ids, vectors]).unwrap()
+            })
+            .collect();
+
+        (schema, batches)
+    }
+
     struct VectorIndexTestContext {
         stats_json: String,
         stats: serde_json::Value,
@@ -2164,6 +2240,59 @@ mod tests {
             .unwrap();
         assert!(result.num_rows() > 0);
     }
+
+    #[tokio::test]
+    async fn test_distributed_ivf_sq_worker_training_respects_fragment_filter() {
+        const ROWS_PER_FRAGMENT: usize = 64;
+        const FRAGMENT_OFFSETS: [f32; 2] = [0.0, 1000.0];
+
+        let test_dir = TempStrDir::default();
+        let dataset_uri = format!("{}/distributed_sq_fragment_filter", test_dir.as_str());
+        let (schema, batches) = make_fragment_offset_batches(ROWS_PER_FRAGMENT, &FRAGMENT_OFFSETS);
+        let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema);
+        let mut dataset = Dataset::write(
+            batches,
+            &dataset_uri,
+            Some(WriteParams {
+                max_rows_per_file: ROWS_PER_FRAGMENT,
+                mode: WriteMode::Overwrite,
+                ..Default::default()
+            }),
+        )
+        .await
+        .unwrap();
+
+        let fragments = dataset.get_fragments();
+        assert_eq!(fragments.len(), FRAGMENT_OFFSETS.len());
+
+        let ivf_params =
+            IvfBuildParams::try_with_centroids(2, build_centroids_for_offsets(&FRAGMENT_OFFSETS))
+                .unwrap();
+        let params = VectorIndexParams::with_ivf_sq_params(
+            DistanceType::L2,
+            ivf_params,
+            SQBuildParams::default(),
+        );
+
+        let segment = dataset
+            .create_index_builder(&["vector"], IndexType::Vector, &params)
+            .name("sq_fragment_filter".to_string())
+            .fragments(vec![fragments[0].id() as u32])
+            .execute_uncommitted()
+            .await
+            .unwrap();
+
+        let scheduler = ScanScheduler::new(
+            Arc::new(dataset.object_store().clone()),
+            SchedulerConfig::default_for_testing(),
+        );
+        let sq_meta = get_sq_metadata(&dataset, scheduler, &segment.uuid.to_string()).await;
+
+        assert_eq!(sq_meta.bounds.start, 0.0);
+        assert_eq!(sq_meta.bounds.end, (DIM - 1) as f64);
+        assert_lt!(sq_meta.bounds.end, FRAGMENT_OFFSETS[1] as f64);
+    }
+
     async fn test_index(
         params: VectorIndexParams,
         nlist: usize,