Add JNI layer and core library improvements for OpenSearch integration by model-collapse · Pull Request #13 · opensearch-project/neural-sparse-cpp

model-collapse · 2026-06-09T06:05:02Z

Summary

Adds JNI bindings (jni/) for integrating the native sparse vector library with OpenSearch's neural-search plugin
Core library improvements for production-scale operation at 138M+ documents across multiple shards

Changes

JNI Layer (`jni/nsparse_jni.cpp`)

Index lifecycle: createIndex, addVectors, addVectorsWithIds, buildIndex, buildAndSaveIndex, saveIndex, loadIndex, deleteIndex
Search: search (SEISMIC) and searchSQ (scalar-quantized SEISMIC)
Memory management: release_build_memory() + malloc_trim() to return RSS after build
Page cache eviction: evictPageCache using madvise(MADV_DONTNEED) via /proc/self/maps scanning with dynamic buffer (getline) and error diagnostics

Core Library

int64 offsets (offset_t): handles >2B non-zeros in CSR indptr arrays (required at 45M+ docs per shard)
Streaming build_and_save: writes clusters batch-by-batch to reduce peak RSS from ~104 GB to ~50 GB
Scalar quantizer: configurable quantization_ceiling parameter for ingest and search paths
IDMapIndex: add_with_ids support for external document ID mapping
K-means clustering: OpenMP parallelization (32-thread builds take 12 min vs 50+ hours single-threaded)
Distance functions: unified SIMD dispatch across AVX2/AVX512/NEON/SVE with prefetch hints

Tests

Extended IDMapIndex tests for add_with_ids and save/load round-trip
Scalar quantizer boundary condition tests
SEISMIC common component tests

Benchmark Results (138M docs, 3 nodes × 46M docs/shard, uint8 SQ)

Force Merge (Index Build)

