perf(ivfpq): speed up transposed PQ8 scans

core/Cargo.toml

@@ -58,6 +58,10 @@ harness = false
 name = "ivfpq_batch_reuse_bench"
 harness = false
 
+[[bench]]
+name = "ivfpq_filter_scan_bench"
+harness = false
+
 [[bench]]
 name = "ivfpq_train_bench"
 harness = false

core/benches/ivfpq_filter_scan_bench.rs

@@ -0,0 +1,157 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use paimon_vindex_core::distance::MetricType;
+use paimon_vindex_core::io::{write_index, IVFPQIndexReader, PosWriter};
+use paimon_vindex_core::ivfpq::{search_batch_reader_filter, IVFPQIndex, RowIdFilter};
+use rand::rngs::StdRng;
+use rand::seq::SliceRandom;
+use rand::{Rng, SeedableRng};
+use std::hint::black_box;
+use std::io::Cursor;
+use std::time::{Duration, Instant};
+
+const D: usize = 768;
+const M: usize = 96;
+const NLIST: usize = 128;
+const NPROBE: usize = 64;
+const NQ: usize = 64;
+const K: usize = 3;
+const ROWS_PER_LIST: usize = 6_800;
+const WARMUPS: usize = 1;
+const ROUNDS: usize = 3;
+
+struct DensityFilter<'a> {
+    row_ranks: &'a [usize],
+    max_rank: usize,
+}
+
+impl RowIdFilter for DensityFilter<'_> {
+    fn contains(&self, id: i64) -> bool {
+        self.row_ranks
+            .get(id as usize)
+            .is_some_and(|&rank| rank < self.max_rank)
+    }
+}
+
+fn randomized_row_ranks() -> Vec<usize> {
+    let mut rng = StdRng::seed_from_u64(43);
+    let mut row_ranks = vec![0; NLIST * ROWS_PER_LIST];
+    let mut row_order = (0..ROWS_PER_LIST).collect::<Vec<_>>();
+    for list_id in 0..NLIST {
+        row_order.shuffle(&mut rng);
+        let base = list_id * ROWS_PER_LIST;
+        for (rank, &row) in row_order.iter().enumerate() {
+            row_ranks[base + row] = rank;
+        }
+    }
+    row_ranks
+}
+
+fn search(
+    reader: &mut IVFPQIndexReader<Cursor<Vec<u8>>>,
+    queries: &[f32],
+    filter: &DensityFilter,
+) -> Duration {
+    let started = Instant::now();
+    let result = search_batch_reader_filter(reader, queries, NQ, K, NPROBE, Some(filter)).unwrap();
+    let elapsed = started.elapsed();
+    black_box(result);
+    elapsed
+}
+
+fn median(samples: &mut [Duration]) -> Duration {
+    samples.sort_unstable();
+    samples[samples.len() / 2]
+}
+
+fn main() {
+    assert_eq!((NQ, NPROBE), (64, 64), "benchmark production query shape");
+    for name in [
+        "PAIMON_VINDEX_LOG_IVFPQ_BATCH_TIMING",
+        "PAIMON_VINDEX_LOG_IVFPQ_BATCH_REUSE",
+    ] {
+        assert!(
+            std::env::var_os(name).is_none(),
+            "unset {name} for this benchmark"
+        );
+    }
+    assert_eq!(
+        rayon::current_num_threads(),
+        1,
+        "run with RAYON_NUM_THREADS=1"
+    );
+    const { assert!(NQ >= 64, "keep production query-table reuse enabled") };
+
+    let mut rng = StdRng::seed_from_u64(42);
+    let mut index = IVFPQIndex::new(D, NLIST, M, MetricType::InnerProduct, false);
+    index.quantizer_centroids = (0..NLIST * D)
+        .map(|_| rng.gen_range(-1.0f32..1.0))
+        .collect();
+    index.pq.centroids = (0..M * index.pq.ksub * index.pq.dsub)
+        .map(|_| rng.gen_range(-1.0f32..1.0))
+        .collect();
+    for list_id in 0..NLIST {
+        let first_id = list_id * ROWS_PER_LIST;
+        index.ids[list_id] = (first_id..first_id + ROWS_PER_LIST)
+            .map(|id| id as i64)
+            .collect();
+        index.codes[list_id] = (0..ROWS_PER_LIST * M).map(|_| rng.gen()).collect();
+    }
+    let queries = (0..NQ * D)
+        .map(|_| rng.gen_range(-1.0f32..1.0))
+        .collect::<Vec<_>>();
+    let row_ranks = randomized_row_ranks();
+    let mut bytes = Vec::new();
+    write_index(&index, &mut PosWriter::new(&mut bytes)).unwrap();
+    let densities = [
+        (1, "6.25"),
+        (2, "12.50"),
+        (3, "18.75"),
+        (4, "25.00"),
+        (8, "50.00"),
+        (16, "100.00"),
+    ];
+
+    println!(
+        "shape: d={D} m={M} nlist={NLIST} nprobe={NPROBE} nq={NQ} k={K} rows_per_list={ROWS_PER_LIST} threads=1 warmups={WARMUPS} rounds={ROUNDS}"
+    );
+    println!("density_percent,p50_ms");
+    for (matching_sixteenths, density) in densities {
+        let filter = DensityFilter {
+            row_ranks: &row_ranks,
+            max_rank: ROWS_PER_LIST * matching_sixteenths / 16,
+        };
+        assert_eq!(
+            (0..ROWS_PER_LIST)
+                .filter(|&id| filter.contains(id as i64))
+                .count(),
+            filter.max_rank
+        );
+        let mut reader = IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap();
+        for _ in 0..WARMUPS {
+            search(&mut reader, &queries, &filter);
+        }
+        let mut samples = (0..ROUNDS)
+            .map(|_| search(&mut reader, &queries, &filter))
+            .collect::<Vec<_>>();
+        println!(
+            "{density},{:.3}",
+            median(&mut samples).as_secs_f64() * 1_000.0
+        );
+    }
+}

