This repo contains the code corresponding to the SIGIR 2026 short paper Scalable K-Means Guided Partitioning for Block-based Sparse Document Retrieval by Parker Carlson*, Sammy Lesner*, Antonio Mallia, and Tao Yang.
*Equal Contribution
If you use this code in your own work or research, please consider citing our work:
@inproceedings{carlson2026guidekp,
author = {Parker Carlson and Sammy Lesner and Antonio Mallia and Tao Yang},
title = {Scalable K-Means Guided Partitioning for Block-based Sparse Document Retrieval},
booktitle = {The 49th International ACM SIGIR Conference on Research and Development in Information Retrieval ({SIGIR '26})},
publisher = {ACM},
year = {2026}
}
The paper can be found at the following DOI: http://doi.org/10.1145/3805712.3809971
This work was built on previous work and the codebase from J. Mackenzie et al: Faster Recursive Graph Bisection, Code Repository.
Our codebase is based on the implementation found in the PISA search engine, which corresponds to the reproducibility study discussed above. The codebase works with the Common Index File Format, an open-source index exchange format for information retrieval experimentation.
You can build the code using Cargo, along with the specific gain function to use, passed in at compile time via an environment variable. Options include:
# <<< Faster Recursive Graph Bisection - loggap based BP >>>
GAIN=approx_1 cargo build --release
GAIN=approx_2 cargo build --release
GAIN=default cargo build --release
# <<< NEW with GuideKP -- kmeans-like BP >>>
GAIN=kmeans cargo build --release
GAIN=BP_kmeans cargo build --release
The valid options from faster-graph-bisection for the implementations of BP are default, approx_1, or approx_2. In this work we add kmeans for GuideKP, and BP_kmeans for a hybrid option between the standard loggap based BP and GuideKP.
This specific hybrid option simply sums the two types of pairwise swap gain in equal parts.
To use this implementation, you will need a CIFF index corresponding to the data you wish to reorder.
For a quick start, a few pre-generated indexes can be found in the official CIFF repository here, or, for CIFF indexes built from MS MARCO, check out the PISA CIFF Hub here.
For this walkthrough we'll reuse faster-recursive-BP's starting example and use Robust04 as found from the official CIFF repo:
mkdir example
cd example
wget https://www.dropbox.com/s/rph6udiqs2k7bfo/robust04-complete-20200306.ciff.gz?dl=0 -O robust-ciff.gz
gunzip robust-ciff.gz
mv robust-ciff robust.ciff
Because many .ciff files are published with BP-ordering, we also included a script to randomly shuffle documents stored in the .ciff file format for a fair comparison between BP and GuideKP. This can be seeded for reproducibility.
The following shows example usage on Robust04:
./target/release/reshuffle -i example/robust.ciff -o example/shuffled_robust.ciff -s 42
Input CIFF:
Postings: 923436
Documents: 528030
Total Terms: 174540872
Avg Doc Length: 330.5510520235593
Generating random mapping with seed 42...
Reading posting lists...
[00:00:10] [========================================] 923436/923436 (0s)
Reading document records...
[00:00:04] [========================================] 528030/528030 (0s)
Applying mapping to posting lists...
[00:00:09] [========================================] 923436/923436 (0s)
Reordering document records...
Writing output CIFF...
[00:00:10] [========================================] 923436/923436 (0s)
[00:00:04] [========================================] 528030/528030 (0s)
✓ Successfully shuffled CIFF file!
Output: "example/shuffled_robust.ciff"
With an index in place, you can run the GuideKP algorithm.
At minimum, GuideKP takes a starting CIFF index and produces a reordered CIFF index.
Adding the --loggap flag will include the computation for the original log-gap metrics of faster-recursive-BP alongside an added set of metrics to evaluate block-based reordering quality.
To reorder the Robust04 index you downloaded above, run:
../target/release/create-rgb --input robust.ciff --output-ciff robust-reordered.ciff --loggap
For which the output might look like:
01:01:41 [INFO] Using the `kmeans` gain function.
01:01:41 [INFO] Opt { input: "robust.ciff", output_ciff: Some("robust-reordered.ciff"), min_len: 4096, cutoff_frequency: 0.1, recursion_stop: 16, swap_iterations: 20, loggap: true, sort_leaf: false, max_depth: 100, input_fidx: None, output_fidx: None, output_mapping: None, metrics_only: false, swap_gain: false, mg_adj: false, block_aware: false, max_iter: 100, restarts: 10, bisecting_km: false, cooling: false, quickselect: false, use_inputblocks: false, blocksize_txt: "blocks.txt" }
01:01:41 [INFO] (1) building forward index
01:01:44 [INFO] forward index stats:
01:01:44 [INFO] total terms: 923436
01:01:44 [INFO] discarded frequent terms: 384
01:01:44 [INFO] discarded infrequent terms: 920126
01:01:44 [INFO] remaining terms: 2927
01:01:44 [INFO] (2) sort empty docs to the back
01:01:44 [INFO] fwd duration: 3.07 secs
01:01:44 [INFO] docs 528030 non_empty 527908
01:01:44 [INFO] put docs back into default order...
01:01:44 [INFO] (3) perform graph bisection
01:02:03 [INFO] rgb duration: 18.64 secs
01:02:03 [INFO] (4) clear forward index
01:02:03 [INFO] (5) starting output operations...
01:02:03 [INFO] --> (5.2) write new ciff file
01:02:03 [INFO] writing to ciff file: robust-reordered.ciff
01:02:13 [INFO] write duration: 9.81 secs
01:02:13 [INFO] (6) compute loggap cost
01:02:13 [INFO]
01:02:13 [INFO] ===== BEFORE =====
01:02:37 [INFO] BPI: 3.975
01:02:37 [INFO]
01:02:37 [INFO] avg block-8 avg cosine distance: 0.879
01:02:37 [INFO] avg block-8 max cosine distance: 0.971
01:02:37 [INFO]
01:02:37 [INFO] avg block-8 avg dot product: 140.856
01:02:37 [INFO] avg block-8 max dot product: 817.753
01:02:37 [INFO]
01:02:37 [INFO] max block-8 max-avg difference dot product: 507865.571
01:02:37 [INFO] avg block-8 max-avg difference dot product: 676.897
01:02:37 [INFO]
01:02:37 [INFO] max block-8 max/avg ratio dot product: 25.543
01:02:37 [INFO] avg block-8 max/avg ratio dot product: 5.254
01:02:37 [INFO]
01:02:37 [INFO] max block-8 max-avg difference cosine to bmax: 0.819
01:02:37 [INFO] avg block-8 max-avg difference cosine to bmax: 0.606
01:02:37 [INFO] avg block-8 avg centroid document distance: 5308.680
01:02:37 [INFO]
01:02:37 [INFO] *** SSE ***: 42469.415
01:02:37 [INFO] *** EBT *** 0.161
01:02:37 [INFO] std cluster sizes (SDCS): 5.322
01:02:37 [INFO] normalized cluster entropy (NCE): 5.039
01:02:37 [INFO]
01:02:37 [INFO] avg terms per block-8: 1087.092
01:02:37 [INFO] max terms per block-8: 38849
01:02:37 [INFO] min terms per block-8: 141
01:02:37 [INFO]
01:02:37 [INFO] avg blocks per term: 77.702
01:02:37 [INFO]
01:02:37 [INFO] ===== AFTER =====
01:03:00 [INFO] BPI: 3.169
01:03:00 [INFO]
01:03:00 [INFO] avg block-8 avg cosine distance: 0.718
01:03:00 [INFO] avg block-8 max cosine distance: 0.844
01:03:00 [INFO]
01:03:00 [INFO] avg block-8 avg dot product: 1355.169
01:03:00 [INFO] avg block-8 max dot product: 8954.782
01:03:00 [INFO]
01:03:00 [INFO] max block-8 max-avg difference dot product: 334387199.214
01:03:00 [INFO] avg block-8 max-avg difference dot product: 7599.613
01:03:00 [INFO]
01:03:00 [INFO] max block-8 max/avg ratio dot product: 27.998
01:03:00 [INFO] avg block-8 max/avg ratio dot product: 2.445
01:03:00 [INFO]
01:03:00 [INFO] max block-8 max-avg difference cosine to bmax: 0.587
01:03:00 [INFO] avg block-8 max-avg difference cosine to bmax: 0.493
01:03:00 [INFO] avg block-8 avg centroid document distance: 4246.156
01:03:00 [INFO]
01:03:00 [INFO] *** SSE ***: 33969.226
01:03:00 [INFO] *** EBT *** 0.196
01:03:00 [INFO] std cluster sizes (SDCS): 5.322
01:03:00 [INFO] normalized cluster entropy (NCE): 5.039
01:03:00 [INFO]
01:03:00 [INFO] avg terms per block-8: 907.693
01:03:00 [INFO] max terms per block-8: 53849
01:03:00 [INFO] min terms per block-8: 7
01:03:00 [INFO]
01:03:00 [INFO] avg blocks per term: 64.879
01:03:00 [INFO]
01:03:00 [INFO] ALL DONE! duration: 78.30 secs
So, with this configuration, we can see that we are using:
- the
kmeansgain function, - minimum postings length of 4096,
- maximum postings length of 0.1 * N (where N is the number of documents in the collection),
- 20 iterations per level, and
- the recursion depth fixed by only recursing while there are more than 16 elements within each partition.
And from the BEFORE and AFTER metrics, we can observe that the RGB process with GuideKP took about 18 seconds, increased the average and max dot product similarity between documents within blocks, improved EBT (estimated boundsum tightness) from 0.161 to 0.196, shifted the distribution of terms per block, decreasing the average but increasing the min/max thresholds, and so on.
If you only want to compute the metrics of a given CIFF file, including the flag --metrics-only will skip the reordering step and simply compute the metrics on the --input CIFF index.
For comparison to other baselines, we also reused the bisecting step structure in this codebase to replicate the bisecting K-means algorithm from Steinbach et al., 2000. To run this other implementation, add the --bisecting-km flag.
../target/release/create-rgb --input robust.ciff --output-ciff robust-reordered.ciff --loggap --bisecting-km
Bisecting-km has additional parameters, restarts and max-iter, to set the number of times k-means is run with a randomly selected starting centroid, and how many recursive splits should be made. Unless the --bisecting-km flag is included, these parameters will have no effect.
A full suite of settings can be found using the --help flag, and is listed below:
Reorders documents using recursive graph bisection and ciff files.
USAGE:
create-rgb [FLAGS] [OPTIONS] --input <input>
FLAGS:
--bisecting-km Enable bisecting_km <<<< NEW
--block-aware Enable block-aware mode <<<< NEW
--cooling Enable Cooling
-h, --help Prints help information
-l, --loggap Show loggap cost
--metrics-only Only compute metrics on the input CIFF, no reordering <<<< NEW
--mg-adj Enable adjusted distances <<<< NEW
--quickselect Enable QuickSelect
--sort-leaf Sort leaf by identifier
--swap-gain Enable swap gain <<<< NEW
--use-inputblocks Enable use_inputblocks <<<< NEW
-V, --version Prints version information
OPTIONS:
-b, --blocksize-txt <blocksize-txt> input blocksizes [default: blocks.txt] <<<< NEW
-c, --cutoff-frequency <cutoff-frequency>
Maximum length to consider in percentage of documents in the index [default: 0.1]
-i, --input <input> Input file ciff file
--input-fidx <input-fidx> Read forward index
--max-depth <max-depth> Maximum depth [default: 100]
--max-iter <max-iter> BisectingKM parameter: Maximum iterations [default: 100] <<<< NEW
-m, --min-len <min-len> Minimum number of occurrences to consider [default: 4096]
-o, --output-ciff <output-ciff> Output ciff file
--output-fidx <output-fidx> Output forward index
--output-mapping <output-mapping> Dump the document map
-r, --recursion-stop <recursion-stop> Min partition size [default: 16]
--restarts <restarts> BisectingKM parameter: # of Restarts [default: 10] <<<< NEW
-s, --swap-iterations <swap-iterations> Swap iterations [default: 20]
Note: cooling and quickselect were originally hardcoded options of Faster Recursive Graph Bisection BP algorithms. For ease of use, we turned them into flags to toggle them on and off from the command line.