DiskSeismic: add scalar-quantized index (disk_seismic_sq) by zirui-song-18 · Pull Request #37 · opensearch-project/neural-sparse-cpp

zirui-song-18 · 2026-08-27T08:16:13Z

Description

This PR is to implement the scalar quantization of DiskSeismicIndex.

  ┌─────────┬────────────────────────────────┬───────────────────┬─────────────────┬───────────────────┬───────────┐  
  │         │           index size           │    build total    │ p50 / p99 (ms)  │    bytes read     │ recall@10 │
  ├─────────┼────────────────────────────────┼───────────────────┼─────────────────┼───────────────────┼───────────┤
  │ float32 │ 122,609,780,704 B (114.19 GiB) │ 254.19 s          │ 0.1393 / 0.2112 │ 57.54 GB          │ 0.900530  │
  ├─────────┼────────────────────────────────┼───────────────────┼─────────────────┼───────────────────┼───────────┤     
  │ 16-bit  │ 92,078,408,368 B (0.7510×)     │ 216.62 s (0.852×) │ 0.1379 / 0.2092 │ 43.46 GB (0.755×) │ 0.900530  │
  ├─────────┼────────────────────────────────┼───────────────────┼─────────────────┼───────────────────┼───────────┤
  │ 8-bit   │ 77,254,839,584 B (0.6301×)     │ 194.84 s (0.767×) │ 0.1348 / 0.2050 │ 35.10 GB (0.610×) │ 0.893696  │
  └─────────┴────────────────────────────────┴───────────────────┴─────────────────┴───────────────────┴───────────┘

Issues Resolved

RFC: [RFC] Disk-resident forward index for SEISMIC, with an inline on-disk layout #21

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

Signed-off-by: Zirui Song <zrsong@amazon.com>

chishui · 2026-08-27T08:47:03Z

