diff --git a/.bumpversion.toml b/.bumpversion.toml index 0909579cb..7dcefa69c 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "0.41.0-beta.13" +current_version = "0.41.0-beta.14" parse = """(?x) (?P0|[1-9]\\d*)\\. (?P0|[1-9]\\d*)\\. diff --git a/.github/workflows/pypi-publish.yml b/.github/workflows/pypi-publish.yml index e71564473..1784c4c8a 100644 --- a/.github/workflows/pypi-publish.yml +++ b/.github/workflows/pypi-publish.yml @@ -15,6 +15,9 @@ on: - ci/check_cargo_lock.py - ci/check_sdist.py - python/pyproject.toml + - python/Cargo.toml + - python/vendor/** + - rust/lancedb/Cargo.toml - Cargo.toml # Change in dependency frequently breaks builds - Cargo.lock diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 930ae67e2..17fe3a7b4 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -29,6 +29,63 @@ env: RUST_BACKTRACE: "1" jobs: + dependency-overrides: + name: Downstream Lance override + runs-on: ubuntu-24.04 + timeout-minutes: 15 + steps: + - uses: actions/checkout@v6 + - uses: actions-rust-lang/setup-rust-toolchain@v1 + - name: Read the pinned Lance tag + id: lance + run: | + python3 - <<'PYTHON' + import os + import tomllib + from pathlib import Path + + manifest = tomllib.loads(Path("Cargo.toml").read_text()) + tag = manifest["workspace"]["dependencies"]["lance"]["tag"] + with open(os.environ["GITHUB_OUTPUT"], "a") as output: + output.write(f"tag={tag}\n") + PYTHON + - uses: actions/checkout@v6 + with: + repository: lance-format/lance + ref: ${{ steps.lance.outputs.tag }} + path: local-lance + - name: Resolve a downstream workspace with a local Lance override + run: | + mkdir -p downstream/src + touch downstream/src/lib.rs + cat > downstream/Cargo.toml <<'TOML' + [package] + name = "lancedb-local-lance-consumer" + version = "0.1.0" + edition = "2024" + + [workspace] + + [dependencies] + lancedb = { path = "../rust/lancedb", default-features = false } + lance-arrow = { path = "../local-lance/rust/lance-arrow" } + + [patch."https://github.com/lance-format/lance.git"] + lance-arrow = { path = "../local-lance/rust/lance-arrow" } + TOML + cargo metadata --manifest-path downstream/Cargo.toml --format-version 1 > downstream/graph.json + python3 - <<'PYTHON' + import json + from pathlib import Path + + graph = json.loads(Path("downstream/graph.json").read_text()) + packages = [p for p in graph["packages"] if p["name"] == "lance-arrow"] + assert len(packages) == 1, [p["id"] for p in packages] + assert Path(packages[0]["manifest_path"]).resolve() == Path( + "local-lance/rust/lance-arrow/Cargo.toml" + ).resolve() + PYTHON + lint: timeout-minutes: 30 runs-on: ubuntu-24.04 diff --git a/Cargo.lock b/Cargo.lock index fac89c4da..03f4b0864 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5700,7 +5700,7 @@ dependencies = [ [[package]] name = "lancedb" -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" dependencies = [ "ahash", "anyhow", @@ -5798,7 +5798,7 @@ dependencies = [ [[package]] name = "lancedb-nodejs" -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" dependencies = [ "arrow-array", "arrow-buffer", @@ -5823,7 +5823,7 @@ dependencies = [ [[package]] name = "lancedb-python" -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" dependencies = [ "arc-swap", "arrow", diff --git a/ci/set_lance_version.py b/ci/set_lance_version.py index c7bf41f55..bef97f8ad 100644 --- a/ci/set_lance_version.py +++ b/ci/set_lance_version.py @@ -2,6 +2,8 @@ import argparse import re import sys import json +import tomllib +from pathlib import Path LANCE_GIT_URL = "https://github.com/lance-format/lance.git" @@ -148,6 +150,18 @@ def update_cargo_toml(line_updater): with open("Cargo.toml", "w") as f: f.writelines(new_lines) + core_manifest = Path("rust/lancedb/Cargo.toml") + if core_manifest.exists(): + lines = core_manifest.read_text().splitlines(keepends=True) + for index, line in enumerate(lines): + if line.startswith("lance-arrow ="): + updated = tomllib.loads(line_updater(line))["lance-arrow"] + config = {"version": updated} if isinstance(updated, str) else updated + if "path" in config: + config["path"] = "../../" + config["path"] + lines[index] = dict_to_toml_line("lance-arrow", config) + core_manifest.write_text("".join(lines)) + def set_stable_version(version: str): """ diff --git a/ci/tests/test_set_lance_version.py b/ci/tests/test_set_lance_version.py index 1493fc59d..1173eab84 100644 --- a/ci/tests/test_set_lance_version.py +++ b/ci/tests/test_set_lance_version.py @@ -4,6 +4,7 @@ import subprocess import sys import tempfile import textwrap +import tomllib import unittest from pathlib import Path @@ -83,6 +84,13 @@ lance_datafusion = {{ "version" = "=8.2.0-beta.4", "features" = ["substrait"], " with self.subTest(version=version), tempfile.TemporaryDirectory() as tmp: workdir = Path(tmp) (workdir / "Cargo.toml").write_text(CARGO_TOML) + core_manifest = workdir / "rust/lancedb/Cargo.toml" + core_manifest.parent.mkdir(parents=True) + core_manifest.write_text( + "[dependencies]\nlance-arrow = { " + 'version = "=1.0.0", git = "https://github.com/lance-format/lance.git", ' + 'tag = "v1.0.0" }\narrow.workspace = true\n' + ) command_log = workdir / "commands.log" fake_bin = workdir / "bin" fake_bin.mkdir() @@ -111,6 +119,26 @@ lance_datafusion = {{ "version" = "=8.2.0-beta.4", "features" = ["substrait"], " + updated_dependencies + UNTOUCHED_DEPENDENCIES, ) + core = tomllib.loads(core_manifest.read_text())["dependencies"] + dependency = core["lance-arrow"] + if isinstance(dependency, str): + dependency = {"version": dependency} + if version == "local": + self.assertEqual( + dependency["path"], "../../../lance/rust/lance-arrow" + ) + else: + expected_version = { + "stable": "9.9.9", + "preview": "10.0.0-beta.3", + }.get(version, version.removeprefix("v")) + self.assertEqual(dependency["version"], f"={expected_version}") + if "beta" in expected_version: + self.assertEqual(dependency["git"], LANCE_GIT_URL) + self.assertEqual(dependency["tag"], f"v{expected_version}") + else: + self.assertNotIn("git", dependency) + self.assertEqual(core["arrow"], {"workspace": True}) commands = command_log.read_text().splitlines() for command in expected_commands: self.assertTrue( diff --git a/deny.toml b/deny.toml index 77412709f..e4982c149 100644 --- a/deny.toml +++ b/deny.toml @@ -177,7 +177,9 @@ wildcards = "deny" # `[workspace.dependencies]` table: any crate used by more than one member must # go through `workspace = true`, and entries nothing uses are an error. This # keeps versions from drifting between the core crate and the bindings. -workspace-dependencies = { duplicates = "deny", unused = "deny" } +# Python bundles a patched path crate while the Rust SDK uses the overridable +# upstream source. Check shared external declarations, not these path bindings. +workspace-dependencies = { duplicates = "deny", unused = "deny", include-path-dependencies = false } # Internal workspace crates reference each other via `path = "..."`, which # cargo-deny sees as a wildcard version. That's fine for private workspace # members (not published to crates.io), so allow it specifically for paths. diff --git a/docs/src/java/java.md b/docs/src/java/java.md index 0c887f66f..9dbccbc81 100644 --- a/docs/src/java/java.md +++ b/docs/src/java/java.md @@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`: com.lancedb lancedb-core - 0.41.0-beta.13 + 0.41.0-beta.14 ``` diff --git a/docs/src/js/interfaces/IvfRqOptions.md b/docs/src/js/interfaces/IvfRqOptions.md index 6b48cf653..298565a62 100644 --- a/docs/src/js/interfaces/IvfRqOptions.md +++ b/docs/src/js/interfaces/IvfRqOptions.md @@ -58,7 +58,8 @@ Number of bits per dimension for residual quantization. This value controls how much each residual component is compressed. The more bits, the more accurate the index will be but the slower search. Typical values -are small integers; the default is 1 bit per dimension. +are small integers. If unset, uses Lance's default (currently 5 bits per +dimension). Set to 1 for maximum compression. *** diff --git a/docs/src/python/python.md b/docs/src/python/python.md index bd8ee486b..5465eaaf2 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -297,6 +297,25 @@ listing a storage directory. ## Tables (Synchronous) +Local and remote tables use the same arguments for index creation, deletion, +index statistics, and search. For example, `table.delete(where="id = 1")`, +`table.create_index("cosine", 2)` (the deprecated vector API), and +`table.search([1.0, 1.0], fast_search=True)` work with either backend. +`fast_search=True` excludes rows that have not been indexed. + +`create_scalar_index` defaults to `replace=True` on both backends. Both scalar +and full-text index creation accept `wait_timeout`. Remote maintenance methods +accept the local options: cleanup and compaction warn that they are no-ops, +while `optimize` reports an unsupported operation. Removed Tantivy options +produce the same explanatory errors on both backends. + +The old remote keyword names `predicate` for `delete`, `index_name` for +`drop_index`, `index_uuid` for `index_stats`, and `column` for `create_fts_index` +remain available with deprecation warnings. Use `where`, `name`, `index_name`, +and `field_names`, respectively. For positional calls, use the argument order +documented by `Table` below; specify index options by keyword when migrating +code that used the former remote-only argument order. + ::: lancedb.table.Table ::: lancedb.table.CompactionOptions diff --git a/docs/web/indexing/vector-index.mdx b/docs/web/indexing/vector-index.mdx index b26e0c295..037d75f8d 100644 --- a/docs/web/indexing/vector-index.mdx +++ b/docs/web/indexing/vector-index.mdx @@ -58,7 +58,7 @@ Lance with `index_type`. | If your priority is... | Use this index | Why | Approx. compression ratio | Python config class | | :--- | :--- | :--- | :--- | :--- | -| Maximum compression | `IVF_RQ` | IVF indexing with RaBitQ quantization | Around `1/32` of raw size | `IvfRq` | +| Maximum compression | `IVF_RQ` | IVF indexing with RaBitQ quantization | Around `1/32` of raw size with `num_bits=1` | `IvfRq` | | Higher accuracy at small dimensions (`dimension <= 256`) | `IVF_PQ` | IVF indexing with product quantization | Usually `1/64` to `1/16` of raw size (depends on `num_sub_vectors`) | `IvfPq` | | Best recall/latency trade-off | `IVF_HNSW_SQ` | IVF-HNSW indexing with scalar | Typically a little larger than `1/4` of raw size | `IvfHnswSq` | | Highest recall / no quantization | `IVF_HNSW_FLAT` | IVF-HNSW indexing with no quantization; embedded library only | Around raw vector size plus HNSW graph overhead | `IvfHnswFlat` | @@ -377,7 +377,7 @@ which has been compressed to $4$ chunks $\times 8$-bit integers $= 32$ bits of q RaBitQ is an advanced quantization technique that outperforms PQ in several ways. It needs no codebook to train, estimates distances very quickly at query-time, -and crucially, quantizes each vector in (with some small overhead) just **one bit per dimension!** +and can quantize each vector in (with some small overhead) as little as **one bit per dimension!** In practice, RaBitQ compresses a $1024$-dimensional `float32` vector into just a few thousand bits, while maintaining good recall. The inner workings of RaBitQ quantization are rather mathematically dense. It generates a quantization codebook by @@ -385,14 +385,19 @@ applying a uniformly random, approximately distance-preserving orthogonal transf where $d$ is the dimensionality of the dataset. We defer the details, and an elegant theoretical error bound, to [the original paper.](https://arxiv.org/pdf/2405.12497) ##### Using RaBitQ {#using-rabitq} -Use RaBitQ quantization by selecting quantized index types ending in the suffix `RQ`. For example, call `create_index` with `index_type="IVF_RQ"`. +Use RaBitQ quantization by selecting quantized index types ending in the suffix `RQ`. For example, import `IvfRq` from `lancedb.index` and call `table.create_index("vector", config=IvfRq())`. Note that when using `IVF_RQ`, the dimension of the dataset must be a multiple of `8`. `num_bits` determines how many bits are used to quantize each dimension. -`1` is the standard RaBitQ setting. Increase to `2`, `4`, or `8` bits to achieve better recall for additional storage and query-time compute. +By default, Python, Rust, and Node leave this unset and inherit Lance's default, currently `5` bits. +Higher values improve recall at the cost of additional storage and query-time compute. +To retain the previous one-bit default and maximum compression, explicitly set `num_bits=1` (for example, `IvfRq(num_bits=1)` in Python). +Five-bit quantized vector codes use roughly five times the storage of one-bit codes; total index size also includes other overhead. -RaBitQ-quantized indexes computed with `num_bits >= 2` use a newer on-disk layout, and cannot be read by some older LanceDB versions. +RaBitQ-quantized indexes computed with `num_bits >= 2` use a newer on-disk layout and cannot be read by some older LanceDB versions. +The new five-bit default therefore affects reader compatibility as well as index size. +Upgrade readers to versions that support multi-bit IVF_RQ before creating indexes with the default settings, or explicitly set `num_bits=1` when sharing indexes with readers that support only one-bit IVF_RQ. See this [blog post](https://lancedb.com/blog/feature-rabitq-quantization/) for further discussion and benchmarking of LanceDB's RaBitQ implementation. @@ -410,7 +415,7 @@ among all integral multiples $\{j \times 10^5: -128 \leq j \leq 127 \}$. | Parameter | Description | | :--- | :--- | -| `num_bits` | Bits per dimension for quantization. Only applies to `IVF_PQ`/`IVF_HNSW_PQ` (default `8`) and `IVF_RQ` (default `1`, RaBitQ); not used by `IVF_FLAT`/`IVF_SQ`/`IVF_HNSW_FLAT`/`IVF_HNSW_SQ`. Higher values improve accuracy at the cost of additional storage. | +| `num_bits` | Bits per dimension for quantization. Only applies to `IVF_PQ`/`IVF_HNSW_PQ` (default `8`) and `IVF_RQ` (`IvfRq()` inherits Lance's default, currently `5`, RaBitQ); not used by `IVF_FLAT`/`IVF_SQ`/`IVF_HNSW_FLAT`/`IVF_HNSW_SQ`. Higher values improve accuracy at the cost of additional storage. | `max_iterations` and `sample_rate` also affect quantizer training, but since they apply to every IVF/HNSW index type (not just quantized ones), they're documented as general [Build-time Parameters](/indexing/vector-index#build-time-parameters) instead. All three are passed as keyword arguments to `create_index`, alongside `index_type`: diff --git a/docs/web/ingestion/create-table.mdx b/docs/web/ingestion/create-table.mdx index 7d4363be7..0921de22e 100644 --- a/docs/web/ingestion/create-table.mdx +++ b/docs/web/ingestion/create-table.mdx @@ -148,9 +148,11 @@ non-nullable column, or provide actual nulls for it, ingestion fails; nullable c omitted or written with null values. Without an explicit schema, Python infers list-like vector values as fixed-size `float32` vector fields from the observed dimension. -For Python ingest, malformed vector values fail by default. If you expect occasional wrong-length, -null, or NaN vectors, choose an `on_bad_vectors` policy: `"drop"` removes those rows, `"fill"` writes -`fill_value`, and `"null"` writes nulls. +For Python ingestion, malformed vectors fail by default (`on_bad_vectors="error"`). This includes +vectors with incorrect sizes, NaNs, or null components, including pandas NaNs converted to Arrow +nulls. Use `"drop"` to remove affected rows, or `"null"` to replace affected vectors with nulls. +Use `"fill"` to replace NaN or null components with `fill_value` while preserving valid components. +This also pads short vectors with `fill_value` and truncates long vectors to the expected dimension. ### From an Arrow Table {#from-an-arrow-table} You can also create LanceDB tables directly from Arrow tables. diff --git a/java/lancedb-core/pom.xml b/java/lancedb-core/pom.xml index 87c5d4cb8..7c34a801a 100644 --- a/java/lancedb-core/pom.xml +++ b/java/lancedb-core/pom.xml @@ -8,7 +8,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.13 + 0.41.0-beta.14 ../pom.xml diff --git a/java/pom.xml b/java/pom.xml index 1f3e0efb0..2dc959887 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -6,7 +6,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.13 + 0.41.0-beta.14 pom ${project.artifactId} LanceDB Java SDK Parent POM diff --git a/nodejs/Cargo.toml b/nodejs/Cargo.toml index 1b9d47e46..e343d68ca 100644 --- a/nodejs/Cargo.toml +++ b/nodejs/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "lancedb-nodejs" edition.workspace = true -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" publish = false license.workspace = true description.workspace = true diff --git a/nodejs/lancedb/indices.ts b/nodejs/lancedb/indices.ts index 43d1cf1b5..0e383e984 100644 --- a/nodejs/lancedb/indices.ts +++ b/nodejs/lancedb/indices.ts @@ -131,7 +131,8 @@ export interface IvfRqOptions { * * This value controls how much each residual component is compressed. The more * bits, the more accurate the index will be but the slower search. Typical values - * are small integers; the default is 1 bit per dimension. + * are small integers. If unset, uses Lance's default (currently 5 bits per + * dimension). Set to 1 for maximum compression. */ numBits?: number; diff --git a/nodejs/npm/darwin-arm64/package.json b/nodejs/npm/darwin-arm64/package.json index bd4b95463..343483a8e 100644 --- a/nodejs/npm/darwin-arm64/package.json +++ b/nodejs/npm/darwin-arm64/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-darwin-arm64", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["darwin"], "cpu": ["arm64"], "main": "lancedb.darwin-arm64.node", diff --git a/nodejs/npm/linux-arm64-gnu/package.json b/nodejs/npm/linux-arm64-gnu/package.json index 6e4e57684..4c3b93c81 100644 --- a/nodejs/npm/linux-arm64-gnu/package.json +++ b/nodejs/npm/linux-arm64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-gnu", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-gnu.node", diff --git a/nodejs/npm/linux-arm64-musl/package.json b/nodejs/npm/linux-arm64-musl/package.json index f77c50ee2..7679a422d 100644 --- a/nodejs/npm/linux-arm64-musl/package.json +++ b/nodejs/npm/linux-arm64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-musl", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-musl.node", diff --git a/nodejs/npm/linux-x64-gnu/package.json b/nodejs/npm/linux-x64-gnu/package.json index b671b7ebf..6a62f8e8a 100644 --- a/nodejs/npm/linux-x64-gnu/package.json +++ b/nodejs/npm/linux-x64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-gnu", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-gnu.node", diff --git a/nodejs/npm/linux-x64-musl/package.json b/nodejs/npm/linux-x64-musl/package.json index b75738ff1..d68341382 100644 --- a/nodejs/npm/linux-x64-musl/package.json +++ b/nodejs/npm/linux-x64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-musl", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-musl.node", diff --git a/nodejs/npm/win32-arm64-msvc/package.json b/nodejs/npm/win32-arm64-msvc/package.json index 7158748f9..42a517310 100644 --- a/nodejs/npm/win32-arm64-msvc/package.json +++ b/nodejs/npm/win32-arm64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-arm64-msvc", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": [ "win32" ], diff --git a/nodejs/npm/win32-x64-msvc/package.json b/nodejs/npm/win32-x64-msvc/package.json index 83d0fb6f2..b116f2819 100644 --- a/nodejs/npm/win32-x64-msvc/package.json +++ b/nodejs/npm/win32-x64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-x64-msvc", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["win32"], "cpu": ["x64"], "main": "lancedb.win32-x64-msvc.node", diff --git a/nodejs/package.json b/nodejs/package.json index 21f903235..f746f6daf 100644 --- a/nodejs/package.json +++ b/nodejs/package.json @@ -11,7 +11,7 @@ "ann" ], "private": false, - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "main": "dist/index.js", "exports": { ".": "./dist/index.js", diff --git a/python/Cargo.toml b/python/Cargo.toml index 076299a7d..265e4d270 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb-python" -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" publish = false edition.workspace = true description = "Python bindings for LanceDB" @@ -22,7 +22,6 @@ bytes.workspace = true lancedb.workspace = true datafusion-common.workspace = true lance-core.workspace = true -# The workspace path dependency includes the patched crate in Maturin sources. lance-arrow.workspace = true lance-namespace.workspace = true lance-namespace-impls.workspace = true diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 1558dd52a..9ed772e1c 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -428,10 +428,13 @@ class DBConnection(EnforceOverrides): it will not add the provided data but will validate against any schema that's specified. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. - One of "error", "drop", "fill". + How to handle vectors with incorrect sizes, NaNs, or null elements. + One of "error", "drop", "fill", "null". fill_value: float The value to use when filling vectors. Only used if on_bad_vectors="fill". + NaN or null elements and missing trailing elements are replaced with this + value. Valid elements are preserved, and vectors longer than the expected + dimension are truncated. storage_options: dict, optional Additional options for the storage backend. Options already set on the connection will be inherited by the table, but can be overridden here. @@ -2550,10 +2553,13 @@ class AsyncConnection(object): it will not add the provided data but will validate against any schema that's specified. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. - One of "error", "drop", "fill". + How to handle vectors with incorrect sizes, NaNs, or null elements. + One of "error", "drop", "fill", "null". fill_value: float The value to use when filling vectors. Only used if on_bad_vectors="fill". + NaN or null elements and missing trailing elements are replaced with this + value. Valid elements are preserved, and vectors longer than the expected + dimension are truncated. storage_options: dict, optional Additional options for the storage backend. Options already set on the connection will be inherited by the table, but can be overridden here. diff --git a/python/python/lancedb/embeddings/open_clip.py b/python/python/lancedb/embeddings/open_clip.py index 8b3f0eb06..6199ef271 100644 --- a/python/python/lancedb/embeddings/open_clip.py +++ b/python/python/lancedb/embeddings/open_clip.py @@ -49,6 +49,9 @@ class OpenClipEmbeddings(EmbeddingFunction): self.name, pretrained=self.pretrained ) model.to(self.device) + # open_clip returns the model in training mode, which keeps dropout, + # stochastic depth and batch norm in their training behaviour. + model.eval() self._model, self._preprocess = model, preprocess self._tokenizer = open_clip.get_tokenizer(self.name) self._ndims = None diff --git a/python/python/lancedb/index.py b/python/python/lancedb/index.py index cf3bf8e4f..80179f4a7 100644 --- a/python/python/lancedb/index.py +++ b/python/python/lancedb/index.py @@ -882,8 +882,12 @@ class IvfRq: num_partitions: int, default sqrt(num_rows) Number of IVF partitions to create. - num_bits: int, default 1 + num_bits: int, optional Number of bits to encode each dimension in the RabitQ codebook. + If None, uses Lance's default (currently 5), matching Rust and Node. + Set to 1 to retain the previous Python default and maximum compression. + Five-bit vector codes use roughly five times the storage of one-bit codes, + in addition to other index overhead. max_iterations: int, default 50 Max iterations to train kmeans when computing IVF partitions. @@ -897,7 +901,7 @@ class IvfRq: distance_type: Literal["l2", "cosine", "dot"] = "l2" num_partitions: Optional[int] = None - num_bits: int = 1 + num_bits: Optional[int] = None max_iterations: int = 50 sample_rate: int = 256 target_partition_size: Optional[int] = None diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 62d00f284..c11e01ff4 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -569,10 +569,10 @@ class RemoteDBConnection(DBConnection): name: str, data: DATA = None, schema: Optional[Union[pa.Schema, LanceModel]] = None, - on_bad_vectors: str = "error", - fill_value: float = 0.0, mode: Optional[str] = None, exist_ok: bool = False, + on_bad_vectors: str = "error", + fill_value: float = 0.0, embedding_functions: Optional[List[EmbeddingFunctionConfig]] = None, *, namespace_path: Optional[List[str]] = None, diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index a01203a15..c07da79f8 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -4,7 +4,7 @@ from datetime import timedelta import deprecation import logging -from functools import cached_property +from functools import cached_property, wraps import os from typing import ( TYPE_CHECKING, @@ -42,19 +42,12 @@ from lancedb.index import ( FTS, BTree, Bitmap, - HnswFlat, - HnswPq, - HnswSq, - IvfFlat, - IvfPq, - IvfRq, - IvfSq, LabelList, ) from lancedb.job import Job from lancedb.functions import FunctionApplication, RefreshColumnResult from lancedb.remote.db import LOOP -from lancedb.table import IndexConfigType, KNOWN_METRICS +from lancedb.table import IndexConfigType, _maybe_add_fts_error_note import pyarrow as pa from lancedb.common import DATA, VEC, VECTOR_COLUMN_NAME @@ -69,7 +62,6 @@ from ..query import ( FullTextQuery, LanceQueryBuilder, LanceTakeQueryBuilder, - LanceVectorQueryBuilder, ) from ..table import ( AsyncTable, @@ -81,7 +73,7 @@ from ..table import ( Table, Tags, ) -from ..types import BaseTokenizerType +from ..types import BaseTokenizerType, ScalarIndexType from ..util import infer_vector_column_name if TYPE_CHECKING: @@ -89,6 +81,28 @@ if TYPE_CHECKING: import polars as pl +def _deprecated_keyword(old_name: str, new_name: str): + """Keep old remote-only keyword names without changing the Table signature.""" + + def decorate(method): + @wraps(method) + def wrapped(self, *args, **kwargs): + if old_name in kwargs: + if new_name in kwargs: + raise TypeError(f"Cannot specify both {old_name} and {new_name}") + warnings.warn( + f"{old_name} is deprecated; use {new_name} instead", + DeprecationWarning, + stacklevel=2, + ) + kwargs[new_name] = kwargs.pop(old_name) + return method(self, *args, **kwargs) + + return wrapped + + return decorate + + class RemoteTable(Table): def __init__( self, @@ -269,7 +283,7 @@ class RemoteTable(Table): """ return LOOP.run(self._table.to_lance(**kwargs)) - def to_polars(self, **kwargs) -> "pl.LazyFrame": + def to_polars(self, batch_size=None) -> "pl.LazyFrame": """Use ``table.search().to_polars()`` to read remote query results.""" raise NotImplementedError( "to_polars() is not supported for remote tables. " @@ -316,9 +330,10 @@ class RemoteTable(Table): self._table.tokenize(query, column=column, index_name=index_name) ) - def index_stats(self, index_uuid: str) -> Optional[IndexStatistics]: + @_deprecated_keyword("index_uuid", "index_name") + def index_stats(self, index_name: str) -> Optional[IndexStatistics]: """List all the stats of a specified index""" - return LOOP.run(self._table.index_stats(index_uuid)) + return LOOP.run(self._table.index_stats(index_name)) @deprecation.deprecated( deprecated_in="0.25.0", @@ -328,43 +343,38 @@ class RemoteTable(Table): def create_scalar_index( self, column: str, - index_type: Literal["BTREE", "BITMAP", "LABEL_LIST", "scalar"] = "scalar", *, - replace: bool = False, + replace: bool = True, + index_type: ScalarIndexType = "BTREE", wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, ): - """Creates a scalar index. + """Create a scalar index on a column. .. deprecated:: 0.25.0 Use :meth:`create_index` with a BTree, Bitmap, or LabelList config instead. Example: ``table.create_index("column", config=BTree())`` - - Parameters - ---------- - column : str - The column to be indexed. Must be a boolean, integer, float, - or string column. - index_type : str - The index type of the scalar index. Must be "scalar" (BTREE), - "BTREE", "BITMAP", or "LABEL_LIST", - replace : bool - If True, replace the existing index with the new one. """ - if index_type == "scalar" or index_type == "BTREE": + if index_type == "scalar": + warnings.warn( + 'index_type="scalar" is deprecated; use "BTREE" instead', + DeprecationWarning, + stacklevel=2, + ) + index_type = "BTREE" + if index_type == "BTREE": config = BTree() elif index_type == "BITMAP": config = Bitmap() elif index_type == "LABEL_LIST": config = LabelList() else: - raise ValueError(f"Unknown index type: {index_type}") - - LOOP.run( + raise ValueError(f"Unknown index type {index_type}") + return LOOP.run( self._table.create_index( column, - config=config, replace=replace, + config=config, wait_timeout=wait_timeout, name=name, ) @@ -375,12 +385,16 @@ class RemoteTable(Table): current_version=__version__, details="Use create_index() with config=FTS() instead.", ) + @_deprecated_keyword("column", "field_names") def create_fts_index( self, - column: str, + field_names: Union[str, List[str]], *, + ordering_field_names: Optional[Union[str, List[str]]] = None, replace: bool = False, - wait_timeout: Optional[timedelta] = None, + writer_heap_size: Optional[int] = 1024 * 1024 * 1024, + use_tantivy: bool = False, + tokenizer_name: Optional[str] = None, with_position: bool = False, # tokenizer configs: base_tokenizer: BaseTokenizerType = "simple", @@ -396,6 +410,7 @@ class RemoteTable(Table): prefix_only: bool = False, block_size: int = 128, document_granularity: DocumentGranularity = DocumentGranularity.ROW, + wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, ): """Create a full-text search index on a column. @@ -404,31 +419,70 @@ class RemoteTable(Table): Use :meth:`create_index` with an FTS config instead. Example: ``table.create_index("text_column", config=FTS())`` """ + if use_tantivy: + raise ValueError( + "Tantivy-based FTS has been removed. " + "Remove use_tantivy and recreate the index with native FTS." + ) + if ordering_field_names is not None: + raise ValueError( + "ordering_field_names was only supported by the removed " + "Tantivy-based FTS implementation." + ) + if writer_heap_size != 1024 * 1024 * 1024: + raise ValueError( + "writer_heap_size was only supported by the removed " + "Tantivy-based FTS implementation." + ) + if not isinstance(field_names, str): + raise ValueError( + "Native FTS indexes can only be created on a single field " + "at a time. To search over multiple text fields, create a " + "separate FTS index for each field." + ) + + if tokenizer_name is None: + tokenizer_configs = { + "base_tokenizer": base_tokenizer, + "language": language, + "with_position": with_position, + "max_token_length": max_token_length, + "lower_case": lower_case, + "stem": stem, + "remove_stop_words": remove_stop_words, + "custom_stop_words": custom_stop_words, + "ascii_folding": ascii_folding, + "ngram_min_length": ngram_min_length, + "ngram_max_length": ngram_max_length, + "prefix_only": prefix_only, + } + else: + tokenizer_configs = self.infer_tokenizer_configs(tokenizer_name) + tokenizer_configs["custom_stop_words"] = custom_stop_words + config = FTS( - with_position=with_position, - base_tokenizer=base_tokenizer, - language=language, - max_token_length=max_token_length, - lower_case=lower_case, - stem=stem, - remove_stop_words=remove_stop_words, - custom_stop_words=custom_stop_words, - ascii_folding=ascii_folding, - ngram_min_length=ngram_min_length, - ngram_max_length=ngram_max_length, - prefix_only=prefix_only, block_size=block_size, document_granularity=document_granularity, + **tokenizer_configs, ) - LOOP.run( - self._table.create_index( - column, - config=config, - replace=replace, - wait_timeout=wait_timeout, - name=name, + + try: + LOOP.run( + self._table.create_index( + field_names, + replace=replace, + config=config, + name=name, + wait_timeout=wait_timeout, + ) ) - ) + except (ValueError, RuntimeError) as e: + _maybe_add_fts_error_note( + e, + base_tokenizer=config.base_tokenizer, + language=config.language, + ) + raise e # New unified API overload @overload @@ -438,6 +492,7 @@ class RemoteTable(Table): /, *, config: IndexConfigType, + replace: bool = ..., wait_timeout: Optional[timedelta] = ..., name: Optional[str] = ..., train: bool = ..., @@ -448,38 +503,56 @@ class RemoteTable(Table): def create_index( self, metric: Literal["l2", "cosine", "dot", "hamming"] = ..., - vector_column_name: str = ..., - index_cache_size: Optional[int] = ..., num_partitions: Optional[int] = ..., num_sub_vectors: Optional[int] = ..., - replace: Optional[bool] = ..., + vector_column_name: str = ..., + replace: bool = ..., accelerator: Optional[str] = ..., - index_type: Literal[ - "VECTOR", "IVF_FLAT", "IVF_SQ", "IVF_PQ", "IVF_HNSW_SQ", "IVF_HNSW_PQ" - ] = ..., - wait_timeout: Optional[timedelta] = ..., - *, + index_cache_size: Optional[int] = ..., num_bits: int = ..., + index_type: Literal[ + "IVF_FLAT", "IVF_SQ", "IVF_PQ", "IVF_RQ", "IVF_HNSW_SQ", "IVF_HNSW_PQ" + ] = ..., + max_iterations: int = ..., + sample_rate: int = ..., + m: int = ..., + ef_construction: int = ..., + *, + wait_timeout: Optional[timedelta] = ..., name: Optional[str] = ..., train: bool = ..., + target_partition_size: Optional[int] = ..., ) -> None: ... def create_index( self, metric: str = "l2", - vector_column_name: str = VECTOR_COLUMN_NAME, - index_cache_size: Optional[int] = None, num_partitions: Optional[int] = None, num_sub_vectors: Optional[int] = None, - replace: Optional[bool] = None, + vector_column_name: str = VECTOR_COLUMN_NAME, + replace: bool = True, accelerator: Optional[str] = None, - index_type="vector", - wait_timeout: Optional[timedelta] = None, - *, + index_cache_size: Optional[int] = None, num_bits: int = 8, + index_type: Literal[ + "IVF_FLAT", + "IVF_SQ", + "IVF_PQ", + "IVF_RQ", + "IVF_HNSW_SQ", + "IVF_HNSW_PQ", + "IVF_HNSW_FLAT", + ] = "IVF_PQ", + max_iterations: int = 50, + sample_rate: int = 256, + m: int = 20, + ef_construction: int = 300, + *, config: Optional[IndexConfigType] = None, + wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, train: bool = True, + target_partition_size: Optional[int] = None, ): """Create an index on a column. @@ -490,6 +563,24 @@ class RemoteTable(Table): argument plus separate ``vector_column_name`` / ``num_partitions`` / etc. parameters, and emits a ``DeprecationWarning``. + Parameters + ---------- + metric : str + For new API: the column name to index. + For legacy API: the distance metric ("l2", "cosine", "dot", "hamming"). + config : IndexConfigType, optional + The index configuration object. If provided, uses the new unified API. + Can be one of: IvfFlat, IvfPq, IvfSq, IvfRq, HnswPq, HnswSq, + BTree, Bitmap, LabelList, Fm, FTS. + replace : bool, default True + Whether to replace an existing index on this column. + wait_timeout : timedelta, optional + Timeout to wait for async indexing to complete. + name : str, optional + Custom name for the index. + train : bool, default True + Whether to train the index with existing data. + Examples -------- New API (recommended): @@ -515,7 +606,6 @@ class RemoteTable(Table): vector_column_name, accelerator, index_cache_size, - replace, ) if is_legacy: @@ -531,6 +621,7 @@ class RemoteTable(Table): stacklevel=2, ) + # Legacy API: first arg is the distance metric column = vector_column_name if accelerator is not None: @@ -539,45 +630,27 @@ class RemoteTable(Table): "and will be ignored." ) - idx_type = index_type.upper() - if idx_type == "VECTOR" or idx_type == "IVF_PQ": - config = IvfPq( - distance_type=metric, - num_partitions=num_partitions, - num_sub_vectors=num_sub_vectors, - num_bits=num_bits, - ) - elif idx_type == "IVF_RQ": - config = IvfRq( - distance_type=metric, - num_partitions=num_partitions, - num_bits=num_bits, - ) - elif idx_type == "IVF_SQ": - config = IvfSq(distance_type=metric, num_partitions=num_partitions) - elif idx_type == "IVF_HNSW_PQ": - config = HnswPq( - distance_type=metric, - num_partitions=num_partitions, - num_sub_vectors=num_sub_vectors, - num_bits=num_bits, - ) - elif idx_type == "IVF_HNSW_SQ": - config = HnswSq(distance_type=metric, num_partitions=num_partitions) - elif idx_type == "IVF_HNSW_FLAT": - config = HnswFlat(distance_type=metric, num_partitions=num_partitions) - elif idx_type == "IVF_FLAT": - config = IvfFlat(distance_type=metric, num_partitions=num_partitions) - else: - raise ValueError( - f"Unknown vector index type: {idx_type}. Valid options are" - " 'IVF_FLAT', 'IVF_PQ', 'IVF_RQ', 'IVF_SQ'," - " 'IVF_HNSW_PQ', 'IVF_HNSW_SQ', 'IVF_HNSW_FLAT'" - ) + # Build config from legacy parameters + config = self._build_vector_config_from_legacy_params( + metric=metric, + index_type="IVF_PQ" + if index_type.upper() == "VECTOR" + else index_type.upper(), + num_partitions=num_partitions, + num_sub_vectors=num_sub_vectors, + num_bits=num_bits, + max_iterations=max_iterations, + sample_rate=sample_rate, + m=m, + ef_construction=ef_construction, + target_partition_size=target_partition_size, + accelerator=None, + ) + else: column = metric - LOOP.run( + return LOOP.run( self._table.create_index( column, replace=replace, @@ -616,37 +689,6 @@ class RemoteTable(Table): ) ) - def _is_legacy_create_index_call( - self, - first_arg: str, - config: Optional[IndexConfigType], - num_partitions: Optional[int], - num_sub_vectors: Optional[int], - vector_column_name: str, - accelerator: Optional[str], - index_cache_size: Optional[int], - replace: Optional[bool], - ) -> bool: - """Detect if this is a legacy create_index call.""" - if config is not None: - return False - if any( - x is not None - for x in ( - num_partitions, - num_sub_vectors, - accelerator, - index_cache_size, - replace, - ) - ): - return True - if vector_column_name != VECTOR_COLUMN_NAME: - return True - if first_arg.lower() in KNOWN_METRICS: - return True - return False - def add( self, data: DATA, @@ -718,9 +760,10 @@ class RemoteTable(Table): query: Union[VEC, str] = None, vector_column_name: Optional[str] = None, query_type="auto", + ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, fast_search: Optional[bool] = None, - ) -> LanceVectorQueryBuilder: + ) -> LanceQueryBuilder: """Create a search query to find the nearest neighbors of the given query vector. We currently support [vector search](https://lancedb.com/docs/search/vector-search/) @@ -811,6 +854,7 @@ class RemoteTable(Table): query, query_type, vector_column_name=vector_column_name, + ordering_field_name=ordering_field_name, fts_columns=fts_columns or [], fast_search=fast_search, ) @@ -870,7 +914,8 @@ class RemoteTable(Table): self._table._do_merge(merge, new_data, on_bad_vectors, fill_value) ) - def delete(self, predicate: str) -> DeleteResult: + @_deprecated_keyword("predicate", "where") + def delete(self, where: Union[str, Expr]) -> DeleteResult: """Delete rows from the table. This can be used to delete a single row, many rows, all rows, or @@ -878,7 +923,7 @@ class RemoteTable(Table): Parameters ---------- - predicate: str + where: Union[str, Expr] The SQL where clause to use when deleting rows. - For example, 'x = 2' or 'x IN (1, 2, 3)'. @@ -922,7 +967,7 @@ class RemoteTable(Table): x vector _distance # doctest: +SKIP 0 2 [3.0, 4.0] 85.0 # doctest: +SKIP """ - return LOOP.run(self._table.delete(predicate)) + return LOOP.run(self._table.delete(where)) def update( self, @@ -984,7 +1029,12 @@ class RemoteTable(Table): self._table.update(where=where, updates=values, updates_sql=values_sql) ) - def cleanup_old_versions(self, *_): + def cleanup_old_versions( + self, + older_than: Optional[timedelta] = None, + *, + delete_unverified: bool = False, + ): """ cleanup_old_versions() is a no-op for remote tables. """ @@ -994,7 +1044,7 @@ class RemoteTable(Table): ) pass - def compact_files(self, *_): + def compact_files(self, *args, **kwargs): """ compact_files() is a no-op for remote tables. """ @@ -1009,6 +1059,7 @@ class RemoteTable(Table): *, cleanup_older_than: Optional[timedelta] = None, delete_unverified: bool = False, + retrain: bool = False, compaction_options: Optional[CompactionOptions] = None, ): """optimize() is not supported for remote tables.""" @@ -1016,6 +1067,7 @@ class RemoteTable(Table): self._table.optimize( cleanup_older_than=cleanup_older_than, delete_unverified=delete_unverified, + retrain=retrain, ) ) @@ -1123,8 +1175,9 @@ class RemoteTable(Table): """No-op for remote tables (no local shard writers).""" return LOOP.run(self._table.close_rte_writers()) - def drop_index(self, index_name: str): - return LOOP.run(self._table.drop_index(index_name)) + @_deprecated_keyword("index_name", "name") + def drop_index(self, name: str) -> None: + return LOOP.run(self._table.drop_index(name)) def prewarm_index(self, name: str) -> None: """Prewarm an index in the table. diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 11ec5c8b5..c577fb58e 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -266,10 +266,8 @@ if TYPE_CHECKING: OnBadVectorsType, AddMode, CreateMode, - VectorIndexType, ScalarIndexType, BaseTokenizerType, - DistanceType, ) # Type alias for index configuration objects @@ -633,10 +631,11 @@ def _sanitize_data( metadata : Optional[dict], default None The embedding metadata to add to the schema. on_bad_vectors : Literal["error", "drop", "fill", "null"], default "error" - What to do if any of the vectors are not the same size or contains NaNs. + How to handle vectors with incorrect sizes, NaNs, or null elements. fill_value : float, default 0.0 The value to use when filling vectors. Only used if on_bad_vectors="fill". - All entries in the vector will be set to this value. + NaN or null elements and missing trailing elements are replaced with this + value. Valid elements are preserved. """ # At this point, the table might not match the schema we are targeting: # 1. There might be embedding columns missing that will be added @@ -1574,12 +1573,18 @@ class Table(ABC): """ raise NotImplementedError("to_lance is not supported for this table type") - def to_polars(self, **kwargs) -> "pl.LazyFrame": + def to_polars(self, batch_size=None) -> "pl.LazyFrame": """Return the table as a Polars LazyFrame. For remote tables, use ``table.search().to_polars()`` to convert query results to a Polars DataFrame. + Parameters + ---------- + batch_size: int, optional + Passed to Polars as the maximum row count for scanned PyArrow + record batches. + Note ---- The Polars streaming engine is not supported because it does not currently @@ -1592,6 +1597,193 @@ class Table(ABC): """ raise NotImplementedError("to_polars is not supported for this table type") + def _is_legacy_create_index_call( + self, + first_arg: str, + config: Optional[IndexConfigType], + num_partitions: Optional[int], + num_sub_vectors: Optional[int], + vector_column_name: str, + accelerator: Optional[str], + index_cache_size: Optional[int], + ) -> bool: + """Detect if this is a legacy create_index call.""" + # If config is provided, it's definitely the new API + if config is not None: + return False + + # If old-style parameters were explicitly set, it's legacy + if any( + x is not None + for x in (num_partitions, num_sub_vectors, accelerator, index_cache_size) + ): + return True + + # If vector_column_name differs from default, it's legacy + if vector_column_name != VECTOR_COLUMN_NAME: + return True + + # If first arg is a known metric, assume legacy + if first_arg.lower() in KNOWN_METRICS: + return True + + # Otherwise assume new API + return False + + def _build_vector_config_from_legacy_params( + self, + metric: str, + index_type: str, + num_partitions: Optional[int], + num_sub_vectors: Optional[int], + num_bits: int, + max_iterations: int, + sample_rate: int, + m: int, + ef_construction: int, + target_partition_size: Optional[int], + accelerator: Optional[str], + ) -> IndexConfigType: + """Build an index config object from legacy parameters.""" + if index_type == "IVF_FLAT": + return IvfFlat( + distance_type=metric, + num_partitions=num_partitions, + max_iterations=max_iterations, + sample_rate=sample_rate, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_SQ": + return IvfSq( + distance_type=metric, + num_partitions=num_partitions, + max_iterations=max_iterations, + sample_rate=sample_rate, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_PQ": + return IvfPq( + distance_type=metric, + num_partitions=num_partitions, + num_sub_vectors=num_sub_vectors, + num_bits=num_bits, + max_iterations=max_iterations, + sample_rate=sample_rate, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_RQ": + return IvfRq( + distance_type=metric, + num_partitions=num_partitions, + num_bits=num_bits, + max_iterations=max_iterations, + sample_rate=sample_rate, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_HNSW_PQ": + return HnswPq( + distance_type=metric, + num_partitions=num_partitions, + num_sub_vectors=num_sub_vectors, + num_bits=num_bits, + max_iterations=max_iterations, + sample_rate=sample_rate, + m=m, + ef_construction=ef_construction, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_HNSW_SQ": + return HnswSq( + distance_type=metric, + num_partitions=num_partitions, + max_iterations=max_iterations, + sample_rate=sample_rate, + m=m, + ef_construction=ef_construction, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_HNSW_FLAT": + return HnswFlat( + distance_type=metric, + num_partitions=num_partitions, + max_iterations=max_iterations, + sample_rate=sample_rate, + m=m, + ef_construction=ef_construction, + target_partition_size=target_partition_size, + ) + else: + raise ValueError(f"Unknown index type {index_type}") + + @staticmethod + def infer_tokenizer_configs(tokenizer_name: str) -> dict: + if tokenizer_name == "default": + return { + "base_tokenizer": "simple", + "language": "English", + "max_token_length": 40, + "lower_case": True, + "stem": False, + "remove_stop_words": False, + "ascii_folding": False, + "ngram_min_length": 3, + "ngram_max_length": 3, + "prefix_only": False, + } + elif tokenizer_name == "raw": + return { + "base_tokenizer": "raw", + "language": "English", + "max_token_length": None, + "lower_case": False, + "stem": False, + "remove_stop_words": False, + "ascii_folding": False, + "ngram_min_length": 3, + "ngram_max_length": 3, + "prefix_only": False, + } + elif tokenizer_name == "whitespace": + return { + "base_tokenizer": "whitespace", + "language": "English", + "max_token_length": None, + "lower_case": False, + "stem": False, + "remove_stop_words": False, + "ascii_folding": False, + "ngram_min_length": 3, + "ngram_max_length": 3, + "prefix_only": False, + } + + # or it's with language stemming with pattern like "en_stem" + if len(tokenizer_name) != 7: + raise ValueError(f"Invalid tokenizer name {tokenizer_name}") + lang = tokenizer_name[:2] + if tokenizer_name[-5:] != "_stem": + raise ValueError(f"Invalid tokenizer name {tokenizer_name}") + if lang not in lang_mapping: + raise ValueError(f"Invalid language code {lang}") + return { + "base_tokenizer": "simple", + "language": lang_mapping[lang], + "max_token_length": 40, + "lower_case": True, + "stem": True, + "remove_stop_words": False, + "ascii_folding": False, + "ngram_min_length": 3, + "ngram_max_length": 3, + "prefix_only": False, + } + # New unified API overload @overload def create_index( @@ -1617,14 +1809,16 @@ class Table(ABC): replace: bool = ..., accelerator: Optional[str] = ..., index_cache_size: Optional[int] = ..., - *, - index_type: VectorIndexType = ..., - wait_timeout: Optional[timedelta] = ..., num_bits: int = ..., + index_type: Literal[ + "IVF_FLAT", "IVF_SQ", "IVF_PQ", "IVF_RQ", "IVF_HNSW_SQ", "IVF_HNSW_PQ" + ] = ..., max_iterations: int = ..., sample_rate: int = ..., m: int = ..., ef_construction: int = ..., + *, + wait_timeout: Optional[timedelta] = ..., name: Optional[str] = ..., train: bool = ..., target_partition_size: Optional[int] = ..., @@ -1632,22 +1826,30 @@ class Table(ABC): def create_index( self, - metric: DistanceType = "l2", + metric: str = "l2", num_partitions: Optional[int] = None, num_sub_vectors: Optional[int] = None, vector_column_name: str = VECTOR_COLUMN_NAME, replace: bool = True, accelerator: Optional[str] = None, index_cache_size: Optional[int] = None, - *, - index_type: VectorIndexType = "IVF_PQ", - wait_timeout: Optional[timedelta] = None, num_bits: int = 8, + index_type: Literal[ + "IVF_FLAT", + "IVF_SQ", + "IVF_PQ", + "IVF_RQ", + "IVF_HNSW_SQ", + "IVF_HNSW_PQ", + "IVF_HNSW_FLAT", + ] = "IVF_PQ", max_iterations: int = 50, sample_rate: int = 256, m: int = 20, ef_construction: int = 300, + *, config: Optional[IndexConfigType] = None, + wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, train: bool = True, target_partition_size: Optional[int] = None, @@ -1984,8 +2186,8 @@ class Table(ABC): "append" and "overwrite". Overwrite adopts explicitly declared blob layouts, while inferred bytes retain the existing blob v2 fields. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. - One of "error", "drop", "fill". + How to handle vectors with incorrect sizes, NaNs, or null elements. + One of "error", "drop", "fill", "null". fill_value: float, default 0. The value to use when filling vectors. Only used if on_bad_vectors="fill". progress: bool, callable, or tqdm-like, optional @@ -2104,6 +2306,7 @@ class Table(ABC): query_type: QueryType = "auto", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceQueryBuilder: """Create a search query to find the nearest neighbors of the given query vector. We currently support [vector search](https://lancedb.com/docs/search/vector-search/) @@ -3845,6 +4048,11 @@ class LanceTable(Table): cfg_type = type(config).__name__ lance_index_type = index_type_map.get(cfg_type, "IVF_PQ") + quantization_kwargs = {} + num_bits = getattr(config, "num_bits", 8) + if num_bits is not None: + quantization_kwargs["num_bits"] = num_bits + dataset.create_index( column=column, index_type=lance_index_type, @@ -3853,12 +4061,12 @@ class LanceTable(Table): num_sub_vectors=getattr(config, "num_sub_vectors", None), replace=replace, accelerator=acc, - num_bits=getattr(config, "num_bits", 8), m=getattr(config, "m", 20), ef_construction=getattr(config, "ef_construction", 300), target_partition_size=getattr( config, "target_partition_size", None ), + **quantization_kwargs, ) self.checkout_latest() return @@ -3902,130 +4110,6 @@ class LanceTable(Table): ) ) - def _is_legacy_create_index_call( - self, - first_arg: str, - config: Optional[IndexConfigType], - num_partitions: Optional[int], - num_sub_vectors: Optional[int], - vector_column_name: str, - accelerator: Optional[str], - index_cache_size: Optional[int], - ) -> bool: - """Detect if this is a legacy create_index call.""" - # If config is provided, it's definitely the new API - if config is not None: - return False - - # If old-style parameters were explicitly set, it's legacy - if any( - x is not None - for x in (num_partitions, num_sub_vectors, accelerator, index_cache_size) - ): - return True - - # If vector_column_name differs from default, it's legacy - if vector_column_name != VECTOR_COLUMN_NAME: - return True - - # If first arg is a known metric, assume legacy - if first_arg.lower() in KNOWN_METRICS: - return True - - # Otherwise assume new API - return False - - def _build_vector_config_from_legacy_params( - self, - metric: str, - index_type: str, - num_partitions: Optional[int], - num_sub_vectors: Optional[int], - num_bits: int, - max_iterations: int, - sample_rate: int, - m: int, - ef_construction: int, - target_partition_size: Optional[int], - accelerator: Optional[str], - ) -> IndexConfigType: - """Build an index config object from legacy parameters.""" - if index_type == "IVF_FLAT": - return IvfFlat( - distance_type=metric, - num_partitions=num_partitions, - max_iterations=max_iterations, - sample_rate=sample_rate, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_SQ": - return IvfSq( - distance_type=metric, - num_partitions=num_partitions, - max_iterations=max_iterations, - sample_rate=sample_rate, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_PQ": - return IvfPq( - distance_type=metric, - num_partitions=num_partitions, - num_sub_vectors=num_sub_vectors, - num_bits=num_bits, - max_iterations=max_iterations, - sample_rate=sample_rate, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_RQ": - return IvfRq( - distance_type=metric, - num_partitions=num_partitions, - num_bits=num_bits, - max_iterations=max_iterations, - sample_rate=sample_rate, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_HNSW_PQ": - return HnswPq( - distance_type=metric, - num_partitions=num_partitions, - num_sub_vectors=num_sub_vectors, - num_bits=num_bits, - max_iterations=max_iterations, - sample_rate=sample_rate, - m=m, - ef_construction=ef_construction, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_HNSW_SQ": - return HnswSq( - distance_type=metric, - num_partitions=num_partitions, - max_iterations=max_iterations, - sample_rate=sample_rate, - m=m, - ef_construction=ef_construction, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_HNSW_FLAT": - return HnswFlat( - distance_type=metric, - num_partitions=num_partitions, - max_iterations=max_iterations, - sample_rate=sample_rate, - m=m, - ef_construction=ef_construction, - target_partition_size=target_partition_size, - ) - else: - raise ValueError(f"Unknown index type {index_type}") - def drop_index(self, name: str) -> None: """ Drops an index from the table @@ -4136,6 +4220,7 @@ class LanceTable(Table): *, replace: bool = True, index_type: ScalarIndexType = "BTREE", + wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, ): """Create a scalar index on a column. @@ -4144,6 +4229,13 @@ class LanceTable(Table): Use :meth:`create_index` with a BTree, Bitmap, or LabelList config instead. Example: ``table.create_index("column", config=BTree())`` """ + if index_type == "scalar": + warnings.warn( + 'index_type="scalar" is deprecated; use "BTREE" instead', + DeprecationWarning, + stacklevel=2, + ) + index_type = "BTREE" if index_type == "BTREE": config = BTree() elif index_type == "BITMAP": @@ -4153,7 +4245,13 @@ class LanceTable(Table): else: raise ValueError(f"Unknown index type {index_type}") return LOOP.run( - self._table.create_index(column, replace=replace, config=config, name=name) + self._table.create_index( + column, + replace=replace, + config=config, + wait_timeout=wait_timeout, + name=name, + ) ) @deprecation.deprecated( @@ -4185,6 +4283,7 @@ class LanceTable(Table): prefix_only: bool = False, block_size: int = 128, document_granularity: DocumentGranularity = DocumentGranularity.ROW, + wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, ): """Create a full-text search index on a column. @@ -4249,6 +4348,7 @@ class LanceTable(Table): replace=replace, config=config, name=name, + wait_timeout=wait_timeout, ) ) except (ValueError, RuntimeError) as e: @@ -4259,69 +4359,6 @@ class LanceTable(Table): ) raise e - @staticmethod - def infer_tokenizer_configs(tokenizer_name: str) -> dict: - if tokenizer_name == "default": - return { - "base_tokenizer": "simple", - "language": "English", - "max_token_length": 40, - "lower_case": True, - "stem": False, - "remove_stop_words": False, - "ascii_folding": False, - "ngram_min_length": 3, - "ngram_max_length": 3, - "prefix_only": False, - } - elif tokenizer_name == "raw": - return { - "base_tokenizer": "raw", - "language": "English", - "max_token_length": None, - "lower_case": False, - "stem": False, - "remove_stop_words": False, - "ascii_folding": False, - "ngram_min_length": 3, - "ngram_max_length": 3, - "prefix_only": False, - } - elif tokenizer_name == "whitespace": - return { - "base_tokenizer": "whitespace", - "language": "English", - "max_token_length": None, - "lower_case": False, - "stem": False, - "remove_stop_words": False, - "ascii_folding": False, - "ngram_min_length": 3, - "ngram_max_length": 3, - "prefix_only": False, - } - - # or it's with language stemming with pattern like "en_stem" - if len(tokenizer_name) != 7: - raise ValueError(f"Invalid tokenizer name {tokenizer_name}") - lang = tokenizer_name[:2] - if tokenizer_name[-5:] != "_stem": - raise ValueError(f"Invalid tokenizer name {tokenizer_name}") - if lang not in lang_mapping: - raise ValueError(f"Invalid language code {lang}") - return { - "base_tokenizer": "simple", - "language": lang_mapping[lang], - "max_token_length": 40, - "lower_case": True, - "stem": True, - "remove_stop_words": False, - "ascii_folding": False, - "ngram_min_length": 3, - "ngram_max_length": 3, - "prefix_only": False, - } - def add( self, data: DATA, @@ -4346,7 +4383,7 @@ class LanceTable(Table): "append" and "overwrite". Overwrite adopts explicitly declared blob layouts, while inferred bytes retain the existing blob v2 fields. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. + How to handle vectors with incorrect sizes, NaNs, or null elements. One of "error", "drop", "fill", "null". fill_value: float, default 0. The value to use when filling vectors. Only used if on_bad_vectors="fill". @@ -4472,6 +4509,7 @@ class LanceTable(Table): query_type: Literal["vector"] = "vector", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceVectorQueryBuilder: ... @overload @@ -4482,6 +4520,7 @@ class LanceTable(Table): query_type: Literal["fts"] = "fts", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceFtsQueryBuilder: ... @overload @@ -4494,6 +4533,7 @@ class LanceTable(Table): query_type: Literal["hybrid"] = "hybrid", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceHybridQueryBuilder: ... @overload @@ -4504,6 +4544,7 @@ class LanceTable(Table): query_type: QueryType = "auto", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceEmptyQueryBuilder: ... def search( @@ -4515,6 +4556,7 @@ class LanceTable(Table): query_type: QueryType = "auto", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceQueryBuilder: """Create a search query to find the nearest neighbors of the given query vector. We currently support [vector search](https://lancedb.com/docs/search/vector-search/) @@ -4599,6 +4641,7 @@ class LanceTable(Table): vector_column_name=vector_column_name, ordering_field_name=ordering_field_name, fts_columns=fts_columns or [], + fast_search=fast_search, ) @classmethod @@ -4663,7 +4706,7 @@ class LanceTable(Table): otherwise just open the table, it will not add the provided data but will validate against any schema that's specified. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. + How to handle vectors with incorrect sizes, NaNs, or null elements. One of "error", "drop", "fill", "null". fill_value: float, default 0. The value to use when filling vectors. Only used if on_bad_vectors="fill". @@ -5359,7 +5402,7 @@ def _handle_bad_vector_column( vector_column_name: str The name of the vector column. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. + How to handle vectors with incorrect sizes, NaNs, or null elements. One of "error", "drop", "fill", "null". fill_value: float, default 0.0 The value to use when filling vectors. Only used if on_bad_vectors="fill". @@ -5378,7 +5421,8 @@ def _handle_bad_vector_column( data = data.set_column(position, vector_column_name, vec_arr) if expected_value_type is not None and ( - pa.types.is_integer(vec_arr.type.value_type) + pa.types.is_null(vec_arr.type.value_type) + or pa.types.is_integer(vec_arr.type.value_type) or pa.types.is_unsigned_integer(vec_arr.type.value_type) ): vec_arr = pa.array(vec_arr.to_pylist(), type=pa.list_(expected_value_type)) @@ -5422,7 +5466,7 @@ def _handle_bad_vector_column( ) else: raise ValueError( - f"Vector column '{vector_column_name}' has NaNs. " + f"Vector column '{vector_column_name}' has NaNs or null values. " "Set on_bad_vectors='drop' to remove them, " "set on_bad_vectors='fill' and fill_value= to replace them, " "or set on_bad_vectors='null' to replace them with null." @@ -5458,29 +5502,11 @@ def _fill_bad_vector_values( arr = arr.combine_chunks() # A fixed-size slice truncates long vectors and pads short vectors with nulls. - # Slice an array marking the original child nulls in parallel so padding nulls - # can be distinguished from null values that were already present. + # Both padding and existing null elements need to be filled. sliced = pc.list_slice(arr, 0, dim, return_fixed_size_list=True) - child_nulls = pc.is_null(arr.values) - parent_nulls = pc.is_null(arr) - if pa.types.is_list(arr.type): - original_child_nulls = pa.ListArray.from_arrays( - arr.offsets, child_nulls, mask=parent_nulls - ) - elif pa.types.is_large_list(arr.type): - original_child_nulls = pa.LargeListArray.from_arrays( - arr.offsets, child_nulls, mask=parent_nulls - ) - else: - original_child_nulls = pa.FixedSizeListArray.from_arrays( - child_nulls, arr.type.list_size, mask=parent_nulls - ) - sliced_child_nulls = pc.list_slice( - original_child_nulls, 0, dim, return_fixed_size_list=True - ) - needs_fill = pc.is_null(sliced_child_nulls.values) values = sliced.values + needs_fill = pc.is_null(values) if pa.types.is_floating(values.type): values_for_nan_check = ( values.cast(pa.float32()) if pa.types.is_float16(values.type) else values @@ -5494,6 +5520,7 @@ def _fill_bad_vector_values( def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray: + """Find vectors with NaN or null components, including pandas-converted NaNs.""" if isinstance(arr, pa.ChunkedArray): values = pa.chunked_array([chunk.flatten() for chunk in arr.chunks]) else: @@ -5504,6 +5531,9 @@ def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray values_has_nan = pc.is_nan(values.cast(pa.float32())) else: values_has_nan = pc.is_nan(values) + # Pandas represents NaN components as Arrow nulls. is_nan returns null + # for these values, so count them as invalid before filtering parent indices. + values_has_nan = pc.fill_null(values_has_nan, True) values_indices = pc.list_parent_indices(arr) if arr.null_count: # Match flatten(), which excludes child spans under null parent lists. @@ -6417,7 +6447,7 @@ class AsyncTable: "append" and "overwrite". Overwrite adopts explicitly declared blob layouts, while inferred bytes retain the existing blob v2 fields. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. + How to handle vectors with incorrect sizes, NaNs, or null elements. One of "error", "drop", "fill", "null". fill_value: float, default 0. The value to use when filling vectors. Only used if on_bad_vectors="fill". diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 207a3d253..3c968a821 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -1039,6 +1039,39 @@ def test_open_clip_opens_percent_encoded_file_uri(tmp_path): assert image.size == (4, 4) +def test_open_clip_model_in_eval_mode(): + """ + open_clip.create_model_and_transforms returns the model in training mode, + which leaves dropout, stochastic depth and batch norm in their training + behaviour. OpenCLIP must switch the model to eval mode before embedding. + """ + from lancedb.embeddings.open_clip import OpenClipEmbeddings + + class FakeModel: + training = True + + def to(self, device): + return self + + def eval(self): + self.training = False + return self + + open_clip = MagicMock() + open_clip.create_model_and_transforms.return_value = (FakeModel(), None, None) + + with patch( + "lancedb.embeddings.open_clip.attempt_import_or_raise", + return_value=open_clip, + ): + func = OpenClipEmbeddings(name="ViTamin-B", pretrained="datacomp1b") + + open_clip.create_model_and_transforms.assert_called_once_with( + "ViTamin-B", pretrained="datacomp1b" + ) + assert func._model.training is False + + def test_jina_generate_image_input_dict_local_path(tmp_path): """ JinaEmbeddings._generate_image_input_dict must accept a local image path diff --git a/python/python/tests/test_index.py b/python/python/tests/test_index.py index 0252a09ad..2ec14f6f8 100644 --- a/python/python/tests/test_index.py +++ b/python/python/tests/test_index.py @@ -421,13 +421,23 @@ async def test_create_4bit_ivfpq_index(some_table: AsyncTable): @pytest.mark.asyncio -async def test_create_ivfrq_index(some_table: AsyncTable): - await some_table.create_index("vector", config=IvfRq(num_bits=1)) +@pytest.mark.parametrize( + "config, expected_bits", + [ + pytest.param(IvfRq(), 5, id="default"), + pytest.param(IvfRq(num_bits=None), 5, id="none"), + pytest.param(IvfRq(num_bits=1), 1, id="one-bit"), + pytest.param(IvfRq(num_bits=4), 4, id="multi-bit"), + ], +) +async def test_create_ivfrq_index(some_table: AsyncTable, config, expected_bits): + await some_table.create_index("vector", config=config) indices = await some_table.list_indices() assert len(indices) == 1 assert indices[0].index_type == "IvfRq" assert indices[0].columns == ["vector"] assert indices[0].name == "vector_idx" + assert indices[0].index_details["compression"]["num_bits"] == expected_bits @pytest.mark.asyncio diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index c1db92948..db3394735 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -5,6 +5,7 @@ import contextlib from datetime import timedelta import http.server import io +import inspect import json import multiprocessing as mp import pickle @@ -765,6 +766,15 @@ def test_create_table_embedding_functions(with_schema): ] +def test_create_table_positional_parameters_match_base_contract(): + from lancedb.db import DBConnection + from lancedb.remote.db import RemoteDBConnection + + expected = list(inspect.signature(DBConnection.create_table).parameters)[:9] + actual = list(inspect.signature(RemoteDBConnection.create_table).parameters)[:9] + assert actual == expected + + def test_create_table_exist_ok(): def handler(request): if request.path == "/v1/table/test/create/?mode=exist_ok": @@ -1145,8 +1155,9 @@ def test_table_create_indices(): # Test create_fts_index with custom name (legacy method) with pytest.warns(DeprecationWarning, match="create_fts_index"): table.create_fts_index( - "text", + field_names="text", wait_timeout=timedelta(seconds=2), + tokenizer_name="en_stem", block_size=256, custom_stop_words=["cloud"], name="custom_fts_idx", @@ -1155,6 +1166,8 @@ def test_table_create_indices(): # Test create_index with custom name (legacy form: vector_column_name kwarg) with pytest.warns(DeprecationWarning, match="create_index"): table.create_index( + "cosine", + 2, vector_column_name="vector", wait_timeout=timedelta(seconds=10), name="custom_vector_idx", @@ -1167,7 +1180,8 @@ def test_table_create_indices(): scalar_req = received_requests[0] assert "name" in scalar_req assert scalar_req["name"] == "custom_scalar_idx" - assert scalar_req["replace"] is False + # The wire protocol omits replace=True (the server default). + assert "replace" not in scalar_req # Check FTS index request has custom name fts_req = received_requests[1] @@ -1176,18 +1190,39 @@ def test_table_create_indices(): assert fts_req["replace"] is False assert fts_req["block_size"] == 256 assert fts_req["custom_stop_words"] == ["cloud"] + assert fts_req["language"] == "English" + assert fts_req["stem"] is True # Check vector index request has custom name vector_req = received_requests[2] assert "name" in vector_req assert vector_req["name"] == "custom_vector_idx" assert "replace" not in vector_req + assert vector_req["metric_type"] == "cosine" + assert vector_req["num_partitions"] == 2 + + with pytest.warns(DeprecationWarning, match="create_index"): + table.create_index( + index_type="IVF_HNSW_SQ", + num_partitions=1, + m=10, + ef_construction=50, + max_iterations=12, + sample_rate=64, + ) + tuned_req = received_requests[-1] + assert tuned_req["index_type"] == "IVF_HNSW_SQ" + assert tuned_req["num_partitions"] == 1 + assert tuned_req["m"] == 10 + assert tuned_req["ef_construction"] == 50 + assert tuned_req["max_iterations"] == 12 + assert tuned_req["sample_rate"] == 64 table.wait_for_index(["custom_scalar_idx"], timedelta(seconds=2)) table.wait_for_index( ["custom_fts_idx", "custom_vector_idx"], timedelta(seconds=2) ) - table.drop_index("custom_vector_idx") + table.drop_index(name="custom_vector_idx") table.drop_index("custom_scalar_idx") table.drop_index("custom_fts_idx") @@ -1583,7 +1618,10 @@ def test_remote_json_fts_config(server_version): assert request["block_size"] == 256 -def test_remote_create_index_new_api(caplog): +@pytest.mark.parametrize( + "rq_kwargs", [{}, {"num_bits": None}, {"num_bits": 1}, {"num_bits": 4}] +) +def test_remote_create_index_new_api(caplog, rq_kwargs): received_requests = [] def handler(request): @@ -1656,7 +1694,7 @@ def test_remote_create_index_new_api(caplog): table.create_index("category", config=BTree()) table.create_index("text", config=FTS(block_size=256)) # IvfRq via new API - table.create_index("vector", config=IvfRq(distance_type="l2")) + table.create_index("vector", config=IvfRq(distance_type="l2", **rq_kwargs)) table.create_index( "vector", config=IvfPq(distance_type="l2"), replace=False ) @@ -1681,6 +1719,10 @@ def test_remote_create_index_new_api(caplog): "vector", ] assert received_requests[2]["block_size"] == 256 + if rq_kwargs.get("num_bits") is None: + assert "num_bits" not in received_requests[3] + else: + assert received_requests[3]["num_bits"] == rq_kwargs["num_bits"] assert received_requests[4]["replace"] is False assert received_requests[5]["replace"] is False assert "replace is not supported" not in caplog.text diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 2fc6a01d5..48fe675d5 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -1751,6 +1751,28 @@ def test_create_index_async_returns_done_job(mem_db: DBConnection): job.cancel() +@pytest.mark.parametrize( + "rq_kwargs", [{}, {"num_bits": None}, {"num_bits": 1}, {"num_bits": 4}] +) +@patch("lancedb.table.LanceTable.to_lance") +def test_create_ivfrq_index_accelerator(mock_to_lance, mem_db: DBConnection, rq_kwargs): + from lancedb.index import IvfRq + + table = mem_db.create_table("test", [{"vector": [1.0] * 8}]) + torch = MagicMock() + torch.cuda.is_available.return_value = True + with patch.dict(sys.modules, {"torch": torch}): + table.create_index("vector", config=IvfRq(accelerator="cuda", **rq_kwargs)) + + kwargs = mock_to_lance.return_value.create_index.call_args.kwargs + assert kwargs["index_type"] == "IVF_RQ" + assert kwargs["accelerator"] == "cuda" + if rq_kwargs.get("num_bits") is None: + assert "num_bits" not in kwargs + else: + assert kwargs["num_bits"] == rq_kwargs["num_bits"] + + @pytest.mark.parametrize("method", ["create_index", "create_index_async"]) @pytest.mark.parametrize("wait_timeout", ["invalid", 1, []]) def test_create_index_invalid_wait_timeout(mem_db: DBConnection, method, wait_timeout): @@ -2170,6 +2192,71 @@ def test_add_with_nans(mem_db: DBConnection): assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0])) +@pytest.mark.parametrize("operation", ["create", "append", "overwrite"]) +@pytest.mark.parametrize("input_type", ["pandas", "rows", "arrow"]) +@pytest.mark.parametrize("on_bad_vectors", [None, "error", "drop", "fill", "null"]) +def test_write_vectors_with_null_elements( + mem_db: DBConnection, operation, input_type, on_bad_vectors +): + schema = pa.schema({"vector": pa.list_(pa.float32(), 2), "id": pa.int64()}) + vectors = [[3.0, 4.0], [np.nan, 1.0], [np.nan, np.nan]] + if input_type == "pandas": + pd = pytest.importorskip("pandas") + + data = pd.DataFrame( + { + "vector": [np.array(v, dtype=np.float32) for v in vectors], + "id": [0, 1, 2], + } + ) + else: + # Pandas turns the NaNs above into null child elements. Explicit nulls + # in other input formats must follow the same bad-vector policy. + data = [ + {"vector": [3.0, 4.0], "id": 0}, + {"vector": [None, 1.0], "id": 1}, + {"vector": [None, None], "id": 2}, + ] + if input_type == "arrow": + data = pa.Table.from_pylist(data, schema=schema) + + options = {} if on_bad_vectors is None else {"on_bad_vectors": on_bad_vectors} + if on_bad_vectors == "fill": + options["fill_value"] = 42.0 + + table = None + if operation != "create": + table = mem_db.create_table("test", schema=schema) + + def write(): + if operation == "create": + return mem_db.create_table("test", data=data, schema=schema, **options) + table.add(data, mode=operation, **options) + return table + + if on_bad_vectors in (None, "error"): + # Create/overwrite can wrap errors from lazy Arrow readers in RuntimeError. + expected_error = ( + ValueError if operation == "append" else (ValueError, RuntimeError) + ) + with pytest.raises(expected_error, match="NaN|null"): + write() + if table is not None: + assert table.count_rows() == 0 + return + + table = write() + if on_bad_vectors == "drop": + expected = [{"vector": [3.0, 4.0], "id": 0}] + else: + expected = [ + {"vector": [3.0, 4.0], "id": 0}, + {"vector": [42.0, 1.0] if on_bad_vectors == "fill" else None, "id": 1}, + {"vector": [42.0, 42.0] if on_bad_vectors == "fill" else None, "id": 2}, + ] + assert table.to_arrow().to_pylist() == expected + + def test_add_with_empty_fixed_size_list_drops_bad_rows(mem_db: DBConnection): class Schema(LanceModel): text: str diff --git a/python/python/tests/test_table_signatures.py b/python/python/tests/test_table_signatures.py new file mode 100644 index 000000000..346935b2b --- /dev/null +++ b/python/python/tests/test_table_signatures.py @@ -0,0 +1,253 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +from datetime import timedelta +import inspect +from unittest.mock import AsyncMock, MagicMock + +import lancedb +from lancedb.expr import col +from lancedb.index import BTree, FTS, HnswSq, IvfPq +from lancedb.remote.table import RemoteTable +from lancedb.table import LanceTable, Table +import pyarrow as pa +import pytest + + +@pytest.mark.parametrize( + "method", + [ + "delete", + "drop_index", + "index_stats", + "create_index", + "create_scalar_index", + "create_fts_index", + "search", + "cleanup_old_versions", + "compact_files", + "optimize", + "to_polars", + ], +) +def test_sync_table_signatures(method): + """A Table call must bind the same arguments on either backend.""" + + def parameters(cls): + return [ + (p.name, p.kind, p.default) + for p in inspect.signature(getattr(cls, method)).parameters.values() + ] + + assert parameters(Table) == parameters(LanceTable) == parameters(RemoteTable) + + +@pytest.fixture(params=[LanceTable, RemoteTable], ids=["local", "remote"]) +def sync_table(request, monkeypatch): + """Exercise real sync wrappers with a recording async backend.""" + handle = MagicMock() + handle.name = "test" + handle.create_index = AsyncMock() + handle.delete = AsyncMock() + handle.drop_index = AsyncMock() + handle.index_stats = AsyncMock() + if request.param is RemoteTable: + table = RemoteTable(handle, "test_db") + else: + table = LanceTable.__new__(LanceTable) + table._table = handle + monkeypatch.setattr(table, "_ensure_no_legacy_fts_index", lambda: None) + return table, handle + + +def test_sync_keyword_names(sync_table): + table, handle = sync_table + predicate = col("id") == 1 + table.delete(where=predicate) + handle.delete.assert_awaited_once() + expected = predicate._inner if isinstance(table, LanceTable) else predicate + assert handle.delete.await_args.args[0] is expected + table.drop_index(name="vector_idx") + handle.drop_index.assert_awaited_once_with("vector_idx") + table.index_stats(index_name="vector_idx") + handle.index_stats.assert_awaited_once_with("vector_idx") + + +def test_legacy_vector_positional_arguments(sync_table): + table, handle = sync_table + with pytest.warns(DeprecationWarning): + table.create_index("cosine", 2) + args, kwargs = handle.create_index.await_args + assert args == ("vector",) + assert isinstance(kwargs["config"], IvfPq) + assert kwargs["config"].distance_type == "cosine" + assert kwargs["config"].num_partitions == 2 + assert kwargs["replace"] is True + + +def test_vector_tuning_options(sync_table): + table, handle = sync_table + timeout = timedelta(seconds=2) + with pytest.warns(DeprecationWarning): + table.create_index( + index_type="IVF_HNSW_SQ", + num_partitions=1, + m=10, + ef_construction=50, + max_iterations=12, + sample_rate=64, + target_partition_size=1000, + wait_timeout=timeout, + name="custom_idx", + replace=False, + train=False, + ) + kwargs = handle.create_index.await_args.kwargs + config = kwargs.pop("config") + assert isinstance(config, HnswSq) + for name, value in { + "num_partitions": 1, + "m": 10, + "ef_construction": 50, + "max_iterations": 12, + "sample_rate": 64, + "target_partition_size": 1000, + }.items(): + assert getattr(config, name) == value + assert kwargs == dict( + wait_timeout=timeout, name="custom_idx", replace=False, train=False + ) + + +def test_scalar_defaults_and_wait_timeout(sync_table): + table, handle = sync_table + timeout = timedelta(seconds=2) + with pytest.warns(DeprecationWarning): + table.create_scalar_index("id", wait_timeout=timeout) + kwargs = handle.create_index.await_args.kwargs + assert isinstance(kwargs["config"], BTree) + assert kwargs["replace"] is True + assert kwargs["wait_timeout"] == timeout + + +def test_fts_tokenizer_and_wait_timeout(sync_table): + table, handle = sync_table + timeout = timedelta(seconds=2) + with pytest.warns(DeprecationWarning): + table.create_fts_index( + field_names="text", + tokenizer_name="en_stem", + replace=True, + wait_timeout=timeout, + custom_stop_words=["example"], + ) + args, kwargs = handle.create_index.await_args + assert args == ("text",) + assert kwargs["replace"] is True + assert kwargs["wait_timeout"] == timeout + config = kwargs["config"] + assert isinstance(config, FTS) + assert config.language == "English" + assert config.stem is True + assert config.custom_stop_words == ["example"] + + +@pytest.mark.parametrize( + "kwargs, message", + [ + ({"use_tantivy": True}, "Tantivy-based FTS has been removed"), + ({"ordering_field_names": "id"}, "ordering_field_names"), + ({"writer_heap_size": 123}, "writer_heap_size"), + ({"tokenizer_name": "invalid"}, "Invalid tokenizer name"), + ], +) +def test_fts_unsupported_options(sync_table, kwargs, message): + table, handle = sync_table + with pytest.warns(DeprecationWarning), pytest.raises(ValueError, match=message): + table.create_fts_index("text", **kwargs) + handle.create_index.assert_not_awaited() + + +@pytest.mark.parametrize( + "method, kwargs", + [ + ( + "cleanup_old_versions", + {"older_than": timedelta(days=1), "delete_unverified": True}, + ), + ("compact_files", {"target_rows_per_fragment": 1000}), + ], +) +def test_remote_maintenance_options_warn(method, kwargs): + table = RemoteTable(MagicMock(name="test"), "test_db") + with pytest.warns(UserWarning, match="no-op"): + getattr(table, method)(**kwargs) + + +def test_remote_optimize_options(): + handle = MagicMock() + handle.optimize = AsyncMock(side_effect=NotImplementedError("remote optimize")) + table = RemoteTable(handle, "test_db") + age = timedelta(days=1) + with pytest.raises(NotImplementedError, match="remote optimize"): + table.optimize(cleanup_older_than=age, delete_unverified=True, retrain=True) + handle.optimize.assert_awaited_once_with( + cleanup_older_than=age, delete_unverified=True, retrain=True + ) + + +@pytest.mark.parametrize( + "method, old, new, value", + [ + ("delete", "predicate", "where", "id = 1"), + ("drop_index", "index_name", "name", "vector_idx"), + ("index_stats", "index_uuid", "index_name", "vector_idx"), + ("create_fts_index", "column", "field_names", "text"), + ], +) +def test_deprecated_remote_keywords(method, old, new, value): + handle = MagicMock() + handle.name = "test" + target = "create_index" if method == "create_fts_index" else method + setattr(handle, target, AsyncMock()) + table = RemoteTable(handle, "test_db") + with pytest.warns(DeprecationWarning, match=f"{old} is deprecated"): + getattr(table, method)(**{old: value}) + assert getattr(handle, target).await_args.args == (value,) + with pytest.raises(TypeError, match="Cannot specify both"): + getattr(table, method)(**{old: value, new: value}) + + +@pytest.mark.parametrize("remote_wrapper", [False, True], ids=["local", "remote"]) +def test_sync_search_and_polars(tmp_path, remote_wrapper): + """Use the native engine through both sync wrappers, without a cloud service.""" + local = lancedb.connect(tmp_path).create_table( + "test", + [{"id": 1, "text": "quick fox", "vector": [1.0, 1.0]}], + ) + table = RemoteTable(local._table, "test_db") if remote_wrapper else local + table.create_index("text", config=FTS()) + assert table.search("quick", query_type="fts").to_arrow()["id"].to_pylist() == [1] + # Both backends must accept the deprecated ordering argument in the same slot. + with pytest.warns(DeprecationWarning, match="ordering_field_name"): + table.search("quick", None, "fts", "id", ["text"]) + fast = table.search([1.0, 1.0], fast_search=True).to_arrow() + assert len(fast) == 0 # No vector index: fast_search excludes unindexed rows. + assert table.search([1.0, 1.0]).to_arrow()["id"].to_pylist() == [1] + if remote_wrapper: + with pytest.raises(NotImplementedError, match=r"table.search\(\).to_polars"): + table.to_polars(batch_size=1) + assert table.search().to_polars()["id"].to_list() == [1] + else: + assert table.to_polars(batch_size=1).collect()["id"].to_list() == [1] + table.delete(where="id = 1") + assert table.count_rows() == 0 + + +def test_index_replacement_defaults(tmp_path): + local = lancedb.connect(tmp_path).create_table("test", pa.table({"id": [1]})) + table = RemoteTable(local._table, "test_db") + for sync in [local, table]: + with pytest.warns(DeprecationWarning): + sync.create_scalar_index("id") + assert len(list(table.list_indices())) == 1 diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index 4fb16bb31..e85afb5bc 100644 --- a/python/python/tests/test_util.py +++ b/python/python/tests/test_util.py @@ -406,7 +406,7 @@ def test_handle_bad_vectors_jagged(on_bad_vectors): ).read_all() if on_bad_vectors == "drop": - expected = pa.array([[1.0, 2.0], [4.0, 5.0], [None, 9.0]]) + expected = pa.array([[1.0, 2.0], [4.0, 5.0]]) elif on_bad_vectors == "fill": expected = pa.array( [ @@ -414,12 +414,12 @@ def test_handle_bad_vectors_jagged(on_bad_vectors): [3.0, 42.0], [4.0, 5.0], [6.0, 7.0], - [None, 9.0], + [42.0, 9.0], [42.0, 42.0], ] ) elif on_bad_vectors == "null": - expected = pa.array([[1.0, 2.0], None, [4.0, 5.0], None, [None, 9.0], None]) + expected = pa.array([[1.0, 2.0], None, [4.0, 5.0], None, None, None]) assert output["vector"].combine_chunks() == expected @@ -430,17 +430,17 @@ def test_handle_bad_vectors_jagged(on_bad_vectors): ( pa.list_(pa.float64()), [[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]], - [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]], + [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [42.0, 3.0], [4.0, 5.0]], ), ( pa.large_list(pa.float64()), [[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]], - [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]], + [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [42.0, 3.0], [4.0, 5.0]], ), ( pa.list_(pa.float64(), 2), [[1.0, float("nan")], None, [None, 3.0]], - [[1.0, 42.0], [42.0, 42.0], [None, 3.0]], + [[1.0, 42.0], [42.0, 42.0], [42.0, 3.0]], ), ], ) @@ -467,7 +467,7 @@ def test_handle_bad_vectors_nan(on_bad_vectors): ).read_all() output = exception_output(e) assert output == ( - "ValueError: Vector column 'vector' has NaNs. Set " + "ValueError: Vector column 'vector' has NaNs or null values. Set " "on_bad_vectors='drop' to remove them, set on_bad_vectors='fill' " "and fill_value= to replace them, or set on_bad_vectors='null' " "to replace them with null." @@ -490,6 +490,46 @@ def test_handle_bad_vectors_nan(on_bad_vectors): assert output["vector"].combine_chunks() == expected +@pytest.mark.parametrize("value_type", [pa.float16(), pa.float32(), pa.float64()]) +@pytest.mark.parametrize( + "list_type", [pa.list_, pa.large_list, lambda t: pa.list_(t, 2)] +) +@pytest.mark.parametrize("chunked", [False, True]) +def test_has_nan_values_with_null_elements(value_type, list_type, chunked): + arr = pa.array( + [[None, None], [1.0, 2.0], [None, 3.0], None, [4.0, float("nan")], [5.0, 6.0]], + type=list_type(value_type), + ).slice(1, 5) + if chunked: + arr = pa.chunked_array([arr.slice(0, 2), arr.slice(2)]) + # Whole null vectors have no components to inspect. Their handling is + # separate from missing or NaN components in a valid vector. + assert has_nan_values(arr).to_pylist() == [False, True, False, True, False] + + +@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"]) +def test_sanitize_vectors_with_only_null_elements(on_bad_vectors): + data = pa.table({"vector": [[None, None]]}) + schema = pa.schema({"vector": pa.list_(pa.float32(), 2)}) + + def sanitize(): + return _sanitize_data( + data, schema, on_bad_vectors=on_bad_vectors, fill_value=42.0 + ).read_all() + + if on_bad_vectors == "error": + with pytest.raises(ValueError, match="null"): + sanitize() + return + + expected = { + "drop": [], + "fill": [{"vector": [42.0, 42.0]}], + "null": [{"vector": None}], + } + assert sanitize().to_pylist() == expected[on_bad_vectors] + + @pytest.mark.parametrize("dtype", [pa.float16(), pa.float32(), pa.float64()]) @pytest.mark.parametrize("large", [False, True]) @pytest.mark.parametrize("chunked", [False, True]) diff --git a/python/src/index.rs b/python/src/index.rs index d4e16a67e..c6c3fe17d 100644 --- a/python/src/index.rs +++ b/python/src/index.rs @@ -136,8 +136,10 @@ pub fn extract_index_params(source: &Option>) -> PyResult, - num_bits: u32, + num_bits: Option, max_iterations: u32, sample_rate: u32, target_partition_size: Option, diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index 486eb1e8e..c18db11fe 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb" -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" edition.workspace = true description = "LanceDB: A serverless, low-latency vector database for AI applications" license.workspace = true @@ -46,7 +46,7 @@ lance-index = { workspace = true, features = ["tokenizer-jieba", "tokenizer-lind lance-table = { workspace = true } lance-linalg = { workspace = true } lance-encoding = { workspace = true } -lance-arrow = { workspace = true } +lance-arrow = { version = "=15.0.0-beta.1", tag = "v15.0.0-beta.1", git = "https://github.com/lance-format/lance.git" } lance-namespace = { workspace = true } lance-namespace-impls = { workspace = true } metrics = { workspace = true, optional = true } diff --git a/rust/lancedb/src/index/vector.rs b/rust/lancedb/src/index/vector.rs index bce77a610..3e6794b54 100644 --- a/rust/lancedb/src/index/vector.rs +++ b/rust/lancedb/src/index/vector.rs @@ -362,6 +362,9 @@ impl IvfRqIndexBuilder { impl_distance_type_setter!(); impl_ivf_params_setter!(); + /// Set the number of bits used to encode each vector dimension. + /// + /// If unset, uses Lance's default (currently 5). Set to 1 for maximum compression. pub fn num_bits(mut self, num_bits: u32) -> Self { self.num_bits = Some(num_bits); self diff --git a/rust/lancedb/src/table/add_data.rs b/rust/lancedb/src/table/add_data.rs index 12404ec73..445f29a18 100644 --- a/rust/lancedb/src/table/add_data.rs +++ b/rust/lancedb/src/table/add_data.rs @@ -45,10 +45,12 @@ pub struct AddResult { #[derive(Debug, Default, Clone, Copy)] pub enum NaNVectorBehavior { - /// Reject any vectors containing NaN values (the default) + /// Reject any vectors containing NaN values or null elements (the default). + /// Whole null vectors are allowed. #[default] Error, - /// Allow NaN values to be added, but they will not be indexed for search + /// Allow NaN values and null elements to be added, but they will not be indexed + /// for search. Keep, } @@ -104,11 +106,12 @@ impl AddDataBuilder { self } - /// Configure how to handle NaN values in vector columns. + /// Configure how to handle NaN values and null elements in vector columns. /// - /// By default, any vectors containing NaN values will be rejected with an - /// error, since NaNs cannot be indexed for search. Setting this to `Keep` - /// will allow NaN values to be added to the table, but they will not be + /// By default, any vectors containing NaN values or null elements will be + /// rejected with an error, since these vectors cannot be indexed for search. + /// Whole null vectors are allowed. Setting this to `Keep` will allow NaN values + /// and null elements to be added to the table, but these vectors will not be /// indexed and will not be searchable. pub fn on_nan_vectors(mut self, behavior: NaNVectorBehavior) -> Self { self.on_nan_vectors = behavior; @@ -953,8 +956,11 @@ mod tests { ); } + #[rstest::rstest] + #[case::nan(Some(f32::NAN))] + #[case::null(None)] #[tokio::test] - async fn test_add_rejects_nan_vectors() { + async fn test_add_rejects_nan_or_null_vectors(#[case] bad_value: Option) { let schema = Arc::new(Schema::new(vec![Field::new( "embedding", DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4), @@ -974,7 +980,12 @@ mod tests { FixedSizeListArray::try_new( Arc::new(Field::new("item", DataType::Float32, true)), 4, - Arc::new(Float32Array::from(vec![0.1, 0.2, f32::NAN, 0.4])), + Arc::new(Float32Array::from(vec![ + Some(0.1), + Some(0.2), + bad_value, + Some(0.4), + ])), None, ) .unwrap(), @@ -984,8 +995,8 @@ mod tests { let res = table.add(batch.clone()).execute().await; let err = res.unwrap_err(); assert!( - err.to_string().contains("NaN"), - "Expected error mentioning NaN values, but got: {err:?}" + err.to_string().contains("NaN or null"), + "Expected error mentioning NaN or null values, but got: {err:?}" ); table diff --git a/rust/lancedb/src/table/create_index.rs b/rust/lancedb/src/table/create_index.rs index 3ac759bdd..1795b778e 100644 --- a/rust/lancedb/src/table/create_index.rs +++ b/rust/lancedb/src/table/create_index.rs @@ -355,7 +355,10 @@ impl NativeTable { index.sample_rate, index.max_iterations, ); - let rq_params = RQBuildParams::new(index.num_bits.unwrap_or(1) as u8); + let rq_params = index + .num_bits + .map(|num_bits| RQBuildParams::new(num_bits as u8)) + .unwrap_or_default(); let lance_idx_params = VectorIndexParams::with_ivf_rq_params( index.distance_type.into(), ivf_params, diff --git a/rust/lancedb/src/table/datafusion.rs b/rust/lancedb/src/table/datafusion.rs index 39f3cd4c6..e6a76c3f5 100644 --- a/rust/lancedb/src/table/datafusion.rs +++ b/rust/lancedb/src/table/datafusion.rs @@ -25,10 +25,11 @@ use datafusion_physical_plan::{ }; use futures::{TryFutureExt, TryStreamExt}; use lance::dataset::{WriteMode, WriteParams}; +use lance_core::{ROW_ADDR, ROW_ADDR_FIELD, ROW_ID, ROW_ID_FIELD}; use super::{AnyQuery, BaseTable}; use crate::{ - Result, + Error, Result, query::{QueryExecutionOptions, QueryFilter, QueryRequest, Select}, }; use arrow_schema::{DataType, Field}; @@ -186,6 +187,46 @@ impl BaseTableAdapter { }) } + /// Expose the requested Lance system columns to DataFusion. + /// + /// Supports [`ROW_ID`] (`_rowid`) and [`ROW_ADDR`] (`_rowaddr`), both nullable + /// `UInt64` fields. Columns are appended in the requested order, skipping any + /// already exposed by this adapter. Other names return [`Error::InvalidInput`]. + /// + /// System columns are hidden by default so existing `SELECT *` queries keep + /// their data-only schema. After opting in, `SELECT *` includes these columns. + /// + /// ``` + /// use lance_core::{ROW_ADDR, ROW_ID}; + /// use lancedb::table::datafusion::BaseTableAdapter; + /// + /// # async fn adapter(table: &lancedb::Table) -> lancedb::Result { + /// BaseTableAdapter::try_new(table.base_table().clone()) + /// .await? + /// .with_system_columns(&[ROW_ID, ROW_ADDR]) + /// # } + /// ``` + pub fn with_system_columns(mut self, columns: &[impl AsRef]) -> Result { + let mut fields = self.schema.fields().to_vec(); + for column in columns { + let column = column.as_ref(); + let field = match column { + ROW_ID => &*ROW_ID_FIELD, + ROW_ADDR => &*ROW_ADDR_FIELD, + _ => { + return Err(Error::InvalidInput { + message: format!("Unsupported system column: {column}"), + }); + } + }; + if !fields.iter().any(|field| field.name() == column) { + fields.push(Arc::new(field.clone())); + } + } + self.schema = Arc::new(ArrowSchema::new(fields)); + Ok(self) + } + /// Set the RTE read routing of every scan; see [`QueryRequest::use_rte`]. pub fn with_use_rte(mut self, use_rte: Option) -> Self { self.use_rte = use_rte; @@ -242,6 +283,15 @@ impl TableProvider for BaseTableAdapter { .map(|i| self.schema.field(*i).name().clone()) .collect(); query.select = Select::Columns(field_names); + } else { + // Select::All excludes system columns, which may be part of our schema. + query.select = Select::Columns( + self.schema + .fields() + .iter() + .map(|field| field.name().clone()) + .collect(), + ); } if !filters.is_empty() { let first = filters.first().unwrap().clone(); @@ -319,6 +369,7 @@ pub mod tests { BinaryArray, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray, UInt32Array, }; use arrow_schema::{DataType, Field, Schema}; + use arrow_select::concat::concat_batches; use datafusion::{ datasource::provider_as_source, prelude::{SessionConfig, SessionContext}, @@ -328,11 +379,13 @@ pub mod tests { use datafusion_execution::SendableRecordBatchStream; use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, col, lit}; use futures::{StreamExt, TryStreamExt}; + use lance_core::{ROW_ADDR, ROW_ADDR_FIELD, ROW_ID, ROW_ID_FIELD}; use tempfile::tempdir; use crate::{ connect, index::{Index, scalar::BTreeIndexBuilder}, + query::{ExecutableQuery, Query, QueryBase, Select}, table::datafusion::BaseTableAdapter, }; @@ -509,6 +562,173 @@ pub mod tests { } } + #[tokio::test] + async fn test_system_columns_opt_in() { + let fixture = TestFixture::new().await; + let ctx = SessionContext::new(); + ctx.register_table("foo", fixture.adapter.clone()).unwrap(); + + let batches = ctx + .sql("SELECT * FROM foo") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches[0].num_columns(), 2); + assert_eq!(batches[0].schema(), fixture.adapter.schema()); + for column in [ROW_ID, ROW_ADDR] { + assert!(ctx.sql(&format!("SELECT {column} FROM foo")).await.is_err()); + } + + let adapter = BaseTableAdapter::try_new(fixture.adapter.table.clone()) + .await + .unwrap() + .with_system_columns(&[ROW_ADDR, ROW_ADDR]) + .unwrap(); + assert_eq!(adapter.schema().fields().len(), 3); + assert_eq!(adapter.schema().field(2), &*ROW_ADDR_FIELD); + assert!(adapter.schema().index_of(ROW_ID).is_err()); + + let adapter = adapter + .with_system_columns(&[ROW_ID.to_string(), ROW_ADDR.to_string()]) + .unwrap(); + assert_eq!(adapter.schema().fields().len(), 4); + assert_eq!(adapter.schema().field(3), &*ROW_ID_FIELD); + assert!(adapter.schema().metadata().is_empty()); + + let err = adapter.with_system_columns(&["_unknown"]).unwrap_err(); + assert!(matches!(err, crate::Error::InvalidInput { .. })); + + // Default adapters must remain compatible with INSERT ... SELECT *. + ctx.sql("INSERT INTO foo SELECT * FROM foo") + .await + .unwrap() + .collect() + .await + .unwrap(); + let batches = ctx + .sql("SELECT COUNT(*) FROM foo") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + batches[0] + .column(0) + .as_primitive::() + .value(0), + 20 + ); + } + + #[tokio::test] + async fn test_system_columns_sql() { + let fixture = TestFixture::new().await; + let adapter = BaseTableAdapter::try_new(fixture.adapter.table.clone()) + .await + .unwrap() + .with_system_columns(&[ROW_ID, ROW_ADDR]) + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table("foo", Arc::new(adapter)).unwrap(); + + // Compare SQL results with the native query's row IDs and row addresses. + let native = Query::new(fixture.adapter.table.clone()) + .select(Select::columns(&["i", "indexed", ROW_ADDR])) + .with_row_id() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let schema = Arc::new( + native[0] + .schema() + .as_ref() + .clone() + .with_metadata(HashMap::new()), + ); + let expected = concat_batches(&schema, &native).unwrap(); + + for columns in [ + vec!["i", ROW_ID], + vec![ROW_ID, "i"], + vec!["i", ROW_ADDR], + vec![ROW_ADDR, "i"], + vec![ROW_ADDR, ROW_ID], + vec![ROW_ID], + vec![ROW_ADDR], + vec!["i", "indexed", ROW_ID, ROW_ADDR], + ] { + let sql = format!("SELECT {} FROM foo ORDER BY i", columns.join(", ")); + let batches = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + let actual = concat_batches(&batches[0].schema(), &batches).unwrap(); + let indices = columns + .iter() + .map(|name| schema.index_of(name).unwrap()) + .collect::>(); + assert_eq!(actual, expected.project(&indices).unwrap(), "{sql}"); + } + + let batches = ctx + .sql("SELECT * FROM foo ORDER BY i") + .await + .unwrap() + .collect() + .await + .unwrap(); + let indices = ["i", "indexed", ROW_ID, ROW_ADDR].map(|name| schema.index_of(name).unwrap()); + assert_eq!(batches[0], expected.project(&indices).unwrap()); + + // Filtering on the ID must work even when the ID is not in the output. + let row_id = expected + .column_by_name(ROW_ID) + .unwrap() + .as_primitive::() + .value(3); + let batches = ctx + .sql(&format!("SELECT i FROM foo WHERE _rowid = {row_id}")) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches[0], expected.project(&[0]).unwrap().slice(3, 1)); + } + + #[tokio::test] + async fn test_system_columns_scan_without_projection() { + let fixture = TestFixture::new().await; + let adapter = BaseTableAdapter::try_new(fixture.adapter.table.clone()) + .await + .unwrap() + .with_system_columns(&[ROW_ADDR, ROW_ID]) + .unwrap(); + let ctx = SessionContext::new(); + let plan = adapter.scan(&ctx.state(), None, &[], None).await.unwrap(); + assert_eq!(plan.schema(), adapter.schema()); + let batches = datafusion_physical_plan::collect(plan, ctx.task_ctx()) + .await + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 10); + for batch in batches { + assert_eq!(batch.schema(), adapter.schema()); + assert_eq!( + batch + .column(2) + .as_primitive::() + .values(), + batch + .column(3) + .as_primitive::() + .values() + ); + } + } + #[tokio::test] async fn test_batch_size() { let fixture = TestFixture::new().await; diff --git a/rust/lancedb/src/table/datafusion/reject_nan.rs b/rust/lancedb/src/table/datafusion/reject_nan.rs index 084c9722b..462216929 100644 --- a/rust/lancedb/src/table/datafusion/reject_nan.rs +++ b/rust/lancedb/src/table/datafusion/reject_nan.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -//! A DataFusion projection that rejects vectors containing NaN values. +//! A DataFusion projection that rejects vectors containing NaN or null elements. use std::sync::{Arc, LazyLock}; @@ -31,10 +31,12 @@ fn is_vector_field(field: &Field) -> bool { } } -/// Wraps the input plan with a projection that checks vector columns for NaN values. +/// Wraps the input plan with a projection that checks vector columns for NaN or +/// null elements. /// /// Non-vector columns pass through unchanged. Vector columns are wrapped with a -/// UDF that returns the column as-is if no NaNs are present, or errors otherwise. +/// UDF that returns the column as-is if no NaN or null elements are present, or +/// errors otherwise. Whole null vectors pass through unchanged. pub fn reject_nan_vectors(input: Arc) -> Result> { let schema = input.schema(); let config = Arc::new(ConfigOptions::default()); @@ -70,7 +72,7 @@ pub fn reject_nan_vectors(input: Arc) -> Result datafusion_common::Result<()> { ) })?; - // Only inspect elements that are both in a valid parent row and non-null - // themselves. Values backing null parent rows or null child elements may - // contain garbage (including NaN) per the Arrow spec. - let has_nan = (0..fsl.len()).filter(|i| fsl.is_valid(*i)).any(|i| { + // Values backing null parent rows may contain garbage per the Arrow spec. + // A null child in a valid row is an invalid vector component, however: + // pandas can represent NaN components as Arrow nulls. + let has_bad_values = (0..fsl.len()).filter(|i| fsl.is_valid(*i)).any(|i| { let row = fsl.value(i); + if row.null_count() > 0 { + return true; + } match row.data_type() { DataType::Float16 => row .as_any() @@ -150,10 +155,10 @@ fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> { } }); - if has_nan { + if has_bad_values { return Err(datafusion_common::DataFusionError::ArrowError( Box::new(arrow_schema::ArrowError::ComputeError( - "Vector column contains NaN values".to_string(), + "Vector column contains NaN or null values".to_string(), )), None, )); @@ -165,7 +170,7 @@ fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> { #[cfg(test)] mod tests { use super::*; - use arrow_array::Float32Array; + use arrow_array::{Float16Array, Float32Array, Float64Array}; #[test] fn test_passes_clean_vectors() { @@ -209,23 +214,43 @@ mod tests { } #[test] - fn test_skips_null_elements_within_valid_row() { - // A valid row with null child elements: the underlying buffer may hold - // NaN but the null bitmap says they're absent — should not reject. - let values = Float32Array::from(vec![ - Some(1.0), - None, // null element — buffer may contain NaN - Some(3.0), - None, // null element - ]); + fn test_rejects_null_elements_within_valid_row() { + let arrays: Vec> = vec![ + Arc::new(Float16Array::from(vec![Some(half::f16::ONE), None])), + Arc::new(Float32Array::from(vec![Some(1.0), None])), + Arc::new(Float64Array::from(vec![Some(1.0), None])), + ]; + for values in arrays { + let fsl = FixedSizeListArray::try_new( + Arc::new(Field::new("item", values.data_type().clone(), true)), + 2, + values, + None, + ) + .unwrap(); + let err = check_no_nans(&fsl).unwrap_err(); + assert!(err.to_string().contains("null")); + } + } + + #[test] + fn test_skips_null_elements_in_null_rows_and_outside_slice() { let fsl = FixedSizeListArray::try_new( Arc::new(Field::new("item", DataType::Float32, true)), 2, - Arc::new(values), - None, // both rows are valid + Arc::new(Float32Array::from(vec![ + None, + Some(1.0), + Some(2.0), + Some(3.0), + None, + None, + ])), + Some(vec![true, true, false].into()), ) .unwrap(); - assert!(check_no_nans(&fsl).is_ok()); + assert!(check_no_nans(&fsl.slice(1, 2)).is_ok()); + assert!(check_no_nans(&fsl.slice(0, 2)).is_err()); } #[test]