Metric	Value
Build time	12.7 min per shard (32 threads)
Peak RSS	59-61 GB (on 128 GB r6i.4xlarge nodes)
Settled RSS	17 GB (JVM only, native index mmap'd on demand)
.nsparse file size	~33 GB per shard

Search Latency (3903 queries, k=10, top_n=3)

heap_factor	p50 (ms)	p95 (ms)	p99 (ms)	QPS
1.03	11	15	18	64
1.08	3	6	8	136
1.15	3	6	8	136
1.25	3	6	8	133
2.00	8	15	18	80
4.00	22	28	29	39

Test plan

cmake -S . -B build -DNSPARSE_ENABLE_TESTS=ON && cmake --build build -j && ctest --test-dir build
JNI build: cmake -S . -B build -DNSPARSE_ENABLE_JNI=ON && cmake --build build --target nsparse_jni
Integration validated at 138M docs across 3-node OpenSearch cluster

🤖 Generated with Claude Code

chishui · 2026-06-30T06:39:08Z

The repo is a pure c++ library for sparse index, logic of specific business and language client should go to separate repos. Could you move JNI related logic out and make the changes in neural-search plugin?
Do you have index building and query benchmark comparison before and after this PR? How is this PR's performance comparing to existing logic under different SIMD envs?

chishui · 2026-06-30T06:39:55Z

 namespace nsparse {
 namespace {

-/**


keep the comments?

Restored the summarize_ doxygen block (fixed the 'Generage' typo and documented the offsets param that the current signature added). Pushed in 7bfa0b0.

chishui · 2026-06-30T06:48:38Z

+
+    // Dense buffer for max-pooling (reused across clusters)
+    std::vector<T> dense_max;
+    if (dimension > 0) {


IIRC, dimension now is required, better to ensure it's non-null and just throw exception at the beginning if (dimension == 0) to make logic simplified

Added if (dimension == 0) throw std::runtime_error(...) at the entry of both builders in seismic_common.cpp (5250b82), since the native engine always supplies the true dimension. I left the dimension == 0 fallback branches inside summarize_ and map_docs_to_clusters_sparse_invindex in place, because those key off vectors->get_dimension() (the vectors payload) rather than the config dimension guarded at build entry — removing them is a separate, larger change. Let me know if you want those torn out too and I can do it in a follow-up.

chishui · 2026-06-30T07:54:07Z

+    const term_t* indices = vectors->indices_data();
+    const auto* values = reinterpret_cast<const T*>(vectors->values_data());
+    size_t dimension = vectors->get_dimension();
+    if (dimension == 0) {


same here, move dimension zero handling to entry point to avoid check spread to many places

Agreed — I'll hoist this out of the helpers. With the explicit strategy parameter (see the thread on line 315), the path decision is made once at the entry point (map_docs_to_clusters) and passed down, so map_docs_to_clusters_sparse_invindex and the other helpers no longer each re-check dimension == 0. That removes the duplicated guard entirely rather than just relocating it.

chishui · 2026-06-30T07:58:28Z

    if (n_clusters == 0 || n_docs == 0) {
        return;
    }
+    if (vectors->get_dimension() > 0) {


why using dimension equal 0 as a switch between your added logic and existing logic? If we want to make it a static choice, we can use compilation option.

Agreed that dimension is the wrong thing to switch on — it's a data property of the vectors, not a strategy flag, and overloading it this way makes the control flow unreadable (you have to know the convention "dimension==0 means old path" to follow the code). I'll fix that.

On the compile option suggestion specifically, I'd push back: I don't think the old-vs-new path should be a static/compile-time choice. A #ifdef/CMake option bakes the decision into the binary — the plugin ships a single .so, so a compile flag would force one path for all users and require a rebuild (or shipping multiple artifacts) to change it. The path selection is better made at runtime.

Proposal: replace the dimension == 0 switch with an explicit, self-documenting strategy — e.g. a ClusteringStrategy { Auto, Dense, SparseInvIndex } enum on the build config — decided once at the entry point and threaded down, so no helper re-derives it from dimension. That addresses both the readability problem and the "check spread to many places" comment on the other thread.

One open question back to you: do we actually expect both paths to coexist long-term, or is the new sparse-invindex path meant to fully replace the old one? If it strictly dominates for our use, the cleaner move is to delete the old path entirely rather than parameterize the choice — happy to go that way instead if that's the intent.

Resolved by removing the dead path rather than parameterizing it.

Digging in, the two implementations selected by dimension == 0 were:

dimension > 0 → the sparse inverted-index clustering path (added in this PR)

dimension == 0 → the older dense-centroid AVX512 path

The dense path turned out to be unreachable and inconsistent with the tests:

In the plugin, createIndex always passes a real dimension (65536), so get_dimension() is never 0 and the sparse inverted-index path is always taken.

Every kmeans/seismic test uses dimension = 10 (the sparse path); none exercise the dense path.

The existing invariant test no_duplicate_and_no_missing_docs_when_centroids_in_docs encodes the sparse path's centroid-skip behavior as required; the dense path (which re-clusters centroid docs via argmax) would violate it.

So instead of a compile option or a runtime strategy flag, I deleted the dense-centroid path entirely (map_docs_to_clusters_avx512 + its centroids_to_dense / initialize_cluster_representatives / get_dense_vector_max_dimension helpers, plus an unused dot_product_typed_dense). map_docs_to_clusters now dispatches purely by element_size — no more dimension == 0 switch. Net -166 lines; full suite (365 tests) still passes.

Pushed as b414549.

chishui · 2026-07-01T02:30:06Z

+
+    // Build offset table (prefix sum)
+    std::vector<uint32_t> offsets(dimension + 1);
+    offsets[0] = 0;


nit: it's already initialized to 0

Removed — the vector is value-initialized to 0. Fixed in the latest push.

chishui · 2026-07-01T05:59:14Z

+    if (n_batches < 1) n_batches = 1;
+    size_t batch_size = (dimension + n_batches - 1) / n_batches;
+
+    fprintf(stderr, "[nsparse] build_inverted_lists: n_docs=%zu, dimension=%zu, "


Firstly, are these logs useful for all users of this library? They seems to me only for debugging. If so, we'd better introduce a verbose parameter similar to https://github.com/facebookresearch/faiss/blob/main/faiss/Index.h#L107 and only print logs when it's true

Done in 5250b82. Added a verbose flag to SeismicClusterParameters (default false, same spirit as faiss Index::verbose) and gated all 7 per-batch fprintf(stderr, ...) diagnostics behind it. It is parsed from the index description (verbose=1) in index_factory for both seismic and seismic_sq, so production builds stay quiet by default.

chishui · 2026-07-01T06:04:11Z

+    size_t mem_budget = 0;
+#ifdef __linux__
+    // Use MemAvailable from /proc/meminfo
+    FILE* meminfo = fopen("/proc/meminfo", "r");


make it a util function

Done in 4896001 — extracted to detail::query_available_memory_bytes() (the /proc/meminfo reader) and detail::resolve_build_mem_budget() (the budget policy), both unit-tested.

chishui · 2026-07-01T06:20:18Z

+    // Choose batch count to fit inverted lists within available memory.
+    // Reserve 4 GB for clustering overhead and OS.
+    constexpr size_t kReserve = 4ULL * 1024 * 1024 * 1024;
+    size_t mem_budget = 0;


Can we not use mem budget to calculate batch but simply have a hard-coded batch size and let users to control thread number used below to tune the memory usage?

A hard-coded batch size is what I'd like to avoid, because the batch count is the dominant control over build-phase peak RSS, and the right value differs by an order of magnitude across deployments — a 128 GB node vs a 256 GB node want very different batching for the same index. Thread count doesn't govern this: the inverted-list memory (total_nnz * (sizeof(idx_t) + element_size)) is materialized per batch regardless of how many threads process it, so tuning threads down doesn't cap the peak, only the parallelism.

That said, your two concerns are both valid and I'll fix them:

The Linux-only path leaving mem_budget = 0 elsewhere is a real bug — on non-Linux it currently falls through to the 16 GB fallback, which is accidental, not intended.

/proc/meminfo parsing should be a util function (noted on the other thread too).

Proposed compromise: keep budget-driven batching but make the budget an explicit parameter — default = auto-detect (Linux) with a documented fixed fallback on other platforms, and allow the caller (the plugin) to override it. That keeps the adaptivity that actually controls memory while removing the platform-specific silent behavior and giving users a real knob. Does an explicit mem_budget_bytes parameter (0 = auto) work for you?

Implemented as proposed, in 4896001. mem_budget_bytes is now an explicit field on SeismicClusterParameters (0 = auto), threaded from the index description string ("...|mem_budget=") — no JNI change needed. Auto-detect stays as the default (with the non-Linux fallback fixed), but callers/the plugin can now set an exact budget to control build-phase peak RSS instead of relying on thread count. Added tests for the explicit and auto paths.

chishui · 2026-07-01T06:27:15Z

@@ -129,30 +133,236 @@ inline int calculate_beta(int beta, int lambda) {
 inline std::vector<InvertedListClusters> build_inverted_lists_clusters(


as inline functions now build_inverted_lists_clusters and build_and_save_inverted_lists_clusters are too big, better create a seismic_common.cpp and move implementation there and not mark them inline

Done in 5250b82. Created nsparse/seismic_common.cpp and moved both build_inverted_lists_clusters and build_and_save_inverted_lists_clusters there (declarations remain in the header, no longer inline). Added the file to NSPARSE_SRC so it is still compiled once per SIMD variant (generic/avx2/avx512/sve) and keeps its per-arch specialization.

chishui · 2026-07-01T06:36:02Z

    for (const size_t& cluster_id : cluster_order) {
        const auto& cluster_score = summary_scores[cluster_id];
-        if (heap.full() && (cluster_score * heap_factor < heap.peek_score())) {
-            if (first_list) {


why removed first_list check? That's an optimization

The first_list check is not always making positive effect, I am considering to make it as an option to the user.

Adds JNI bindings for the native sparse vector library to integrate with OpenSearch's neural-search plugin. Key capabilities: - Index lifecycle: create, add vectors (with IDs), build, save, load, delete - Search: dense-scoring SEISMIC and SEISMIC-SQ (scalar quantized) search - Memory management: streaming build_and_save, release_build_memory, malloc_trim - Page cache eviction: madvise(MADV_DONTNEED) via /proc/self/maps scanning Core library changes for production-scale operation (138M+ documents): - int64 offsets (offset_t) to handle >2B non-zeros in CSR indptr - Streaming build_and_save: writes clusters batch-by-batch to reduce peak RSS - Scalar quantizer improvements: configurable quantization ceilings - IDMapIndex: add_with_ids support for external document ID mapping - K-means clustering: OpenMP parallelization for 32-thread builds - Distance functions: unified SIMD dispatch across AVX2/AVX512/NEON/SVE Validated at 138M docs (3 shards × 46M), uint8 SQ SEISMIC: - Build time: 12.7 min per shard (32 threads) - Peak RSS: 59-61 GB during force merge (on 128 GB nodes) - Search: 3ms p50 at heap_factor=1.08, top_n=3 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Per upstream review feedback on PR opensearch-project#13: language-client / integration glue should not live in the core C++ library. The JNI bindings (jni/nsparse_jni.cpp + its CMake target) move to the OpenSearch neural-search plugin, which compiles them against this library. - Delete jni/ (nsparse_jni.cpp, CMakeLists.txt) - Remove NSPARSE_BUILD_JNI option and add_subdirectory(jni) from the top-level CMakeLists; leave a pointer comment. - .gitignore: also ignore out-of-source build-*/ build_*/ __pycache__/ dirs. Library + full test suite (365 tests) build and pass unchanged. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

map_docs_to_clusters had two implementations selected by `dimension == 0`: - dimension > 0 -> sparse inverted-index path (Path B) - dimension == 0 -> dense-centroid AVX512 path (Path A), else Path B anyway Path A is unreachable in practice and inconsistent with the tests: - The plugin always builds with a real dimension (createIndex passes 65536), so get_dimension() is never 0 in production and Path B is always taken. - Every kmeans/seismic test uses dimension=10 (Path B). None exercise Path A. - The existing invariant test `no_duplicate_and_no_missing_docs_when_ centroids_in_docs` encodes Path B's centroid-skip behavior as required; Path A (which re-clusters centroid docs via argmax) would violate it. Addresses the PR opensearch-project#13 review comment about using `dimension == 0` as a code-path switch: rather than parameterize the choice, remove the dead path. - Delete map_docs_to_clusters_avx512 + its helpers (centroids_to_dense, initialize_cluster_representatives, get_dense_vector_max_dimension) and the AVX512-only include block. - Delete the unused dot_product_typed_dense helper. - map_docs_to_clusters now dispatches purely by element_size to the sparse inverted-index path. Net -166 lines. Full test suite (365 tests) builds and passes unchanged (Release, AVX-512 host). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

The swignsparse_avx2/avx512/sve.swig files are generated at build time from swignsparse.swig via configure_file(COPYONLY). They were removed from the tree in b414549 but the .gitignore rule to keep them out wasn't committed, so a build + `git add` would re-track them. Commit the ignore rule. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…allback Addresses the PR opensearch-project#13 review: `mem_budget` was computed inline from /proc/meminfo and left at 0 on non-Linux, silently falling through to a hardcoded 16 GB. Two problems: the value was undetectable/opaque off Linux, and callers had no way to control build-phase peak RSS. Changes: - SeismicClusterParameters gains `mem_budget_bytes` (default 0 = auto). This rides the existing index description string, so no JNI signature change: "seismic[_sq],...|mem_budget=<bytes>". - New detail::query_available_memory_bytes() — the /proc/meminfo reader extracted into a documented util (returns 0 where unsupported). - New detail::resolve_build_mem_budget(requested): explicit budget wins; else auto-detect minus a 4 GB reserve; else a documented 16 GB fallback. Never returns 0 (guards the batch divisor). - Both build_inverted_lists_clusters and build_and_save_inverted_lists_clusters honor an explicit budget; the streaming path keeps its 8 GB auto-cap when none is given. - index_factory parses `mem_budget` for seismic and seismic_sq. - Tests: 4 cases covering explicit-verbatim, small-explicit, auto-nonzero, and auto == query/fallback. Full suite 369/369; plugin JNI rebuild verified. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

std::vector<uint32_t> offsets(dimension + 1) already value-initializes all elements to 0, so the explicit assignment was redundant. Addresses a PR opensearch-project#13 review nit. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

@return

The PR's rewrite of summarize_() accidentally removed its @brief/@param/@return block. Restore it (fix the "Generage" typo, document the offsets param that the current signature added). Addresses PR opensearch-project#13 "keep the comments?". Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Addresses review comments on PR opensearch-project#13: - Gate the 7 per-batch progress fprintf(stderr) diagnostics behind a new SeismicClusterParameters.verbose flag (default false), so production builds stay quiet. Parsed from the index description ("verbose=1") in index_factory for both seismic and seismic_sq. - Move build_inverted_lists_clusters and build_and_save_inverted_lists_clusters out of the header into a new seismic_common.cpp (declarations remain in the header). Added to NSPARSE_SRC so every SIMD variant still specializes it. - Split the two large builders into shared helpers (compute_batch_plan, build_batch_invlists, count_non_empty_lists). Both distinct mem_budget policies (resolve_build_mem_budget vs explicit-or-8GB-cap) and all OMP pragmas are preserved byte-for-byte. - Guard against dimension == 0 at the entry of both builders. All 369 tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

chishui

Please fix DCO for each commit
Run clang-tidy and clang-format

chishui · 2026-07-15T03:40:25Z

                     const float* values) = 0;
+    virtual void reserve(size_t num_vectors, size_t total_nnz) {}
    virtual void build();
+    virtual void build_and_save(const char* path);


Please remove build_and_save and release_build_memory as we were aligned offline

chishui · 2026-07-15T03:42:42Z


    for (const size_t& cluster_id : cluster_order) {
        const auto& cluster_score = summary_scores[cluster_id];
-        if (heap.full() && (cluster_score * heap_factor < heap.peek_score())) {


As we discussed offline, let's keep first_list optimization and have the change in a dedicated PR when there are strong evidence showing that it's benefit to remove

chishui · 2026-07-15T03:44:09Z

    add_subdirectory(nsparse/python)
 endif()

+# Note: JNI bindings for the OpenSearch neural-search plugin live in that


this comment is no longer needed?

chishui · 2026-07-15T03:44:25Z

+__pycache__/
+
+# SWIG SIMD variants are generated from swignsparse.swig at build time (configure_file COPYONLY)
+nsparse/python/swignsparse_avx2.swig


chishui · 2026-07-15T03:45:16Z

+        active_terms.clear();
+        std::vector<std::pair<term_t, T>> summary_vec;
+
+        if (dimension > 0) {


as we discussed, dimension check should be removed?

chishui · 2026-07-15T06:06:02Z

+        total_nnz * (sizeof(idx_t) + element_size);
+
+    // Cap per-batch inverted list memory to keep peak RSS under physical RAM.
+    // The CSR (~46 GB for float32 at 46M docs) remains throughout; each batch


the comments are specific to certain dataset and should be removed.

chishui · 2026-07-15T06:08:58Z

+    // interface it falls back to a fixed default. Set explicitly (e.g. from
+    // the native-engine index description, "mem_budget=<bytes>") to control
+    // build-phase peak RSS deterministically.
+    size_t mem_budget_bytes = 0;


We'd better make it mem_budget_in_megabytes or even mem_budget_in_gigabytes to make it simple for user.

nit: you can ask LLM to make comments precise and concise

chishui · 2026-07-15T06:20:10Z

+        constexpr size_t kMaxBatchBytes = 8ULL * 1024 * 1024 * 1024;  // 8 GB
+        mem_budget = std::min(kMaxBatchBytes, invlist_bytes_full / 3);
+    }
+    if (mem_budget == 0) mem_budget = invlist_bytes_full;


if mem_budget is 0 here, then invlist_bytes_full must be 0

chishui · 2026-07-15T06:24:29Z

+    return clustered_inverted_lists;
+}
+
+void build_and_save_inverted_lists_clusters(


can we only have one responsibility for one function? already have build_inverted_lists_clusters above, not necessary to have a build_and_save_inverted_lists_clusters with duplicate code plus save to file. any places which need to build then save, they can simply call build_inverted_lists_clusters then save to whatever they want.

chishui · 2026-07-15T06:28:10Z

+    return clustered_inverted_lists;
+}
+
+void build_and_save_inverted_lists_clusters(


also, this function has quite different implementation with build_inverted_lists_clusters

model-collapse requested review from chishui, yuye-aws and zirui-song-18 as code owners June 9, 2026 06:05

chishui reviewed Jul 1, 2026

View reviewed changes

model-collapse force-pushed the feature/jni-opensearch-integration branch from 57e5763 to 0f20144 Compare July 7, 2026 05:20

wanglovesyang and others added 7 commits July 7, 2026 06:19

chishui reviewed Jul 15, 2026

View reviewed changes

		@@ -129,30 +133,236 @@ inline int calculate_beta(int beta, int lambda) {
		inline std::vector<InvertedListClusters> build_inverted_lists_clusters(

Uh oh!

Conversation

model-collapse commented Jun 9, 2026

Summary

Changes

JNI Layer (jni/nsparse_jni.cpp)

Core Library

Tests

Benchmark Results (138M docs, 3 nodes × 46M docs/shard, uint8 SQ)

Force Merge (Index Build)

Search Latency (3903 queries, k=10, top_n=3)

Test plan

Uh oh!

chishui commented Jun 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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

JNI Layer (`jni/nsparse_jni.cpp`)

chishui commented Jun 30, 2026 •

edited

Loading