Skip to content

fix(python): preserve PQ num_bits in model training#7583

Open
ddupg wants to merge 2 commits into
lance-format:mainfrom
ddupg:fix/pq-num-bits
Open

fix(python): preserve PQ num_bits in model training#7583
ddupg wants to merge 2 commits into
lance-format:mainfrom
ddupg:fix/pq-num-bits

Conversation

@ddupg

@ddupg ddupg commented Jul 2, 2026

Copy link
Copy Markdown
Contributor

Background

While building vector indexes through Ray-based distributed indexing, we found that the PQ num_bits option could be ignored by the global PQ training path. The codebook is parameterized by this value: for example, 4-bit PQ should train 16 centroids per sub-vector, while 8-bit PQ trains 256.

Before this change, Python/PyO3 PQ training always used 8 bits internally. That can produce a PQ codebook whose bit width does not match the later index build settings, which can make pre-trained PQ artifacts inconsistent with segment construction and hurt quantization/index quality.

Summary

  • expose num_bits on Python PQ training helpers and keep the default at 8 bits
  • pass num_bits through the PyO3 PQ training path instead of hard-coding 8
  • persist num_bits in saved PqModel metadata and reuse it when building from pre-trained PQ models

Tests

  • make build PYTHON=3.12
  • uv run --frozen pytest python/tests/test_indices.py::test_gen_pq python/tests/test_indices.py::test_indices_builder_multivector_distributed_dimensions
  • uv run --frozen ruff check python/tests/test_indices.py
  • uv run --frozen ruff format --check python/tests/test_indices.py

@ddupg ddupg force-pushed the fix/pq-num-bits branch from c4a6a9e to 23e6d9e Compare July 2, 2026 12:15
@ddupg ddupg marked this pull request as ready for review July 2, 2026 12:19
@wjones127 wjones127 self-assigned this Jul 2, 2026
@wjones127 wjones127 self-requested a review July 2, 2026 21:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

A-python Python bindings bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants