rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-27 14:06:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge main into gatekeeper/fix-2325-1

Browse Source 
Retain both retrain and compaction_options on RemoteTable.optimize, preserving upstream signature alignment and the existing compaction API.

Validation: rebuilt the Python extension; 66 signature and optimize tests plus the Python tuning example pass. Python lint and Python/Rust formatting checks pass. Root Ruff retains 49 inherited documentation diagnostics in files identical to main.
This commit is contained in:
Gatefixer
2026-10-26 12:01:47 +00:00
parent 2a85c0eb9b 3dd579785e
commit 3280acb667
44 changed files with 1422 additions and 466 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/Cargo.toml
+1 -2
View File
@@ -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
python/python/lancedb/db.py
+10 -4
View File
@@ -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.
python/python/lancedb/embeddings/open_clip.py
+3
View File
@@ -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
python/python/lancedb/index.py
+6 -2
View File
@@ -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
python/python/lancedb/remote/db.py
+2 -2
View File
@@ -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,
python/python/lancedb/remote/table.py
+198 -145
View File
@@ -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.
python/python/lancedb/table.py
+259 -229
View File
@@ -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=<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".
python/python/tests/test_embeddings.py
+33
View File
@@ -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
python/python/tests/test_index.py
+12 -2
View File
@@ -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
python/python/tests/test_remote_db.py
+47 -5
View File
@@ -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
python/python/tests/test_table.py
+87
View File
@@ -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
python/python/tests/test_table_signatures.py
+253
View File
@@ -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
python/python/tests/test_util.py
+47 -7
View File
@@ -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=<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])
python/src/index.rs
+5 -3
View File
@@ -136,8 +136,10 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
let mut ivf_rq_builder = IvfRqIndexBuilder::default()
.distance_type(distance_type)
.max_iterations(params.max_iterations)
.sample_rate(params.sample_rate)
.num_bits(params.num_bits);
.sample_rate(params.sample_rate);
if let Some(num_bits) = params.num_bits {
ivf_rq_builder = ivf_rq_builder.num_bits(num_bits);
}
if let Some(num_partitions) = params.num_partitions {
ivf_rq_builder = ivf_rq_builder.num_partitions(num_partitions);
}
@@ -265,7 +267,7 @@ struct IvfSqParams {
struct IvfRqParams {
distance_type: String,
num_partitions: Option<u32>,
num_bits: u32,
num_bits: Option<u32>,
max_iterations: u32,
sample_rate: u32,
target_partition_size: Option<u32>,