mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-14 15:52:26 +00:00
lancedb-gatefixer[bot]
85b0eba176
## Summary - infer fixed-size float32 Arrow vectors from `VectorField` metadata when a model uses the standard `list[float]` annotation - document the mypy-compatible embedding-field declaration while preserving the existing `Vector(dim)` runtime API - add regression coverage for schema conversion and model construction ## Root cause `Vector(dim)` is a runtime constrained-type factory, but Python's static typing rules do not permit call expressions in annotation positions. Mypy's suggested `Vector[...]` spelling also cannot represent a dimension computed at runtime and fails because `Vector` is a function. `VectorField` already carries the embedding function and its dimension, but schema conversion previously ignored that metadata for a statically valid `list[float]` annotation. The runtime factory remains available for dimension-constrained Pydantic validation. Embedding-backed models that require mypy compatibility can now use `list[float] = embedding.VectorField()`; LanceDB derives the same nullable fixed-size float32 Arrow type from the field metadata. ## Validation - `cd python && uv run --extra tests pytest python/tests/test_pydantic.py -q` (23 passed) - focused mypy check reproducing the original error and accepting the new `LanceModel` declaration - isolated schema checks with Pydantic 1.10.24 and 2.11.7 - `ruff format .` - `ruff check .` Fixes #1982 <!-- lance-gatekeeper-fix:v1 agent=877facbc659cdc9424bd367e12426366 generation=1 --> Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
736 lines
20 KiB
Python
736 lines
20 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
import json
|
|
from datetime import date, datetime
|
|
from enum import Enum
|
|
from typing import List, Optional, Tuple
|
|
|
|
import pyarrow as pa
|
|
import pydantic
|
|
import pytest
|
|
from lancedb.conftest import MockTextEmbeddingFunction
|
|
from lancedb.pydantic import (
|
|
LanceModel,
|
|
MultiVector,
|
|
Vector,
|
|
pydantic_to_schema,
|
|
)
|
|
from pydantic import BaseModel
|
|
from pydantic import Field
|
|
|
|
|
|
def test_pydantic_to_arrow():
|
|
class StructModel(pydantic.BaseModel):
|
|
a: str
|
|
b: Optional[float]
|
|
|
|
class TestModel(pydantic.BaseModel):
|
|
id: int
|
|
s: str
|
|
vec: list[float]
|
|
li: list[int]
|
|
lili: list[list[float]]
|
|
litu: list[tuple[float, float]]
|
|
opt: Optional[str] = None
|
|
st: StructModel
|
|
dt: date
|
|
dtt: datetime
|
|
dt_with_tz: datetime = Field(json_schema_extra={"tz": "Asia/Shanghai"})
|
|
# d: dict
|
|
|
|
# TODO: test we can actually convert the model into data.
|
|
# m = TestModel(
|
|
# id=1,
|
|
# s="hello",
|
|
# vec=[1.0, 2.0, 3.0],
|
|
# li=[2, 3, 4],
|
|
# lili=[[2.5, 1.5], [3.5, 4.5], [5.5, 6.5]],
|
|
# litu=[(2.5, 1.5), (3.5, 4.5), (5.5, 6.5)],
|
|
# st=StructModel(a="a", b=1.0),
|
|
# dt=date.today(),
|
|
# dtt=datetime.now(),
|
|
# dt_with_tz=datetime.now(pytz.timezone("Asia/Shanghai")),
|
|
# )
|
|
|
|
schema = pydantic_to_schema(TestModel)
|
|
|
|
expect_schema = pa.schema(
|
|
[
|
|
pa.field("id", pa.int64(), False),
|
|
pa.field("s", pa.utf8(), False),
|
|
pa.field("vec", pa.list_(pa.float64()), False),
|
|
pa.field("li", pa.list_(pa.int64()), False),
|
|
pa.field("lili", pa.list_(pa.list_(pa.float64())), False),
|
|
pa.field("litu", pa.list_(pa.list_(pa.float64())), False),
|
|
pa.field("opt", pa.utf8(), True),
|
|
pa.field(
|
|
"st",
|
|
pa.struct(
|
|
[pa.field("a", pa.utf8(), False), pa.field("b", pa.float64(), True)]
|
|
),
|
|
False,
|
|
),
|
|
pa.field("dt", pa.date32(), False),
|
|
pa.field("dtt", pa.timestamp("us"), False),
|
|
pa.field("dt_with_tz", pa.timestamp("us", tz="Asia/Shanghai"), False),
|
|
]
|
|
)
|
|
assert schema == expect_schema
|
|
|
|
|
|
def test_optional_types_py310():
|
|
class TestModel(pydantic.BaseModel):
|
|
a: str | None
|
|
b: None | str
|
|
c: Optional[str]
|
|
|
|
schema = pydantic_to_schema(TestModel)
|
|
|
|
expect_schema = pa.schema(
|
|
[
|
|
pa.field("a", pa.utf8(), True),
|
|
pa.field("b", pa.utf8(), True),
|
|
pa.field("c", pa.utf8(), True),
|
|
]
|
|
)
|
|
assert schema == expect_schema
|
|
|
|
|
|
def test_optional_structs():
|
|
class SplitInfo(pydantic.BaseModel):
|
|
start_frame: int
|
|
end_frame: int
|
|
|
|
class TestModel(pydantic.BaseModel):
|
|
id: str
|
|
split: SplitInfo | None = None
|
|
|
|
schema = pydantic_to_schema(TestModel)
|
|
|
|
expect_schema = pa.schema(
|
|
[
|
|
pa.field("id", pa.utf8(), False),
|
|
pa.field(
|
|
"split",
|
|
pa.struct(
|
|
[
|
|
pa.field("start_frame", pa.int64(), False),
|
|
pa.field("end_frame", pa.int64(), False),
|
|
]
|
|
),
|
|
True,
|
|
),
|
|
]
|
|
)
|
|
assert schema == expect_schema
|
|
|
|
|
|
def test_optional_struct_list_py310():
|
|
class SplitInfo(pydantic.BaseModel):
|
|
start_frame: int
|
|
end_frame: int
|
|
|
|
class TestModel(pydantic.BaseModel):
|
|
id: str
|
|
splits: list[SplitInfo] | None = None
|
|
|
|
schema = pydantic_to_schema(TestModel)
|
|
|
|
expect_schema = pa.schema(
|
|
[
|
|
pa.field("id", pa.utf8(), False),
|
|
pa.field(
|
|
"splits",
|
|
pa.list_(
|
|
pa.struct(
|
|
[
|
|
pa.field("start_frame", pa.int64(), False),
|
|
pa.field("end_frame", pa.int64(), False),
|
|
]
|
|
)
|
|
),
|
|
True,
|
|
),
|
|
]
|
|
)
|
|
assert schema == expect_schema
|
|
|
|
|
|
def test_nested_struct_list():
|
|
class SplitInfo(pydantic.BaseModel):
|
|
start_frame: int
|
|
end_frame: int
|
|
|
|
class TestModel(pydantic.BaseModel):
|
|
id: str
|
|
splits: list[SplitInfo]
|
|
|
|
schema = pydantic_to_schema(TestModel)
|
|
|
|
expect_schema = pa.schema(
|
|
[
|
|
pa.field("id", pa.utf8(), False),
|
|
pa.field(
|
|
"splits",
|
|
pa.list_(
|
|
pa.struct(
|
|
[
|
|
pa.field("start_frame", pa.int64(), False),
|
|
pa.field("end_frame", pa.int64(), False),
|
|
]
|
|
)
|
|
),
|
|
False,
|
|
),
|
|
]
|
|
)
|
|
assert schema == expect_schema
|
|
|
|
|
|
def test_bare_generic_raises_type_error():
|
|
# A bare, unparameterised List/Tuple has no element type to map to Arrow.
|
|
# It should raise a clear TypeError, not crash with AttributeError: __args__.
|
|
for bare in (List, Tuple):
|
|
|
|
class TestModel(pydantic.BaseModel):
|
|
items: bare
|
|
|
|
with pytest.raises(TypeError, match="unsupported type"):
|
|
pydantic_to_schema(TestModel)
|
|
|
|
|
|
def test_nested_struct_list_optional():
|
|
class SplitInfo(pydantic.BaseModel):
|
|
start_frame: int
|
|
end_frame: int
|
|
|
|
class TestModel(pydantic.BaseModel):
|
|
id: str
|
|
splits: Optional[list[SplitInfo]] = None
|
|
|
|
schema = pydantic_to_schema(TestModel)
|
|
|
|
expect_schema = pa.schema(
|
|
[
|
|
pa.field("id", pa.utf8(), False),
|
|
pa.field(
|
|
"splits",
|
|
pa.list_(
|
|
pa.struct(
|
|
[
|
|
pa.field("start_frame", pa.int64(), False),
|
|
pa.field("end_frame", pa.int64(), False),
|
|
]
|
|
)
|
|
),
|
|
True,
|
|
),
|
|
]
|
|
)
|
|
assert schema == expect_schema
|
|
|
|
|
|
def test_nested_struct_list_optional_items():
|
|
class SplitInfo(pydantic.BaseModel):
|
|
start_frame: int
|
|
end_frame: int
|
|
|
|
class TestModel(pydantic.BaseModel):
|
|
id: str
|
|
splits: list[Optional[SplitInfo]]
|
|
|
|
schema = pydantic_to_schema(TestModel)
|
|
|
|
expect_schema = pa.schema(
|
|
[
|
|
pa.field("id", pa.utf8(), False),
|
|
pa.field(
|
|
"splits",
|
|
pa.list_(
|
|
pa.field(
|
|
"item",
|
|
pa.struct(
|
|
[
|
|
pa.field("start_frame", pa.int64(), False),
|
|
pa.field("end_frame", pa.int64(), False),
|
|
]
|
|
),
|
|
True,
|
|
)
|
|
),
|
|
False,
|
|
),
|
|
]
|
|
)
|
|
assert schema == expect_schema
|
|
|
|
|
|
def test_nested_struct_list_optional_container_and_items():
|
|
class SplitInfo(pydantic.BaseModel):
|
|
start_frame: int
|
|
end_frame: int
|
|
|
|
class TestModel(pydantic.BaseModel):
|
|
id: str
|
|
splits: Optional[list[Optional[SplitInfo]]] = None
|
|
|
|
schema = pydantic_to_schema(TestModel)
|
|
|
|
expect_schema = pa.schema(
|
|
[
|
|
pa.field("id", pa.utf8(), False),
|
|
pa.field(
|
|
"splits",
|
|
pa.list_(
|
|
pa.field(
|
|
"item",
|
|
pa.struct(
|
|
[
|
|
pa.field("start_frame", pa.int64(), False),
|
|
pa.field("end_frame", pa.int64(), False),
|
|
]
|
|
),
|
|
True,
|
|
)
|
|
),
|
|
True,
|
|
),
|
|
]
|
|
)
|
|
assert schema == expect_schema
|
|
|
|
|
|
def test_nested_struct_list_optional_items_pep604():
|
|
class SplitInfo(pydantic.BaseModel):
|
|
start_frame: int
|
|
end_frame: int
|
|
|
|
class TestModel(pydantic.BaseModel):
|
|
id: str
|
|
splits: list[SplitInfo | None]
|
|
|
|
schema = pydantic_to_schema(TestModel)
|
|
|
|
expect_schema = pa.schema(
|
|
[
|
|
pa.field("id", pa.utf8(), False),
|
|
pa.field(
|
|
"splits",
|
|
pa.list_(
|
|
pa.field(
|
|
"item",
|
|
pa.struct(
|
|
[
|
|
pa.field("start_frame", pa.int64(), False),
|
|
pa.field("end_frame", pa.int64(), False),
|
|
]
|
|
),
|
|
True,
|
|
)
|
|
),
|
|
False,
|
|
),
|
|
]
|
|
)
|
|
assert schema == expect_schema
|
|
|
|
|
|
def test_pydantic_to_arrow_py38():
|
|
class StructModel(pydantic.BaseModel):
|
|
a: str
|
|
b: Optional[float]
|
|
|
|
class TestModel(pydantic.BaseModel):
|
|
id: int
|
|
s: str
|
|
vec: List[float]
|
|
li: List[int]
|
|
lili: List[List[float]]
|
|
litu: List[Tuple[float, float]]
|
|
opt: Optional[str] = None
|
|
st: StructModel
|
|
dt: date
|
|
dtt: datetime
|
|
dt_with_tz: datetime = Field(json_schema_extra={"tz": "Asia/Shanghai"})
|
|
# d: dict
|
|
|
|
# TODO: test we can actually convert the model to Arrow data.
|
|
# m = TestModel(
|
|
# id=1,
|
|
# s="hello",
|
|
# vec=[1.0, 2.0, 3.0],
|
|
# li=[2, 3, 4],
|
|
# lili=[[2.5, 1.5], [3.5, 4.5], [5.5, 6.5]],
|
|
# litu=[(2.5, 1.5), (3.5, 4.5), (5.5, 6.5)],
|
|
# st=StructModel(a="a", b=1.0),
|
|
# dt=date.today(),
|
|
# dtt=datetime.now(),
|
|
# dt_with_tz=datetime.now(pytz.timezone("Asia/Shanghai")),
|
|
# )
|
|
|
|
schema = pydantic_to_schema(TestModel)
|
|
|
|
expect_schema = pa.schema(
|
|
[
|
|
pa.field("id", pa.int64(), False),
|
|
pa.field("s", pa.utf8(), False),
|
|
pa.field("vec", pa.list_(pa.float64()), False),
|
|
pa.field("li", pa.list_(pa.int64()), False),
|
|
pa.field("lili", pa.list_(pa.list_(pa.float64())), False),
|
|
pa.field("litu", pa.list_(pa.list_(pa.float64())), False),
|
|
pa.field("opt", pa.utf8(), True),
|
|
pa.field(
|
|
"st",
|
|
pa.struct(
|
|
[pa.field("a", pa.utf8(), False), pa.field("b", pa.float64(), True)]
|
|
),
|
|
False,
|
|
),
|
|
pa.field("dt", pa.date32(), False),
|
|
pa.field("dtt", pa.timestamp("us"), False),
|
|
pa.field("dt_with_tz", pa.timestamp("us", tz="Asia/Shanghai"), False),
|
|
]
|
|
)
|
|
assert schema == expect_schema
|
|
|
|
|
|
def test_nullable_vector():
|
|
class NullableModel(pydantic.BaseModel):
|
|
vec: Vector(16, nullable=False)
|
|
|
|
schema = pydantic_to_schema(NullableModel)
|
|
assert schema == pa.schema([pa.field("vec", pa.list_(pa.float32(), 16), False)])
|
|
|
|
class DefaultModel(pydantic.BaseModel):
|
|
vec: Vector(16)
|
|
|
|
schema = pydantic_to_schema(DefaultModel)
|
|
assert schema == pa.schema([pa.field("vec", pa.list_(pa.float32(), 16), True)])
|
|
|
|
class NotNullableModel(pydantic.BaseModel):
|
|
vec: Vector(16)
|
|
|
|
schema = pydantic_to_schema(NotNullableModel)
|
|
assert schema == pa.schema([pa.field("vec", pa.list_(pa.float32(), 16), True)])
|
|
|
|
|
|
def test_bare_vector_raises_clear_error():
|
|
namespace = {
|
|
"__name__": "test_model_without_pyarrow",
|
|
"LanceModel": LanceModel,
|
|
"Vector": Vector,
|
|
}
|
|
|
|
with pytest.raises(TypeError, match=r"Vector must be parameterized.*Vector\(128\)"):
|
|
exec("class TestModel(LanceModel):\n vector: Vector", namespace)
|
|
|
|
|
|
def test_embedding_vector_list_annotation():
|
|
embedding = MockTextEmbeddingFunction.create()
|
|
|
|
class StaticTypingModel(LanceModel):
|
|
text: str = embedding.SourceField()
|
|
vector: list[float] = embedding.VectorField()
|
|
|
|
schema = pydantic_to_schema(StaticTypingModel)
|
|
assert schema == pa.schema(
|
|
[
|
|
pa.field("text", pa.utf8(), False),
|
|
pa.field("vector", pa.list_(pa.float32(), embedding.ndims()), True),
|
|
]
|
|
)
|
|
|
|
model = StaticTypingModel(text="hello", vector=[0.0] * embedding.ndims())
|
|
assert model.vector == [0.0] * embedding.ndims()
|
|
|
|
|
|
def test_fixed_size_list_field():
|
|
class TestModel(pydantic.BaseModel):
|
|
vec: Vector(16)
|
|
li: List[int]
|
|
|
|
data = TestModel(vec=list(range(16)), li=[1, 2, 3])
|
|
assert json.loads(data.model_dump_json()) == {
|
|
"vec": list(range(16)),
|
|
"li": [1, 2, 3],
|
|
}
|
|
|
|
schema = pydantic_to_schema(TestModel)
|
|
assert schema == pa.schema(
|
|
[
|
|
pa.field("vec", pa.list_(pa.float32(), 16)),
|
|
pa.field("li", pa.list_(pa.int64()), False),
|
|
]
|
|
)
|
|
|
|
json_schema = TestModel.model_json_schema()
|
|
|
|
assert json_schema == {
|
|
"properties": {
|
|
"vec": {
|
|
"items": {"type": "number"},
|
|
"maxItems": 16,
|
|
"minItems": 16,
|
|
"title": "Vec",
|
|
"type": "array",
|
|
},
|
|
"li": {"items": {"type": "integer"}, "title": "Li", "type": "array"},
|
|
},
|
|
"required": ["vec", "li"],
|
|
"title": "TestModel",
|
|
"type": "object",
|
|
}
|
|
|
|
|
|
def test_fixed_size_list_validation():
|
|
class TestModel(pydantic.BaseModel):
|
|
vec: Vector(8)
|
|
|
|
with pytest.raises(pydantic.ValidationError):
|
|
TestModel(vec=range(9))
|
|
|
|
with pytest.raises(pydantic.ValidationError):
|
|
TestModel(vec=range(7))
|
|
|
|
TestModel(vec=range(8))
|
|
|
|
|
|
def test_lance_model():
|
|
class TestModel(LanceModel):
|
|
vector: Vector(16) = Field(default=[0.0] * 16)
|
|
li: List[int] = Field(default=[1, 2, 3])
|
|
|
|
schema = pydantic_to_schema(TestModel)
|
|
assert schema == TestModel.to_arrow_schema()
|
|
assert TestModel.field_names() == ["vector", "li"]
|
|
|
|
t = TestModel()
|
|
assert t == TestModel(vec=[0.0] * 16, li=[1, 2, 3])
|
|
|
|
|
|
def test_optional_nested_model():
|
|
class WAMedia(BaseModel):
|
|
url: str
|
|
mimetype: str
|
|
filename: Optional[str]
|
|
error: Optional[str]
|
|
data: bytes
|
|
|
|
class WALocation(BaseModel):
|
|
description: Optional[str]
|
|
latitude: str
|
|
longitude: str
|
|
|
|
class ReplyToMessage(BaseModel):
|
|
id: str
|
|
participant: str
|
|
body: str
|
|
|
|
class Message(BaseModel):
|
|
id: str
|
|
timestamp: int
|
|
from_: str
|
|
fromMe: bool
|
|
to: str
|
|
body: str
|
|
hasMedia: Optional[bool]
|
|
media: WAMedia
|
|
mediaUrl: Optional[str]
|
|
ack: Optional[int]
|
|
ackName: Optional[str]
|
|
author: Optional[str]
|
|
location: Optional[WALocation]
|
|
vCards: Optional[List[str]]
|
|
replyTo: Optional[ReplyToMessage]
|
|
|
|
class AnyEvent(LanceModel):
|
|
id: str
|
|
session: str
|
|
metadata: Optional[str] = None
|
|
engine: str
|
|
event: str
|
|
|
|
class MessageEvent(AnyEvent):
|
|
payload: Message
|
|
|
|
schema = pydantic_to_schema(MessageEvent)
|
|
|
|
payload = schema.field("payload")
|
|
assert payload.type == pa.struct(
|
|
[
|
|
pa.field("id", pa.utf8(), False),
|
|
pa.field("timestamp", pa.int64(), False),
|
|
pa.field("from_", pa.utf8(), False),
|
|
pa.field("fromMe", pa.bool_(), False),
|
|
pa.field("to", pa.utf8(), False),
|
|
pa.field("body", pa.utf8(), False),
|
|
pa.field("hasMedia", pa.bool_(), True),
|
|
pa.field(
|
|
"media",
|
|
pa.struct(
|
|
[
|
|
pa.field("url", pa.utf8(), False),
|
|
pa.field("mimetype", pa.utf8(), False),
|
|
pa.field("filename", pa.utf8(), True),
|
|
pa.field("error", pa.utf8(), True),
|
|
pa.field("data", pa.binary(), False),
|
|
]
|
|
),
|
|
False,
|
|
),
|
|
pa.field("mediaUrl", pa.utf8(), True),
|
|
pa.field("ack", pa.int64(), True),
|
|
pa.field("ackName", pa.utf8(), True),
|
|
pa.field("author", pa.utf8(), True),
|
|
pa.field(
|
|
"location",
|
|
pa.struct(
|
|
[
|
|
pa.field("description", pa.utf8(), True),
|
|
pa.field("latitude", pa.utf8(), False),
|
|
pa.field("longitude", pa.utf8(), False),
|
|
]
|
|
),
|
|
True, # Optional
|
|
),
|
|
pa.field("vCards", pa.list_(pa.utf8()), True),
|
|
pa.field(
|
|
"replyTo",
|
|
pa.struct(
|
|
[
|
|
pa.field("id", pa.utf8(), False),
|
|
pa.field("participant", pa.utf8(), False),
|
|
pa.field("body", pa.utf8(), False),
|
|
]
|
|
),
|
|
True,
|
|
),
|
|
]
|
|
)
|
|
|
|
|
|
def test_multi_vector():
|
|
class TestModel(pydantic.BaseModel):
|
|
vec: MultiVector(8)
|
|
|
|
schema = pydantic_to_schema(TestModel)
|
|
assert schema == pa.schema(
|
|
[pa.field("vec", pa.list_(pa.list_(pa.float32(), 8)), True)]
|
|
)
|
|
|
|
with pytest.raises(pydantic.ValidationError):
|
|
TestModel(vec=[[1.0] * 7])
|
|
|
|
with pytest.raises(pydantic.ValidationError):
|
|
TestModel(vec=[[1.0] * 9])
|
|
|
|
TestModel(vec=[[1.0] * 8])
|
|
TestModel(vec=[[1.0] * 8, [2.0] * 8])
|
|
|
|
TestModel(vec=[])
|
|
|
|
|
|
def test_multi_vector_nullable():
|
|
class NullableModel(pydantic.BaseModel):
|
|
vec: MultiVector(16, nullable=False)
|
|
|
|
schema = pydantic_to_schema(NullableModel)
|
|
assert schema == pa.schema(
|
|
[pa.field("vec", pa.list_(pa.list_(pa.float32(), 16)), False)]
|
|
)
|
|
|
|
class DefaultModel(pydantic.BaseModel):
|
|
vec: MultiVector(16)
|
|
|
|
schema = pydantic_to_schema(DefaultModel)
|
|
assert schema == pa.schema(
|
|
[pa.field("vec", pa.list_(pa.list_(pa.float32(), 16)), True)]
|
|
)
|
|
|
|
|
|
def test_multi_vector_in_lance_model():
|
|
class TestModel(LanceModel):
|
|
id: int
|
|
vectors: MultiVector(16) = Field(default=[[0.0] * 16])
|
|
|
|
schema = pydantic_to_schema(TestModel)
|
|
assert schema == TestModel.to_arrow_schema()
|
|
assert TestModel.field_names() == ["id", "vectors"]
|
|
|
|
t = TestModel(id=1)
|
|
assert t.vectors == [[0.0] * 16]
|
|
|
|
|
|
def test_aliases_in_lance_model(mem_db):
|
|
data = [
|
|
{"vector": [3.1, 4.1], "item": "foo", "price": 10.0},
|
|
{"vector": [5.9, 6.5], "item": "bar", "price": 20.0},
|
|
]
|
|
tbl = mem_db.create_table("items", data=data)
|
|
|
|
class TestModel(LanceModel):
|
|
name: str = Field(alias="item")
|
|
price: float
|
|
distance: float = Field(alias="_distance")
|
|
|
|
model = (
|
|
tbl.search([5.9, 6.5])
|
|
.distance_type("cosine")
|
|
.limit(1)
|
|
.to_pydantic(TestModel)[0]
|
|
)
|
|
assert hasattr(model, "name")
|
|
assert hasattr(model, "distance")
|
|
assert model.distance < 0.01
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aliases_in_lance_model_async(mem_db_async):
|
|
data = [
|
|
{"vector": [8.3, 2.5], "item": "foo", "price": 12.0},
|
|
{"vector": [7.7, 3.9], "item": "bar", "price": 11.2},
|
|
]
|
|
tbl = await mem_db_async.create_table("items", data=data)
|
|
|
|
class TestModel(LanceModel):
|
|
name: str = Field(alias="item")
|
|
price: float
|
|
distance: float = Field(alias="_distance")
|
|
|
|
model = (
|
|
await tbl.vector_search([7.7, 3.9])
|
|
.distance_type("cosine")
|
|
.limit(1)
|
|
.to_pydantic(TestModel)
|
|
)[0]
|
|
assert hasattr(model, "name")
|
|
assert hasattr(model, "distance")
|
|
assert model.distance < 0.01
|
|
|
|
|
|
def test_enum_types():
|
|
"""Enum fields should map to the Arrow type of their value (issue #1846)."""
|
|
|
|
class StrStatus(str, Enum):
|
|
PENDING = "pending"
|
|
RUNNING = "running"
|
|
DONE = "done"
|
|
|
|
class IntPriority(int, Enum):
|
|
LOW = 1
|
|
MEDIUM = 2
|
|
HIGH = 3
|
|
|
|
class TestModel(pydantic.BaseModel):
|
|
status: StrStatus
|
|
priority: IntPriority
|
|
opt_status: Optional[StrStatus] = None
|
|
|
|
schema = pydantic_to_schema(TestModel)
|
|
|
|
assert schema.field("status").type == pa.dictionary(pa.int32(), pa.utf8())
|
|
assert schema.field("priority").type == pa.int64()
|
|
assert schema.field("opt_status").type == pa.dictionary(pa.int32(), pa.utf8())
|
|
assert schema.field("opt_status").nullable
|