core/src/ivfpq.rs

@@ -1219,44 +1219,18 @@ fn scan_codes_transposed_with_scratch(
     }
 
     dists.resize(count, 0.0);
-    let table = &sim_table[..ksub];
-    let column = &codes[..count];
-    let mut row = 0usize;
-    while row + 8 <= count {
-        dists[row] = dis0 + table[column[row] as usize];
-        dists[row + 1] = dis0 + table[column[row + 1] as usize];
-        dists[row + 2] = dis0 + table[column[row + 2] as usize];
-        dists[row + 3] = dis0 + table[column[row + 3] as usize];
-        dists[row + 4] = dis0 + table[column[row + 4] as usize];
-        dists[row + 5] = dis0 + table[column[row + 5] as usize];
-        dists[row + 6] = dis0 + table[column[row + 6] as usize];
-        dists[row + 7] = dis0 + table[column[row + 7] as usize];
-        row += 8;
-    }
-    while row < count {
-        dists[row] = dis0 + table[column[row] as usize];
-        row += 1;
-    }
+    transposed_column_init(
+        &mut dists[..count],
+        &codes[..count],
+        &sim_table[..ksub],
+        dis0,
+    );
     for sub in 1..m {
-        let table = &sim_table[sub * ksub..(sub + 1) * ksub];
-        let col_base = sub * count;
-        let column = &codes[col_base..col_base + count];
-        let mut row = 0usize;
-        while row + 8 <= count {
-            dists[row] += table[column[row] as usize];
-            dists[row + 1] += table[column[row + 1] as usize];
-            dists[row + 2] += table[column[row + 2] as usize];
-            dists[row + 3] += table[column[row + 3] as usize];
-            dists[row + 4] += table[column[row + 4] as usize];
-            dists[row + 5] += table[column[row + 5] as usize];
-            dists[row + 6] += table[column[row + 6] as usize];
-            dists[row + 7] += table[column[row + 7] as usize];
-            row += 8;
-        }
-        while row < count {
-            dists[row] += table[column[row] as usize];
-            row += 1;
-        }
+        transposed_column_add(
+            &mut dists[..count],
+            &codes[sub * count..(sub + 1) * count],
+            &sim_table[sub * ksub..(sub + 1) * ksub],
+        );
     }
 
     if let Some(rows) = matching_rows {
@@ -1270,6 +1244,63 @@ fn scan_codes_transposed_with_scratch(
     }
 }
 
