fix(python): preserve PQ num_bits in model training

python/python/lance/indices/builder.py

@@ -164,6 +164,7 @@ def train_pq(
         *,
         sample_rate: int = 256,
         max_iters: int = 50,
+        num_bits: int = 8,
         fragment_ids: Optional[list[int]] = None,
     ) -> PqModel:
         """
@@ -195,14 +196,16 @@ def train_pq(
             This parameter is used in the same way as in the IVF model.
         max_iters: int
             This parameter is used in the same way as in the IVF model.
+        num_bits: int
+            The number of bits used to encode each PQ centroid.
         fragment_ids: list[int], optional
             If provided, train using only the specified fragments from the dataset.
         """
         from lance.lance import indices
 
         num_rows = self._count_rows(fragment_ids)
         num_subvectors = self._normalize_pq_params(num_subvectors, self.dimension)
-        self._verify_pq_sample_rate(num_rows, sample_rate)
+        self._verify_pq_sample_rate(num_rows, sample_rate, num_bits)
         distance_type = ivf_model.distance_type
         pq_codebook = indices.train_pq_model(
             self.dataset._ds,
@@ -214,8 +217,9 @@ def train_pq(
             max_iters,
             ivf_model.centroids,
             fragment_ids,
+            num_bits,
         )
-        return PqModel(num_subvectors, pq_codebook)
+        return PqModel(num_subvectors, pq_codebook, num_bits=num_bits)
 
     def prepare_global_ivf_pq(
         self,
@@ -226,6 +230,7 @@ def prepare_global_ivf_pq(
         accelerator: Optional[Union[str, "torch.Device"]] = None,
         sample_rate: int = 256,
         max_iters: int = 50,
+        num_bits: int = 8,
         fragment_ids: Optional[list[int]] = None,
     ) -> dict:
         """
@@ -267,6 +272,7 @@ def prepare_global_ivf_pq(
             num_subvectors,
             sample_rate=sample_rate,
             max_iters=max_iters,
+            num_bits=num_bits,
             fragment_ids=fragment_ids,
         )
 
@@ -381,6 +387,7 @@ def transform_vectors(
             dest_uri,
             fragments,
             partition_ds_uri,
+            pq.num_bits,
         )
 
     def shuffle_transformed_vectors(
@@ -471,6 +478,7 @@ def load_shuffled_vectors(
                 num_subvectors,
                 distance_type,
                 index_name,
+                pq.num_bits,
             )
         else:
             raise ValueError("filenames must be a list of strings")
@@ -526,13 +534,17 @@ def _verify_base_sample_rate(self, sample_rate: int):
                 f"The sample_rate must be an int greater than 1, got {sample_rate}"
             )
 
-    def _verify_pq_sample_rate(self, num_rows: int, sample_rate: int):
+    def _verify_pq_sample_rate(
+        self, num_rows: int, sample_rate: int, num_bits: int = 8
+    ):
         self._verify_base_sample_rate(sample_rate)
-        if 256 * sample_rate > num_rows:
+        required_rows = (2**num_bits) * sample_rate
+        if required_rows > num_rows:
             raise ValueError(
                 "There are not enough rows in the dataset to create PQ"
-                f" codebook with a sample rate of {sample_rate}.  {sample_rate * 256}"
-                f" rows needed and there are {num_rows}"
+                f" codebook with a sample rate of {sample_rate} and num_bits"
+                f" of {num_bits}.  {required_rows} rows needed and there are"
+                f" {num_rows}"
             )
 
     def _verify_ivf_sample_rate(

python/python/lance/indices/pq.py

@@ -14,9 +14,13 @@ class PqModel:
     Can be saved / loaded to checkpoint progress.
     """
 
-    def __init__(self, num_subvectors: int, codebook: pa.FixedSizeListArray):
+    def __init__(
+        self, num_subvectors: int, codebook: pa.FixedSizeListArray, *, num_bits: int = 8
+    ):
         self.num_subvectors = num_subvectors
         """The number of subvectors to divide source vectors into"""
+        self.num_bits = num_bits
+        """The number of bits used to encode each PQ centroid"""
         self.codebook = codebook
         """The centroids of the PQ clusters"""
 
@@ -42,7 +46,10 @@ def save(self, uri: str, *, storage_options: Optional[Dict[str, str]] = None):
             uri,
             pa.schema(
                 [pa.field("codebook", self.codebook.type)],
-                metadata={b"num_subvectors": str(self.num_subvectors).encode()},
+                metadata={
+                    b"num_subvectors": str(self.num_subvectors).encode(),
+                    b"num_bits": str(self.num_bits).encode(),
+                },
             ),
             storage_options=storage_options,
         ) as writer:
@@ -65,9 +72,10 @@ def load(cls, uri: str, *, storage_options: Optional[Dict[str, str]] = None):
         """
         reader = LanceFileReader(uri, storage_options=storage_options)
         num_rows = reader.metadata().num_rows
-        metadata = reader.metadata().schema.metadata
+        metadata = reader.metadata().schema.metadata or {}
         num_subvectors = int(metadata[b"num_subvectors"].decode())
+        num_bits = int(metadata.get(b"num_bits", b"8").decode())
         codebook = (
             reader.read_all(batch_size=num_rows).to_table().column("codebook").chunk(0)
         )
-        return cls(num_subvectors, codebook)
+        return cls(num_subvectors, codebook, num_bits=num_bits)

python/python/lance/lance/indices/__init__.pyi

@@ -48,6 +48,7 @@ def train_pq_model(
     max_iters: int,
     ivf_model: pa.Array,
     fragment_ids: Optional[list[int]] = None,
+    num_bits: int = 8,
 ) -> pa.Array: ...
 def transform_vectors(
     dataset,
@@ -58,6 +59,9 @@ def transform_vectors(
     ivf_centroids: pa.Array,
     pq_codebook: pa.Array,
     dst_uri: str,
+    fragments: list,
+    partitions_ds_uri: Optional[str] = None,
+    num_bits: int = 8,
 ): ...
 def build_rq_model(
     dimension: int,

python/python/tests/test_indices.py

@@ -209,6 +209,18 @@ def test_gen_pq(tmpdir, rand_dataset, rand_ivf):
     assert pq.dimension == reloaded.dimension
     assert pq.codebook == reloaded.codebook
 
+    pq_4bit = IndicesBuilder(rand_dataset, "vectors").train_pq(
+        rand_ivf,
+        sample_rate=2,
+        num_bits=4,
+    )
+    assert pq_4bit.num_bits == 4
+    assert len(pq_4bit.codebook) == 16
+
+    pq_4bit.save(str(tmpdir / "pq_4bit"))
+    reloaded = PqModel.load(str(tmpdir / "pq_4bit"))
+    assert reloaded.num_bits == 4
+
 
 def test_ivf_centroids_fragment_ids(tmpdir):
     rows_per_fragment = 32
@@ -300,7 +312,7 @@ def test_indices_builder_multivector_distributed_dimensions(tmpdir, monkeypatch)
 
     captured_dimensions = {}
 
-    def train_pq_model(*args):
+    def train_pq_model(*args, **kwargs):
         captured_dimensions["train_pq"] = args[2]
         return codebook
 

python/src/indices.rs

@@ -232,14 +232,15 @@ async fn do_train_pq_model(
     distance_type: &str,
     sample_rate: u32,
     max_iters: u32,
+    num_bits: u32,
     ivf_model: IvfModel,
     fragment_ids: Option<Vec<u32>>,
 ) -> PyResult<ArrayData> {
     // We verify distance_type earlier so can unwrap here
     let distance_type = DistanceType::try_from(distance_type).unwrap();
     let params = PQBuildParams {
         num_sub_vectors: num_subvectors as usize,
-        num_bits: 8,
+        num_bits: num_bits as usize,
         max_iters: max_iters as usize,
         sample_rate: sample_rate as usize,
         ..Default::default()
@@ -260,7 +261,7 @@ async fn do_train_pq_model(
 
 #[pyfunction]
 #[allow(clippy::too_many_arguments)]
-#[pyo3(signature=(dataset, column, dimension, num_subvectors, distance_type, sample_rate, max_iters, ivf_centroids, fragment_ids=None))]
+#[pyo3(signature=(dataset, column, dimension, num_subvectors, distance_type, sample_rate, max_iters, ivf_centroids, fragment_ids=None, num_bits=8))]
 fn train_pq_model<'py>(
     py: Python<'py>,
     dataset: &Dataset,
@@ -272,6 +273,7 @@ fn train_pq_model<'py>(
     max_iters: u32,
     ivf_centroids: PyArrowType<ArrayData>,
     fragment_ids: Option<Vec<u32>>,
+    num_bits: u32,
 ) -> PyResult<Bound<'py, PyAny>> {
     let ivf_centroids = ivf_centroids.0;
     let ivf_centroids = FixedSizeListArray::from(ivf_centroids);
@@ -291,6 +293,7 @@ fn train_pq_model<'py>(
             distance_type,
             sample_rate,
             max_iters,
+            num_bits,
             ivf_model,
             fragment_ids,
         ),
@@ -398,7 +401,7 @@ async fn do_transform_vectors(
 
 #[pyfunction]
 #[allow(clippy::too_many_arguments)]
-#[pyo3(signature=(dataset, column, dimension, num_subvectors, distance_type, ivf_centroids, pq_codebook, dst_uri, fragments, partitions_ds_uri=None))]
+#[pyo3(signature=(dataset, column, dimension, num_subvectors, distance_type, ivf_centroids, pq_codebook, dst_uri, fragments, partitions_ds_uri=None, num_bits=8))]
 pub fn transform_vectors(
     py: Python<'_>,
     dataset: &Dataset,
@@ -411,6 +414,7 @@ pub fn transform_vectors(
     dst_uri: &str,
     fragments: Vec<FileFragment>,
     partitions_ds_uri: Option<&str>,
+    num_bits: u32,
 ) -> PyResult<()> {
     let ivf_centroids = ivf_centroids.0;
     let ivf_centroids = FixedSizeListArray::from(ivf_centroids);
@@ -419,7 +423,7 @@ pub fn transform_vectors(
     let distance_type = DistanceType::try_from(distance_type).unwrap();
     let pq = ProductQuantizer::new(
         num_subvectors as usize,
-        /*num_bits=*/ 8,
+        num_bits,
         dimension,
         codebook,
         distance_type,
@@ -561,7 +565,7 @@ async fn do_load_shuffled_vectors(
 }
 
 #[pyfunction]
-#[pyo3(signature=(filenames, dir_path, dataset, column, ivf_centroids, pq_codebook, pq_dimension, num_subvectors, distance_type, index_name=None))]
+#[pyo3(signature=(filenames, dir_path, dataset, column, ivf_centroids, pq_codebook, pq_dimension, num_subvectors, distance_type, index_name=None, num_bits=8))]
 #[allow(clippy::too_many_arguments)]
 pub fn load_shuffled_vectors(
     filenames: Vec<String>,
@@ -574,6 +578,7 @@ pub fn load_shuffled_vectors(
     num_subvectors: u32,
     distance_type: &str,
     index_name: Option<&str>,
+    num_bits: u32,
 ) -> PyResult<()> {
     let mut default_idx_name = column.to_string();
     default_idx_name.push_str("_idx");
@@ -595,7 +600,7 @@ pub fn load_shuffled_vectors(
     let distance_type = DistanceType::try_from(distance_type).unwrap();
     let pq_model = ProductQuantizer::new(
         num_subvectors as usize,
-        /*num_bits=*/ 8,
+        num_bits,
         pq_dimension,
         codebook,
         distance_type,