feat: build ivf partition using disk based shuffler #1312

eddyxu · 2023-09-22T02:30:52Z

Reduce the memory consumption to run IVF_PQ
Building block for distributed IVF_PQ indexing
Better multi-thread support during IVF assignment

westonpace

Some thoughts from a partial review.

westonpace · 2023-09-27T18:25:03Z

rust/lance-linalg/src/kmeans.rs

@@ -190,64 +185,38 @@ pub struct KMeanMembership {
 impl KMeanMembership {
    /// Reconstruct a KMeans model from the membership.
    async fn to_kmeans(&self) -> Result<KMeans> {
+        let time = std::time::Instant::now();


Maybe #[instrument] the method instead?

westonpace · 2023-09-27T18:34:11Z

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

+                        .iter()
+                        .zip(centroid.values().iter())
+                        .map(|(v, c)| *v - *c)
+                        .collect::<Vec<_>>() // How to avoid one memory allocation here?


flat_map should accept a function that returns an iterator. Can you just remove the collect? Maybe you need to do .copied() to convert from Iterator<Item = &f32> to Iterator<Item = f32>?

It complaints returns a value referencing data owned by the current function (centroid) tho.

Ok, fixed in another way: using foreach to append to pre-allocated vector.

westonpace · 2023-09-27T18:35:55Z

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

@@ -649,6 +635,7 @@ pub async fn build_ivf_pq_index(
    #[cfg(not(feature = "opq"))]
    let transforms: Vec<Box<dyn Transformer>> = vec![];

+    let start = std::time::Instant::now();


Can we use a span?

westonpace · 2023-09-27T20:15:15Z

rust/lance-linalg/src/kmeans.rs

+                for i in 0..dimension {
+                    new_centroids[*c as usize * dimension + i] += v[i];
+                }


This might be a bit faster (especially for large dimension) as it can avoid some bounds checking.

Suggested change

for i in 0..dimension {

new_centroids[*c as usize * dimension + i] += v[i];

}

for (old, new) in new_centroids[(*c * dimension)..((*c + 1) * dimension)]

.iter_mut()

.zip(v)

{

*old += new;

}

wjones127 · 2023-09-27T21:14:10Z

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

+        .and_then(|batch_and_range| async move {
+            // Split batch into per-partition batches
+            let (batch, range) = batch_and_range;
+            Ok(stream::iter(range).map(move |part_id| {
+                let predictions = BooleanArray::from_unary(
+                    batch
+                        .column_by_name(PARTITION_ID_COLUMN)
+                        .unwrap()
+                        .as_primitive::<UInt32Type>(),
+                    |pid| pid == part_id,
+                );
+                let parted_batch =
+                    filter_record_batch(&batch, &predictions)?.drop_column(PARTITION_ID_COLUMN)?;
+                Ok::<(u32, RecordBatch), Error>((part_id, parted_batch))
+            }))
+        })


This is CPU bound, so I think we need to put this into tokio::task::spawn_blocking() to get any concurrency.

I changed this piece of code.

wjones127 · 2023-09-27T21:20:23Z

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

+#[cfg(test)]
+mod tests {}


This going to be filled in? Or removed?

wjones127 · 2023-09-27T21:26:04Z

rust/lance-linalg/src/kmeans.rs

+            .values()
+            .chunks_exact(dimension)
+            .zip(self.cluster_ids.iter())
+            .for_each(|(v, c)| {


nit: using single variable names makes this harder to read. (I had to cross reference different parts of the file to figure out what they meant.)

Suggested change

.for_each(|(v, c)| {

.for_each(|(vector, cluster_id)| {

chebbyChefNEQ · 2023-09-27T22:02:16Z

rust/lance-linalg/src/kmeans.rs

+            .for_each(|(v, c)| {
+                cluster_cnts[*c as usize] += 1;
+                for i in 0..dimension {
+                    new_centroids[*c as usize * dimension + i] += v[i];


this feels like it can be an overflow hazard. We probably haven't hit it because of sampling I'm guessing

Changed to use iterator

chebbyChefNEQ · 2023-09-27T22:08:52Z

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

+        0..num_partitions,
+    )
+    .await?;
+    println!("Building partitions: {}s", start.elapsed().as_secs_f32());


s/building/built

westonpace · 2023-09-28T03:33:53Z

rust/lance-linalg/src/kmeans.rs

+                warn!("KMeans: cluster {} is empty", i);
+            } else {
+                for j in 0..dimension {
+                    new_centroids[i * dimension + j] /= cluster_cnts[i] as f32;


Is there any possibility that cluster_cnts[i] is zero here? Maybe if K is larger than the # of data points?

Yea. the training algorithm does not work that well to make sure no cluster at the moment.

Part of #1312

parent 8775776 author Lei Xu <[email protected]> 1695332018 -0700 committer Lei Xu <[email protected]> 1695926481 -0700 Ivf::compute_partitions use instructment avoid memory copy during residual computation

wjones127 · 2023-10-02T19:50:37Z

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

+        let indices = UInt32Array::from(row_ids.clone());
+        // Use `take` to select rows.
+        let str_arr = take(&struct_arr, &indices, None)?;
+        let parted_batch: RecordBatch = str_arr.as_struct().into();


Don't we have that new take() method now?

You are right. Fixed.

eddyxu force-pushed the lei/ivf_build_partition branch from 8486ce2 to edb8e9c Compare September 25, 2023 16:41

eddyxu self-assigned this Sep 25, 2023

eddyxu added benchmark vector Vector Search labels Sep 25, 2023

eddyxu requested review from wjones127, chebbyChefNEQ and westonpace September 26, 2023 16:41

eddyxu force-pushed the lei/ivf_build_partition branch from 7ee145f to 63ffe33 Compare September 26, 2023 20:06

eddyxu mentioned this pull request Sep 27, 2023

feat: fused argmin and argmax to avoid unnecessary memory allocation #1325

Merged

eddyxu force-pushed the lei/ivf_build_partition branch from 9df4538 to 72d7c5c Compare September 27, 2023 05:00

eddyxu marked this pull request as ready for review September 27, 2023 05:15

eddyxu requested a review from QianZhu September 27, 2023 05:24

westonpace reviewed Sep 27, 2023

View reviewed changes

wjones127 requested changes Sep 27, 2023

View reviewed changes

chebbyChefNEQ reviewed Sep 27, 2023

View reviewed changes

eddyxu force-pushed the lei/ivf_build_partition branch from e4dd67f to c36e76d Compare September 27, 2023 22:25

westonpace reviewed Sep 28, 2023

View reviewed changes

This was referenced Sep 28, 2023

chore(rust): improve kmean training time complexity #1330

Merged

chore: convenience methods for reading and writing local files #1329

Merged

eddyxu added a commit that referenced this pull request Sep 28, 2023

chore(rust): improve kmean training time complexity (#1330)

8775776

Part of #1312

eddyxu force-pushed the lei/ivf_build_partition branch from 5777884 to e6fb3b3 Compare September 28, 2023 18:43

eddyxu marked this pull request as draft September 28, 2023 18:49

eddyxu force-pushed the lei/ivf_build_partition branch from e6fb3b3 to 91dd61d Compare September 28, 2023 19:02

eddyxu mentioned this pull request Sep 28, 2023

chore: ivf model improvement #1333

Closed

eddyxu force-pushed the lei/ivf_build_partition branch 2 times, most recently from 953f8c1 to efb6167 Compare September 29, 2023 20:26

eddyxu and others added 3 commits October 2, 2023 11:09

t # This is a combination of 10 commits.

c04e8ce

parent 8775776 author Lei Xu <[email protected]> 1695332018 -0700 committer Lei Xu <[email protected]> 1695926481 -0700 Ivf::compute_partitions use instructment avoid memory copy during residual computation

use path

2c85274

move read_batch as debug

7869cdd

eddyxu and others added 7 commits October 2, 2023 11:09

filter

a2ae8aa

batches

662bd30

sla

b2d96a3

sample pq code before running residual

f18a3d3

increase prefetch batch size

74cfa8c

fix write shuffler

4b9ac26

shuffler

27c863a

eddyxu force-pushed the lei/ivf_build_partition branch from b468322 to 27c863a Compare October 2, 2023 18:10

eddyxu added 2 commits October 2, 2023 11:17

try windows fix

c32770c

remove unused tests

55b80dd

eddyxu marked this pull request as ready for review October 2, 2023 18:37

wjones127 approved these changes Oct 2, 2023

View reviewed changes

use record take

065be04

eddyxu merged commit fbe5798 into main Oct 2, 2023
15 checks passed

eddyxu deleted the lei/ivf_build_partition branch October 2, 2023 21:02

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: build ivf partition using disk based shuffler #1312

feat: build ivf partition using disk based shuffler #1312

eddyxu commented Sep 22, 2023 •

edited

Loading

westonpace left a comment

westonpace Sep 27, 2023

eddyxu Sep 27, 2023

westonpace Sep 27, 2023

eddyxu Sep 27, 2023

eddyxu Sep 28, 2023

westonpace Sep 27, 2023

westonpace Sep 27, 2023

eddyxu Sep 28, 2023

wjones127 Sep 27, 2023

eddyxu Oct 2, 2023

wjones127 Sep 27, 2023

wjones127 Sep 27, 2023

eddyxu Sep 28, 2023

chebbyChefNEQ Sep 27, 2023 •

edited

Loading

eddyxu Sep 28, 2023

chebbyChefNEQ Sep 27, 2023

eddyxu Sep 28, 2023

westonpace Sep 28, 2023

eddyxu Sep 28, 2023

wjones127 Oct 2, 2023

eddyxu Oct 2, 2023

-                for i in 0..dimension {
-                    new_centroids[*c as usize * dimension + i] += v[i];
-                }
+            for (old, new) in new_centroids[(*c * dimension)..((*c + 1) * dimension)]
+                .iter_mut()
+                .zip(v)
+            {
+                *old += new;
+            }

		#[cfg(test)]
		mod tests {}

feat: build ivf partition using disk based shuffler #1312

feat: build ivf partition using disk based shuffler #1312

Conversation

eddyxu commented Sep 22, 2023 • edited Loading

westonpace left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chebbyChefNEQ Sep 27, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eddyxu commented Sep 22, 2023 •

edited

Loading

chebbyChefNEQ Sep 27, 2023 •

edited

Loading