+// A u8 code cannot index out of a 256-entry table, so converting the LUT to a
+// fixed-size array reference lets the compiler drop the per-lookup bounds
+// checks that otherwise dominate this hot loop for 8-bit scans.
+#[inline]
+fn transposed_column_init(dists: &mut [f32], column: &[u8], table: &[f32], dis0: f32) {
+    debug_assert_eq!(dists.len(), column.len());
+    if let Ok(table) = <&[f32; 256]>::try_from(table) {
+        let mut dist_chunks = dists.chunks_exact_mut(8);
+        let mut code_chunks = column.chunks_exact(8);
+        for (dist8, code8) in (&mut dist_chunks).zip(&mut code_chunks) {
+            let dist8: &mut [f32; 8] = dist8.try_into().unwrap();
+            let code8: &[u8; 8] = code8.try_into().unwrap();
+            for i in 0..8 {
+                dist8[i] = dis0 + table[code8[i] as usize];
+            }
+        }
+        for (dist, &code) in dist_chunks
+            .into_remainder()
+            .iter_mut()
+            .zip(code_chunks.remainder())
+        {
+            *dist = dis0 + table[code as usize];
+        }
+    } else {
+        for (dist, &code) in dists.iter_mut().zip(column) {
+            *dist = dis0 + table[code as usize];
+        }
+    }
+}
+
+#[inline]
+fn transposed_column_add(dists: &mut [f32], column: &[u8], table: &[f32]) {
+    debug_assert_eq!(dists.len(), column.len());
+    if let Ok(table) = <&[f32; 256]>::try_from(table) {
+        let mut dist_chunks = dists.chunks_exact_mut(8);
+        let mut code_chunks = column.chunks_exact(8);
+        for (dist8, code8) in (&mut dist_chunks).zip(&mut code_chunks) {
+            let dist8: &mut [f32; 8] = dist8.try_into().unwrap();
+            let code8: &[u8; 8] = code8.try_into().unwrap();
+            for i in 0..8 {
+                dist8[i] += table[code8[i] as usize];
+            }
+        }
+        for (dist, &code) in dist_chunks
+            .into_remainder()
+            .iter_mut()
+            .zip(code_chunks.remainder())
+        {
+            *dist += table[code as usize];
+        }
+    } else {
+        for (dist, &code) in dists.iter_mut().zip(column) {
+            *dist += table[code as usize];
+        }
+    }
+}
+
 /// Scan inverted list codes with 4-code batching for ILP (row-major layout).
 fn scan_codes_batched(
     sim_table: &[f32],
@@ -3469,7 +3500,7 @@ mod tests {
 
     #[test]
     fn transposed_scan_matches_scalar_distance_table() {
-        let count = 37;
+        let count = 40;
         let m = 7;
         let ksub = 256;
         let dis0 = 3.25;
@@ -3507,26 +3538,47 @@ mod tests {
         expected.sort_by(|left, right| left.0.total_cmp(&right.0));
         assert_eq!(heap.into_sorted(), expected);
 
-        let matching_positions = (0..count).step_by(5).collect::<Vec<_>>();
-        let matching_rows = MatchingRows::Sparse(matching_positions);
-        let mut filtered_heap = TopKHeap::new(matching_rows.len());
-        scan_codes_transposed_with_scratch(
-            &table,
-            &codes,
-            &ids,
+        for matching_count in [4, 5, 6] {
+            let matching_rows = MatchingRows::Sparse((0..matching_count).collect());
+            let mut filtered_heap = TopKHeap::new(matching_rows.len());
+            scan_codes_transposed_with_scratch(
+                &table,
+                &codes,
+                &ids,
+                count,
+                m,
+                ksub,
+                dis0,
+                Some(&matching_rows),
+                &mut filtered_heap,
+                &mut scratch,
+            );
+            let filtered_expected = expected
+                .iter()
+                .copied()
+                .filter(|(_, id)| *id < ids[0] + matching_count as i64)
+                .collect::<Vec<_>>();
+            assert_eq!(filtered_heap.into_sorted(), filtered_expected);
+        }
+    }
+
+    #[test]
+    fn transposed_sparse_scan_uses_configured_crossover() {
+        let count = 40;
+        let boundary = count / TRANSPOSED_SPARSE_SCAN_DIVISOR;
+        let at_boundary = MatchingRows::Sparse((0..boundary).collect());
+        let above_boundary = MatchingRows::Sparse((0..boundary + 1).collect());
+
+        assert!(should_scan_sparse(
             count,
-            m,
-            ksub,
-            dis0,
-            Some(&matching_rows),
-            &mut filtered_heap,
-            &mut scratch,
-        );
-        let filtered_expected = expected
-            .into_iter()
-            .filter(|(_, id)| (id - ids[0]) % 5 == 0)
-            .collect::<Vec<_>>();
-        assert_eq!(filtered_heap.into_sorted(), filtered_expected);
+            &at_boundary,
+            TRANSPOSED_SPARSE_SCAN_DIVISOR
+        ));
+        assert!(!should_scan_sparse(
+            count,
+            &above_boundary,
+            TRANSPOSED_SPARSE_SCAN_DIVISOR
+        ));
     }
 
     #[test]
@@ -3630,14 +3682,14 @@ mod tests {
     fn ivfpq_batch_timing_distinguishes_sparse_and_dense_scan_work() {
         let mut timing = IvfpqBatchTiming::default();
         let sparse_rows = MatchingRows::Sparse((0..10).collect());
-        let dense_rows = MatchingRows::Sparse((0..20).collect());
+        let dense_rows = MatchingRows::Sparse((0..26).collect());
 
         timing.record_scan_work(100, Some(&sparse_rows), 4, 8, true);
         timing.record_scan_work(100, Some(&dense_rows), 3, 8, true);
 
         assert_eq!(timing.unique_list_rows, 200);
-        assert_eq!(timing.matched_rows, 30);
-        assert_eq!(timing.pq_codes_evaluated, 100);
+        assert_eq!(timing.matched_rows, 36);
+        assert_eq!(timing.pq_codes_evaluated, 118);
         assert_eq!(timing.sparse_query_list_pairs, 4);
         assert_eq!(timing.dense_query_list_pairs, 3);
         assert_eq!(timing.actual_pq_codes_evaluated, 340);