mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-27 16:38:31 +00:00
Compare commits
15 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 670bda8725 | |||
| 10151b2dc8 | |||
| 53b4c3b715 | |||
| 638430fdb4 | |||
| 2221b8df6a | |||
| 14eeae4bd4 | |||
| 320755ed55 | |||
| e55c2da7b1 | |||
| d33b05328c | |||
| 82f5355b71 | |||
| 40cff9b644 | |||
| edf95e53fc | |||
| 0b5eba085d | |||
| 21bf859c0b | |||
| e0499de959 |
@@ -42,6 +42,8 @@ listing a storage directory.
|
||||
|
||||
::: lancedb.table.Table
|
||||
|
||||
::: lancedb.table.CompactionOptions
|
||||
|
||||
::: lancedb.table.FragmentStatistics
|
||||
|
||||
::: lancedb.table.FragmentSummaryStats
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
import * as fs from "node:fs";
|
||||
import * as vm from "node:vm";
|
||||
import * as arrow15 from "apache-arrow-15";
|
||||
import * as arrow16 from "apache-arrow-16";
|
||||
import * as arrow17 from "apache-arrow-17";
|
||||
@@ -42,41 +40,6 @@ function sampleRecords(): Array<Record<string, any>> {
|
||||
];
|
||||
}
|
||||
|
||||
it("serializes an Arrow Table created in another JavaScript realm", async () => {
|
||||
const context = vm.createContext({
|
||||
TextDecoder,
|
||||
TextEncoder,
|
||||
console,
|
||||
setTimeout,
|
||||
clearTimeout,
|
||||
});
|
||||
vm.runInContext(
|
||||
fs.readFileSync(
|
||||
require.resolve("apache-arrow-15/Arrow.es2015.min"),
|
||||
"utf8",
|
||||
),
|
||||
context,
|
||||
);
|
||||
const foreignTable: unknown = vm.runInContext(
|
||||
"Arrow.tableFromArrays({ id: new Int32Array([1, 2, 3]), text: ['foo', 'bar', 'baz'] })",
|
||||
context,
|
||||
);
|
||||
|
||||
const foreignMetadata = (
|
||||
foreignTable as { schema: { metadata: Map<string, string> } }
|
||||
).schema.metadata;
|
||||
expect(foreignMetadata).not.toBeInstanceOf(Map);
|
||||
|
||||
const buf = await fromDataToBuffer(
|
||||
foreignTable as Parameters<typeof fromDataToBuffer>[0],
|
||||
);
|
||||
const actual = currentTableFromIPC(buf);
|
||||
|
||||
expect(actual.numRows).toBe(3);
|
||||
expect(actual.getChild("id")?.toJSON()).toEqual([1, 2, 3]);
|
||||
expect(actual.getChild("text")?.toJSON()).toEqual(["foo", "bar", "baz"]);
|
||||
});
|
||||
|
||||
it("preserves field metadata from a provided schema", async function () {
|
||||
const jsonMetadata = new Map([["ARROW:extension:name", "lance.json"]]);
|
||||
const schema = new CurrentSchema([
|
||||
|
||||
@@ -72,7 +72,8 @@ export type FieldLike =
|
||||
};
|
||||
|
||||
export type DataLike =
|
||||
| import("apache-arrow").Data
|
||||
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
|
||||
| import("apache-arrow").Data<Struct<any>>
|
||||
| {
|
||||
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
|
||||
type: any;
|
||||
@@ -81,7 +82,6 @@ export type DataLike =
|
||||
stride: number;
|
||||
nullable: boolean;
|
||||
children: DataLike[];
|
||||
dictionary?: { data: readonly DataLike[] };
|
||||
get nullCount(): number;
|
||||
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
|
||||
values: Buffers<any>[BufferType.DATA];
|
||||
|
||||
@@ -94,24 +94,17 @@ export function sanitizeMetadata(
|
||||
if (metadataLike === undefined || metadataLike === null) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
let entries: IterableIterator<[unknown, unknown]>;
|
||||
try {
|
||||
entries = Map.prototype.entries.call(metadataLike);
|
||||
} catch {
|
||||
if (!(metadataLike instanceof Map)) {
|
||||
throw Error("Expected metadata, if present, to be a Map<string, string>");
|
||||
}
|
||||
|
||||
const metadata = new Map<string, string>();
|
||||
for (const [key, value] of entries) {
|
||||
if (typeof key !== "string" || typeof value !== "string") {
|
||||
for (const item of metadataLike) {
|
||||
if (typeof item[0] !== "string" || typeof item[1] !== "string") {
|
||||
throw Error(
|
||||
"Expected metadata, if present, to be a Map<string, string> but it had non-string keys or values",
|
||||
);
|
||||
}
|
||||
metadata.set(key, value);
|
||||
}
|
||||
return metadata;
|
||||
return metadataLike as Map<string, string>;
|
||||
}
|
||||
|
||||
export function sanitizeInt(typeLike: object) {
|
||||
|
||||
@@ -38,7 +38,7 @@ from .materialized_view import (
|
||||
MaterializedView,
|
||||
MaterializedViewDefinition,
|
||||
)
|
||||
from .table import AsyncTable, Table
|
||||
from .table import AsyncTable, CompactionOptions, Table
|
||||
from .types import BaseTokenizerType
|
||||
from ._lancedb import Session
|
||||
from .namespace import (
|
||||
@@ -545,6 +545,7 @@ __all__ = [
|
||||
"AsyncJob",
|
||||
"AsyncLanceNamespaceDBConnection",
|
||||
"AsyncTable",
|
||||
"CompactionOptions",
|
||||
"FtsToken",
|
||||
"col",
|
||||
"Expr",
|
||||
|
||||
@@ -270,8 +270,7 @@ def _iter_projection_pairs(
|
||||
if isinstance(expr, str):
|
||||
yield name, expr
|
||||
elif isinstance(expr, Expr):
|
||||
source = expr._column_name()
|
||||
yield name, source if source is not None else expr.to_sql()
|
||||
yield name, expr.to_sql()
|
||||
return
|
||||
for column in projection:
|
||||
if isinstance(column, str):
|
||||
@@ -281,8 +280,7 @@ def _iter_projection_pairs(
|
||||
if isinstance(expr, str):
|
||||
yield name, expr
|
||||
elif isinstance(expr, Expr):
|
||||
source = expr._column_name()
|
||||
yield name, source if source is not None else expr.to_sql()
|
||||
yield name, expr.to_sql()
|
||||
|
||||
|
||||
def _set_blob_column(tbl: pa.Table, output_name: str, blobs: pa.Array) -> pa.Table:
|
||||
|
||||
@@ -87,7 +87,6 @@ class PyExpr:
|
||||
def contains(self, substr: "PyExpr") -> "PyExpr": ...
|
||||
def isin(self, values: List["PyExpr"]) -> "PyExpr": ...
|
||||
def cast(self, data_type: pa.DataType) -> "PyExpr": ...
|
||||
def column_name(self) -> Optional[str]: ...
|
||||
def to_sql(self) -> str: ...
|
||||
|
||||
def expr_col(name: str) -> PyExpr: ...
|
||||
@@ -374,6 +373,7 @@ class Table:
|
||||
*,
|
||||
cleanup_since_ms: Optional[int] = None,
|
||||
delete_unverified: Optional[bool] = None,
|
||||
compaction_options: Optional[Dict[str, Any]] = None,
|
||||
) -> OptimizeStats: ...
|
||||
async def uri(self) -> str: ...
|
||||
async def initial_storage_options(self) -> Optional[Dict[str, str]]: ...
|
||||
@@ -609,7 +609,6 @@ class PyQueryRequest:
|
||||
filter: Optional[Union[str, bytes]]
|
||||
full_text_search: Optional[FullTextQuery]
|
||||
select: Optional[Union[str, List[str]]]
|
||||
select_source_columns: Optional[Dict[str, str]]
|
||||
fast_search: Optional[bool]
|
||||
with_row_id: Optional[bool]
|
||||
use_lsm: Optional[bool]
|
||||
|
||||
@@ -249,10 +249,6 @@ class Expr:
|
||||
|
||||
# ── utilities ────────────────────────────────────────────────────────────
|
||||
|
||||
def _column_name(self) -> str | None:
|
||||
"""Return the source name when this is a bare column expression."""
|
||||
return self._inner.column_name()
|
||||
|
||||
def to_sql(self) -> str:
|
||||
"""Render the expression as a SQL string (useful for debugging)."""
|
||||
return self._inner.to_sql()
|
||||
@@ -316,7 +312,7 @@ def func(name: str, *args: ExprLike) -> Expr:
|
||||
--------
|
||||
>>> from lancedb.expr import col, func
|
||||
>>> func("lower", col("name"))
|
||||
Expr(lower(`name`))
|
||||
Expr(lower(name))
|
||||
"""
|
||||
inner_args = [_coerce(a)._inner for a in args]
|
||||
return Expr(expr_func(name, inner_args))
|
||||
|
||||
@@ -167,12 +167,6 @@ def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]:
|
||||
return {"columns": projection}
|
||||
|
||||
|
||||
def _query_request_projection(req: "PyQueryRequest") -> QueryProjection:
|
||||
if req.select_source_columns is not None:
|
||||
return req.select_source_columns
|
||||
return req.select
|
||||
|
||||
|
||||
def _scanner_kwargs_for_query(
|
||||
query: Query,
|
||||
blob_mode: BlobMode,
|
||||
@@ -2805,16 +2799,15 @@ class AsyncQueryBase(object):
|
||||
|
||||
req = self._inner.to_query_request()
|
||||
schema = await self._table.schema()
|
||||
projection = _query_request_projection(req)
|
||||
self._blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||
schema,
|
||||
projection,
|
||||
req.select,
|
||||
with_row_id=self._with_row_id,
|
||||
)
|
||||
if not self._blob_auto_row_id:
|
||||
self._blob_paths = ()
|
||||
return
|
||||
self._blob_paths = tuple(blob_v2_projection_sources(schema, projection).keys())
|
||||
self._blob_paths = tuple(blob_v2_projection_sources(schema, req.select).keys())
|
||||
self._inner.with_row_id()
|
||||
|
||||
def select(self, columns: Union[List[str], dict[str, str]]) -> Self:
|
||||
@@ -3901,15 +3894,14 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||||
blob_paths: tuple[str, ...] = ()
|
||||
if self._table is not None:
|
||||
schema = await self._table.schema()
|
||||
projection = _query_request_projection(req)
|
||||
blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||
schema,
|
||||
projection,
|
||||
req.select,
|
||||
with_row_id=self._with_row_id,
|
||||
)
|
||||
if blob_auto_row_id:
|
||||
blob_paths = tuple(
|
||||
blob_v2_projection_sources(schema, projection).keys()
|
||||
blob_v2_projection_sources(schema, req.select).keys()
|
||||
)
|
||||
self._blob_auto_row_id = blob_auto_row_id
|
||||
self._blob_paths = blob_paths
|
||||
|
||||
@@ -36,7 +36,6 @@ from lancedb._lancedb import (
|
||||
UpdateResult,
|
||||
)
|
||||
from lancedb.embeddings.base import EmbeddingFunctionConfig
|
||||
from lancedb.expr import Expr
|
||||
from lancedb.index import (
|
||||
FTS,
|
||||
BTree,
|
||||
@@ -67,7 +66,16 @@ from ..query import (
|
||||
LanceTakeQueryBuilder,
|
||||
LanceVectorQueryBuilder,
|
||||
)
|
||||
from ..table import AsyncTable, BlobMode, Branches, IndexStatistics, Query, Table, Tags
|
||||
from ..table import (
|
||||
AsyncTable,
|
||||
BlobMode,
|
||||
Branches,
|
||||
CompactionOptions,
|
||||
IndexStatistics,
|
||||
Query,
|
||||
Table,
|
||||
Tags,
|
||||
)
|
||||
from ..types import BaseTokenizerType
|
||||
|
||||
|
||||
@@ -864,7 +872,7 @@ class RemoteTable(Table):
|
||||
|
||||
def update(
|
||||
self,
|
||||
where: Optional[Union[str, Expr]] = None,
|
||||
where: Optional[str] = None,
|
||||
values: Optional[dict] = None,
|
||||
*,
|
||||
values_sql: Optional[Dict[str, str]] = None,
|
||||
@@ -875,11 +883,9 @@ class RemoteTable(Table):
|
||||
|
||||
Parameters
|
||||
----------
|
||||
where: str or [Expr][lancedb.expr.Expr], optional
|
||||
The filter condition. Can be a SQL string or a type-safe
|
||||
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
|
||||
[lit][lancedb.expr.lit]. The filter must not be empty, or it will
|
||||
error.
|
||||
where: str, optional
|
||||
The SQL where clause to use when updating rows. For example, 'x = 2'
|
||||
or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error.
|
||||
values: dict, optional
|
||||
The values to update. The keys are the column names and the values
|
||||
are the values to set.
|
||||
@@ -953,6 +959,7 @@ class RemoteTable(Table):
|
||||
*,
|
||||
cleanup_older_than: Optional[timedelta] = None,
|
||||
delete_unverified: bool = False,
|
||||
compaction_options: Optional[CompactionOptions] = None,
|
||||
):
|
||||
"""
|
||||
optimize() is a no-op on LanceDB Cloud.
|
||||
|
||||
+119
-26
@@ -22,6 +22,7 @@ from typing import (
|
||||
Optional,
|
||||
Sequence,
|
||||
Tuple,
|
||||
TypedDict,
|
||||
Union,
|
||||
overload,
|
||||
)
|
||||
@@ -222,6 +223,88 @@ IndexConfigType = Union[
|
||||
FTS,
|
||||
]
|
||||
|
||||
|
||||
class CompactionOptions(TypedDict, total=False):
|
||||
"""Options that control file compaction during table optimization.
|
||||
|
||||
Unspecified options use Lance's defaults.
|
||||
|
||||
Compaction planning is row based. Lowering ``target_rows_per_fragment``
|
||||
based on the expected row size can bound later compaction passes once
|
||||
oversized fragments have been rewritten. It does not split an existing
|
||||
fragment, so the first pass over an oversized fragment is not subject to
|
||||
that bound. ``max_bytes_per_file`` limits output file size, not compaction
|
||||
memory. Source budgets keep whole planned tasks; if the first task exceeds
|
||||
a budget, that run performs no compaction work.
|
||||
|
||||
Examples
|
||||
--------
|
||||
Derive a steady-state row target from the expected row size:
|
||||
|
||||
>>> desired_fragment_bytes = 750 * 1024 * 1024
|
||||
>>> average_row_bytes = 1_500_000
|
||||
>>> options: CompactionOptions = {
|
||||
... "target_rows_per_fragment": max(
|
||||
... 1, desired_fragment_bytes // average_row_bytes
|
||||
... ),
|
||||
... }
|
||||
>>> await table.optimize(compaction_options=options) # doctest: +SKIP
|
||||
"""
|
||||
|
||||
target_rows_per_fragment: int
|
||||
"""Target rows per fragment; existing oversized fragments are not split."""
|
||||
|
||||
max_rows_per_group: int
|
||||
"""Maximum number of rows per row group (default: 1,024)."""
|
||||
|
||||
max_bytes_per_file: Optional[int]
|
||||
"""Maximum output data-file size; this does not bound compaction memory."""
|
||||
|
||||
materialize_deletions: bool
|
||||
"""Whether to rewrite fragments containing deleted rows (default: True)."""
|
||||
|
||||
materialize_deletions_threshold: float
|
||||
"""Minimum deleted-row fraction that makes a fragment eligible (default: 0.1)."""
|
||||
|
||||
num_threads: Optional[int]
|
||||
"""Number of compaction tasks to run in parallel."""
|
||||
|
||||
batch_size: Optional[int]
|
||||
"""Number of rows per input scan batch."""
|
||||
|
||||
io_buffer_size: Optional[int]
|
||||
"""Maximum number of bytes queued in the input scan I/O buffer."""
|
||||
|
||||
defer_index_remap: bool
|
||||
"""Whether to defer index remapping during compaction (default: False)."""
|
||||
|
||||
index_remap_mode: Literal["direct", "compact"]
|
||||
"""How to construct the old-to-new row-address mapping."""
|
||||
|
||||
compaction_mode: Optional[
|
||||
Literal["reencode", "try_binary_copy", "force_binary_copy"]
|
||||
]
|
||||
"""Whether compaction re-encodes data or uses binary copying."""
|
||||
|
||||
binary_copy_read_batch_bytes: Optional[int]
|
||||
"""Number of bytes read per batch during binary-copy compaction."""
|
||||
|
||||
max_source_fragments: Optional[int]
|
||||
"""Maximum number of source fragments compacted in one run."""
|
||||
|
||||
max_source_rows: Optional[int]
|
||||
"""Maximum live source rows per run, applied to whole planned tasks."""
|
||||
|
||||
max_source_bytes: Optional[int]
|
||||
"""Maximum source bytes per run, applied to whole planned tasks."""
|
||||
|
||||
excluded_fragment_ids: List[int]
|
||||
"""Fragment IDs to leave unchanged and use as planning boundaries."""
|
||||
|
||||
max_overlays_per_fragment: Optional[int]
|
||||
"""Maximum overlays before a fragment is fully compacted."""
|
||||
|
||||
|
||||
# Known distance metrics for legacy API detection
|
||||
KNOWN_METRICS = {"l2", "cosine", "dot", "hamming"}
|
||||
|
||||
@@ -1744,7 +1827,7 @@ class Table(ABC):
|
||||
@abstractmethod
|
||||
def update(
|
||||
self,
|
||||
where: Optional[Union[str, Expr]] = None,
|
||||
where: Optional[str] = None,
|
||||
values: Optional[dict] = None,
|
||||
*,
|
||||
values_sql: Optional[Dict[str, str]] = None,
|
||||
@@ -1759,11 +1842,9 @@ class Table(ABC):
|
||||
|
||||
Parameters
|
||||
----------
|
||||
where: str or [Expr][lancedb.expr.Expr], optional
|
||||
The filter condition. Can be a SQL string or a type-safe
|
||||
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
|
||||
[lit][lancedb.expr.lit]. The filter must not be empty, or it will
|
||||
error.
|
||||
where: str, optional
|
||||
The SQL where clause to use when updating rows. For example, 'x = 2'
|
||||
or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error.
|
||||
values: dict, optional
|
||||
The values to update. The keys are the column names and the values
|
||||
are the values to set.
|
||||
@@ -1781,7 +1862,6 @@ class Table(ABC):
|
||||
Examples
|
||||
--------
|
||||
>>> import lancedb
|
||||
>>> from lancedb.expr import col
|
||||
>>> import pandas as pd
|
||||
>>> data = pd.DataFrame({"x": [1, 2, 3], "vector": [[1.0, 2], [3, 4], [5, 6]]})
|
||||
>>> db = lancedb.connect("./.lancedb")
|
||||
@@ -1791,7 +1871,7 @@ class Table(ABC):
|
||||
0 1 [1.0, 2.0]
|
||||
1 2 [3.0, 4.0]
|
||||
2 3 [5.0, 6.0]
|
||||
>>> table.update(where=col("x") == 2, values={"vector": [10.0, 10]})
|
||||
>>> table.update(where="x = 2", values={"vector": [10.0, 10]})
|
||||
UpdateResult(rows_updated=1, version=2)
|
||||
>>> table.to_pandas()
|
||||
x vector
|
||||
@@ -1875,6 +1955,7 @@ class Table(ABC):
|
||||
cleanup_older_than: Optional[timedelta] = None,
|
||||
delete_unverified: bool = False,
|
||||
retrain: bool = False,
|
||||
compaction_options: Optional[CompactionOptions] = None,
|
||||
):
|
||||
"""
|
||||
Optimize the on-disk data and indices for better performance.
|
||||
@@ -1907,6 +1988,11 @@ class Table(ABC):
|
||||
|
||||
retrain: bool, default False
|
||||
This parameter is no longer used and is deprecated.
|
||||
compaction_options: CompactionOptions, optional
|
||||
Options that control file compaction. For large rows, derive a lower
|
||||
``target_rows_per_fragment`` from the expected row size to bound later
|
||||
passes. This does not cap the first pass over an existing oversized
|
||||
fragment; see [CompactionOptions][lancedb.table.CompactionOptions].
|
||||
|
||||
Notes
|
||||
-----
|
||||
@@ -3844,7 +3930,7 @@ class LanceTable(Table):
|
||||
|
||||
def update(
|
||||
self,
|
||||
where: Optional[Union[str, Expr]] = None,
|
||||
where: Optional[str] = None,
|
||||
values: Optional[dict] = None,
|
||||
*,
|
||||
values_sql: Optional[Dict[str, str]] = None,
|
||||
@@ -3855,11 +3941,9 @@ class LanceTable(Table):
|
||||
|
||||
Parameters
|
||||
----------
|
||||
where: str or [Expr][lancedb.expr.Expr], optional
|
||||
The filter condition. Can be a SQL string or a type-safe
|
||||
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
|
||||
[lit][lancedb.expr.lit]. The filter must not be empty, or it will
|
||||
error.
|
||||
where: str, optional
|
||||
The SQL where clause to use when updating rows. For example, 'x = 2'
|
||||
or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error.
|
||||
values: dict, optional
|
||||
The values to update. The keys are the column names and the values
|
||||
are the values to set.
|
||||
@@ -3877,7 +3961,6 @@ class LanceTable(Table):
|
||||
Examples
|
||||
--------
|
||||
>>> import lancedb
|
||||
>>> from lancedb.expr import col
|
||||
>>> import pandas as pd
|
||||
>>> data = pd.DataFrame({"x": [1, 2, 3], "vector": [[1.0, 2], [3, 4], [5, 6]]})
|
||||
>>> db = lancedb.connect("./.lancedb")
|
||||
@@ -3887,7 +3970,7 @@ class LanceTable(Table):
|
||||
0 1 [1.0, 2.0]
|
||||
1 2 [3.0, 4.0]
|
||||
2 3 [5.0, 6.0]
|
||||
>>> table.update(where=col("x") == 2, values={"vector": [10.0, 10]})
|
||||
>>> table.update(where="x = 2", values={"vector": [10.0, 10]})
|
||||
UpdateResult(rows_updated=1, version=2)
|
||||
>>> table.to_pandas()
|
||||
x vector
|
||||
@@ -4025,6 +4108,7 @@ class LanceTable(Table):
|
||||
cleanup_older_than: Optional[timedelta] = None,
|
||||
delete_unverified: bool = False,
|
||||
retrain: bool = False,
|
||||
compaction_options: Optional[CompactionOptions] = None,
|
||||
):
|
||||
"""
|
||||
Optimize the on-disk data and indices for better performance.
|
||||
@@ -4057,6 +4141,11 @@ class LanceTable(Table):
|
||||
|
||||
retrain: bool, default False
|
||||
This parameter is no longer used and is deprecated.
|
||||
compaction_options: CompactionOptions, optional
|
||||
Options that control file compaction. For large rows, derive a lower
|
||||
``target_rows_per_fragment`` from the expected row size to bound later
|
||||
passes. This does not cap the first pass over an existing oversized
|
||||
fragment; see [CompactionOptions][lancedb.table.CompactionOptions].
|
||||
|
||||
Notes
|
||||
-----
|
||||
@@ -4071,6 +4160,7 @@ class LanceTable(Table):
|
||||
cleanup_older_than=cleanup_older_than,
|
||||
delete_unverified=delete_unverified,
|
||||
retrain=retrain,
|
||||
compaction_options=compaction_options,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -6001,7 +6091,7 @@ class AsyncTable:
|
||||
self,
|
||||
updates: Optional[Dict[str, Any]] = None,
|
||||
*,
|
||||
where: Optional[Union[str, Expr]] = None,
|
||||
where: Optional[str] = None,
|
||||
updates_sql: Optional[Dict[str, str]] = None,
|
||||
) -> UpdateResult:
|
||||
"""
|
||||
@@ -6016,11 +6106,9 @@ class AsyncTable:
|
||||
The updates to apply. The keys should be the name of the column to
|
||||
update. The values should be the new values to assign. This is
|
||||
required unless updates_sql is supplied.
|
||||
where: str or [Expr][lancedb.expr.Expr], optional
|
||||
The filter condition. Can be a SQL string or a type-safe
|
||||
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
|
||||
[lit][lancedb.expr.lit]. Only rows that satisfy this filter will
|
||||
be updated.
|
||||
where: str, optional
|
||||
An SQL filter that controls which rows are updated. For example, 'x = 2'
|
||||
or 'x IN (1, 2, 3)'. Only rows that satisfy this filter will be udpated.
|
||||
updates_sql: dict, optional
|
||||
The updates to apply, expressed as SQL expression strings. The keys should
|
||||
be column names. The values should be SQL expressions. These can be SQL
|
||||
@@ -6038,14 +6126,13 @@ class AsyncTable:
|
||||
--------
|
||||
>>> import asyncio
|
||||
>>> import lancedb
|
||||
>>> from lancedb.expr import col
|
||||
>>> import pandas as pd
|
||||
>>> async def demo_update():
|
||||
... data = pd.DataFrame({"x": [1, 2], "vector": [[1, 2], [3, 4]]})
|
||||
... db = await lancedb.connect_async("./.lancedb")
|
||||
... table = await db.create_table("my_table", data)
|
||||
... # x is [1, 2], vector is [[1, 2], [3, 4]]
|
||||
... await table.update({"vector": [10, 10]}, where=col("x") == 2)
|
||||
... await table.update({"vector": [10, 10]}, where="x = 2")
|
||||
... # x is [1, 2], vector is [[1, 2], [10, 10]]
|
||||
... await table.update(updates_sql={"x": "x + 1"})
|
||||
... # x is [2, 3], vector is [[1, 2], [10, 10]]
|
||||
@@ -6059,8 +6146,7 @@ class AsyncTable:
|
||||
if updates is not None:
|
||||
updates_sql = {k: value_to_sql(v) for k, v in updates.items()}
|
||||
|
||||
predicate = where.to_sql() if isinstance(where, Expr) else where
|
||||
return await self._inner.update(updates_sql, predicate)
|
||||
return await self._inner.update(updates_sql, where)
|
||||
|
||||
async def add_columns(
|
||||
self,
|
||||
@@ -6467,6 +6553,7 @@ class AsyncTable:
|
||||
cleanup_older_than: Optional[timedelta] = None,
|
||||
delete_unverified: bool = False,
|
||||
retrain=False,
|
||||
compaction_options: Optional[CompactionOptions] = None,
|
||||
) -> OptimizeStats:
|
||||
"""
|
||||
Optimize the on-disk data and indices for better performance.
|
||||
@@ -6499,6 +6586,11 @@ class AsyncTable:
|
||||
|
||||
retrain: bool, default False
|
||||
This parameter is no longer used and is deprecated.
|
||||
compaction_options: CompactionOptions, optional
|
||||
Options that control file compaction. For large rows, derive a lower
|
||||
``target_rows_per_fragment`` from the expected row size to bound later
|
||||
passes. This does not cap the first pass over an existing oversized
|
||||
fragment; see [CompactionOptions][lancedb.table.CompactionOptions].
|
||||
|
||||
Notes
|
||||
-----
|
||||
@@ -6525,6 +6617,7 @@ class AsyncTable:
|
||||
return await self._inner.optimize(
|
||||
cleanup_since_ms=cleanup_since_ms,
|
||||
delete_unverified=delete_unverified,
|
||||
compaction_options=compaction_options,
|
||||
)
|
||||
|
||||
async def list_indices(self) -> Iterable[IndexConfig]:
|
||||
|
||||
@@ -8,12 +8,7 @@ import pyarrow.compute as pc
|
||||
import pytest
|
||||
|
||||
import lancedb
|
||||
from lancedb._blob import (
|
||||
blob_v2_projection_sources,
|
||||
read_row_ids_from_hits,
|
||||
stash_auto_row_ids,
|
||||
)
|
||||
from lancedb.expr import col
|
||||
from lancedb._blob import read_row_ids_from_hits, stash_auto_row_ids
|
||||
from lancedb.index import FTS
|
||||
from lancedb.schema import blob_column_paths, blob_v2_column_paths
|
||||
|
||||
@@ -75,14 +70,6 @@ def test_blob_v2_column_paths_include_list_children():
|
||||
]
|
||||
|
||||
|
||||
def test_blob_v2_projection_sources_use_typed_column_name():
|
||||
schema = pa.schema([lancedb.blob("blob")])
|
||||
|
||||
assert blob_v2_projection_sources(schema, {"blob_alias": col("blob")}) == {
|
||||
"blob_alias": "blob"
|
||||
}
|
||||
|
||||
|
||||
def _legacy_v1_table(name):
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema(
|
||||
@@ -179,20 +166,6 @@ async def test_async_table_to_pandas_descriptions_mode_omits_row_id():
|
||||
assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_typed_blob_projection_preserves_source_column():
|
||||
db = await lancedb.connect_async("memory:///typed_blob_projection")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")])
|
||||
table = await db.create_table("typed_blob_projection", schema=schema)
|
||||
await table.add([{"id": 1, "blob": b"alpha"}])
|
||||
|
||||
hits = await table.query().select({"blob_alias": col("blob")}).to_arrow()
|
||||
|
||||
assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
|
||||
blobs = await table.fetch_blobs("blob", hits)
|
||||
assert blobs.to_pylist() == [b"alpha"]
|
||||
|
||||
|
||||
def test_fetch_blobs_round_trip():
|
||||
table = _blob_table(
|
||||
"round_trip",
|
||||
@@ -430,50 +403,6 @@ async def test_blob_v2_hybrid_fetch_blobs_async():
|
||||
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_hybrid_typed_blob_projection_preserves_source_column():
|
||||
db = await lancedb.connect_async("memory:///hybrid_typed_blob")
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("text", pa.utf8()),
|
||||
pa.field("vector", pa.list_(pa.float32(), list_size=2)),
|
||||
lancedb.blob("blob"),
|
||||
]
|
||||
)
|
||||
table = await db.create_table("hybrid_typed_blob", schema=schema)
|
||||
await table.add(
|
||||
[
|
||||
{
|
||||
"id": 1,
|
||||
"text": "hello alpha",
|
||||
"vector": [1.0, 0.0],
|
||||
"blob": b"alpha",
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"text": "hello beta",
|
||||
"vector": [0.9, 0.1],
|
||||
"blob": b"beta",
|
||||
},
|
||||
]
|
||||
)
|
||||
await table.create_index("text", config=FTS(with_position=False))
|
||||
|
||||
hits = await (
|
||||
table.query()
|
||||
.nearest_to([1.0, 0.0])
|
||||
.nearest_to_text("hello")
|
||||
.select({"blob_alias": col("blob")})
|
||||
.limit(2)
|
||||
.to_arrow()
|
||||
)
|
||||
|
||||
assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
|
||||
blobs = await table.fetch_blobs("blob", hits)
|
||||
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}
|
||||
|
||||
|
||||
def test_blob_file_seek_read_and_read_range():
|
||||
payload = _identifiable_payload(1024)
|
||||
table = _blob_table("seek_read", [{"id": 1, "image": payload}])
|
||||
|
||||
@@ -52,7 +52,7 @@ class TestExprConstruction:
|
||||
def test_func(self):
|
||||
e = func("lower", col("name"))
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "lower(`name`)"
|
||||
assert e.to_sql() == "lower(name)"
|
||||
|
||||
def test_func_unknown_raises(self):
|
||||
with pytest.raises(Exception):
|
||||
@@ -115,7 +115,7 @@ class TestExprOperators:
|
||||
def test_and_operator(self):
|
||||
e = (col("age") > lit(18)) & (col("status") == lit("active"))
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "((age > 18) AND (`status` = 'active'))"
|
||||
assert e.to_sql() == "((age > 18) AND (status = 'active'))"
|
||||
|
||||
def test_or_operator(self):
|
||||
e = (col("a") == lit(1)) | (col("b") == lit(2))
|
||||
@@ -166,7 +166,7 @@ class TestExprOperators:
|
||||
def test_coerce_plain_str(self):
|
||||
e = col("name") == "alice"
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(`name` = 'alice')"
|
||||
assert e.to_sql() == "(name = 'alice')"
|
||||
|
||||
def test_reflexive_comparisons(self):
|
||||
# 10 < col("age") swaps to col("age") > 10
|
||||
@@ -198,85 +198,85 @@ class TestExprBytesLiteral:
|
||||
|
||||
def test_bytes_equality_expr_sql(self):
|
||||
e = col("data") == lit(b"\xca\xfe")
|
||||
assert e.to_sql() == "(`data` = X'CAFE')"
|
||||
assert e.to_sql() == "(data = X'CAFE')"
|
||||
|
||||
def test_bytes_ne_expr_sql(self):
|
||||
e = col("data") != lit(b"\xff")
|
||||
assert e.to_sql() == "(`data` <> X'FF')"
|
||||
assert e.to_sql() == "(data <> X'FF')"
|
||||
|
||||
def test_bytes_compound_expr_sql(self):
|
||||
e = (col("data") == lit(b"\x01")) & (col("id") > lit(5))
|
||||
assert e.to_sql() == "((`data` = X'01') AND (id > 5))"
|
||||
assert e.to_sql() == "((data = X'01') AND (id > 5))"
|
||||
|
||||
def test_bytes_in_function_call(self):
|
||||
# Regression test: binary literals inside scalar function calls
|
||||
# used to fail because DataFusion's unparser does not support Binary
|
||||
# scalars. Now handled via a placeholder-substitution rewrite.
|
||||
e = func("contains", col("data"), lit(b"\xff"))
|
||||
assert e.to_sql() == "contains(`data`, X'FF')"
|
||||
assert e.to_sql() == "contains(data, X'FF')"
|
||||
|
||||
def test_bytes_in_not(self):
|
||||
e = ~(col("data") == lit(b"\xff"))
|
||||
assert e.to_sql() == "NOT (`data` = X'FF')"
|
||||
assert e.to_sql() == "NOT (data = X'FF')"
|
||||
|
||||
|
||||
class TestExprStringMethods:
|
||||
def test_lower(self):
|
||||
e = col("name").lower()
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "lower(`name`)"
|
||||
assert e.to_sql() == "lower(name)"
|
||||
|
||||
def test_upper(self):
|
||||
e = col("name").upper()
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "upper(`name`)"
|
||||
assert e.to_sql() == "upper(name)"
|
||||
|
||||
def test_contains(self):
|
||||
e = col("text").contains(lit("hello"))
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "contains(`text`, 'hello')"
|
||||
assert e.to_sql() == "contains(text, 'hello')"
|
||||
|
||||
def test_contains_with_str_coerce(self):
|
||||
e = col("text").contains("hello")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "contains(`text`, 'hello')"
|
||||
assert e.to_sql() == "contains(text, 'hello')"
|
||||
|
||||
def test_chained_lower_eq(self):
|
||||
e = col("name").lower() == lit("alice")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(lower(`name`) = 'alice')"
|
||||
assert e.to_sql() == "(lower(name) = 'alice')"
|
||||
|
||||
|
||||
class TestExprCast:
|
||||
def test_cast_string(self):
|
||||
e = col("id").cast("string")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "arrow_cast(id, 'Utf8')"
|
||||
assert e.to_sql() == "CAST(id AS VARCHAR)"
|
||||
|
||||
def test_cast_int32(self):
|
||||
e = col("score").cast("int32")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "arrow_cast(score, 'Int32')"
|
||||
assert e.to_sql() == "CAST(score AS INTEGER)"
|
||||
|
||||
def test_cast_float64(self):
|
||||
e = col("val").cast("float64")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "arrow_cast(val, 'Float64')"
|
||||
assert e.to_sql() == "CAST(val AS DOUBLE)"
|
||||
|
||||
def test_cast_pyarrow_type(self):
|
||||
e = col("score").cast(pa.int32())
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "arrow_cast(score, 'Int32')"
|
||||
assert e.to_sql() == "CAST(score AS INTEGER)"
|
||||
|
||||
def test_cast_pyarrow_float64(self):
|
||||
e = col("val").cast(pa.float64())
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "arrow_cast(val, 'Float64')"
|
||||
assert e.to_sql() == "CAST(val AS DOUBLE)"
|
||||
|
||||
def test_cast_pyarrow_string(self):
|
||||
e = col("id").cast(pa.string())
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "arrow_cast(id, 'Utf8')"
|
||||
assert e.to_sql() == "CAST(id AS VARCHAR)"
|
||||
|
||||
def test_cast_pyarrow_and_string_equivalent(self):
|
||||
# pa.int32() and "int32" should produce equivalent SQL
|
||||
@@ -597,14 +597,14 @@ class TestExprIsin:
|
||||
def test_isin_strs(self):
|
||||
assert (
|
||||
col("status").isin(["active", "pending"]).to_sql()
|
||||
== "`status` IN ('active', 'pending')"
|
||||
== "status IN ('active', 'pending')"
|
||||
)
|
||||
|
||||
def test_isin_coerces_and_mixes(self):
|
||||
assert col("id").isin([lit(1), 2]).to_sql() == "id IN (1, 2)"
|
||||
|
||||
def test_isin_empty(self):
|
||||
assert col("id").isin([]).to_sql() == "false"
|
||||
assert col("id").isin([]).to_sql() == "id IN ()"
|
||||
|
||||
def test_isin_filter(self, simple_table):
|
||||
result = simple_table.search().where(col("id").isin([1, 3, 5])).to_arrow()
|
||||
|
||||
@@ -11,9 +11,8 @@ import warnings
|
||||
import weakref
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import date, datetime, timedelta
|
||||
from decimal import Decimal
|
||||
from time import sleep
|
||||
from typing import List
|
||||
from typing import Any, List
|
||||
from unittest.mock import patch
|
||||
|
||||
import lancedb
|
||||
@@ -337,21 +336,6 @@ async def test_update_async(mem_db_async: AsyncConnection):
|
||||
assert await table.count_rows("id == 10") == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_expr_filter_literals_async(mem_db_async: AsyncConnection):
|
||||
values = ["5", "4.66e-84", "it's"]
|
||||
table = await mem_db_async.create_table(
|
||||
"update_expr_literals",
|
||||
data=[{"field": value, "result": "original"} for value in values],
|
||||
)
|
||||
|
||||
for value in values:
|
||||
update_res = await table.update({"result": value}, where=col("field") == value)
|
||||
assert update_res.rows_updated == 1
|
||||
|
||||
assert (await table.to_arrow())["result"].to_pylist() == values
|
||||
|
||||
|
||||
def test_create_table(mem_db: DBConnection):
|
||||
schema = pa.schema(
|
||||
{
|
||||
@@ -2359,148 +2343,6 @@ def test_update(mem_db: DBConnection):
|
||||
assert np.allclose(v, np.array([[1.2, 1.9], [1.1, 1.1]]))
|
||||
|
||||
|
||||
def test_update_expr_filter_literals(mem_db: DBConnection):
|
||||
values = ["5", "4.66e-84", "it's"]
|
||||
table = mem_db.create_table(
|
||||
"update_expr_literals",
|
||||
data=[{"field": value, "result": "original"} for value in values],
|
||||
)
|
||||
|
||||
for value in values:
|
||||
update_res = table.update(where=col("field") == value, values={"result": value})
|
||||
assert update_res.rows_updated == 1
|
||||
|
||||
assert table.to_arrow()["result"].to_pylist() == values
|
||||
|
||||
|
||||
def test_update_expr_filter_preserves_typed_semantics(mem_db: DBConnection):
|
||||
low = Decimal("1.234567890123456789")
|
||||
high = Decimal("1.234567890123456790")
|
||||
decimal_schema = pa.schema(
|
||||
[("val", pa.decimal128(19, 18)), ("result", pa.string())]
|
||||
)
|
||||
decimal_table = mem_db.create_table(
|
||||
"update_expr_decimal",
|
||||
pa.table(
|
||||
{"val": [low, high], "result": ["old", "old"]},
|
||||
schema=decimal_schema,
|
||||
),
|
||||
)
|
||||
predicate = col("val") < lit(high)
|
||||
assert decimal_table.search().where(predicate).to_arrow().num_rows == 1
|
||||
result = decimal_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 1
|
||||
|
||||
keyword_table = mem_db.create_table(
|
||||
"update_expr_keyword", [{"null": 1, "result": "old"}]
|
||||
)
|
||||
predicate = col("null") == 1
|
||||
assert keyword_table.search().where(predicate).to_arrow().num_rows == 1
|
||||
result = keyword_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 1
|
||||
|
||||
empty_in_table = mem_db.create_table(
|
||||
"update_expr_empty_in", [{"id": 1, "result": "old"}]
|
||||
)
|
||||
predicate = col("id").isin([])
|
||||
assert empty_in_table.search().where(predicate).to_arrow().num_rows == 0
|
||||
result = empty_in_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 0
|
||||
|
||||
marker = "__lancedb_binary_placeholder_0__"
|
||||
binary_schema = pa.schema(
|
||||
[("payload", pa.binary()), ("text", pa.string()), ("result", pa.string())]
|
||||
)
|
||||
binary_table = mem_db.create_table(
|
||||
"update_expr_binary",
|
||||
pa.table(
|
||||
{
|
||||
"payload": [b"\x01", b"\x02"],
|
||||
"text": ["other", marker],
|
||||
"result": ["old", "old"],
|
||||
},
|
||||
schema=binary_schema,
|
||||
),
|
||||
)
|
||||
predicate = (col("payload") == lit(b"\x01")) | (col("text") == marker)
|
||||
assert binary_table.search().where(predicate).to_arrow().num_rows == 2
|
||||
result = binary_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 2
|
||||
|
||||
nonfinite_table = mem_db.create_table(
|
||||
"update_expr_nonfinite",
|
||||
[{"x": 1.0, "result": "old"}, {"x": 2.0, "result": "old"}],
|
||||
)
|
||||
predicate = col("x") < float("inf")
|
||||
assert nonfinite_table.search().where(predicate).to_arrow().num_rows == 2
|
||||
result = nonfinite_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 2
|
||||
|
||||
float16_table = mem_db.create_table(
|
||||
"update_expr_float16",
|
||||
[{"x": 1.0, "result": "old"}, {"x": 3.0, "result": "old"}],
|
||||
)
|
||||
predicate = col("x").cast(pa.float16()) < 2.0
|
||||
assert float16_table.search().where(predicate).to_arrow().num_rows == 1
|
||||
result = float16_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 1
|
||||
|
||||
string_cast_table = mem_db.create_table(
|
||||
"update_expr_string_cast",
|
||||
[{"x": 1, "result": "old"}, {"x": 2, "result": "old"}],
|
||||
)
|
||||
predicate = col("x").cast("string") == "1"
|
||||
assert string_cast_table.search().where(predicate).to_arrow().num_rows == 1
|
||||
result = string_cast_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 1
|
||||
|
||||
quoted_identifier_schema = pa.schema(
|
||||
[("payload", pa.binary()), ("odd'name", pa.int64()), ("result", pa.string())]
|
||||
)
|
||||
quoted_identifier_table = mem_db.create_table(
|
||||
"update_expr_quoted_identifier",
|
||||
pa.table(
|
||||
{"payload": [b"\x01"], "odd'name": [1], "result": ["old"]},
|
||||
schema=quoted_identifier_schema,
|
||||
),
|
||||
)
|
||||
predicate = (col("payload") == lit(b"\x01")) & (col("odd'name") == 1)
|
||||
assert quoted_identifier_table.search().where(predicate).to_arrow().num_rows == 1
|
||||
result = quoted_identifier_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 1
|
||||
|
||||
decimal256_schema = pa.schema(
|
||||
[("val", pa.decimal256(40, 2)), ("result", pa.string())]
|
||||
)
|
||||
decimal256_table = mem_db.create_table(
|
||||
"update_expr_decimal256",
|
||||
pa.table(
|
||||
{
|
||||
"val": [Decimal("1.00"), Decimal("3.00")],
|
||||
"result": ["old", "old"],
|
||||
},
|
||||
schema=decimal256_schema,
|
||||
),
|
||||
)
|
||||
predicate = col("val") < lit(Decimal("2.00")).cast(pa.decimal256(40, 2))
|
||||
assert decimal256_table.search().where(predicate).to_arrow().num_rows == 1
|
||||
result = decimal256_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 1
|
||||
|
||||
binary_empty_table = mem_db.create_table(
|
||||
"update_expr_binary_empty",
|
||||
pa.table(
|
||||
{"payload": [b"\x01", b"\x02"], "result": ["old", "old"]},
|
||||
schema=pa.schema([("payload", pa.binary()), ("result", pa.string())]),
|
||||
),
|
||||
)
|
||||
predicate = (col("payload") == lit(b"\x01")).isin([])
|
||||
assert binary_empty_table.search().where(predicate).to_arrow().num_rows == 0
|
||||
assert predicate.to_sql() == "false"
|
||||
result = binary_empty_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 0
|
||||
|
||||
|
||||
def test_update_with_arrow_scalar(mem_db: DBConnection):
|
||||
schema = pa.schema({"id": pa.int64(), "vector": pa.list_(pa.float32(), 4)})
|
||||
table = mem_db.create_table("my_table", schema=schema)
|
||||
@@ -3876,6 +3718,100 @@ async def test_optimize(mem_db_async: AsyncConnection):
|
||||
assert await table.query().to_arrow() == pa.table({"x": [[1], [2]]})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_optimize_compaction_options(mem_db_async: AsyncConnection):
|
||||
table = await mem_db_async.create_table("test", data=[{"x": 1}])
|
||||
await table.add([{"x": 2}])
|
||||
|
||||
stats = await table.optimize(
|
||||
compaction_options={
|
||||
"target_rows_per_fragment": 1,
|
||||
"batch_size": 1,
|
||||
"num_threads": 1,
|
||||
}
|
||||
)
|
||||
assert stats.compaction.fragments_removed == 0
|
||||
assert stats.compaction.fragments_added == 0
|
||||
|
||||
stats = await table.optimize(compaction_options={"target_rows_per_fragment": 3})
|
||||
assert stats.compaction.fragments_removed == 2
|
||||
assert stats.compaction.fragments_added == 1
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid compaction option: unknown"):
|
||||
await table.optimize(compaction_options={"unknown": 1})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("option", ["max_source_rows", "max_source_bytes"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_optimize_compaction_source_limits(
|
||||
mem_db_async: AsyncConnection, option: str
|
||||
):
|
||||
table = await mem_db_async.create_table("test", data=[{"x": 1}])
|
||||
await table.add([{"x": 2}])
|
||||
|
||||
stats = await table.optimize(
|
||||
compaction_options={
|
||||
"target_rows_per_fragment": 3,
|
||||
option: 1,
|
||||
}
|
||||
)
|
||||
assert stats.compaction.fragments_removed == 0
|
||||
assert stats.compaction.fragments_added == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_optimize_compaction_excluded_fragments(mem_db_async: AsyncConnection):
|
||||
table = await mem_db_async.create_table("test", data=[{"x": 1}])
|
||||
await table.add([{"x": 2}])
|
||||
|
||||
stats = await table.optimize(
|
||||
compaction_options={
|
||||
"target_rows_per_fragment": 3,
|
||||
"excluded_fragment_ids": [0],
|
||||
}
|
||||
)
|
||||
assert stats.compaction.fragments_removed == 0
|
||||
assert stats.compaction.fragments_added == 0
|
||||
|
||||
stats = await table.optimize(compaction_options={"target_rows_per_fragment": 3})
|
||||
assert stats.compaction.fragments_removed == 2
|
||||
assert stats.compaction.fragments_added == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("option", "value", "message"),
|
||||
[
|
||||
("target_rows_per_fragment", 0, "must be between 1 and 4294967295"),
|
||||
("max_rows_per_group", 0, "must be between 1 and 4294967295"),
|
||||
("batch_size", 0, "must be between 1 and 4294967295"),
|
||||
("num_threads", 0, "must be greater than 0"),
|
||||
("target_rows_per_fragment", 2**32, "must be between 1 and 4294967295"),
|
||||
("max_rows_per_group", 2**32, "must be between 1 and 4294967295"),
|
||||
("batch_size", 2**32, "must be between 1 and 4294967295"),
|
||||
("io_buffer_size", 2**63, "must be at most 9223372036854775807"),
|
||||
("max_source_rows", 0, "must be greater than 0"),
|
||||
("max_source_bytes", 0, "must be greater than 0"),
|
||||
(
|
||||
"excluded_fragment_ids",
|
||||
[-1],
|
||||
"must contain values between 0 and 4294967295",
|
||||
),
|
||||
(
|
||||
"excluded_fragment_ids",
|
||||
[2**32],
|
||||
"must contain values between 0 and 4294967295",
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_optimize_compaction_options_validation(
|
||||
mem_db_async: AsyncConnection, option: str, value: Any, message: str
|
||||
):
|
||||
table = await mem_db_async.create_table("test", data=[{"x": 1}])
|
||||
with pytest.raises(ValueError, match=message):
|
||||
await table.optimize(compaction_options={option: value})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_optimize_delete_unverified(tmp_db_async: AsyncConnection, tmp_path):
|
||||
table = await tmp_db_async.create_table(
|
||||
|
||||
@@ -130,14 +130,6 @@ impl PyExpr {
|
||||
|
||||
// ── utilities ────────────────────────────────────────────────────────────
|
||||
|
||||
/// Return the referenced column name for a bare column expression.
|
||||
fn column_name(&self) -> Option<String> {
|
||||
match &self.0 {
|
||||
DfExpr::Column(column) if column.relation.is_none() => Some(column.name.clone()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Render the expression as a SQL string (useful for debugging).
|
||||
fn to_sql(&self) -> PyResult<String> {
|
||||
lancedb::expr::expr_to_sql_string(&self.0).map_err(|e| PyValueError::new_err(e.to_string()))
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -326,7 +325,6 @@ pub struct PyQueryRequest {
|
||||
pub filter: Option<PyQueryFilter>,
|
||||
pub full_text_search: Option<PyLanceDB<FtsQuery>>,
|
||||
pub select: PySelect,
|
||||
pub select_source_columns: Option<HashMap<String, String>>,
|
||||
pub fast_search: Option<bool>,
|
||||
pub with_row_id: Option<bool>,
|
||||
pub use_lsm: Option<bool>,
|
||||
@@ -357,7 +355,6 @@ impl From<AnyQuery> for PyQueryRequest {
|
||||
full_text_search: query_request
|
||||
.full_text_search
|
||||
.map(|fts| PyLanceDB(fts.query)),
|
||||
select_source_columns: PySelect::source_columns(&query_request.select),
|
||||
select: PySelect(query_request.select),
|
||||
fast_search: Some(query_request.fast_search),
|
||||
with_row_id: Some(query_request.with_row_id),
|
||||
@@ -383,7 +380,6 @@ impl From<AnyQuery> for PyQueryRequest {
|
||||
offset: vector_query.base.offset,
|
||||
filter: vector_query.base.filter.map(PyQueryFilter),
|
||||
full_text_search: None,
|
||||
select_source_columns: PySelect::source_columns(&vector_query.base.select),
|
||||
select: PySelect(vector_query.base.select),
|
||||
fast_search: Some(vector_query.base.fast_search),
|
||||
with_row_id: Some(vector_query.base.with_row_id),
|
||||
@@ -416,25 +412,6 @@ impl From<AnyQuery> for PyQueryRequest {
|
||||
#[derive(Clone)]
|
||||
pub struct PySelect(Select);
|
||||
|
||||
impl PySelect {
|
||||
fn source_columns(select: &Select) -> Option<HashMap<String, String>> {
|
||||
match select {
|
||||
Select::Expr(pairs) => Some(
|
||||
pairs
|
||||
.iter()
|
||||
.filter_map(|(output, expr)| match expr {
|
||||
lancedb::expr::DfExpr::Column(column) if column.relation.is_none() => {
|
||||
Some((output.clone(), column.name.clone()))
|
||||
}
|
||||
_ => None,
|
||||
})
|
||||
.collect(),
|
||||
),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'py> IntoPyObject<'py> for PySelect {
|
||||
type Target = PyAny;
|
||||
type Output = Bound<'py, Self::Target>;
|
||||
|
||||
+132
-7
@@ -20,8 +20,9 @@ use arrow::{
|
||||
use lancedb::blob::{BlobFile, BlobRangeRequest};
|
||||
use lancedb::index::scalar::FtsIndexBuilder;
|
||||
use lancedb::table::{
|
||||
AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken,
|
||||
NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
AddDataMode, ColumnAlteration, CompactionMode, CompactionOptions, Duration,
|
||||
FieldMetadataUpdate, FtsToken as LanceDbFtsToken, IndexRemapMode, NewColumnTransform,
|
||||
OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
};
|
||||
use lancedb::tokenize as lancedb_tokenize;
|
||||
use pyo3::{
|
||||
@@ -100,6 +101,128 @@ enum PredicateArg {
|
||||
Sql(String),
|
||||
}
|
||||
|
||||
fn validate_positive_u32(value: u64, name: &str) -> PyResult<usize> {
|
||||
if !(1..=u32::MAX as u64).contains(&value) {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"{name} must be between 1 and {}",
|
||||
u32::MAX
|
||||
)));
|
||||
}
|
||||
Ok(value as usize)
|
||||
}
|
||||
|
||||
fn positive_u32(value: &Bound<'_, PyAny>, name: &str) -> PyResult<usize> {
|
||||
validate_positive_u32(value.extract()?, name)
|
||||
}
|
||||
|
||||
fn optional_positive_u32(value: &Bound<'_, PyAny>, name: &str) -> PyResult<Option<usize>> {
|
||||
value
|
||||
.extract::<Option<u64>>()?
|
||||
.map(|value| validate_positive_u32(value, name))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn optional_positive_usize(value: &Bound<'_, PyAny>, name: &str) -> PyResult<Option<usize>> {
|
||||
let value: Option<usize> = value.extract()?;
|
||||
if value == Some(0) {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"{name} must be greater than 0"
|
||||
)));
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
fn optional_positive_u64(value: &Bound<'_, PyAny>, name: &str) -> PyResult<Option<u64>> {
|
||||
let value: Option<u64> = value.extract()?;
|
||||
if value == Some(0) {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"{name} must be greater than 0"
|
||||
)));
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
fn u32_list(value: &Bound<'_, PyAny>, name: &str) -> PyResult<Vec<u32>> {
|
||||
value
|
||||
.extract::<Vec<i64>>()?
|
||||
.into_iter()
|
||||
.map(|value| {
|
||||
u32::try_from(value).map_err(|_| {
|
||||
PyValueError::new_err(format!(
|
||||
"{name} must contain values between 0 and {}",
|
||||
u32::MAX
|
||||
))
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn optional_i64_bounded_u64(value: &Bound<'_, PyAny>, name: &str) -> PyResult<Option<u64>> {
|
||||
let value: Option<u64> = value.extract()?;
|
||||
if value.is_some_and(|value| value > i64::MAX as u64) {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"{name} must be at most {}",
|
||||
i64::MAX
|
||||
)));
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
fn parse_compaction_options(options: Option<&Bound<'_, PyDict>>) -> PyResult<CompactionOptions> {
|
||||
let mut parsed = CompactionOptions::default();
|
||||
let Some(options) = options else {
|
||||
return Ok(parsed);
|
||||
};
|
||||
|
||||
for (key, value) in options.iter() {
|
||||
let key: String = key.extract()?;
|
||||
match key.as_str() {
|
||||
"target_rows_per_fragment" => {
|
||||
parsed.target_rows_per_fragment = positive_u32(&value, &key)?
|
||||
}
|
||||
"max_rows_per_group" => parsed.max_rows_per_group = positive_u32(&value, &key)?,
|
||||
"max_bytes_per_file" => parsed.max_bytes_per_file = value.extract()?,
|
||||
"materialize_deletions" => parsed.materialize_deletions = value.extract()?,
|
||||
"materialize_deletions_threshold" => {
|
||||
parsed.materialize_deletions_threshold = value.extract()?
|
||||
}
|
||||
"num_threads" => parsed.num_threads = optional_positive_usize(&value, &key)?,
|
||||
"batch_size" => parsed.batch_size = optional_positive_u32(&value, &key)?,
|
||||
"io_buffer_size" => parsed.io_buffer_size = optional_i64_bounded_u64(&value, &key)?,
|
||||
"defer_index_remap" => parsed.defer_index_remap = value.extract()?,
|
||||
"index_remap_mode" => {
|
||||
let mode: String = value.extract()?;
|
||||
parsed.index_remap_mode = IndexRemapMode::try_from(mode.as_str())
|
||||
.map_err(|err| PyValueError::new_err(err.to_string()))?;
|
||||
}
|
||||
"compaction_mode" => {
|
||||
let mode: Option<String> = value.extract()?;
|
||||
parsed.compaction_mode = mode
|
||||
.map(|mode| {
|
||||
CompactionMode::try_from(mode.as_str())
|
||||
.map_err(|err| PyValueError::new_err(err.to_string()))
|
||||
})
|
||||
.transpose()?;
|
||||
}
|
||||
"binary_copy_read_batch_bytes" => {
|
||||
parsed.binary_copy_read_batch_bytes = value.extract()?
|
||||
}
|
||||
"max_source_fragments" => parsed.max_source_fragments = value.extract()?,
|
||||
"max_source_rows" => parsed.max_source_rows = optional_positive_usize(&value, &key)?,
|
||||
"max_source_bytes" => parsed.max_source_bytes = optional_positive_u64(&value, &key)?,
|
||||
"excluded_fragment_ids" => parsed.excluded_fragment_ids = u32_list(&value, &key)?,
|
||||
"max_overlays_per_fragment" => parsed.max_overlays_per_fragment = value.extract()?,
|
||||
_ => {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"Invalid compaction option: {key}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(parsed)
|
||||
}
|
||||
|
||||
/// Statistics about a compaction operation.
|
||||
#[pyclass(get_all, from_py_object)]
|
||||
#[derive(Clone, Debug)]
|
||||
@@ -1340,13 +1463,15 @@ impl Table {
|
||||
}
|
||||
|
||||
/// Optimize the on-disk data by compacting and pruning old data, for better performance.
|
||||
#[pyo3(signature = (cleanup_since_ms=None, delete_unverified=None))]
|
||||
pub fn optimize(
|
||||
self_: PyRef<'_, Self>,
|
||||
#[pyo3(signature = (cleanup_since_ms=None, delete_unverified=None, compaction_options=None))]
|
||||
pub fn optimize<'py>(
|
||||
self_: PyRef<'py, Self>,
|
||||
cleanup_since_ms: Option<u64>,
|
||||
delete_unverified: Option<bool>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
compaction_options: Option<&Bound<'py, PyDict>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
let compaction_options = parse_compaction_options(compaction_options)?;
|
||||
let older_than = if let Some(ms) = cleanup_since_ms {
|
||||
if ms > i64::MAX as u64 {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
@@ -1362,7 +1487,7 @@ impl Table {
|
||||
future_into_py(self_.py(), async move {
|
||||
let compaction_stats = inner
|
||||
.optimize(OptimizeAction::Compact {
|
||||
options: lancedb::table::CompactionOptions::default(),
|
||||
options: compaction_options,
|
||||
remap_options: None,
|
||||
})
|
||||
.await
|
||||
|
||||
+4
-120
@@ -157,7 +157,7 @@ mod tests {
|
||||
use datafusion_common::ScalarValue;
|
||||
let expr = col("data").eq(lit(ScalarValue::Binary(Some(vec![0xca, 0xfe]))));
|
||||
let sql = expr_to_sql_string(&expr).unwrap();
|
||||
assert_eq!(sql, "(`data` = X'CAFE')");
|
||||
assert_eq!(sql, "(data = X'CAFE')");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -167,7 +167,7 @@ mod tests {
|
||||
let int_expr = col("id").gt(lit(5i64));
|
||||
let combined = bin_expr.and(int_expr);
|
||||
let sql = expr_to_sql_string(&combined).unwrap();
|
||||
assert_eq!(sql, "((`data` = X'01') AND (id > 5))");
|
||||
assert_eq!(sql, "((data = X'01') AND (id > 5))");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -185,7 +185,7 @@ mod tests {
|
||||
// serialized correctly (regression test for placeholder rewrite path).
|
||||
let expr = contains(col("data"), lit(ScalarValue::Binary(Some(vec![0xff]))));
|
||||
let sql = expr_to_sql_string(&expr).unwrap();
|
||||
assert_eq!(sql, "contains(`data`, X'FF')");
|
||||
assert_eq!(sql, "contains(data, X'FF')");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -196,7 +196,7 @@ mod tests {
|
||||
.eq(lit(ScalarValue::Binary(Some(vec![0xab, 0xcd]))))
|
||||
.not();
|
||||
let sql = expr_to_sql_string(&expr).unwrap();
|
||||
assert_eq!(sql, "NOT (`data` = X'ABCD')");
|
||||
assert_eq!(sql, "NOT (data = X'ABCD')");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -206,122 +206,6 @@ mod tests {
|
||||
assert!(sql.contains("IN"), "expected IN in: {}", sql);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_is_in() {
|
||||
let expr = is_in(col("id"), vec![]);
|
||||
assert_eq!(expr_to_sql_string(&expr).unwrap(), "false");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_is_in_discards_binary_children() {
|
||||
use datafusion_common::ScalarValue;
|
||||
|
||||
let expr = is_in(
|
||||
col("payload").eq(lit(ScalarValue::Binary(Some(vec![0x01])))),
|
||||
vec![],
|
||||
);
|
||||
assert_eq!(expr_to_sql_string(&expr).unwrap(), "false");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_keyword_identifier() {
|
||||
let expr = col("null").eq(lit(1i64));
|
||||
assert_eq!(expr_to_sql_string(&expr).unwrap(), "(`null` = 1)");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_decimal_literal_preserves_type() {
|
||||
use datafusion_common::ScalarValue;
|
||||
|
||||
let expr = col("val").lt(lit(ScalarValue::Decimal128(
|
||||
Some(1_234_567_890_123_456_790),
|
||||
19,
|
||||
18,
|
||||
)));
|
||||
let sql = expr_to_sql_string(&expr).unwrap();
|
||||
assert_eq!(
|
||||
sql,
|
||||
"(val < arrow_cast('1.234567890123456790', 'Decimal128(19, 18)'))"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_non_finite_float_literal_preserves_type() {
|
||||
let expr = col("x").lt(lit(f64::INFINITY));
|
||||
assert_eq!(
|
||||
expr_to_sql_string(&expr).unwrap(),
|
||||
"(x < arrow_cast('inf', 'Float64'))"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cast_uses_arrow_type_name() {
|
||||
let string = expr_cast(col("x"), DataType::Utf8);
|
||||
assert_eq!(
|
||||
expr_to_sql_string(&string).unwrap(),
|
||||
"arrow_cast(x, 'Utf8')"
|
||||
);
|
||||
|
||||
let int32 = expr_cast(col("x"), DataType::Int32);
|
||||
assert_eq!(
|
||||
expr_to_sql_string(&int32).unwrap(),
|
||||
"arrow_cast(x, 'Int32')"
|
||||
);
|
||||
|
||||
let expr = expr_cast(col("x"), DataType::Float16).lt(lit(2.0));
|
||||
assert_eq!(
|
||||
expr_to_sql_string(&expr).unwrap(),
|
||||
"(arrow_cast(x, 'Float16') < 2.0)"
|
||||
);
|
||||
|
||||
let decimal = expr_cast(lit("2.00"), DataType::Decimal256(40, 2));
|
||||
assert_eq!(
|
||||
expr_to_sql_string(&decimal).unwrap(),
|
||||
"arrow_cast('2.00', 'Decimal256(40, 2)')"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_binary_placeholder_does_not_rewrite_user_string() {
|
||||
use datafusion_common::ScalarValue;
|
||||
|
||||
let marker = "__lancedb_binary_placeholder_0__";
|
||||
let expr = col("payload")
|
||||
.eq(lit(ScalarValue::Binary(Some(vec![0x01]))))
|
||||
.or(col("text").eq(lit(marker)));
|
||||
assert_eq!(
|
||||
expr_to_sql_string(&expr).unwrap(),
|
||||
"((payload = X'01') OR (`text` = '__lancedb_binary_placeholder_0__'))"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_binary_binding_skips_quoted_identifiers() {
|
||||
use datafusion_common::ScalarValue;
|
||||
|
||||
let expr = col("payload")
|
||||
.eq(lit(ScalarValue::Binary(Some(vec![0x01]))))
|
||||
.and(col("odd'name").eq(lit(1i64)))
|
||||
.and(col("odd`'name").eq(lit(2i64)));
|
||||
assert_eq!(
|
||||
expr_to_sql_string(&expr).unwrap(),
|
||||
"(((payload = X'01') AND (`odd'name` = 1)) AND (`odd``'name` = 2))"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_binary_placeholder_collision_search_is_linear() {
|
||||
use datafusion_common::ScalarValue;
|
||||
|
||||
let collision_shaped = format!("__lancedb_binary_placeholder_0__{}", "_".repeat(64_000));
|
||||
let expr = col("payload")
|
||||
.eq(lit(ScalarValue::Binary(Some(vec![0x01]))))
|
||||
.and(col("text").eq(lit(collision_shaped.clone())));
|
||||
let sql = expr_to_sql_string(&expr).unwrap();
|
||||
assert!(sql.contains("X'01'"));
|
||||
assert!(sql.contains(&format!("'{collision_shaped}'")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multiple_binary_literals() {
|
||||
use datafusion_common::ScalarValue;
|
||||
|
||||
+42
-220
@@ -1,24 +1,13 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::{
|
||||
any::TypeId,
|
||||
collections::{HashMap, HashSet},
|
||||
};
|
||||
use std::any::TypeId;
|
||||
|
||||
use arrow_array::types::{
|
||||
Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, DecimalType,
|
||||
};
|
||||
use arrow_schema::DataType;
|
||||
use datafusion_common::ScalarValue;
|
||||
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
|
||||
use datafusion_expr::Expr;
|
||||
use datafusion_functions::core::expr_fn::{
|
||||
arrow_cast as datafusion_arrow_cast, arrow_try_cast as datafusion_arrow_try_cast,
|
||||
};
|
||||
use datafusion_sql::sqlparser::{
|
||||
dialect::{Dialect as SqlParserDialect, GenericDialect},
|
||||
keywords::ALL_KEYWORDS,
|
||||
tokenizer::{Token, Tokenizer},
|
||||
};
|
||||
use datafusion_sql::unparser::{self, dialect::Dialect as UnparserDialect};
|
||||
@@ -38,13 +27,11 @@ struct LanceSqlDialect;
|
||||
|
||||
impl UnparserDialect for LanceSqlDialect {
|
||||
fn identifier_quote_style(&self, identifier: &str) -> Option<char> {
|
||||
let identifier_upper = identifier.to_ascii_uppercase();
|
||||
let needs_quote =
|
||||
(identifier_upper != "ID" && ALL_KEYWORDS.contains(&identifier_upper.as_str()))
|
||||
|| identifier.chars().any(|c| c.is_ascii_uppercase())
|
||||
|| !identifier.chars().enumerate().all(|(i, c)| {
|
||||
c == '_' || c.is_ascii_alphabetic() || (i > 0 && c.is_ascii_digit())
|
||||
});
|
||||
let needs_quote = identifier.chars().any(|c| c.is_ascii_uppercase())
|
||||
|| !identifier
|
||||
.chars()
|
||||
.enumerate()
|
||||
.all(|(i, c)| c == '_' || c.is_ascii_alphabetic() || (i > 0 && c.is_ascii_digit()));
|
||||
if needs_quote { Some('`') } else { None }
|
||||
}
|
||||
}
|
||||
@@ -113,128 +100,24 @@ fn bytes_to_hex_sql(bytes: &[u8]) -> String {
|
||||
format!("X'{hex}'")
|
||||
}
|
||||
|
||||
fn string_literals(expr: &Expr) -> HashSet<String> {
|
||||
let mut literals = HashSet::new();
|
||||
/// Returns true if *expr* contains a `Binary` or `LargeBinary` scalar literal
|
||||
/// anywhere in its subtree. DataFusion's SQL unparser cannot serialize those
|
||||
/// variants, so we route such expressions through a placeholder-substitution
|
||||
/// path that emits SQL `X'...'` byte-string literals.
|
||||
fn has_binary_literal(expr: &Expr) -> bool {
|
||||
let mut found = false;
|
||||
let _ = expr.apply(&mut |e: &Expr| {
|
||||
if let Expr::Literal(
|
||||
ScalarValue::Utf8(Some(value))
|
||||
| ScalarValue::LargeUtf8(Some(value))
|
||||
| ScalarValue::Utf8View(Some(value)),
|
||||
_,
|
||||
) = e
|
||||
{
|
||||
literals.insert(value.clone());
|
||||
}
|
||||
Ok(TreeNodeRecursion::Continue)
|
||||
});
|
||||
literals
|
||||
}
|
||||
|
||||
fn typed_string_literal(value: String, data_type: DataType) -> Expr {
|
||||
datafusion_arrow_cast(
|
||||
Expr::Literal(ScalarValue::Utf8(Some(value)), None),
|
||||
Expr::Literal(ScalarValue::Utf8(Some(data_type.to_string())), None),
|
||||
)
|
||||
}
|
||||
|
||||
fn next_binary_placeholder(user_strings: &HashSet<String>, next_id: &mut usize) -> String {
|
||||
loop {
|
||||
let placeholder = format!("{BINARY_PLACEHOLDER_PREFIX}{}__", *next_id);
|
||||
*next_id += 1;
|
||||
if !user_strings.contains(&placeholder) {
|
||||
return placeholder;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn bind_binary_literals(
|
||||
sql: &str,
|
||||
mut bindings: HashMap<String, Vec<u8>>,
|
||||
) -> crate::Result<String> {
|
||||
let bytes = sql.as_bytes();
|
||||
let mut output = Vec::with_capacity(bytes.len());
|
||||
let mut index = 0;
|
||||
|
||||
// Walk SQL string tokens once. Placeholders are plain, unescaped string
|
||||
// literals, so this remains linear even when user strings are large or
|
||||
// deliberately resemble the placeholder prefix.
|
||||
while index < bytes.len() {
|
||||
if bytes[index] == b'`' {
|
||||
let identifier_start = index;
|
||||
index += 1;
|
||||
let mut identifier_end = None;
|
||||
while index < bytes.len() {
|
||||
if bytes[index] == b'`' {
|
||||
if index + 1 < bytes.len() && bytes[index + 1] == b'`' {
|
||||
index += 2;
|
||||
} else {
|
||||
index += 1;
|
||||
identifier_end = Some(index);
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
index += 1;
|
||||
}
|
||||
}
|
||||
|
||||
let Some(identifier_end) = identifier_end else {
|
||||
return Err(crate::Error::InvalidInput {
|
||||
message: "unterminated identifier while binding binary literal".to_string(),
|
||||
});
|
||||
};
|
||||
output.extend_from_slice(&bytes[identifier_start..identifier_end]);
|
||||
continue;
|
||||
}
|
||||
|
||||
if bytes[index] != b'\'' {
|
||||
output.push(bytes[index]);
|
||||
index += 1;
|
||||
continue;
|
||||
}
|
||||
|
||||
let literal_start = index;
|
||||
index += 1;
|
||||
let content_start = index;
|
||||
let mut escaped = false;
|
||||
let mut content_end = None;
|
||||
while index < bytes.len() {
|
||||
if bytes[index] == b'\'' {
|
||||
if index + 1 < bytes.len() && bytes[index + 1] == b'\'' {
|
||||
escaped = true;
|
||||
index += 2;
|
||||
} else {
|
||||
content_end = Some(index);
|
||||
index += 1;
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
index += 1;
|
||||
}
|
||||
}
|
||||
|
||||
let Some(content_end) = content_end else {
|
||||
return Err(crate::Error::InvalidInput {
|
||||
message: "unterminated string while binding binary literal".to_string(),
|
||||
});
|
||||
};
|
||||
|
||||
let placeholder = &sql[content_start..content_end];
|
||||
if !escaped && let Some(value) = bindings.remove(placeholder) {
|
||||
output.extend_from_slice(bytes_to_hex_sql(&value).as_bytes());
|
||||
if matches!(
|
||||
e,
|
||||
Expr::Literal(ScalarValue::Binary(_) | ScalarValue::LargeBinary(_), _)
|
||||
) {
|
||||
found = true;
|
||||
Ok(TreeNodeRecursion::Stop)
|
||||
} else {
|
||||
output.extend_from_slice(&bytes[literal_start..index]);
|
||||
Ok(TreeNodeRecursion::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
if !bindings.is_empty() {
|
||||
return Err(crate::Error::InvalidInput {
|
||||
message: "failed to bind binary literal while serializing expression".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
String::from_utf8(output).map_err(|e| crate::Error::InvalidInput {
|
||||
message: format!("failed to bind binary literal: {e}"),
|
||||
})
|
||||
});
|
||||
found
|
||||
}
|
||||
|
||||
fn run_unparser(expr: &Expr) -> crate::Result<String> {
|
||||
@@ -247,37 +130,25 @@ fn run_unparser(expr: &Expr) -> crate::Result<String> {
|
||||
}
|
||||
|
||||
pub fn expr_to_sql_string(expr: &Expr) -> crate::Result<String> {
|
||||
// DataFusion's unparser needs a few adaptations before its SQL can be
|
||||
// reparsed by Lance without changing the typed expression's semantics:
|
||||
//
|
||||
// * decimal literals need an explicit cast to preserve precision and scale;
|
||||
// * casts need exact Arrow type names rather than SQL type aliases;
|
||||
// * an empty IN list is valid in DataFusion but invalid SQL;
|
||||
// * binary literals are unsupported by the unparser and need placeholders.
|
||||
// Eliminate empty membership expressions before visiting their children.
|
||||
// Otherwise a discarded binary child could leave behind a stale binding.
|
||||
// Fast path: no binary literals — DataFusion's unparser handles everything.
|
||||
if !has_binary_literal(expr) {
|
||||
return run_unparser(expr);
|
||||
}
|
||||
|
||||
// Slow path: DataFusion's unparser cannot serialize `Binary`/`LargeBinary`
|
||||
// scalars, so we rewrite each one to a unique string-literal placeholder,
|
||||
// let the unparser do the rest of the work, then substitute the SQL
|
||||
// `X'...'` byte-string literal back in. This keeps the operator/function
|
||||
// serialization logic centralized in DataFusion and works for every
|
||||
// expression node type the unparser supports.
|
||||
let mut bindings: Vec<Vec<u8>> = Vec::new();
|
||||
let rewritten = expr
|
||||
.clone()
|
||||
.transform(|e: Expr| match e {
|
||||
Expr::InList(in_list) if in_list.list.is_empty() => Ok(Transformed::yes(
|
||||
Expr::Literal(ScalarValue::Boolean(Some(in_list.negated)), None),
|
||||
)),
|
||||
other => Ok(Transformed::no(other)),
|
||||
})
|
||||
.map_err(|e| crate::Error::InvalidInput {
|
||||
message: format!("failed to rewrite expression: {e}"),
|
||||
})?
|
||||
.data;
|
||||
|
||||
let user_strings = string_literals(&rewritten);
|
||||
let mut next_placeholder_id = 0;
|
||||
let mut binary_bindings = HashMap::new();
|
||||
let rewritten = rewritten
|
||||
.transform(|e: Expr| match e {
|
||||
Expr::Literal(ScalarValue::Binary(Some(bytes)), m)
|
||||
| Expr::Literal(ScalarValue::LargeBinary(Some(bytes)), m) => {
|
||||
let placeholder = next_binary_placeholder(&user_strings, &mut next_placeholder_id);
|
||||
binary_bindings.insert(placeholder.clone(), bytes);
|
||||
let placeholder = format!("{}{}__", BINARY_PLACEHOLDER_PREFIX, bindings.len());
|
||||
bindings.push(bytes);
|
||||
Ok(Transformed::yes(Expr::Literal(
|
||||
ScalarValue::Utf8(Some(placeholder)),
|
||||
m,
|
||||
@@ -287,57 +158,6 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result<String> {
|
||||
| Expr::Literal(ScalarValue::LargeBinary(None), m) => {
|
||||
Ok(Transformed::yes(Expr::Literal(ScalarValue::Null, m)))
|
||||
}
|
||||
Expr::Literal(ScalarValue::Decimal32(Some(value), precision, scale), _m) => {
|
||||
let value = Decimal32Type::format_decimal(value, precision, scale);
|
||||
Ok(Transformed::yes(typed_string_literal(
|
||||
value,
|
||||
DataType::Decimal32(precision, scale),
|
||||
)))
|
||||
}
|
||||
Expr::Literal(ScalarValue::Decimal64(Some(value), precision, scale), _m) => {
|
||||
let value = Decimal64Type::format_decimal(value, precision, scale);
|
||||
Ok(Transformed::yes(typed_string_literal(
|
||||
value,
|
||||
DataType::Decimal64(precision, scale),
|
||||
)))
|
||||
}
|
||||
Expr::Literal(ScalarValue::Decimal128(Some(value), precision, scale), _m) => {
|
||||
let value = Decimal128Type::format_decimal(value, precision, scale);
|
||||
Ok(Transformed::yes(typed_string_literal(
|
||||
value,
|
||||
DataType::Decimal128(precision, scale),
|
||||
)))
|
||||
}
|
||||
Expr::Literal(ScalarValue::Decimal256(Some(value), precision, scale), _m) => {
|
||||
let value = Decimal256Type::format_decimal(value, precision, scale);
|
||||
Ok(Transformed::yes(typed_string_literal(
|
||||
value,
|
||||
DataType::Decimal256(precision, scale),
|
||||
)))
|
||||
}
|
||||
Expr::Literal(ScalarValue::Float16(Some(value)), _m) if !value.is_finite() => Ok(
|
||||
Transformed::yes(typed_string_literal(value.to_string(), DataType::Float16)),
|
||||
),
|
||||
Expr::Literal(ScalarValue::Float32(Some(value)), _m) if !value.is_finite() => Ok(
|
||||
Transformed::yes(typed_string_literal(value.to_string(), DataType::Float32)),
|
||||
),
|
||||
Expr::Literal(ScalarValue::Float64(Some(value)), _m) if !value.is_finite() => Ok(
|
||||
Transformed::yes(typed_string_literal(value.to_string(), DataType::Float64)),
|
||||
),
|
||||
Expr::Cast(cast) => Ok(Transformed::yes(datafusion_arrow_cast(
|
||||
*cast.expr,
|
||||
Expr::Literal(
|
||||
ScalarValue::Utf8(Some(cast.field.data_type().to_string())),
|
||||
None,
|
||||
),
|
||||
))),
|
||||
Expr::TryCast(cast) => Ok(Transformed::yes(datafusion_arrow_try_cast(
|
||||
*cast.expr,
|
||||
Expr::Literal(
|
||||
ScalarValue::Utf8(Some(cast.field.data_type().to_string())),
|
||||
None,
|
||||
),
|
||||
))),
|
||||
other => Ok(Transformed::no(other)),
|
||||
})
|
||||
.map_err(|e| crate::Error::InvalidInput {
|
||||
@@ -345,12 +165,14 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result<String> {
|
||||
})?
|
||||
.data;
|
||||
|
||||
let sql = run_unparser(&rewritten)?;
|
||||
if binary_bindings.is_empty() {
|
||||
Ok(sql)
|
||||
} else {
|
||||
bind_binary_literals(&sql, binary_bindings)
|
||||
let mut sql = run_unparser(&rewritten)?;
|
||||
for (i, bytes) in bindings.iter().enumerate() {
|
||||
// The unparser quotes string literals with single quotes, so the
|
||||
// placeholder appears as `'__lancedb_binary_placeholder_<i>__'`.
|
||||
let quoted = format!("'{}{}__'", BINARY_PLACEHOLDER_PREFIX, i);
|
||||
sql = sql.replace("ed, &bytes_to_hex_sql(bytes));
|
||||
}
|
||||
Ok(sql)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -103,7 +103,9 @@ pub use lance::dataset::refs::{BranchContents, Ref, TagContents, Tags as LanceTa
|
||||
pub use lance::dataset::scanner::DatasetRecordBatchStream;
|
||||
pub use lance_index::optimize::OptimizeOptions;
|
||||
pub use lsm_stats::{BucketStats, GenerationStats, LsmStats, MemtableStats};
|
||||
pub use optimize::{CompactionOptions, OptimizeAction, OptimizeStats};
|
||||
pub use optimize::{
|
||||
CompactionMode, CompactionOptions, IndexRemapMode, OptimizeAction, OptimizeStats,
|
||||
};
|
||||
pub use refresh::RefreshColumnResult;
|
||||
pub use schema_evolution::{
|
||||
AddColumnsResult, AlterColumnsResult, DropColumnsResult, FieldMetadataUpdate,
|
||||
|
||||
@@ -15,7 +15,7 @@ use lance_index::optimize::OptimizeOptions;
|
||||
use log::info;
|
||||
|
||||
pub use chrono::Duration;
|
||||
pub use lance::dataset::optimize::CompactionOptions;
|
||||
pub use lance::dataset::optimize::{CompactionMode, CompactionOptions, IndexRemapMode};
|
||||
|
||||
use super::NativeTable;
|
||||
use crate::error::Result;
|
||||
|
||||
Reference in New Issue
Block a user