perf!: store PQ codebook in global buffer instead of scheme metadata (#3829)

storing the PQ codebook in schema metadata is not efficient as it's in
JSON format.
this causes slow cold latency and high memory footprint when opening the
index, so move it to global buffer, and store it as tensor protobuf
format

---------

Signed-off-by: BubbleCal <bubble-cal@outlook.com>
Co-authored-by: Will Jones <willjones127@gmail.com>

Author: BubbleCal

.github/workflows/python.yml

@@ -120,14 +120,16 @@ jobs:
       - uses: ./.github/workflows/run_tests
       - name: Generate forward compatibility files
         run: python python/tests/forward_compat/datagen.py
-      - name: Install old wheel
+      - name: Run forward compatibility tests (pylance 0.16.0)
         run: |
           python -m venv venv
           source venv/bin/activate
           pip install pytest pylance==0.16.0
-      - name: Run forward compatibility tests
+          pytest python/tests/forward_compat --run-forward
+      - name: Run forward compatibility tests (pylance 0.29.1.beta2)
         run: |
           source venv/bin/activate
+          pip install pytest --pre --extra-index-url https://pypi.fury.io/lancedb/ pylance==0.29.1.beta2
           pytest python/tests/forward_compat --run-forward
       # Make sure wheels are not included in the Rust cache
       - name: Delete wheels

Cargo.lock

@@ -0,0 +1 @@
+../../../rust-toolchain.toml

java/core/lance-jni/rust-toolchain.toml

@@ -5,7 +5,9 @@
 
 from pathlib import Path
 
+import lance
 import pyarrow as pa
+import pyarrow.compute as pc
 from lance.file import LanceFileWriter
 
 
@@ -92,6 +94,32 @@ def write_large():
         writer.write_batch(build_large())
 
 
+def write_dataset_pq_buffer():
+    # In https://github.com/lancedb/lance/pull/3829, we started storing the PQ
+    # codebook in a global buffer instead of the schema metadata as JSON.
+
+    ndims = 32
+    nvecs = 512
+
+    data = pa.table(
+        {
+            "id": pa.array(range(nvecs)),
+            "vec": pa.FixedSizeListArray.from_arrays(
+                pc.random(ndims * nvecs).cast(pa.float32()), ndims
+            ),
+        }
+    )
+
+    dataset = lance.write_dataset(data, get_path("pq_in_schema"))
+    dataset.create_index(
+        "vec",
+        "IVF_PQ",
+        num_partitions=1,
+        num_sub_vectors=4,
+    )
+
+
 if __name__ == "__main__":
     write_basic_types()
     write_large()
+    write_dataset_pq_buffer()

python/python/tests/forward_compat/datagen.py

@@ -1,8 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright The Lance Authors
 
+import lance
+import pyarrow as pa
+import pyarrow.compute as pc
 import pytest
 from lance.file import LanceFileReader
+from packaging.version import Version
 
 from .datagen import build_basic_types, build_large, get_path
 
@@ -18,3 +22,21 @@ def test_scans():
     expected_large = build_large()
     actual_large = LanceFileReader(str(get_path("large.lance"))).read_all().to_table()
     assert actual_large.equals(expected_large)
+
+
+@pytest.mark.forward
+@pytest.mark.skipif(
+    Version(lance.__version__) < Version("0.29.1.beta2"),  # at least 0.29.1-beta.2
+    reason="Lance 0.29.1-beta.2 would ignore indices too new",
+)
+def test_pq_buffer():
+    ds = lance.dataset(get_path("pq_in_schema"))
+    # the index should be ignored, still able to query (brute force)
+    q = pc.random(32).cast(pa.float32())
+    ds.to_table(
+        nearest={
+            "q": q,
+            "k": 4,
+            "column": "vec",
+        }
+    )

python/python/tests/forward_compat/test_compat.py

@@ -0,0 +1 @@
+../rust-toolchain.toml

python/rust-toolchain.toml

@@ -191,7 +191,7 @@ impl IndexType {
             | Self::IvfSq
             | Self::IvfPq
             | Self::IvfHnswSq
-            | Self::IvfHnswPq => 0,
+            | Self::IvfHnswPq => 1,
         }
     }
 }

rust/lance-index/src/lib.rs

@@ -208,12 +208,8 @@ impl Quantization for FlatQuantizer {
         }))
     }
 
-    fn metadata(
-        &self,
-        _: Option<crate::vector::quantizer::QuantizationMetadata>,
-    ) -> Result<serde_json::Value> {
-        let metadata = FlatMetadata { dim: self.dim };
-        Ok(serde_json::to_value(metadata)?)
+    fn metadata(&self, _: Option<crate::vector::quantizer::QuantizationMetadata>) -> FlatMetadata {
+        FlatMetadata { dim: self.dim }
     }
 
     fn metadata_key() -> &'static str {
@@ -290,12 +286,8 @@ impl Quantization for FlatBinQuantizer {
         }))
     }
 
-    fn metadata(
-        &self,
-        _: Option<crate::vector::quantizer::QuantizationMetadata>,
-    ) -> Result<serde_json::Value> {
-        let metadata = FlatMetadata { dim: self.dim };
-        Ok(serde_json::to_value(metadata)?)
+    fn metadata(&self, _: Option<crate::vector::quantizer::QuantizationMetadata>) -> FlatMetadata {
+        FlatMetadata { dim: self.dim }
     }
 
     fn metadata_key() -> &'static str {