mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-11 07:42:26 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
894a9e7e6f | ||
|
|
667cf32e78 |
Generated
+298
-260
File diff suppressed because it is too large
Load Diff
+14
-14
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
lance = { "version" = "=10.1.0-beta.1", default-features = false, "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=10.1.0-beta.1", default-features = false, "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=10.1.0-beta.1", default-features = false, "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance = { "version" = "=10.1.0-beta.2", default-features = false, "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=10.1.0-beta.2", default-features = false, "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=10.1.0-beta.2", default-features = false, "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=10.1.0-beta.2", "tag" = "v10.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
ahash = "0.8"
|
||||
# Note that this one does not include pyarrow
|
||||
arrow = { version = "58.0.0", optional = false }
|
||||
|
||||
@@ -10,24 +10,6 @@ Reranks the results using the Reciprocal Rank Fusion (RRF) algorithm.
|
||||
|
||||
## Methods
|
||||
|
||||
### outputSchema()
|
||||
|
||||
```ts
|
||||
outputSchema(inputSchema): Promise<Schema<any>>
|
||||
```
|
||||
|
||||
Declare the RRF output schema for vector-only query execution.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **inputSchema**: `Schema`<`any`>
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`Schema`<`any`>>
|
||||
|
||||
***
|
||||
|
||||
### rerankHybrid()
|
||||
|
||||
```ts
|
||||
|
||||
@@ -8,27 +8,6 @@
|
||||
|
||||
## Methods
|
||||
|
||||
### outputSchema()?
|
||||
|
||||
```ts
|
||||
optional outputSchema(inputSchema): Promise<Schema<any>>
|
||||
```
|
||||
|
||||
Declare the schema returned when reranking a vector-only query.
|
||||
|
||||
This is required for vector-only reranking so query schema introspection
|
||||
and execution agree. Hybrid-only rerankers may omit it.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **inputSchema**: `Schema`<`any`>
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`Schema`<`any`>>
|
||||
|
||||
***
|
||||
|
||||
### rerankHybrid()
|
||||
|
||||
```ts
|
||||
|
||||
+1
-1
@@ -28,7 +28,7 @@
|
||||
<properties>
|
||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||
<arrow.version>15.0.0</arrow.version>
|
||||
<lance-core.version>10.1.0-beta.1</lance-core.version>
|
||||
<lance-core.version>10.1.0-beta.2</lance-core.version>
|
||||
<spotless.skip>false</spotless.skip>
|
||||
<spotless.version>2.30.0</spotless.version>
|
||||
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
|
||||
|
||||
@@ -79,22 +79,6 @@ describe("rerankers", function () {
|
||||
expect(result).toHaveLength(2);
|
||||
});
|
||||
|
||||
it("returns relevance scores when reranking a vector search", async function () {
|
||||
const query = table
|
||||
.vectorSearch([0.1, 0.1])
|
||||
.limit(2)
|
||||
.rerank(await RRFReranker.create());
|
||||
const schema = await query.outputSchema();
|
||||
const result = await query.toArray();
|
||||
|
||||
expect(schema.fields.map((field) => field.name)).toContain(
|
||||
"_relevance_score",
|
||||
);
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0]._relevance_score).toBeCloseTo(1 / 60);
|
||||
expect(result[1]._relevance_score).toBeCloseTo(1 / 61);
|
||||
});
|
||||
|
||||
it("does not keep process alive after rerank query", async function () {
|
||||
const script = `
|
||||
import * as lancedb from "./dist/index.js";
|
||||
|
||||
+13
-25
@@ -5,11 +5,9 @@ import {
|
||||
Table as ArrowTable,
|
||||
type IntoVector,
|
||||
RecordBatch,
|
||||
createEmptyTable,
|
||||
extractVectorBuffer,
|
||||
fromBufferToRecordBatch,
|
||||
fromRecordBatchToBuffer,
|
||||
fromTableToBuffer,
|
||||
tableFromIPC,
|
||||
} from "./arrow";
|
||||
import { type IvfPqOptions } from "./indices";
|
||||
@@ -746,30 +744,20 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
}
|
||||
|
||||
rerank(reranker: Reranker): VectorQuery {
|
||||
super.doCall((inner) => {
|
||||
const outputSchema = reranker.outputSchema?.bind(reranker);
|
||||
inner.rerank(
|
||||
async (args) => {
|
||||
const vecResults = await fromBufferToRecordBatch(args.vecResults);
|
||||
const ftsResults = await fromBufferToRecordBatch(args.ftsResults);
|
||||
const result = await reranker.rerankHybrid(
|
||||
args.query,
|
||||
vecResults as RecordBatch,
|
||||
ftsResults as RecordBatch,
|
||||
);
|
||||
super.doCall((inner) =>
|
||||
inner.rerank(async (args) => {
|
||||
const vecResults = await fromBufferToRecordBatch(args.vecResults);
|
||||
const ftsResults = await fromBufferToRecordBatch(args.ftsResults);
|
||||
const result = await reranker.rerankHybrid(
|
||||
args.query,
|
||||
vecResults as RecordBatch,
|
||||
ftsResults as RecordBatch,
|
||||
);
|
||||
|
||||
const buffer = fromRecordBatchToBuffer(result);
|
||||
return buffer;
|
||||
},
|
||||
outputSchema
|
||||
? async (args) => {
|
||||
const inputSchema = tableFromIPC(args.inputSchema).schema;
|
||||
const result = await outputSchema(inputSchema);
|
||||
return fromTableToBuffer(createEmptyTable(result));
|
||||
}
|
||||
: undefined,
|
||||
);
|
||||
});
|
||||
const buffer = fromRecordBatchToBuffer(result);
|
||||
return buffer;
|
||||
}),
|
||||
);
|
||||
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -1,22 +1,14 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import { RecordBatch, Schema } from "apache-arrow";
|
||||
import { RecordBatch } from "apache-arrow";
|
||||
|
||||
export * from "./rrf";
|
||||
|
||||
// Interface for a reranker. A reranker is used to rerank vector and hybrid
|
||||
// search results. For vector-only searches, query is empty and ftsResults is an
|
||||
// empty batch with the same schema as vecResults.
|
||||
// Interface for a reranker. A reranker is used to rerank the results from a
|
||||
// vector and FTS search. This is useful for combining the results from both
|
||||
// search methods.
|
||||
export interface Reranker {
|
||||
/**
|
||||
* Declare the schema returned when reranking a vector-only query.
|
||||
*
|
||||
* This is required for vector-only reranking so query schema introspection
|
||||
* and execution agree. Hybrid-only rerankers may omit it.
|
||||
*/
|
||||
outputSchema?(inputSchema: Schema): Promise<Schema>;
|
||||
|
||||
rerankHybrid(
|
||||
query: string,
|
||||
vecResults: RecordBatch,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import { Field, Float32, RecordBatch, Schema } from "apache-arrow";
|
||||
import { RecordBatch } from "apache-arrow";
|
||||
import { fromBufferToRecordBatch, fromRecordBatchToBuffer } from "../arrow";
|
||||
import { RrfReranker as NativeRRFReranker } from "../native";
|
||||
|
||||
@@ -24,17 +24,6 @@ export class RRFReranker {
|
||||
);
|
||||
}
|
||||
|
||||
/** Declare the RRF output schema for vector-only query execution. */
|
||||
async outputSchema(inputSchema: Schema): Promise<Schema> {
|
||||
return new Schema(
|
||||
[
|
||||
...inputSchema.fields,
|
||||
new Field("_relevance_score", new Float32(), false),
|
||||
],
|
||||
inputSchema.metadata,
|
||||
);
|
||||
}
|
||||
|
||||
async rerankHybrid(
|
||||
query: string,
|
||||
vecResults: RecordBatch,
|
||||
|
||||
+2
-3
@@ -6,8 +6,8 @@ use std::sync::Arc;
|
||||
use crate::error::NapiErrorExt;
|
||||
use crate::error::convert_error;
|
||||
use crate::iterator::RecordBatchIterator;
|
||||
use crate::rerankers::RerankHybridCallbackArgs;
|
||||
use crate::rerankers::Reranker;
|
||||
use crate::rerankers::{RerankHybridCallbackArgs, RerankOutputSchemaCallbackArgs};
|
||||
use crate::util::{parse_distance_type, schema_to_buffer};
|
||||
use arrow_array::{
|
||||
Array, Float16Array as ArrowFloat16Array, Float32Array as ArrowFloat32Array,
|
||||
@@ -388,9 +388,8 @@ impl VectorQuery {
|
||||
pub fn rerank(
|
||||
&mut self,
|
||||
rerank_hybrid: Function<RerankHybridCallbackArgs, Promise<Buffer>>,
|
||||
output_schema: Option<Function<RerankOutputSchemaCallbackArgs, Promise<Buffer>>>,
|
||||
) -> napi::Result<()> {
|
||||
let reranker = Reranker::new(rerank_hybrid, output_schema)?;
|
||||
let reranker = Reranker::new(rerank_hybrid)?;
|
||||
self.inner = self.inner.clone().rerank(Arc::new(reranker));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
+2
-51
@@ -6,7 +6,7 @@ use async_trait::async_trait;
|
||||
use napi::{bindgen_prelude::*, threadsafe_function::ThreadsafeFunction};
|
||||
use napi_derive::napi;
|
||||
|
||||
use lancedb::ipc::{batches_to_ipc_file, ipc_file_to_schema, schema_to_ipc_file};
|
||||
use lancedb::ipc::batches_to_ipc_file;
|
||||
use lancedb::rerankers::Reranker as LanceDBReranker;
|
||||
use lancedb::{error::Error, ipc::ipc_file_to_batches};
|
||||
|
||||
@@ -21,72 +21,28 @@ type RerankHybridFn = ThreadsafeFunction<
|
||||
true,
|
||||
>;
|
||||
|
||||
type RerankOutputSchemaFn = ThreadsafeFunction<
|
||||
RerankOutputSchemaCallbackArgs,
|
||||
Promise<Buffer>,
|
||||
RerankOutputSchemaCallbackArgs,
|
||||
Status,
|
||||
false,
|
||||
true,
|
||||
>;
|
||||
|
||||
/// Reranker implementation that "wraps" a NodeJS Reranker implementation.
|
||||
/// This contains references to the callbacks that can be used to invoke the
|
||||
/// reranking methods on the NodeJS implementation and handles serializing the
|
||||
/// record batches to Arrow IPC buffers.
|
||||
pub struct Reranker {
|
||||
rerank_hybrid: RerankHybridFn,
|
||||
output_schema: Option<RerankOutputSchemaFn>,
|
||||
}
|
||||
|
||||
impl Reranker {
|
||||
pub fn new(
|
||||
rerank_hybrid: Function<RerankHybridCallbackArgs, Promise<Buffer>>,
|
||||
output_schema: Option<Function<RerankOutputSchemaCallbackArgs, Promise<Buffer>>>,
|
||||
) -> napi::Result<Self> {
|
||||
let rerank_hybrid = rerank_hybrid
|
||||
.build_threadsafe_function()
|
||||
.weak::<true>()
|
||||
.build()?;
|
||||
let output_schema = output_schema
|
||||
.map(|output_schema| {
|
||||
output_schema
|
||||
.build_threadsafe_function()
|
||||
.weak::<true>()
|
||||
.build()
|
||||
})
|
||||
.transpose()?;
|
||||
Ok(Self {
|
||||
rerank_hybrid,
|
||||
output_schema,
|
||||
})
|
||||
Ok(Self { rerank_hybrid })
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl lancedb::rerankers::Reranker for Reranker {
|
||||
async fn output_schema(
|
||||
&self,
|
||||
input: &arrow_schema::SchemaRef,
|
||||
) -> lancedb::error::Result<arrow_schema::SchemaRef> {
|
||||
let output_schema = self.output_schema.as_ref().ok_or(Error::NotSupported {
|
||||
message: "vector rerankers must declare their output schema".to_string(),
|
||||
})?;
|
||||
let callback_args = RerankOutputSchemaCallbackArgs {
|
||||
input_schema: Buffer::from(schema_to_ipc_file(input.as_ref())?),
|
||||
};
|
||||
let promised_buffer: Promise<Buffer> = output_schema
|
||||
.call_async(callback_args)
|
||||
.await
|
||||
.map_err(|e| Error::Runtime {
|
||||
message: format!("napi error status={}, reason={}", e.status, e.reason),
|
||||
})?;
|
||||
let buffer = promised_buffer.await.map_err(|e| Error::Runtime {
|
||||
message: format!("napi error status={}, reason={}", e.status, e.reason),
|
||||
})?;
|
||||
ipc_file_to_schema(buffer.to_vec())
|
||||
}
|
||||
|
||||
async fn rerank_hybrid(
|
||||
&self,
|
||||
query: &str,
|
||||
@@ -130,11 +86,6 @@ pub struct RerankHybridCallbackArgs {
|
||||
pub fts_results: Buffer,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
pub struct RerankOutputSchemaCallbackArgs {
|
||||
pub input_schema: Buffer,
|
||||
}
|
||||
|
||||
fn buffer_to_record_batch(buffer: Buffer) -> Result<RecordBatch> {
|
||||
let mut reader = ipc_file_to_batches(buffer.to_vec()).default_error()?;
|
||||
reader
|
||||
|
||||
+2
-6
@@ -339,9 +339,7 @@ impl Table {
|
||||
let transforms = NewColumnTransform::SqlExpressions(transforms);
|
||||
let res = self
|
||||
.inner_ref()?
|
||||
.add_columns()
|
||||
.transform(transforms)
|
||||
.execute()
|
||||
.add_columns(transforms, None)
|
||||
.await
|
||||
.default_error()?;
|
||||
Ok(res.into())
|
||||
@@ -358,9 +356,7 @@ impl Table {
|
||||
let transforms = NewColumnTransform::AllNulls(schema);
|
||||
let res = self
|
||||
.inner_ref()?
|
||||
.add_columns()
|
||||
.transform(transforms)
|
||||
.execute()
|
||||
.add_columns(transforms, None)
|
||||
.await
|
||||
.default_error()?;
|
||||
Ok(res.into())
|
||||
|
||||
@@ -707,9 +707,6 @@ class LanceDBConnection(DBConnection):
|
||||
self._namespace_client_properties = namespace_client_properties
|
||||
if _inner is not None:
|
||||
self._conn = _inner
|
||||
# Native-derived wrappers resolve this in their async reconstruction
|
||||
# path so construction never synchronously re-enters LOOP.
|
||||
self._read_consistency_interval = read_consistency_interval
|
||||
self._cached_namespace_client = None
|
||||
return
|
||||
|
||||
@@ -759,14 +756,11 @@ class LanceDBConnection(DBConnection):
|
||||
# storage_options. Also, this class really shouldn't be holding any state
|
||||
# beyond _conn.
|
||||
self._conn = AsyncConnection(LOOP.run(do_connect()))
|
||||
# Keep property access synchronous so debugger introspection cannot wait on
|
||||
# the background loop while that thread is suspended at a breakpoint.
|
||||
self._read_consistency_interval = read_consistency_interval
|
||||
self._cached_namespace_client: Optional[LanceNamespace] = None
|
||||
|
||||
@property
|
||||
def read_consistency_interval(self) -> Optional[timedelta]:
|
||||
return self._read_consistency_interval
|
||||
return LOOP.run(self._conn.get_read_consistency_interval())
|
||||
|
||||
@property
|
||||
def session(self) -> Optional[Session]:
|
||||
@@ -777,16 +771,8 @@ class LanceDBConnection(DBConnection):
|
||||
return self._conn.uri
|
||||
|
||||
@classmethod
|
||||
def from_inner(
|
||||
cls,
|
||||
inner: LanceDbConnection,
|
||||
read_consistency_interval: Optional[timedelta],
|
||||
):
|
||||
return cls(
|
||||
None,
|
||||
read_consistency_interval=read_consistency_interval,
|
||||
_inner=inner,
|
||||
)
|
||||
def from_inner(cls, inner: LanceDbConnection):
|
||||
return cls(None, _inner=inner)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}(uri={self._conn.uri!r})"
|
||||
|
||||
@@ -92,10 +92,8 @@ class LanceMergeInsertBuilder(object):
|
||||
self._when_not_matched_by_source_delete = True
|
||||
if isinstance(condition, Expr):
|
||||
self._when_not_matched_by_source_condition_expr = condition._inner
|
||||
self._when_not_matched_by_source_condition = None
|
||||
else:
|
||||
elif condition is not None:
|
||||
self._when_not_matched_by_source_condition = condition
|
||||
self._when_not_matched_by_source_condition_expr = None
|
||||
return self
|
||||
|
||||
def use_index(self, use_index: bool) -> LanceMergeInsertBuilder:
|
||||
|
||||
@@ -226,7 +226,7 @@ class PermutationBuilder:
|
||||
|
||||
async def do_execute():
|
||||
inner_tbl = await self._async.execute()
|
||||
return await LanceTable.from_inner(inner_tbl)
|
||||
return LanceTable.from_inner(inner_tbl)
|
||||
|
||||
return LOOP.run(do_execute())
|
||||
|
||||
|
||||
@@ -2697,7 +2697,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
self._fts_query.phrase_query(True)
|
||||
if self._distance_type:
|
||||
self._vector_query.metric(self._distance_type)
|
||||
if self._minimum_nprobes is not None:
|
||||
if self._minimum_nprobes:
|
||||
self._vector_query.minimum_nprobes(self._minimum_nprobes)
|
||||
if self._maximum_nprobes is not None:
|
||||
self._vector_query.maximum_nprobes(self._maximum_nprobes)
|
||||
|
||||
@@ -1595,8 +1595,8 @@ class Table(ABC):
|
||||
Prefer this over :meth:`fetch_blobs` for large payloads. ``row_ids`` is
|
||||
a ``list[int]`` or a query ``pyarrow.Table`` carrying row identity via
|
||||
``_rowid`` or a ``_lance_row_id`` field on the blob descriptor. Null
|
||||
rows are ``None``. Remote tables require LanceDB Cloud server 0.5.0 or
|
||||
newer.
|
||||
rows are ``None``. Unsupported on LanceDB Cloud, where
|
||||
:meth:`fetch_blobs` returns full bytes instead.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
@@ -2182,15 +2182,11 @@ class LanceTable(Table):
|
||||
return self.name
|
||||
|
||||
@classmethod
|
||||
async def from_inner(cls, tbl: LanceDBTable):
|
||||
from .db import AsyncConnection, LanceDBConnection
|
||||
def from_inner(cls, tbl: LanceDBTable):
|
||||
from .db import LanceDBConnection
|
||||
|
||||
async_tbl = AsyncTable(tbl)
|
||||
inner_conn = tbl.database()
|
||||
read_consistency_interval = await AsyncConnection(
|
||||
inner_conn
|
||||
).get_read_consistency_interval()
|
||||
conn = LanceDBConnection.from_inner(inner_conn, read_consistency_interval)
|
||||
conn = LanceDBConnection.from_inner(tbl.database())
|
||||
return cls(
|
||||
conn,
|
||||
async_tbl.name,
|
||||
|
||||
@@ -77,23 +77,6 @@ def test_sync_repr_does_not_use_background_loop(tmp_path, monkeypatch):
|
||||
assert repr(table) == f"LanceTable(name='test', _conn={db!r})"
|
||||
|
||||
|
||||
def test_read_consistency_interval_does_not_use_background_loop(tmp_path, monkeypatch):
|
||||
from lancedb.background_loop import LOOP
|
||||
from lancedb.db import LanceDBConnection
|
||||
|
||||
consistency_interval = timedelta(seconds=5)
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=consistency_interval)
|
||||
db_from_inner = LanceDBConnection.from_inner(db._inner, consistency_interval)
|
||||
|
||||
def fail_run(*args, **kwargs):
|
||||
raise AssertionError("properties should not use the Python background loop")
|
||||
|
||||
monkeypatch.setattr(LOOP, "run", fail_run)
|
||||
|
||||
assert db.read_consistency_interval == consistency_interval
|
||||
assert db_from_inner.read_consistency_interval == consistency_interval
|
||||
|
||||
|
||||
def test_ingest_pd(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
|
||||
|
||||
@@ -123,19 +123,6 @@ async def test_async_hybrid_query_default_limit(table: AsyncTable):
|
||||
assert texts.count("a") == 1
|
||||
|
||||
|
||||
def test_hybrid_query_minimum_nprobes_zero_raises(sync_table: Table):
|
||||
# minimum_nprobes(0) must raise the same validation error a plain vector
|
||||
# query raises, not silently no-op because 0 is falsy.
|
||||
with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"):
|
||||
(
|
||||
sync_table.search(query_type="hybrid")
|
||||
.vector([0.0, 0.4])
|
||||
.text("dog")
|
||||
.minimum_nprobes(0)
|
||||
.to_arrow()
|
||||
)
|
||||
|
||||
|
||||
def test_hybrid_query_distance_range(sync_table: Table):
|
||||
reranker = RRFReranker(return_score="all")
|
||||
result = (
|
||||
|
||||
@@ -6,7 +6,6 @@ import math
|
||||
import pytest
|
||||
|
||||
from lancedb import DBConnection, Table, connect
|
||||
from lancedb.background_loop import LOOP
|
||||
from lancedb.permutation import Permutation, Permutations, permutation_builder
|
||||
|
||||
|
||||
@@ -32,25 +31,6 @@ def test_split_random_ratios(mem_db):
|
||||
assert 65 <= split_1_count <= 75 # ~70% ± tolerance
|
||||
|
||||
|
||||
def test_execute_does_not_reenter_background_loop(tmp_path, monkeypatch):
|
||||
import threading
|
||||
|
||||
db = connect(tmp_path)
|
||||
tbl = db.create_table("test_table", pa.table({"x": range(10)}))
|
||||
original_run = LOOP.run
|
||||
|
||||
def fail_on_reentry(future):
|
||||
assert threading.current_thread() is not LOOP.thread
|
||||
return original_run(future)
|
||||
|
||||
monkeypatch.setattr(LOOP, "run", fail_on_reentry)
|
||||
|
||||
permutation_tbl = permutation_builder(tbl).execute()
|
||||
|
||||
assert permutation_tbl.count_rows() == 10
|
||||
assert permutation_tbl._conn.read_consistency_interval is None
|
||||
|
||||
|
||||
def test_split_random_counts(mem_db):
|
||||
"""Test random splitting with absolute counts."""
|
||||
tbl = mem_db.create_table(
|
||||
|
||||
@@ -2055,24 +2055,6 @@ def blob_remote_table(*, server_version=Version("0.5.0")):
|
||||
request.send_header("phalanx-version", str(server_version))
|
||||
request.end_headers()
|
||||
request.wfile.write(json.dumps(BLOB_DESCRIBE_RESPONSE).encode())
|
||||
elif request.path.startswith("/v1/table/test/blob/image/"):
|
||||
path = request.path.partition("?")[0]
|
||||
row_id = int(path.split("/")[-2])
|
||||
payload = {10: b"alpha", 20: None, 30: b"gamma"}[row_id]
|
||||
if payload is None:
|
||||
request.send_response(204)
|
||||
request.end_headers()
|
||||
return
|
||||
byte_range = request.headers["Range"].removeprefix("bytes=")
|
||||
start_text, end_text = byte_range.split("-", maxsplit=1)
|
||||
start = int(start_text)
|
||||
end = int(end_text) if end_text else len(payload) - 1
|
||||
chunk = payload[start : end + 1]
|
||||
request.send_response(206)
|
||||
request.send_header("Content-Range", f"bytes {start}-{end}/{len(payload)}")
|
||||
request.send_header("Content-Length", str(len(chunk)))
|
||||
request.end_headers()
|
||||
request.wfile.write(chunk)
|
||||
elif request.path == "/v1/table/test/query/":
|
||||
content_len = int(request.headers.get("Content-Length", 0))
|
||||
body = json.loads(request.rfile.read(content_len))
|
||||
@@ -2110,21 +2092,8 @@ def test_remote_blob_columns_and_fetch():
|
||||
assert table.blob_columns() == ["image"]
|
||||
blobs = table.fetch_blobs("image", [10, 20, 30])
|
||||
assert blobs.to_pylist() == [b"alpha", None, b"gamma"]
|
||||
|
||||
|
||||
def test_remote_blob_files_are_lazy_seekable_handles():
|
||||
with blob_remote_table() as table:
|
||||
files = table.fetch_blob_files("image", [10, 20, 30])
|
||||
|
||||
assert len(files) == 3
|
||||
alpha, null_row, gamma = files
|
||||
assert null_row is None
|
||||
assert alpha is not None
|
||||
assert gamma is not None
|
||||
assert alpha.size() == 5
|
||||
assert alpha.read_range(1, 3) == b"lph"
|
||||
gamma.seek(2)
|
||||
assert gamma.read() == b"mma"
|
||||
with pytest.raises(NotImplementedError, match="Use fetch_blobs for full bytes"):
|
||||
table.fetch_blob_files("image", [10, 20, 30])
|
||||
|
||||
|
||||
def test_remote_blob_fetch_accepts_query_table():
|
||||
|
||||
@@ -6,7 +6,6 @@ import os
|
||||
import sys
|
||||
import threading
|
||||
import warnings
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import date, datetime, timedelta
|
||||
from time import sleep
|
||||
from typing import List
|
||||
@@ -2125,27 +2124,6 @@ def test_delete(mem_db: DBConnection):
|
||||
assert table.to_arrow()["id"].to_pylist() == [1]
|
||||
|
||||
|
||||
def test_concurrent_deletes_are_thread_safe(mem_db: DBConnection):
|
||||
num_workers = 8
|
||||
table = mem_db.create_table(
|
||||
"my_table", data=[{"id": row_id} for row_id in range(num_workers)]
|
||||
)
|
||||
barrier = threading.Barrier(num_workers)
|
||||
|
||||
def delete(row_id: int):
|
||||
barrier.wait()
|
||||
return table.delete(f"id = {row_id}")
|
||||
|
||||
with ThreadPoolExecutor(max_workers=num_workers) as pool:
|
||||
results = list(pool.map(delete, range(num_workers)))
|
||||
|
||||
assert all(result.num_deleted_rows == 1 for result in results)
|
||||
assert sorted(result.version for result in results) == list(
|
||||
range(2, num_workers + 2)
|
||||
)
|
||||
assert table.count_rows() == 0
|
||||
|
||||
|
||||
def test_delete_expr(mem_db: DBConnection):
|
||||
table = mem_db.create_table(
|
||||
"my_table",
|
||||
@@ -2386,29 +2364,6 @@ def test_merge_insert_by_source_delete_expr(mem_db: DBConnection):
|
||||
assert table.to_arrow().sort_by("a") == expected
|
||||
|
||||
|
||||
def test_merge_insert_by_source_delete_reconfigure(mem_db: DBConnection):
|
||||
# Calling when_not_matched_by_source_delete() again with no condition must
|
||||
# widen the delete to unconditional, not keep the earlier condition around.
|
||||
table = mem_db.create_table(
|
||||
"my_table",
|
||||
data=pa.table({"a": [1, 2, 3], "b": ["a", "b", "c"]}),
|
||||
)
|
||||
new_data = pa.table({"a": [2, 4], "b": ["x", "z"]})
|
||||
|
||||
merge_insert_res = (
|
||||
table.merge_insert("a")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.when_not_matched_by_source_delete("a > 2")
|
||||
.when_not_matched_by_source_delete()
|
||||
.execute(new_data)
|
||||
)
|
||||
assert merge_insert_res.num_deleted_rows == 2
|
||||
|
||||
expected = pa.table({"a": [2, 4], "b": ["x", "z"]})
|
||||
assert table.to_arrow().sort_by("a") == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_merge_insert_by_source_delete_expr_async(
|
||||
mem_db_async: AsyncConnection,
|
||||
|
||||
+13
-40
@@ -426,11 +426,9 @@ pub struct PyBlobFile {
|
||||
impl PyBlobFile {
|
||||
fn read_bytes(self_: PyRef<'_, Self>) -> PyResult<Py<PyBytes>> {
|
||||
let inner = self_.inner.clone();
|
||||
let py = self_.py();
|
||||
let bytes = py
|
||||
.detach(move || block_on(async move { inner.read().await }))
|
||||
let bytes = block_on(async move { inner.read().await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
|
||||
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
|
||||
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
|
||||
}
|
||||
|
||||
pub fn read(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
@@ -446,32 +444,24 @@ impl PyBlobFile {
|
||||
|
||||
fn close(self_: PyRef<'_, Self>) -> PyResult<()> {
|
||||
let inner = self_.inner.clone();
|
||||
self_
|
||||
.py()
|
||||
.detach(move || block_on(async move { inner.close().await }))
|
||||
block_on(async move { inner.close().await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob close failed: {e}")))
|
||||
}
|
||||
|
||||
fn is_closed(self_: PyRef<'_, Self>) -> bool {
|
||||
let inner = self_.inner.clone();
|
||||
self_
|
||||
.py()
|
||||
.detach(move || block_on(async move { inner.is_closed().await }))
|
||||
block_on(async move { inner.is_closed().await })
|
||||
}
|
||||
|
||||
fn seek(self_: PyRef<'_, Self>, position: u64) -> PyResult<()> {
|
||||
let inner = self_.inner.clone();
|
||||
self_
|
||||
.py()
|
||||
.detach(move || block_on(async move { inner.seek(position).await }))
|
||||
block_on(async move { inner.seek(position).await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob seek failed: {e}")))
|
||||
}
|
||||
|
||||
fn tell(self_: PyRef<'_, Self>) -> PyResult<u64> {
|
||||
let inner = self_.inner.clone();
|
||||
self_
|
||||
.py()
|
||||
.detach(move || block_on(async move { inner.tell().await }))
|
||||
block_on(async move { inner.tell().await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob tell failed: {e}")))
|
||||
}
|
||||
|
||||
@@ -485,20 +475,16 @@ impl PyBlobFile {
|
||||
.checked_add(length as u64)
|
||||
.ok_or_else(|| PyValueError::new_err("offset + length overflowed"))?;
|
||||
let inner = self_.inner.clone();
|
||||
let py = self_.py();
|
||||
let bytes = py
|
||||
.detach(move || block_on(async move { inner.read_range(offset..end).await }))
|
||||
let bytes = block_on(async move { inner.read_range(offset..end).await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read_range failed: {e}")))?;
|
||||
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
|
||||
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
|
||||
}
|
||||
|
||||
fn read_up_to(self_: PyRef<'_, Self>, length: usize) -> PyResult<Py<PyBytes>> {
|
||||
let inner = self_.inner.clone();
|
||||
let py = self_.py();
|
||||
let bytes = py
|
||||
.detach(move || block_on(async move { inner.read_up_to(length).await }))
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read_up_to failed: {e}")))?;
|
||||
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
|
||||
let bytes = block_on(async move { inner.read_up_to(length).await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
|
||||
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -745,9 +731,6 @@ impl Table {
|
||||
|
||||
#[allow(private_interfaces)]
|
||||
pub fn delete(self_: PyRef<'_, Self>, condition: PredicateArg) -> PyResult<Bound<'_, PyAny>> {
|
||||
// Do not hold the Python borrow across the await. The cloned Rust table
|
||||
// handle is thread-safe and allows deletes on the same Python table to
|
||||
// run concurrently without PyO3 reporting "Already borrowed".
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let result = match &condition {
|
||||
@@ -1378,12 +1361,7 @@ impl Table {
|
||||
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let result = inner
|
||||
.add_columns()
|
||||
.transform(definitions)
|
||||
.execute()
|
||||
.await
|
||||
.infer_error()?;
|
||||
let result = inner.add_columns(definitions, None).await.infer_error()?;
|
||||
Ok(AddColumnsResult::from(result))
|
||||
})
|
||||
}
|
||||
@@ -1397,12 +1375,7 @@ impl Table {
|
||||
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let result = inner
|
||||
.add_columns()
|
||||
.transform(transform)
|
||||
.execute()
|
||||
.await
|
||||
.infer_error()?;
|
||||
let result = inner.add_columns(transform, None).await.infer_error()?;
|
||||
Ok(AddColumnsResult::from(result))
|
||||
})
|
||||
}
|
||||
|
||||
@@ -49,8 +49,8 @@ lance-namespace = { workspace = true }
|
||||
lance-namespace-impls = { workspace = true }
|
||||
metrics = { workspace = true, optional = true }
|
||||
metrics-util = { workspace = true, optional = true }
|
||||
# Pin the transitive GooseFS SDK until the 0.1.6 compile break is fixed upstream.
|
||||
goosefs-sdk = { version = "=0.1.5", optional = true }
|
||||
# Keep the direct dependency aligned with the version required by OpenDAL.
|
||||
goosefs-sdk = { version = "=0.1.8", optional = true }
|
||||
moka = { workspace = true }
|
||||
pin-project = { workspace = true }
|
||||
tokio = { version = "1.23", features = ["rt-multi-thread", "sync"] }
|
||||
|
||||
+2
-198
@@ -9,7 +9,6 @@
|
||||
//!
|
||||
//! Blob tables require Lance file format >= 2.2 and stable row ids at create.
|
||||
|
||||
use std::ops::Range;
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::LargeBinaryArray;
|
||||
@@ -18,202 +17,10 @@ use arrow_schema::{DataType, Field, Schema};
|
||||
use lance::dataset::{BlobRangeRequest as LanceBlobRangeRequest, Dataset, WriteParams};
|
||||
use lance_arrow::FieldExt;
|
||||
use lance_encoding::version::LanceFileVersion;
|
||||
use lance_io::object_store::ObjectStore;
|
||||
use object_store::path::Path;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
/// Seekable handle for one blob value, backed by local storage or a remote
|
||||
/// HTTP byte-range endpoint.
|
||||
#[derive(Debug)]
|
||||
pub struct BlobFile {
|
||||
inner: BlobFileInner,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum BlobFileInner {
|
||||
Native(lance::dataset::BlobFile),
|
||||
#[cfg(feature = "remote")]
|
||||
Remote(Box<crate::remote::table::blobs::RemoteBlobFile>),
|
||||
}
|
||||
|
||||
impl From<lance::dataset::BlobFile> for BlobFile {
|
||||
fn from(value: lance::dataset::BlobFile) -> Self {
|
||||
Self {
|
||||
inner: BlobFileInner::Native(value),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "remote")]
|
||||
impl From<crate::remote::table::blobs::RemoteBlobFile> for BlobFile {
|
||||
fn from(value: crate::remote::table::blobs::RemoteBlobFile) -> Self {
|
||||
Self {
|
||||
inner: BlobFileInner::Remote(Box::new(value)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl BlobFile {
|
||||
/// Inline reader over a data-file slice.
|
||||
pub fn new_inline(
|
||||
object_store: Arc<ObjectStore>,
|
||||
path: Path,
|
||||
position: u64,
|
||||
size: u64,
|
||||
) -> Self {
|
||||
lance::dataset::BlobFile::new_inline(object_store, path, position, size).into()
|
||||
}
|
||||
|
||||
/// Dedicated sidecar-file reader.
|
||||
pub fn new_dedicated(object_store: Arc<ObjectStore>, path: Path, size: u64) -> Self {
|
||||
lance::dataset::BlobFile::new_dedicated(object_store, path, size).into()
|
||||
}
|
||||
|
||||
/// Packed reader for a slice in a shared sidecar.
|
||||
pub fn new_packed(
|
||||
object_store: Arc<ObjectStore>,
|
||||
path: Path,
|
||||
position: u64,
|
||||
size: u64,
|
||||
) -> Self {
|
||||
lance::dataset::BlobFile::new_packed(object_store, path, position, size).into()
|
||||
}
|
||||
|
||||
/// External reader at a resolved object location.
|
||||
pub fn new_external(
|
||||
object_store: Arc<ObjectStore>,
|
||||
path: Path,
|
||||
uri: String,
|
||||
position: u64,
|
||||
size: u64,
|
||||
) -> Self {
|
||||
lance::dataset::BlobFile::new_external(object_store, path, uri, position, size).into()
|
||||
}
|
||||
|
||||
/// Close the handle.
|
||||
pub async fn close(&self) -> lance_core::Result<()> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.close().await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.close().await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether the handle is closed.
|
||||
pub async fn is_closed(&self) -> bool {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.is_closed().await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.is_closed(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Read a range without moving the cursor.
|
||||
pub async fn read_range(&self, range: Range<u64>) -> lance_core::Result<bytes::Bytes> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.read_range(range).await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.read_range(range).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Read ranges without moving the cursor.
|
||||
pub async fn read_ranges(
|
||||
&self,
|
||||
ranges: &[Range<u64>],
|
||||
) -> lance_core::Result<Vec<bytes::Bytes>> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.read_ranges(ranges).await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.read_ranges(ranges).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Read from the cursor to the end.
|
||||
pub async fn read(&self) -> lance_core::Result<bytes::Bytes> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.read().await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.read().await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Read up to `len` bytes and advance the cursor.
|
||||
pub async fn read_up_to(&self, len: usize) -> lance_core::Result<bytes::Bytes> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.read_up_to(len).await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.read_up_to(len).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Move the cursor to `new_cursor`.
|
||||
pub async fn seek(&self, new_cursor: u64) -> lance_core::Result<()> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.seek(new_cursor).await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.seek(new_cursor).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Current cursor position.
|
||||
pub async fn tell(&self) -> lance_core::Result<u64> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.tell().await,
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.tell().await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Blob length in bytes.
|
||||
pub fn size(&self) -> u64 {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.size(),
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(file) => file.size(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Physical byte offset in the data file. `None` on remote handles. The
|
||||
/// Cloud byte-range route does not expose storage layout.
|
||||
pub fn position(&self) -> Option<u64> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => Some(file.position()),
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Path of the data file holding the blob. `None` on remote handles. The
|
||||
/// Cloud byte-range route does not expose storage layout.
|
||||
pub fn data_path(&self) -> Option<&Path> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => Some(file.data_path()),
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Native storage layout. `None` on remote handles. The Cloud byte-range
|
||||
/// route does not expose layout.
|
||||
pub fn kind(&self) -> Option<lance_core::datatypes::BlobKind> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => Some(file.kind()),
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// External URI for native handles. Remote handles do not expose storage URIs.
|
||||
pub fn uri(&self) -> Option<&str> {
|
||||
match &self.inner {
|
||||
BlobFileInner::Native(file) => file.uri(),
|
||||
#[cfg(feature = "remote")]
|
||||
BlobFileInner::Remote(_) => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
pub use lance::dataset::BlobFile;
|
||||
|
||||
/// One row-specific blob range read request.
|
||||
///
|
||||
@@ -457,10 +264,7 @@ pub(crate) async fn take_blob_files_aligned(
|
||||
|
||||
let handles = dataset.take_blobs(row_ids, column).await?;
|
||||
ensure_all_row_ids_resolved(column, row_ids.len(), handles.len())?;
|
||||
Ok(handles
|
||||
.into_iter()
|
||||
.map(|handle| handle.map(Into::into))
|
||||
.collect())
|
||||
Ok(handles)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -8,12 +8,15 @@ use std::fs::create_dir_all;
|
||||
use std::path::Path;
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
|
||||
use futures::TryStreamExt;
|
||||
use lance::dataset::refs::Ref;
|
||||
use lance::dataset::{ReadParams, WriteMode, builder::DatasetBuilder};
|
||||
use lance::io::{ObjectStore, ObjectStoreParams, WrappingObjectStore};
|
||||
use lance_datafusion::utils::StreamingWriteSource;
|
||||
use lance_encoding::version::LanceFileVersion;
|
||||
use lance_io::object_store::{StorageOptionsAccessor, StorageOptionsProvider};
|
||||
use lance_io::object_store::{
|
||||
DirCursor, ReadDirOptions, StorageOptionsAccessor, StorageOptionsProvider,
|
||||
};
|
||||
use lance_table::io::commit::commit_handler_from_url;
|
||||
use object_store::local::LocalFileSystem;
|
||||
use snafu::ResultExt;
|
||||
@@ -718,6 +721,54 @@ impl ListingDatabase {
|
||||
self.namespace_database.clone()
|
||||
}
|
||||
|
||||
/// List up to `limit` table names, resuming after the table named `start_after`.
|
||||
///
|
||||
/// The cursor and the page size go into the object store's list request rather than
|
||||
/// being applied to a full listing, so the cost of a page is set by the size of the
|
||||
/// page and not by the size of the database. Stores with no paginated list API fall
|
||||
/// back to a full listing, which is what this did for every store before.
|
||||
///
|
||||
/// Names come back in the order the store lists the directories in, which is by key:
|
||||
/// `foo-bar` precedes `foo`, because the `-` of `foo-bar.lance` sorts below the `.` of
|
||||
/// `foo.lance`. Pagination has to follow the order the cursor is pushed down in, so
|
||||
/// that is the order both listing methods report and the order `start_after` resumes
|
||||
/// in. It matches sorting by name except between a name and one that extends it.
|
||||
async fn list_table_dirs(
|
||||
&self,
|
||||
start_after: Option<&str>,
|
||||
limit: Option<usize>,
|
||||
) -> Result<Vec<String>> {
|
||||
let dir_suffix = format!(".{}", LANCE_EXTENSION);
|
||||
let options = ReadDirOptions {
|
||||
// An empty name means "from the start": that is how comparing names against it
|
||||
// behaved, and how a client looping on a page token spells its first request.
|
||||
// Built into a cursor it would instead sit after every name below `.lance`.
|
||||
resume_from: start_after
|
||||
.filter(|name| !name.is_empty())
|
||||
.map(|name| DirCursor::after_directory(format!("{name}{dir_suffix}"))),
|
||||
page_size: limit,
|
||||
};
|
||||
let mut entries = self
|
||||
.object_store
|
||||
.read_dir_stream(self.base_path.clone(), options);
|
||||
|
||||
let mut names = Vec::new();
|
||||
while limit.is_none_or(|limit| names.len() < limit) {
|
||||
let Some(entry) = entries.try_next().await? else {
|
||||
break;
|
||||
};
|
||||
// A table is the directory `<name>.lance`; anything else under the database
|
||||
// prefix belongs to something other than a table.
|
||||
if !entry.is_dir() {
|
||||
continue;
|
||||
}
|
||||
if let Some(name) = entry.name.strip_suffix(&dir_suffix) {
|
||||
names.push(name.to_string());
|
||||
}
|
||||
}
|
||||
Ok(names)
|
||||
}
|
||||
|
||||
async fn drop_tables(&self, names: Vec<String>) -> Result<()> {
|
||||
let object_store_params = ObjectStoreParams {
|
||||
storage_options_accessor: if self.storage_options.is_empty() {
|
||||
@@ -959,80 +1010,37 @@ impl Database for ListingDatabase {
|
||||
if !request.namespace_path.is_empty() {
|
||||
return self.namespace_database().table_names(request).await;
|
||||
}
|
||||
let mut f = self
|
||||
.object_store
|
||||
.read_dir(self.base_path.clone())
|
||||
.await?
|
||||
.iter()
|
||||
.map(Path::new)
|
||||
.filter(|path| {
|
||||
let is_lance = path
|
||||
.extension()
|
||||
.and_then(|e| e.to_str())
|
||||
.map(|e| e == LANCE_EXTENSION);
|
||||
is_lance.unwrap_or(false)
|
||||
})
|
||||
.filter_map(|p| p.file_stem().and_then(|s| s.to_str().map(String::from)))
|
||||
.collect::<Vec<String>>();
|
||||
f.sort();
|
||||
if let Some(start_after) = request.start_after {
|
||||
let index = f
|
||||
.iter()
|
||||
.position(|name| name.as_str() > start_after.as_str())
|
||||
.unwrap_or(f.len());
|
||||
f.drain(0..index);
|
||||
}
|
||||
if let Some(limit) = request.limit {
|
||||
f.truncate(limit as usize);
|
||||
}
|
||||
Ok(f)
|
||||
self.list_table_dirs(
|
||||
request.start_after.as_deref(),
|
||||
request.limit.map(|limit| limit as usize),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
|
||||
if request.id.as_ref().map(|v| !v.is_empty()).unwrap_or(false) {
|
||||
return self.namespace_database().list_tables(request).await;
|
||||
}
|
||||
let mut f = self
|
||||
.object_store
|
||||
.read_dir(self.base_path.clone())
|
||||
.await?
|
||||
.iter()
|
||||
.map(Path::new)
|
||||
.filter(|path| {
|
||||
let is_lance = path
|
||||
.extension()
|
||||
.and_then(|e| e.to_str())
|
||||
.map(|e| e == LANCE_EXTENSION);
|
||||
is_lance.unwrap_or(false)
|
||||
})
|
||||
.filter_map(|p| p.file_stem().and_then(|s| s.to_str().map(String::from)))
|
||||
.collect::<Vec<String>>();
|
||||
f.sort();
|
||||
let limit = request.limit.map(|limit| limit as usize);
|
||||
// Reading one past the page is how we learn whether another page follows, without
|
||||
// a second request. The extra name is dropped before the response goes out.
|
||||
let mut tables = self
|
||||
.list_table_dirs(
|
||||
request.page_token.as_deref(),
|
||||
limit.map(|limit| limit.saturating_add(1)),
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Handle pagination with page_token
|
||||
if let Some(ref page_token) = request.page_token {
|
||||
let index = f
|
||||
.iter()
|
||||
.position(|name| name.as_str() > page_token.as_str())
|
||||
.unwrap_or(f.len());
|
||||
f.drain(0..index);
|
||||
}
|
||||
|
||||
// Determine if there's a next page
|
||||
let next_page_token = if let Some(limit) = request.limit {
|
||||
if f.len() > limit as usize {
|
||||
let token = f[limit as usize].clone();
|
||||
f.truncate(limit as usize);
|
||||
Some(token)
|
||||
} else {
|
||||
None
|
||||
let next_page_token = match limit {
|
||||
Some(limit) if tables.len() > limit => {
|
||||
tables.truncate(limit);
|
||||
tables.last().cloned()
|
||||
}
|
||||
} else {
|
||||
None
|
||||
_ => None,
|
||||
};
|
||||
|
||||
Ok(ListTablesResponse {
|
||||
tables: f,
|
||||
tables,
|
||||
page_token: next_page_token,
|
||||
})
|
||||
}
|
||||
@@ -1322,6 +1330,130 @@ mod tests {
|
||||
(tempdir, db)
|
||||
}
|
||||
|
||||
async fn create_tables(db: &ListingDatabase, names: &[&str]) {
|
||||
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
|
||||
for name in names {
|
||||
db.create_table(CreateTableRequest {
|
||||
name: name.to_string(),
|
||||
namespace_path: vec![],
|
||||
data: Box::new(RecordBatch::new_empty(schema.clone())) as Box<dyn Scannable>,
|
||||
mode: CreateTableMode::Create,
|
||||
write_options: Default::default(),
|
||||
location: None,
|
||||
namespace_client: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
/// Paging with the returned token has to visit every table exactly once. The token is
|
||||
/// the last name of the page, which is what `page_token` resumes after.
|
||||
#[tokio::test]
|
||||
async fn test_list_tables_pages_over_every_table_once() {
|
||||
let (_tempdir, db) = setup_database().await;
|
||||
create_tables(&db, &["a", "b", "c", "d", "e"]).await;
|
||||
|
||||
let mut seen = Vec::new();
|
||||
let mut page_token = None;
|
||||
loop {
|
||||
let page = db
|
||||
.list_tables(ListTablesRequest {
|
||||
limit: Some(2),
|
||||
page_token,
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
seen.extend(page.tables);
|
||||
match page.page_token {
|
||||
Some(token) => page_token = Some(token),
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(seen, vec!["a", "b", "c", "d", "e"]);
|
||||
}
|
||||
|
||||
/// The last page reports no token, so a caller paging by token knows to stop without
|
||||
/// asking for an empty page.
|
||||
#[tokio::test]
|
||||
async fn test_list_tables_exhausted_page_has_no_token() {
|
||||
let (_tempdir, db) = setup_database().await;
|
||||
create_tables(&db, &["a", "b"]).await;
|
||||
|
||||
let page = db
|
||||
.list_tables(ListTablesRequest {
|
||||
limit: Some(2),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(page.tables, vec!["a", "b"]);
|
||||
assert_eq!(page.page_token, None);
|
||||
}
|
||||
|
||||
/// Listing follows the order the object store lists directories in, so a name that
|
||||
/// extends another comes first: `-` sorts below the `.` of `.lance`. Pagination pushes
|
||||
/// its cursor into the list request, so it cannot report a different order than the
|
||||
/// one it resumes in.
|
||||
#[tokio::test]
|
||||
async fn test_listing_order_follows_the_store_not_the_table_name() {
|
||||
let (_tempdir, db) = setup_database().await;
|
||||
create_tables(&db, &["users", "users-archive", "users.old"]).await;
|
||||
|
||||
#[allow(deprecated)]
|
||||
let names = db.table_names(TableNamesRequest::default()).await.unwrap();
|
||||
assert_eq!(names, vec!["users-archive", "users", "users.old"]);
|
||||
|
||||
// Resuming after a name skips everything the store lists before it, which is what
|
||||
// paging by the previous page's last name relies on.
|
||||
#[allow(deprecated)]
|
||||
let after = db
|
||||
.table_names(TableNamesRequest {
|
||||
start_after: Some("users-archive".to_string()),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(after, vec!["users", "users.old"]);
|
||||
}
|
||||
|
||||
/// An empty `start_after` means "from the start". A name that sorts below `.lance` is
|
||||
/// what disappears if it is treated as a cursor instead.
|
||||
#[tokio::test]
|
||||
async fn test_empty_start_after_lists_from_the_start() {
|
||||
let (_tempdir, db) = setup_database().await;
|
||||
create_tables(&db, &["-dash", "alpha"]).await;
|
||||
|
||||
#[allow(deprecated)]
|
||||
let names = db
|
||||
.table_names(TableNamesRequest {
|
||||
start_after: Some(String::new()),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(names, vec!["-dash", "alpha"]);
|
||||
}
|
||||
|
||||
/// Only directories named `<name>.lance` are tables; loose files and other directories
|
||||
/// under the database prefix are not.
|
||||
#[tokio::test]
|
||||
async fn test_listing_ignores_non_table_children() {
|
||||
let (tempdir, db) = setup_database().await;
|
||||
create_tables(&db, &["real"]).await;
|
||||
std::fs::write(tempdir.path().join("loose.lance"), b"not a table").unwrap();
|
||||
create_dir_all(tempdir.path().join("scratch")).unwrap();
|
||||
|
||||
#[allow(deprecated)]
|
||||
let names = db.table_names(TableNamesRequest::default()).await.unwrap();
|
||||
|
||||
assert_eq!(names, vec!["real"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_listing_database_root_ops_do_not_create_manifest() {
|
||||
let tempdir = tempdir().unwrap();
|
||||
|
||||
+2
-219
@@ -511,9 +511,7 @@ pub trait QueryBase {
|
||||
|
||||
/// Rerank the results using the specified reranker.
|
||||
///
|
||||
/// For vector-only searches, the reranker receives the vector results and an
|
||||
/// empty full-text result set and must declare its output schema. Reranking
|
||||
/// multiple query vectors in one query is not supported.
|
||||
/// This is currently only supported for Hybrid Search.
|
||||
fn rerank(self, reranker: Arc<dyn Reranker>) -> Self;
|
||||
|
||||
/// The method to normalize the scores. Can be "rank" or "Score". If "Rank",
|
||||
@@ -1140,44 +1138,6 @@ pub struct VectorQuery {
|
||||
}
|
||||
|
||||
impl VectorQuery {
|
||||
fn check_vector_rerank_supported(&self) -> Result<()> {
|
||||
if self.request.query_vector.len() > 1 {
|
||||
return Err(Error::NotSupported {
|
||||
message: "reranking multiple query vectors is not supported; execute one query per vector"
|
||||
.to_string(),
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn vector_rerank_output_schema(&self) -> Result<SchemaRef> {
|
||||
self.check_vector_rerank_supported()?;
|
||||
|
||||
// Rerankers receive row IDs internally. Apply their schema transform to
|
||||
// that exact input and then hide the row ID from the declared public
|
||||
// schema unless it was explicitly requested.
|
||||
let vector_query = self.clone().with_row_id();
|
||||
let plan = vector_query
|
||||
.create_plan(QueryExecutionOptions::default())
|
||||
.await?;
|
||||
let reranker = self
|
||||
.request
|
||||
.base
|
||||
.reranker
|
||||
.as_ref()
|
||||
.expect("vector_rerank_output_schema requires a reranker");
|
||||
let input_schema = plan.schema();
|
||||
let output_schema = reranker.output_schema(&input_schema).await?;
|
||||
|
||||
if self.request.base.with_row_id {
|
||||
Ok(output_schema)
|
||||
} else {
|
||||
Ok(RecordBatch::new_empty(output_schema)
|
||||
.drop_column(ROW_ID)?
|
||||
.schema())
|
||||
}
|
||||
}
|
||||
|
||||
fn new(base: Query) -> Self {
|
||||
Self {
|
||||
parent: base.parent,
|
||||
@@ -1483,61 +1443,6 @@ impl VectorQuery {
|
||||
Ok(single_batch_stream(results, max_batch_length))
|
||||
}
|
||||
|
||||
async fn execute_vector_rerank(
|
||||
&self,
|
||||
options: QueryExecutionOptions,
|
||||
) -> Result<SendableRecordBatchStream> {
|
||||
self.check_vector_rerank_supported()?;
|
||||
|
||||
let max_batch_length = options.max_batch_length as usize;
|
||||
let internal_options = options.without_output_batch_length_limit();
|
||||
|
||||
// RRF needs row IDs to assign and preserve scores. Keep them internal unless
|
||||
// the caller explicitly requested them.
|
||||
let vector_query = self.clone().with_row_id();
|
||||
let vector_results = vector_query
|
||||
.inner_execute_with_options(internal_options)
|
||||
.await?;
|
||||
let schema = vector_results.schema();
|
||||
let vector_results = vector_results.try_collect::<Vec<_>>().await?;
|
||||
let vector_results = concat_batches(&schema, vector_results.iter())?;
|
||||
let vector_schema = vector_results.schema();
|
||||
let fts_results = RecordBatch::new_empty(vector_schema.clone());
|
||||
|
||||
let reranker = self
|
||||
.request
|
||||
.base
|
||||
.reranker
|
||||
.as_ref()
|
||||
.expect("execute_vector_rerank requires a reranker");
|
||||
let expected_schema = reranker.output_schema(&vector_schema).await?;
|
||||
let mut results = reranker
|
||||
.rerank_hybrid("", vector_results, fts_results)
|
||||
.await?;
|
||||
|
||||
check_reranker_result(&results)?;
|
||||
if results.schema() != expected_schema {
|
||||
return Err(Error::Schema {
|
||||
message: format!(
|
||||
"reranker returned schema {:?}, but declared {:?}",
|
||||
results.schema(),
|
||||
expected_schema
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
let limit = self.request.base.limit.unwrap_or(DEFAULT_TOP_K);
|
||||
if results.num_rows() > limit {
|
||||
results = results.slice(0, limit);
|
||||
}
|
||||
|
||||
if !self.request.base.with_row_id {
|
||||
results = results.drop_column(ROW_ID)?;
|
||||
}
|
||||
|
||||
Ok(single_batch_stream(results, max_batch_length))
|
||||
}
|
||||
|
||||
async fn inner_execute_with_options(
|
||||
&self,
|
||||
options: QueryExecutionOptions,
|
||||
@@ -1590,23 +1495,6 @@ impl ExecutableQuery for VectorQuery {
|
||||
return Ok(hybrid_result);
|
||||
}
|
||||
|
||||
if self.request.base.reranker.is_some() {
|
||||
let timeout = options.timeout;
|
||||
let mut rerank_options = options;
|
||||
// A single outer deadline covers planning, candidate collection,
|
||||
// schema declaration, and the complete reranker callback.
|
||||
rerank_options.timeout = None;
|
||||
let execution = self.execute_vector_rerank(rerank_options);
|
||||
return match timeout {
|
||||
Some(timeout) => tokio::time::timeout(timeout, execution)
|
||||
.await
|
||||
.map_err(|_| Error::Timeout {
|
||||
message: format!("Query timeout after {} ms", timeout.as_millis()),
|
||||
})?,
|
||||
None => execution.await,
|
||||
};
|
||||
}
|
||||
|
||||
self.inner_execute_with_options(options).await
|
||||
}
|
||||
|
||||
@@ -1619,15 +1507,6 @@ impl ExecutableQuery for VectorQuery {
|
||||
let query = AnyQuery::VectorQuery(self.request.clone());
|
||||
self.parent.analyze_plan(&query, options).await
|
||||
}
|
||||
|
||||
async fn output_schema(&self) -> Result<SchemaRef> {
|
||||
if self.request.base.full_text_search.is_none() && self.request.base.reranker.is_some() {
|
||||
self.vector_rerank_output_schema().await
|
||||
} else {
|
||||
let plan = self.create_plan(QueryExecutionOptions::default()).await?;
|
||||
Ok(plan.schema())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HasQuery for VectorQuery {
|
||||
@@ -1764,13 +1643,7 @@ impl ExecutableQuery for TakeQuery {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{
|
||||
collections::HashSet,
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
},
|
||||
};
|
||||
use std::{collections::HashSet, sync::Arc};
|
||||
|
||||
use super::*;
|
||||
use arrow::{array::downcast_array, compute::concat_batches, datatypes::Int32Type};
|
||||
@@ -1786,31 +1659,6 @@ mod tests {
|
||||
|
||||
use crate::{Table, connect, database::CreateTableMode, index::Index};
|
||||
|
||||
#[derive(Debug)]
|
||||
struct SlowReranker {
|
||||
invoked: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Reranker for SlowReranker {
|
||||
async fn output_schema(&self, input: &SchemaRef) -> Result<SchemaRef> {
|
||||
RRFReranker::default().output_schema(input).await
|
||||
}
|
||||
|
||||
async fn rerank_hybrid(
|
||||
&self,
|
||||
query: &str,
|
||||
vector_results: RecordBatch,
|
||||
fts_results: RecordBatch,
|
||||
) -> Result<RecordBatch> {
|
||||
self.invoked.store(true, Ordering::SeqCst);
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
RRFReranker::default()
|
||||
.rerank_hybrid(query, vector_results, fts_results)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_setters_getters() {
|
||||
// TODO: Switch back to memory://foo after https://github.com/lancedb/lancedb/issues/1051
|
||||
@@ -2522,71 +2370,6 @@ mod tests {
|
||||
// We don't guarantee order.
|
||||
assert!(query_index.values().contains(&0));
|
||||
assert!(query_index.values().contains(&1));
|
||||
|
||||
let reranked = query.rerank(Arc::new(RRFReranker::default()));
|
||||
let Err(execute_error) = reranked.execute().await else {
|
||||
panic!("multi-vector reranking should be rejected");
|
||||
};
|
||||
assert!(
|
||||
execute_error
|
||||
.to_string()
|
||||
.contains("reranking multiple query vectors is not supported")
|
||||
);
|
||||
let schema_error = reranked.output_schema().await.unwrap_err();
|
||||
assert!(
|
||||
schema_error
|
||||
.to_string()
|
||||
.contains("reranking multiple query vectors is not supported")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_vector_rerank_timeout_covers_reranker() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let table = make_test_table(&tmp_dir).await;
|
||||
let invoked = Arc::new(AtomicBool::new(false));
|
||||
let reranker = SlowReranker {
|
||||
invoked: invoked.clone(),
|
||||
};
|
||||
|
||||
let result = table
|
||||
.vector_search(&[0.1, 0.2, 0.3, 0.4])
|
||||
.unwrap()
|
||||
.limit(1)
|
||||
.rerank(Arc::new(reranker))
|
||||
.execute_with_options(QueryExecutionOptions {
|
||||
timeout: Some(Duration::from_secs(1)),
|
||||
..Default::default()
|
||||
})
|
||||
.await;
|
||||
|
||||
assert!(invoked.load(Ordering::SeqCst));
|
||||
assert!(matches!(result, Err(Error::Timeout { .. })));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_vector_rerank_output_schema_matches_execution() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let table = make_test_table(&tmp_dir).await;
|
||||
let query = table
|
||||
.vector_search(&[0.1, 0.2, 0.3, 0.4])
|
||||
.unwrap()
|
||||
.limit(1)
|
||||
.rerank(Arc::new(RRFReranker::default()));
|
||||
|
||||
let promised = query.output_schema().await.unwrap();
|
||||
let actual = query
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.next()
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.schema();
|
||||
|
||||
assert_eq!(promised, actual);
|
||||
assert!(promised.column_with_name("_relevance_score").is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
pub mod blobs;
|
||||
mod blobs;
|
||||
pub mod insert;
|
||||
|
||||
use self::insert::{RemoteWriteExec, WriteOp};
|
||||
@@ -2939,7 +2939,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
}
|
||||
|
||||
#[derive(Serialize, Clone, Debug)]
|
||||
pub(crate) struct MergeInsertRequest {
|
||||
pub struct MergeInsertRequest {
|
||||
on: String,
|
||||
when_matched_update_all: bool,
|
||||
when_matched_update_all_filt: Option<String>,
|
||||
@@ -3089,12 +3089,10 @@ mod tests {
|
||||
Box::pin(table.delete("false").map_ok(|_| ())),
|
||||
Box::pin(
|
||||
table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"x".into(),
|
||||
"y".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("x".into(), "y".into())]),
|
||||
None,
|
||||
)
|
||||
.map_ok(|_| ()),
|
||||
),
|
||||
Box::pin(async {
|
||||
@@ -4302,9 +4300,32 @@ mod tests {
|
||||
"fetch_blobs",
|
||||
);
|
||||
|
||||
let message = table
|
||||
.fetch_blob_files("image", &[1])
|
||||
.await
|
||||
.unwrap_err()
|
||||
.to_string();
|
||||
assert!(
|
||||
message.contains("fetch_blob_files is not supported on LanceDB Cloud"),
|
||||
"got: {message}"
|
||||
);
|
||||
assert!(
|
||||
!message.contains("Use fetch_blobs"),
|
||||
"old server must not be told to use fetch_blobs, got: {message}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_blob_files_point_at_fetch_blobs_on_a_blob_capable_server() {
|
||||
let table = Table::new_with_handler_version(
|
||||
"my_table",
|
||||
semver::Version::new(0, 5, 0),
|
||||
|_| -> http::Response<String> { panic!("fetch_blob_files must not reach the server") },
|
||||
);
|
||||
|
||||
assert_not_supported_error(
|
||||
table.fetch_blob_files("image", &[1]).await.unwrap_err(),
|
||||
"requires LanceDB Cloud server 0.5.0 or newer",
|
||||
"Use fetch_blobs for full bytes",
|
||||
);
|
||||
}
|
||||
|
||||
@@ -6390,12 +6411,13 @@ mod tests {
|
||||
});
|
||||
|
||||
let result = table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![
|
||||
("b".into(), "a + 1".into()),
|
||||
("x".into(), "cast(NULL as int32)".into()),
|
||||
]))
|
||||
.execute()
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![
|
||||
("b".into(), "a + 1".into()),
|
||||
("x".into(), "cast(NULL as int32)".into()),
|
||||
]),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -7120,12 +7142,10 @@ mod tests {
|
||||
}
|
||||
"add_columns" => {
|
||||
let _ = table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"c".into(),
|
||||
"a + 1".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("c".into(), "a + 1".into())]),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
"drop_columns" => {
|
||||
@@ -9883,12 +9903,10 @@ mod tests {
|
||||
.await
|
||||
.unwrap();
|
||||
branch
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"b".into(),
|
||||
"a + 1".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("b".into(), "a + 1".into())]),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
branch
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -33,7 +33,7 @@ use crate::table::{AddResult, MergeResult};
|
||||
/// same Arrow-IPC streaming body and error side-channel; only the target
|
||||
/// endpoint, query parameters, and parsed result type differ.
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) enum WriteOp {
|
||||
pub enum WriteOp {
|
||||
/// `add`: stream to `/v1/table/{id}/insert/`, optionally overwriting.
|
||||
Insert { overwrite: bool },
|
||||
/// `merge_insert`: stream to `/v1/table/{id}/merge_insert/` with the merge
|
||||
@@ -49,7 +49,7 @@ pub(crate) enum WriteOp {
|
||||
/// The parsed server response for a completed write, discriminated by the
|
||||
/// operation that produced it.
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) enum WriteResult {
|
||||
pub enum WriteResult {
|
||||
Add(AddResult),
|
||||
Merge(MergeResult),
|
||||
}
|
||||
|
||||
@@ -8,7 +8,6 @@ use arrow::{
|
||||
compute::{concat_batches, filter_record_batch},
|
||||
};
|
||||
use arrow_array::{BooleanArray, RecordBatch, UInt64Array};
|
||||
use arrow_schema::SchemaRef;
|
||||
use async_trait::async_trait;
|
||||
use lance::dataset::ROW_ID;
|
||||
|
||||
@@ -48,28 +47,16 @@ impl std::fmt::Display for NormalizeMethod {
|
||||
}
|
||||
}
|
||||
|
||||
/// Interface for a reranker. A reranker is used to rerank vector and hybrid
|
||||
/// search results. This is useful for combining results from multiple search
|
||||
/// methods or assigning a relevance score to vector search results.
|
||||
/// Interface for a reranker. A reranker is used to rerank the results from a
|
||||
/// vector and FTS search. This is useful for combining the results from both
|
||||
/// search methods.
|
||||
#[async_trait]
|
||||
pub trait Reranker: std::fmt::Debug + Sync + Send {
|
||||
/// Declare the schema returned by [`Self::rerank_hybrid`] for a vector-only
|
||||
/// query.
|
||||
///
|
||||
/// Vector reranking validates the returned batch against this schema so
|
||||
/// [`crate::query::ExecutableQuery::output_schema`] and execution cannot
|
||||
/// disagree. Rerankers that only support hybrid search do not need to
|
||||
/// implement this method.
|
||||
async fn output_schema(&self, _input: &SchemaRef) -> Result<SchemaRef> {
|
||||
Err(Error::NotSupported {
|
||||
message: "vector rerankers must declare their output schema".to_string(),
|
||||
})
|
||||
}
|
||||
// TODO support vector reranking and FTS reranking. Currently only hybrid reranking is supported.
|
||||
|
||||
/// Rerank function receives the individual results from the vector and FTS search
|
||||
/// results. You can choose to use any of the results to generate the final results,
|
||||
/// allowing maximum flexibility. For a vector-only search, `query` is empty and
|
||||
/// `fts_results` is an empty batch with the same schema as `vector_results`.
|
||||
/// allowing maximum flexibility.
|
||||
async fn rerank_hybrid(
|
||||
&self,
|
||||
query: &str,
|
||||
|
||||
@@ -9,7 +9,7 @@ use arrow::{
|
||||
compute::{sort_to_indices, take},
|
||||
};
|
||||
use arrow_array::{Float32Array, RecordBatch, UInt64Array};
|
||||
use arrow_schema::{DataType, Field, Schema, SchemaRef, SortOptions};
|
||||
use arrow_schema::{DataType, Field, Schema, SortOptions};
|
||||
use async_trait::async_trait;
|
||||
use lance::dataset::ROW_ID;
|
||||
|
||||
@@ -44,19 +44,6 @@ impl Default for RRFReranker {
|
||||
|
||||
#[async_trait]
|
||||
impl Reranker for RRFReranker {
|
||||
async fn output_schema(&self, input: &SchemaRef) -> Result<SchemaRef> {
|
||||
let mut fields = input.fields().to_vec();
|
||||
fields.push(Arc::new(Field::new(
|
||||
RELEVANCE_SCORE,
|
||||
DataType::Float32,
|
||||
false,
|
||||
)));
|
||||
Ok(Arc::new(Schema::new_with_metadata(
|
||||
fields,
|
||||
input.metadata().clone(),
|
||||
)))
|
||||
}
|
||||
|
||||
async fn rerank_hybrid(
|
||||
&self,
|
||||
_query: &str,
|
||||
@@ -148,9 +135,15 @@ impl Reranker for RRFReranker {
|
||||
.collect();
|
||||
|
||||
// add relevance score to schema
|
||||
let schema = self.output_schema(&combined_results.schema()).await?;
|
||||
let mut fields = combined_results.schema().fields().to_vec();
|
||||
fields.push(Arc::new(Field::new(
|
||||
RELEVANCE_SCORE,
|
||||
DataType::Float32,
|
||||
false,
|
||||
)));
|
||||
let schema = Schema::new(fields);
|
||||
|
||||
let combined_results = RecordBatch::try_new(schema, columns)?;
|
||||
let combined_results = RecordBatch::try_new(Arc::new(schema), columns)?;
|
||||
|
||||
Ok(combined_results)
|
||||
}
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
|
||||
//! LanceDB Table APIs
|
||||
|
||||
use crate::blob::BlobFile;
|
||||
use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader};
|
||||
use arrow_schema::{Schema, SchemaRef};
|
||||
use async_trait::async_trait;
|
||||
@@ -13,6 +12,7 @@ use datafusion_physical_plan::ExecutionPlan;
|
||||
use datafusion_physical_plan::display::DisplayableExecutionPlan;
|
||||
use futures::StreamExt;
|
||||
use futures::stream::FuturesUnordered;
|
||||
use lance::dataset::BlobFile;
|
||||
pub use lance::dataset::ColumnAlteration;
|
||||
pub use lance::dataset::NewColumnTransform;
|
||||
pub use lance::dataset::ReadParams;
|
||||
@@ -65,7 +65,6 @@ use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path};
|
||||
use self::dataset::DatasetConsistencyWrapper;
|
||||
use self::merge::MergeInsertBuilder;
|
||||
|
||||
pub mod add_columns;
|
||||
mod add_data;
|
||||
pub mod branch_merge;
|
||||
mod create_index;
|
||||
@@ -80,7 +79,6 @@ pub mod schema_evolution;
|
||||
pub mod update;
|
||||
pub mod write_progress;
|
||||
use crate::index::waiter::wait_for_index;
|
||||
pub use add_columns::AddColumnsBuilder;
|
||||
#[cfg(feature = "remote")]
|
||||
pub(crate) use add_data::PreprocessingOutput;
|
||||
pub use add_data::{AddDataBuilder, AddDataMode, AddResult, NaNVectorBehavior};
|
||||
@@ -1622,8 +1620,12 @@ impl Table {
|
||||
}
|
||||
|
||||
/// Add new columns to the table, providing values to fill in.
|
||||
pub fn add_columns(&self) -> AddColumnsBuilder {
|
||||
AddColumnsBuilder::new(self.inner.clone())
|
||||
pub async fn add_columns(
|
||||
&self,
|
||||
transforms: NewColumnTransform,
|
||||
read_columns: Option<Vec<String>>,
|
||||
) -> Result<AddColumnsResult> {
|
||||
self.inner.add_columns(transforms, read_columns).await
|
||||
}
|
||||
|
||||
/// Change a column's name or nullability.
|
||||
|
||||
@@ -1,161 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! Builder for adding columns to a table.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use lance::dataset::NewColumnTransform;
|
||||
|
||||
use super::BaseTable;
|
||||
use super::schema_evolution::AddColumnsResult;
|
||||
use crate::{Error, Result};
|
||||
|
||||
/// Adds columns to a table. See [`Table::add_columns`](super::Table::add_columns).
|
||||
pub struct AddColumnsBuilder {
|
||||
parent: Arc<dyn BaseTable>,
|
||||
transform: Option<NewColumnTransform>,
|
||||
read_columns: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for AddColumnsBuilder {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("AddColumnsBuilder")
|
||||
.field("parent", &self.parent)
|
||||
.field("has_transform", &self.transform.is_some())
|
||||
.field("read_columns", &self.read_columns)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl AddColumnsBuilder {
|
||||
pub(crate) fn new(parent: Arc<dyn BaseTable>) -> Self {
|
||||
Self {
|
||||
parent,
|
||||
transform: None,
|
||||
read_columns: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set how the new columns' values are produced. Required.
|
||||
pub fn transform(mut self, transform: NewColumnTransform) -> Self {
|
||||
self.transform = Some(transform);
|
||||
self
|
||||
}
|
||||
|
||||
/// Limit which existing columns a [`NewColumnTransform::BatchUDF`] mapper
|
||||
/// receives. Every other transform determines what it reads, so setting
|
||||
/// this alongside one is an error rather than a silent no-op.
|
||||
pub fn read_columns(mut self, columns: impl IntoIterator<Item = impl Into<String>>) -> Self {
|
||||
self.read_columns = Some(columns.into_iter().map(Into::into).collect());
|
||||
self
|
||||
}
|
||||
|
||||
/// Add the columns.
|
||||
pub async fn execute(self) -> Result<AddColumnsResult> {
|
||||
let Self {
|
||||
parent,
|
||||
transform,
|
||||
read_columns,
|
||||
} = self;
|
||||
|
||||
let Some(transform) = transform else {
|
||||
return Err(Error::InvalidInput {
|
||||
message: "add_columns requires a transform".into(),
|
||||
});
|
||||
};
|
||||
|
||||
if read_columns.is_some() && !matches!(transform, NewColumnTransform::BatchUDF(_)) {
|
||||
return Err(Error::InvalidInput {
|
||||
message: "read_columns applies only to a BatchUDF transform; \
|
||||
every other transform determines what it reads"
|
||||
.into(),
|
||||
});
|
||||
}
|
||||
|
||||
parent.add_columns(transform, read_columns).await
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::{Int32Array, RecordBatch, record_batch};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use lance::dataset::{BatchUDF, NewColumnTransform};
|
||||
|
||||
use crate::Table;
|
||||
use crate::connect;
|
||||
|
||||
async fn table_with_two_columns(name: &str) -> Table {
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
let batch = record_batch!(("x", Int32, [1, 2, 3]), ("y", Int32, [10, 20, 30])).unwrap();
|
||||
conn.create_table(name, batch).execute().await.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_requires_a_transform() {
|
||||
let table = table_with_two_columns("no_transform").await;
|
||||
let err = table.add_columns().execute().await.unwrap_err();
|
||||
assert!(
|
||||
err.to_string().contains("requires a transform"),
|
||||
"got: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_columns_with_sql_expressions_is_rejected() {
|
||||
let table = table_with_two_columns("read_cols_sql").await;
|
||||
let err = table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"doubled".into(),
|
||||
"x * 2".into(),
|
||||
)]))
|
||||
.read_columns(["x"])
|
||||
.execute()
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(err.to_string().contains("BatchUDF"), "got: {err}");
|
||||
|
||||
let schema = table.schema().await.unwrap();
|
||||
assert!(
|
||||
schema.field_with_name("doubled").is_err(),
|
||||
"a rejected call must not commit"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_columns_limits_what_a_batch_udf_sees() {
|
||||
let table = table_with_two_columns("read_cols_udf").await;
|
||||
|
||||
let output_schema = Arc::new(Schema::new(vec![Field::new("sum", DataType::Int32, true)]));
|
||||
let mapper_schema = output_schema.clone();
|
||||
let udf = BatchUDF {
|
||||
mapper: Box::new(move |batch: &RecordBatch| {
|
||||
assert!(batch.column_by_name("x").is_some());
|
||||
assert!(batch.column_by_name("y").is_none(), "y was not requested");
|
||||
let x = batch["x"].as_any().downcast_ref::<Int32Array>().unwrap();
|
||||
let doubled: Int32Array = x.iter().map(|v| v.map(|v| v * 2)).collect();
|
||||
Ok(RecordBatch::try_new(
|
||||
mapper_schema.clone(),
|
||||
vec![Arc::new(doubled)],
|
||||
)?)
|
||||
}),
|
||||
output_schema,
|
||||
result_checkpoint: None,
|
||||
};
|
||||
|
||||
table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::BatchUDF(udf))
|
||||
.read_columns(["x"])
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let schema = table.schema().await.unwrap();
|
||||
assert!(schema.field_with_name("sum").is_ok());
|
||||
}
|
||||
}
|
||||
@@ -576,12 +576,10 @@ mod tests {
|
||||
|
||||
// Add a new physical column AFTER the embedding column.
|
||||
table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"score".into(),
|
||||
"42.0".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("score".into(), "42.0".into())]),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -685,9 +683,7 @@ mod tests {
|
||||
true,
|
||||
)]));
|
||||
table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::AllNulls(nested_schema))
|
||||
.execute()
|
||||
.add_columns(NewColumnTransform::AllNulls(nested_schema), None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
|
||||
@@ -193,12 +193,10 @@ mod tests {
|
||||
|
||||
// Add a computed column
|
||||
let result = table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"doubled".into(),
|
||||
"id * 2".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("doubled".into(), "id * 2".into())]),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -253,12 +251,13 @@ mod tests {
|
||||
|
||||
// Add multiple columns at once
|
||||
table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![
|
||||
("y".into(), "x + 1".into()),
|
||||
("z".into(), "x * x".into()),
|
||||
]))
|
||||
.execute()
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![
|
||||
("y".into(), "x + 1".into()),
|
||||
("z".into(), "x * x".into()),
|
||||
]),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -284,12 +283,10 @@ mod tests {
|
||||
|
||||
// Add a column with a constant value
|
||||
table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"constant".into(),
|
||||
"42".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("constant".into(), "42".into())]),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -662,12 +659,10 @@ mod tests {
|
||||
|
||||
// Add column increments version
|
||||
let add_result = table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"c".into(),
|
||||
"a + b".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("c".into(), "a + b".into())]),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(add_result.version > v1);
|
||||
|
||||
@@ -9,17 +9,14 @@ use arrow_array::{
|
||||
};
|
||||
use arrow_schema::{DataType, Field, Fields, Schema};
|
||||
use futures::TryStreamExt;
|
||||
use lance::Dataset;
|
||||
use lance_encoding::version::LanceFileVersion;
|
||||
use lancedb::{
|
||||
Connection, Error, Result, Table,
|
||||
blob::{BlobRangeRequest, blob},
|
||||
connect, connect_namespace,
|
||||
database::listing::{
|
||||
ListingDatabaseOptions, NewTableConfig, OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
|
||||
},
|
||||
database::listing::OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
|
||||
query::{ExecutableQuery, QueryBase},
|
||||
table::{AddDataMode, CompactionOptions, OptimizeAction, OptimizeStats},
|
||||
table::{AddDataMode, CompactionOptions, OptimizeAction},
|
||||
};
|
||||
use tempfile::tempdir;
|
||||
|
||||
@@ -1078,252 +1075,3 @@ async fn fetch_blob_files_aligns_across_fragments_with_nulls_and_dups() -> Resul
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Rows exercising the null/empty interleavings from
|
||||
/// <https://github.com/lancedb/lancedb/issues/3744>: a payload, a null, a valid
|
||||
/// empty value, then payloads whose descriptors a fragment rewrite used to zero.
|
||||
fn null_empty_input_batch() -> RecordBatch {
|
||||
let owned = [
|
||||
Some(dedicated_blob_bytes(1)),
|
||||
None,
|
||||
Some(Vec::new()),
|
||||
Some(dedicated_blob_bytes(4)),
|
||||
Some(dedicated_blob_bytes(5)),
|
||||
Some(dedicated_blob_bytes(6)),
|
||||
];
|
||||
let payloads: Vec<Option<&[u8]>> = owned.iter().map(|payload| payload.as_deref()).collect();
|
||||
binary_input_batch(&[1, 2, 3, 4, 5, 6], &payloads)
|
||||
}
|
||||
|
||||
/// One `(id, Some((payload length, first byte)))` per live row, or `(id, None)`
|
||||
/// for a null blob. Comparing lengths and first bytes keeps failure output
|
||||
/// readable where comparing whole payloads would not.
|
||||
type BlobSummary = Vec<(i64, Option<(usize, Option<u8>)>)>;
|
||||
|
||||
/// The rows [`null_empty_input_batch`] leaves behind after `id IN (1, 4)` is
|
||||
/// deleted: a null, a valid empty value, and the two payloads that follow them.
|
||||
fn expected_null_empty_survivors() -> BlobSummary {
|
||||
vec![
|
||||
(2, None),
|
||||
(3, Some((0, None))),
|
||||
(5, Some((DEDICATED_BLOB_LEN, Some(5)))),
|
||||
(6, Some((DEDICATED_BLOB_LEN, Some(6)))),
|
||||
]
|
||||
}
|
||||
|
||||
/// `optimize()` only rewrites a fragment when lance's compaction planner selects
|
||||
/// it — here because the delete pushes the fragment past
|
||||
/// `materialize_deletions_threshold` (0.1 by default; these tests delete 2 of 6
|
||||
/// rows). Without this check, a planner or threshold change upstream would leave
|
||||
/// both regression tests green while no rewrite happened at all.
|
||||
fn assert_compacted(stats: &OptimizeStats) {
|
||||
let metrics = stats
|
||||
.compaction
|
||||
.as_ref()
|
||||
.expect("OptimizeAction::All runs compaction");
|
||||
assert!(
|
||||
metrics.fragments_removed >= 1,
|
||||
"optimize() rewrote no fragment, so this test proves nothing: {metrics:?}"
|
||||
);
|
||||
}
|
||||
|
||||
fn summarize(rows: &[(i64, Option<Vec<u8>>)]) -> BlobSummary {
|
||||
rows.iter()
|
||||
.map(|(id, payload)| {
|
||||
(
|
||||
*id,
|
||||
payload
|
||||
.as_ref()
|
||||
.map(|bytes| (bytes.len(), bytes.first().copied())),
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn sorted_id_rowid(table: &Table) -> Result<Vec<(i64, u64)>> {
|
||||
let mut pairs = collect_id_rowid(table).await?;
|
||||
pairs.sort_by_key(|(id, _)| *id);
|
||||
Ok(pairs)
|
||||
}
|
||||
|
||||
/// `{position, size}` descriptors of a legacy v1 blob column, keyed by `id`.
|
||||
async fn v1_blob_descriptors(table: &Table) -> Result<Vec<(i64, Option<(u64, u64)>)>> {
|
||||
let batches = table
|
||||
.query()
|
||||
.execute()
|
||||
.await?
|
||||
.try_collect::<Vec<_>>()
|
||||
.await?;
|
||||
let batch = arrow_select::concat::concat_batches(&batches[0].schema(), &batches).unwrap();
|
||||
let ids = batch
|
||||
.column_by_name("id")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<Int64Array>()
|
||||
.unwrap();
|
||||
let descriptors = batch
|
||||
.column_by_name("image")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<StructArray>()
|
||||
.expect("v1 blob column reads back as a descriptor struct");
|
||||
let position = descriptors
|
||||
.column_by_name("position")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<UInt64Array>()
|
||||
.unwrap();
|
||||
let size = descriptors
|
||||
.column_by_name("size")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<UInt64Array>()
|
||||
.unwrap();
|
||||
let mut rows: Vec<(i64, Option<(u64, u64)>)> = (0..batch.num_rows())
|
||||
.map(|row| {
|
||||
let descriptor =
|
||||
(!descriptors.is_null(row)).then(|| (position.value(row), size.value(row)));
|
||||
(ids.value(row), descriptor)
|
||||
})
|
||||
.collect();
|
||||
rows.sort_by_key(|(id, _)| *id);
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
/// Payload bytes of every live row of a legacy v1 blob column, keyed by `id`.
|
||||
/// [`Table::fetch_blobs`] rejects v1 columns, so read them through lance.
|
||||
async fn v1_blob_payloads(dataset_uri: &str, table: &Table) -> Result<Vec<(i64, Option<Vec<u8>>)>> {
|
||||
let pairs = sorted_id_rowid(table).await?;
|
||||
let row_ids: Vec<u64> = pairs.iter().map(|(_, row_id)| *row_id).collect();
|
||||
let dataset = Arc::new(Dataset::open(dataset_uri).await?);
|
||||
let files = dataset.take_blobs(&row_ids, "image").await?;
|
||||
assert_eq!(
|
||||
files.len(),
|
||||
pairs.len(),
|
||||
"take_blobs returned {} handles for {} live rows",
|
||||
files.len(),
|
||||
pairs.len()
|
||||
);
|
||||
let mut rows = Vec::with_capacity(pairs.len());
|
||||
for ((id, _), file) in pairs.iter().zip(files) {
|
||||
let payload = match file {
|
||||
Some(file) => Some(file.read().await?.to_vec()),
|
||||
None => None,
|
||||
};
|
||||
rows.push((*id, payload));
|
||||
}
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
/// Length and first byte of every live blob v2 value, keyed by `id`.
|
||||
async fn blob_v2_values(table: &Table) -> Result<BlobSummary> {
|
||||
let pairs = sorted_id_rowid(table).await?;
|
||||
let row_ids: Vec<u64> = pairs.iter().map(|(_, row_id)| *row_id).collect();
|
||||
let bytes = table.fetch_blobs("image", &row_ids).await?;
|
||||
Ok(pairs
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(slot, (id, _))| {
|
||||
let value = (!bytes.is_null(slot))
|
||||
.then(|| (bytes.value(slot).len(), bytes.value(slot).first().copied()));
|
||||
(*id, value)
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Regression test for [#3744]: on storage 2.0 (legacy v1 descriptors),
|
||||
/// compaction rewrote every payload following a null or empty value in the same
|
||||
/// fragment as `{position: 0, size: 0}`, so the payload bytes read back as `b""`
|
||||
/// and the new fragment no longer referenced them at all.
|
||||
///
|
||||
/// [#3744]: https://github.com/lancedb/lancedb/issues/3744
|
||||
#[tokio::test]
|
||||
async fn optimize_preserves_v1_blob_payloads_with_null_and_empty() -> Result<()> {
|
||||
let tmp = tempdir().unwrap();
|
||||
let db_uri = tmp.path().to_str().unwrap().to_string();
|
||||
let db = connect(&db_uri)
|
||||
.database_options(&ListingDatabaseOptions {
|
||||
new_table_config: NewTableConfig {
|
||||
data_storage_version: Some(LanceFileVersion::V2_0),
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
})
|
||||
.execute()
|
||||
.await?;
|
||||
let legacy = Field::new("image", DataType::LargeBinary, true).with_metadata(
|
||||
std::collections::HashMap::from([("lance-encoding:blob".to_string(), "true".to_string())]),
|
||||
);
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int64, false),
|
||||
legacy,
|
||||
]));
|
||||
let table = db.create_empty_table("t", schema).execute().await?;
|
||||
table.add(null_empty_input_batch()).execute().await?;
|
||||
assert_eq!(
|
||||
storage_format_version(&table).await,
|
||||
LanceFileVersion::V2_0.resolve(),
|
||||
"v1 blob descriptors only exist below storage 2.2"
|
||||
);
|
||||
let dataset_uri = table.uri().await?;
|
||||
|
||||
// Any rewrite triggers it; deleting rows is the shape from the issue.
|
||||
table.delete("id IN (1, 4)").await?;
|
||||
let descriptors_before = v1_blob_descriptors(&table).await?;
|
||||
let before = v1_blob_payloads(&dataset_uri, &table).await?;
|
||||
assert_eq!(
|
||||
summarize(&before),
|
||||
expected_null_empty_survivors(),
|
||||
"test setup no longer produces the null/empty/payload mix"
|
||||
);
|
||||
|
||||
let stats = table.optimize(OptimizeAction::All).await?;
|
||||
assert_compacted(&stats);
|
||||
|
||||
let descriptors_after = v1_blob_descriptors(&table).await?;
|
||||
let after = v1_blob_payloads(&dataset_uri, &table).await?;
|
||||
assert_eq!(
|
||||
summarize(&after),
|
||||
summarize(&before),
|
||||
"optimize() lost blob payloads; descriptors before={descriptors_before:?} after={descriptors_after:?}"
|
||||
);
|
||||
assert!(after == before, "optimize() changed blob payload bytes");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Regression test for the blob v2 half of [#3744]: compaction rewrote a valid
|
||||
/// empty value as null, destroying the null-vs-empty distinction.
|
||||
///
|
||||
/// [#3744]: https://github.com/lancedb/lancedb/issues/3744
|
||||
#[tokio::test]
|
||||
async fn optimize_preserves_blob_v2_null_and_empty_distinction() -> Result<()> {
|
||||
let tmp = tempdir().unwrap();
|
||||
let db = connect(tmp.path().to_str().unwrap()).execute().await?;
|
||||
let table = db
|
||||
.create_empty_table("t", blob_table_schema())
|
||||
.execute()
|
||||
.await?;
|
||||
table.add(null_empty_input_batch()).execute().await?;
|
||||
assert!(
|
||||
storage_format_version(&table).await >= LanceFileVersion::V2_2,
|
||||
"blob v2 columns require storage >= 2.2"
|
||||
);
|
||||
|
||||
table.delete("id IN (1, 4)").await?;
|
||||
let before = blob_v2_values(&table).await?;
|
||||
assert_eq!(
|
||||
before,
|
||||
expected_null_empty_survivors(),
|
||||
"test setup no longer produces the null/empty/payload mix"
|
||||
);
|
||||
|
||||
let stats = table.optimize(OptimizeAction::All).await?;
|
||||
assert_compacted(&stats);
|
||||
|
||||
assert_eq!(
|
||||
blob_v2_values(&table).await?,
|
||||
before,
|
||||
"optimize() changed blob v2 values"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user