+      sq_(quantizer_type, vmin, vmax),
+      cluster_parameter_(parameter) {}
+
+void DiskSeismicScalarQuantizedIndex::read_csr(const char* file_path,


do you need to override this? I think you only want to handle kMmap use case

Yes — the override exists only to reject the kMmap case: a mapped CSR is borrowed as float, but this index searches over codes, so mapping it would misread. The in-memory case just delegates to the base. So it's scoped to exactly the kMmap use case you mention. DiskSeismicIndex doesn't need it (it's float), and this mirrors SeismicScalarQuantizedIndex::read_csr. Kept.

chishui · 2026-08-27T08:47:40Z

+// A quantizer described by an index file. bytes_per_value() treats anything but
+// QT_8bit as 16-bit, so an undefined type would silently pick an element width
+// rather than be rejected.
+ScalarQuantizer quantizer_from_file(QuantizerType type, float vmin,


while it's named "quantizer from file", but the signature has nothing to do with file

Good catch — renamed to make_scalar_quantizer(type, vmin, vmax);

chishui · 2026-08-27T08:52:35Z

+}
+
+// A candidate block: its summary score and its (posting list, cluster) address.
+struct BlockCandidate {


it's shared with diskseismicindex, reuse

chishui · 2026-08-27T08:53:45Z

+// Score every document of one block against the dense query codes, dedup via
+// `visited` and honor the id selector. Vectors come from the inline forward
+// index (fwd) when loaded, else from the in-RAM CSR of a fresh build.
+void score_block(const detail::InlineForwardIndex* fwd,


seems these helper functions share the logic with DiskSeismicIndex, could you refactor them to make them work for both indices?

chishui · 2026-08-27T08:55:19Z

+    // Resolve `cut` and `k_prime`. A DiskSeismicSearchParameters (including a
+    // DiskSeismicSQSearchParameters) carries k_prime; a plain
+    // SeismicSearchParameters (or null) uses the default budget.
+    const DiskSeismicSearchParameters defaults;


duplicate logic block, make it shared. Check even if search logic can be refactored to a shared function

chishui · 2026-08-27T08:57:01Z

+    query_sq.encode(values, codes.data(), nnz);
+    const uint8_t* query_values = codes.data();
+
+    std::vector<std::vector<float>> result_distances(


share code with line 204,205

chishui · 2026-08-27T08:57:56Z

+
+auto DiskSeismicScalarQuantizedIndex::single_query(
+    std::vector<uint8_t>& dense, absl::flat_hash_set<idx_t>& visited,
+    const term_t* q_idx, const uint8_t* q_val_bytes, size_t q_len,


expand q_xxx to query_xxx

chishui · 2026-08-27T08:59:26Z

+}
+
+void DiskSeismicScalarQuantizedIndex::write_index(IOWriter* io_writer) {
+    write_header(io_writer);


change to write_quantization_header

Renamed to write_quantization_header. One note: the sibling SeismicScalarQuantizedIndex (from #36) calls its equivalent write_quantizer_header — want me to rename that one to match, or should I keep write_quantizer_header here for parity with SESQ? Happy either way; just want them consistent.

sorry, please make them consistent

chishui · 2026-08-27T09:00:36Z

+
+DiskSeismicScalarQuantizedIndex* DiskSeismicScalarQuantizedIndex::mmap_index(
+    int dimension, const char* index_file, size_t pos) {
+    throw_if_null(index_file, "index_file must not be null");


there is a valid file check in check.h, use that for file check, you may also want to apply in DiskSeismicIndex

I'm using throw_if_null(index_file, ...) from checks.h here, same as SeismicIndex/InvertedIndex/DiskSeismicIndex::mmap_index; the file's openability is then enforced by MmapFile's constructor, which throws if it can't open. I didn't spot a stronger file-validity helper in checks.h (it has null/positive/overflow checks). Did you have a specific one in mind — or want me to add a shared throw_if_invalid_index_file and apply it across all the mmap readers? Happy to do that consistently.

interesting, I remember I implemented a check for file pointer and whether file exist, but it's not there, I guess it's removed somehow or in a local branch, never mind.

chishui · 2026-08-27T09:52:13Z

+    return {cut, k_prime};
+}
+
+pair_of_score_id_vectors_t padded_results(idx_t n, int k) {


what about initialize_padded_results

Renamed to initialize_padded_results.

chishui · 2026-08-27T09:52:40Z

+                n, std::vector<idx_t>(k, INVALID_IDX))};
+}
+
+pair_of_score_id_vector_t groc_search_query(


groc is not a common term here, right?

You are right. Renamed the function to block_budget_query and dropped "GroC" from the comments.

chishui · 2026-08-27T09:53:49Z

+            query_values + i * element_size, element_size,
+            dense + static_cast<size_t>(query_indices[i]) * element_size);
+    }
+    visited.clear();


the function will only be called once per query, right?

Yes — one call per query inside the #pragma omp for.

chishui

write_index / read_index / mmap_index / add / build and the whole batch-search loop are still near-verbatim copies of DiskSeismicIndex's, differing only in the value width and the quantizer header. Extract a shared base instead of keeping a third parallel copy.
Same for the tests -- both the C++ helpers and the python file are forks of the disk_seismic ones.
write_quantization_header vs SESQ's write_quantizer_header is still inconsistent.

chishui · 2026-08-27T10:02:58Z

+    // summary (stored at the same width), so scores are comparable across
+    // posting lists for the global ranking.
+    std::vector<BlockCandidate> candidates;
+    std::vector<float> score_scratch;


both allocate per query on the hot path. pass them in as per-thread scratch like dense/visited, or at least reserve

Moved candidates and score_scratch out to per-thread scratch, passed in alongside dense/visited, so nothing allocates per query now.

chishui · 2026-08-27T10:02:58Z

+        "read_index(file, IndexIoFlag::kUseMmap)");
+}
+
+void DiskSeismicScalarQuantizedIndex::write_quantization_header(


SeismicScalarQuantizedIndex still has write_quantizer_header, please rename that one too

Done — renamed SESQ's write_quantizer_header and read_quantizer_header to write_quantization_header/read_quantization_header so both indexes match.

chishui · 2026-08-27T10:02:58Z

+    inv_list_writer.mmap_deserialize(&cursor);
+    detail::InlineForwardIndex forward;
+    forward.mmap_deserialize(&cursor);
+    throw_if_element_size_mismatch(forward, sq);


this only checks fwd. score_summaries_transposed dispatches on the summaries' own element_size_, verify that against the quantizer too?

Good catch — added an element_size() accessor on InvertedListClusters and validate_mapped_payload now checks the summaries' width against the quantizer in addition to the forward index's.

chishui · 2026-08-27T10:02:58Z

+    return results;
+}
+
+void DiskSeismicScalarQuantizedIndex::write_index(IOWriter* io_writer) {


this and read_index/mmap_index/add/build are still copies of DiskSeismicIndex's, differing only in the code width and the quantizer header. shared base?

Done — added DiskSeismicIndexBase (disk_seismic_index_base.{h,cpp}). It owns add/build/search/write_index/read_index and the mmap payload load; both indexes derive from it and supply only the differences via hooks (code_element_size, encode_values, encode_query, decode_scores, write_payload_header, validate_mapped_payload). The two .cpps dropped ~730 lines between them. DiskSeismic's bit-exact parity tests still pass, so behavior is unchanged.

chishui · 2026-08-27T10:02:58Z

+
+void DiskSeismicScalarQuantizedIndex::write_index(IOWriter* io_writer) {
+    write_quantization_header(io_writer);
+    const uint64_t nv = num_vectors_;


nit: drop the const instead of const_cast-ing it away. same in disk_seismic_index.cpp

Dropped the const on nv instead of casting it away, here and in disk_seismic_index.cpp

chishui · 2026-08-27T10:02:58Z

+    query_sq.encode(values, codes.data(), nnz);
+    const uint8_t* query_batch = codes.data();
+
+    pair_of_score_id_vectors_t results = detail::padded_results(n, k);


every row is overwritten below, so the -1.0/INVALID_IDX fill is n*k wasted writes on this path

chishui · 2026-08-27T10:02:58Z

    }

+    if (index_type == "disk_seismic_sq") {
+        std::string quantizer_str = get_param("quantizer", "8bit");


duplicate of the seismic_sq block below, move the quantizer/vmin/vmax/cluster parsing into a function

an unknown quantizer= value silently becomes 8bit, reject it?

FYI the seismic_sq copy has .lambda = lambda = lambda (line 157) -- a shared helper kills that too

Extracted parse_cluster_params and parse_quantizer_config, shared by seismic/disk_seismic/seismic_sq/disk_seismic_sq.

An unrecognized quantizer= now throws instead of defaulting to 8bit (test added).

That also removed the .lambda = lambda = lambda self-assignments.

chishui · 2026-08-27T10:02:58Z

 // Use %factory to enable proper downcasting based on runtime type
 %newobject nsparse::index_factory;
-%factory(nsparse::Index* nsparse::index_factory, nsparse::BrutalIndex, nsparse::SeismicIndex, nsparse::SeismicScalarQuantizedIndex, nsparse::IDMapIndex, nsparse::InvertedIndex);
+%factory(nsparse::Index* nsparse::index_factory, nsparse::BrutalIndex, nsparse::SeismicIndex, nsparse::SeismicScalarQuantizedIndex, nsparse::DiskSeismicScalarQuantizedIndex, nsparse::IDMapIndex, nsparse::InvertedIndex);


nsparse::DiskSeismicIndex is missing from both %factory lists, add it?

Added nsparse::DiskSeismicIndex to both %factory lists.

Signed-off-by: Zirui Song <zrsong@amazon.com>

zirui-song-18 · 2026-08-27T12:52:37Z

write_index / read_index / mmap_index / add / build and the whole batch-search loop are still near-verbatim copies of DiskSeismicIndex's, differing only in the value width and the quantizer header. Extract a shared base instead of keeping a third parallel copy.

Same for the tests -- both the C++ helpers and the python file are forks of the disk_seismic ones.

write_quantization_header vs SESQ's write_quantizer_header is still inconsistent.

Thanks — I think this predated my last push. Here's where each stands on the latest revision:

Shared base — done: DiskSeismicIndexBase (disk_seismic_index_base.{h,cpp}) now owns add/build/search/write_index/read_index/mmap-load and the batch loop; DiskSeismicIndex and DiskSeismicScalarQuantizedIndex are thin, supplying only the width/encode/decode/header via virtual hooks. No third parallel copy.
Tests — now deduped too: shared C++ helpers moved to tests/disk_seismic_test_util.h (both suites use it), and the Python contract tests moved to a DiskSeismicContract base class in disk_seismic_contract.py that both test_disk_seismic_index.py and test_disk_seismic_sq_index.py subclass — the SQ file keeps only its quantization-specific tests.
Naming — done: renamed SESQ's write_quantizer_header/read_quantizer_header to write_quantization_header/read_quantization_header, so both indexes match.

zirui-song-18 requested review from chishui, model-collapse and yuye-aws as code owners August 27, 2026 08:16

DiskSeismic: add scalar-quantized index (disk_seismic_sq)

df101bc

Signed-off-by: Zirui Song <zrsong@amazon.com>

zirui-song-18 force-pushed the disk-seismic-sq branch from 2d3cfce to df101bc Compare August 27, 2026 08:58

chishui reviewed Aug 27, 2026

View reviewed changes

zirui-song-18 force-pushed the disk-seismic-sq branch 2 times, most recently from 9fa616b to a5110f9 Compare August 27, 2026 09:34

zirui-song-18 requested a review from chishui August 27, 2026 09:38

chishui reviewed Aug 27, 2026

View reviewed changes

zirui-song-18 force-pushed the disk-seismic-sq branch from a5110f9 to d5c8204 Compare August 27, 2026 12:26

address Liyun's comments

891d9a0

Signed-off-by: Zirui Song <zrsong@amazon.com>

zirui-song-18 force-pushed the disk-seismic-sq branch from d5c8204 to 891d9a0 Compare August 27, 2026 12:48

zirui-song-18 requested a review from chishui August 27, 2026 12:53

chishui approved these changes Aug 28, 2026

View reviewed changes

chishui merged commit 9a74567 into opensearch-project:main Aug 28, 2026
9 checks passed

Conversation

zirui-song-18 commented Aug 27, 2026

Description

Issues Resolved

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chishui left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!