mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-04 12:38:38 +00:00
Merge remote-tracking branch 'refs/remotes/origin/main' into gatekeeper/fix-2820-1
# Conflicts: # rust/lancedb/src/table/query.rs
This commit is contained in:
@@ -1292,6 +1292,18 @@ abstract updateFieldMetadata(updates): Promise<UpdateFieldMetadataResult>
|
||||
|
||||
Update per-field (column) metadata.
|
||||
|
||||
The following keys are treated specially, by convention, and should be
|
||||
used when appropriate:
|
||||
|
||||
- `lancedb:description`: for a human-readable description of a field.
|
||||
- `lancedb:tag:<name>`: for a user-defined key-value tag, where the suffix
|
||||
names the tag category; e.g. `lancedb:tag:model: "clip"`.
|
||||
- `lancedb:logical-column`: for a column grouping; e.g. `feature_v1` and
|
||||
`feature_v2` might be in the same logical column.
|
||||
- `lancedb:status`: for status options (`production`, `candidate`,
|
||||
`deprecated`, `archived`) to designate the current life cycle state of
|
||||
this column.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **updates**: [`FieldMetadataUpdate`](../interfaces/FieldMetadataUpdate.md)[]
|
||||
|
||||
@@ -17,7 +17,8 @@ metadata: Record<string, null | string>;
|
||||
```
|
||||
|
||||
Metadata key/value pairs. Merged into the field's existing metadata by
|
||||
default; a value of `null` deletes that key.
|
||||
default; a value of `null` deletes that key. See
|
||||
[Table.updateFieldMetadata](../classes/Table.md#updatefieldmetadata) for the conventional `lancedb:*` keys.
|
||||
|
||||
***
|
||||
|
||||
|
||||
@@ -170,7 +170,7 @@ test("basic table examples", async () => {
|
||||
// --8<-- [end:create_index]
|
||||
|
||||
// --8<-- [start:delete_rows]
|
||||
await tbl.delete('item = "fizz"');
|
||||
await tbl.delete("item = 'fizz'");
|
||||
// --8<-- [end:delete_rows]
|
||||
|
||||
// --8<-- [start:drop_table]
|
||||
|
||||
@@ -727,11 +727,11 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* Add a query vector to the search
|
||||
*
|
||||
* This method can be called multiple times to add multiple query vectors
|
||||
* to the search. If multiple query vectors are added, then they will be searched
|
||||
* in parallel, and the results will be concatenated. A column called `query_index`
|
||||
* will be added to indicate the index of the query vector that produced the result.
|
||||
*
|
||||
* Performance wise, this is equivalent to running multiple queries concurrently.
|
||||
* to the search. A column called `query_index` will be added to indicate the index
|
||||
* of the query vector that produced the result. Flat searches share one table scan
|
||||
* across the query vectors, avoiding the scan and memory amplification of running
|
||||
* multiple queries concurrently. Indexed searches may still perform per-vector
|
||||
* index work.
|
||||
*/
|
||||
addQueryVector(vector: IntoVector): VectorQuery {
|
||||
if (vector instanceof Promise) {
|
||||
|
||||
+14
-1
@@ -630,6 +630,18 @@ export abstract class Table {
|
||||
|
||||
/**
|
||||
* Update per-field (column) metadata.
|
||||
*
|
||||
* The following keys are treated specially, by convention, and should be
|
||||
* used when appropriate:
|
||||
*
|
||||
* - `lancedb:description`: for a human-readable description of a field.
|
||||
* - `lancedb:tag:<name>`: for a user-defined key-value tag, where the suffix
|
||||
* names the tag category; e.g. `lancedb:tag:model: "clip"`.
|
||||
* - `lancedb:logical-column`: for a column grouping; e.g. `feature_v1` and
|
||||
* `feature_v2` might be in the same logical column.
|
||||
* - `lancedb:status`: for status options (`production`, `candidate`,
|
||||
* `deprecated`, `archived`) to designate the current life cycle state of
|
||||
* this column.
|
||||
* @param {FieldMetadataUpdate[]} updates One or more per-field updates. Each
|
||||
* update's metadata is merged into the field's existing metadata by default;
|
||||
* a value of `null` deletes that key, and `replace: true` swaps the whole map.
|
||||
@@ -1555,7 +1567,8 @@ export interface FieldMetadataUpdate {
|
||||
path: string;
|
||||
/**
|
||||
* Metadata key/value pairs. Merged into the field's existing metadata by
|
||||
* default; a value of `null` deletes that key.
|
||||
* default; a value of `null` deletes that key. See
|
||||
* {@link Table.updateFieldMetadata} for the conventional `lancedb:*` keys.
|
||||
*/
|
||||
metadata: Record<string, string | null>;
|
||||
/** If true, replace the field's entire metadata map instead of merging. */
|
||||
|
||||
@@ -3407,9 +3407,10 @@ class AsyncQuery(AsyncStandardQuery):
|
||||
pass in multiple vectors. When multiple vectors are passed in, if the vector
|
||||
column is with multivector type, then the vectors will be treated as a single
|
||||
query. Or the vectors will be treated as multiple queries, this can be useful
|
||||
if you want to find the nearest vectors to multiple query vectors.
|
||||
This is not expected to be faster than making multiple queries concurrently;
|
||||
it is just a convenience method. If multiple vectors are passed in then
|
||||
if you want to find the nearest vectors to multiple query vectors. Flat
|
||||
searches share one table scan across the query vectors, avoiding the scan
|
||||
and memory amplification of making multiple queries concurrently. If
|
||||
multiple vectors are passed in then
|
||||
an additional column `query_index` will be added to the results. This column
|
||||
will contain the index of the query vector that the result is nearest to.
|
||||
"""
|
||||
@@ -3538,8 +3539,8 @@ class AsyncFTSQuery(AsyncStandardQuery):
|
||||
|
||||
Typically, a single vector is passed in as the query. However, you can also
|
||||
pass in multiple vectors. This can be useful if you want to find the nearest
|
||||
vectors to multiple query vectors. This is not expected to be faster than
|
||||
making multiple queries concurrently; it is just a convenience method.
|
||||
vectors to multiple query vectors. Flat searches share one table scan across
|
||||
the query vectors instead of issuing concurrent full scans.
|
||||
If multiple vectors are passed in then an additional column `query_index`
|
||||
will be added to the results. This column will contain the index of the
|
||||
query vector that the result is nearest to.
|
||||
|
||||
@@ -2127,12 +2127,25 @@ class Table(ABC):
|
||||
----------
|
||||
updates : dict
|
||||
One or more dicts, each with:
|
||||
|
||||
- "path": str — dot-path to the field (e.g. "embedding" or "a.b.c").
|
||||
- "metadata": dict[str, str | None] — keys to set; a value of ``None``
|
||||
deletes that key.
|
||||
- "replace": bool, optional — replace the field's whole metadata map
|
||||
instead of merging (default False).
|
||||
|
||||
The following keys are treated specially, by convention, and should
|
||||
be used when appropriate:
|
||||
|
||||
- "lancedb:description": for a human-readable description of a field.
|
||||
- ``"lancedb:tag:<name>"`` for a user-defined key-value tag, where the
|
||||
suffix names the tag category; e.g. "lancedb:tag:model": "clip".
|
||||
- "lancedb:logical-column" for a column grouping; e.g. "feature_v1"
|
||||
and "feature_v2" might be in the same logical column.
|
||||
- "lancedb:status" for status options ("production", "candidate",
|
||||
"deprecated", "archived") to designate the current life cycle
|
||||
state of this column.
|
||||
|
||||
Returns
|
||||
-------
|
||||
UpdateFieldMetadataResult
|
||||
|
||||
@@ -105,7 +105,7 @@ def test_quickstart(tmp_path):
|
||||
tbl.create_index(num_sub_vectors=1)
|
||||
# --8<-- [end:create_index]
|
||||
# --8<-- [start:delete_rows]
|
||||
tbl.delete('item = "fizz"')
|
||||
tbl.delete("item = 'fizz'")
|
||||
# --8<-- [end:delete_rows]
|
||||
# --8<-- [start:drop_table]
|
||||
db.drop_table("my_table")
|
||||
@@ -201,7 +201,7 @@ async def test_quickstart_async(tmp_path):
|
||||
await tbl.create_index("vector")
|
||||
# --8<-- [end:create_index_async]
|
||||
# --8<-- [start:delete_rows_async]
|
||||
await tbl.delete('item = "fizz"')
|
||||
await tbl.delete("item = 'fizz'")
|
||||
# --8<-- [end:delete_rows_async]
|
||||
# --8<-- [start:drop_table_async]
|
||||
await db.drop_table("my_table_async")
|
||||
|
||||
@@ -266,7 +266,7 @@ def test_table():
|
||||
tbl.add(pydantic_model_items)
|
||||
# --8<-- [end:add_table_from_pydantic]
|
||||
# --8<-- [start:delete_row]
|
||||
tbl.delete('item = "fizz"')
|
||||
tbl.delete("item = 'fizz'")
|
||||
# --8<-- [end:delete_row]
|
||||
# --8<-- [start:delete_specific_row]
|
||||
data = [
|
||||
@@ -538,7 +538,7 @@ async def test_table_async():
|
||||
await async_tbl.add(pydantic_model_items)
|
||||
# --8<-- [end:add_table_async_from_pydantic]
|
||||
# --8<-- [start:delete_row_async]
|
||||
await async_tbl.delete('item = "fizz"')
|
||||
await async_tbl.delete("item = 'fizz'")
|
||||
# --8<-- [end:delete_row_async]
|
||||
# --8<-- [start:delete_specific_row_async]
|
||||
data = [
|
||||
|
||||
@@ -897,6 +897,23 @@ def test_query_builder_batches(table):
|
||||
assert rs_list["id"][1] == 2
|
||||
|
||||
|
||||
def test_batch_vector_query_shares_filtered_flat_scan(table):
|
||||
query = (
|
||||
table.search([[1.0, 2.0], [3.0, 4.0]])
|
||||
.where("id > 0", prefilter=True)
|
||||
.limit(1)
|
||||
.select(["id"])
|
||||
)
|
||||
|
||||
plan = query.explain_plan(verbose=True)
|
||||
assert "KNNVectorDistance: queries=2" in plan
|
||||
assert "UnionExec" not in plan
|
||||
|
||||
results = query.to_arrow()
|
||||
assert len(results) == 2
|
||||
assert results["query_index"].to_pylist() == [0, 1]
|
||||
|
||||
|
||||
def test_dynamic_projection(table):
|
||||
rs = (
|
||||
LanceVectorQueryBuilder(table, [0, 0], "vector")
|
||||
|
||||
@@ -19,6 +19,7 @@
|
||||
|
||||
mod sql;
|
||||
|
||||
pub(crate) use sql::canonicalize_sql_predicate;
|
||||
pub use sql::expr_to_sql_string;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
@@ -1,10 +1,16 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::any::TypeId;
|
||||
|
||||
use datafusion_common::ScalarValue;
|
||||
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
|
||||
use datafusion_expr::Expr;
|
||||
use datafusion_sql::unparser::{self, dialect::Dialect};
|
||||
use datafusion_sql::sqlparser::{
|
||||
dialect::{Dialect as SqlParserDialect, GenericDialect},
|
||||
tokenizer::{Token, Tokenizer},
|
||||
};
|
||||
use datafusion_sql::unparser::{self, dialect::Dialect as UnparserDialect};
|
||||
|
||||
/// Unparser dialect that matches the quoting style expected by the Lance SQL
|
||||
/// parser. Lance uses backtick (`` ` ``) as the only delimited-identifier
|
||||
@@ -19,7 +25,7 @@ use datafusion_sql::unparser::{self, dialect::Dialect};
|
||||
/// lower-case by the SQL parser, which would break case-sensitive schemas).
|
||||
struct LanceSqlDialect;
|
||||
|
||||
impl Dialect for LanceSqlDialect {
|
||||
impl UnparserDialect for LanceSqlDialect {
|
||||
fn identifier_quote_style(&self, identifier: &str) -> Option<char> {
|
||||
let needs_quote = identifier.chars().any(|c| c.is_ascii_uppercase())
|
||||
|| !identifier
|
||||
@@ -30,6 +36,61 @@ impl Dialect for LanceSqlDialect {
|
||||
}
|
||||
}
|
||||
|
||||
/// Lance's tokenizer dialect with SQL-standard double-quoted identifiers added.
|
||||
///
|
||||
/// Keep this deliberately small: Lance's parser wraps `GenericDialect` and
|
||||
/// delegates only identifier recognition, leaving every other dialect option at
|
||||
/// its default. In particular, `/*! ... */` remains an ordinary block comment.
|
||||
#[derive(Debug, Default)]
|
||||
struct PredicateDialect(GenericDialect);
|
||||
|
||||
impl SqlParserDialect for PredicateDialect {
|
||||
fn dialect(&self) -> TypeId {
|
||||
self.0.dialect()
|
||||
}
|
||||
|
||||
fn is_identifier_start(&self, ch: char) -> bool {
|
||||
self.0.is_identifier_start(ch)
|
||||
}
|
||||
|
||||
fn is_identifier_part(&self, ch: char) -> bool {
|
||||
self.0.is_identifier_part(ch)
|
||||
}
|
||||
|
||||
fn is_delimited_identifier_start(&self, ch: char) -> bool {
|
||||
ch == '"' || ch == '`'
|
||||
}
|
||||
}
|
||||
|
||||
/// Canonicalize a raw SQL predicate for Lance's parser.
|
||||
///
|
||||
/// Lance wraps [`GenericDialect`] for identifier recognition while retaining the
|
||||
/// default dialect behavior for every other lexical option. [`PredicateDialect`]
|
||||
/// mirrors that contract and additionally recognizes `"` as an identifier
|
||||
/// delimiter, allowing this function to rewrite only those identifier tokens.
|
||||
pub fn canonicalize_sql_predicate(predicate: &str) -> crate::Result<String> {
|
||||
let dialect = PredicateDialect::default();
|
||||
let tokens = Tokenizer::new(&dialect, predicate)
|
||||
.with_unescape(false)
|
||||
.tokenize()
|
||||
.map_err(|err| crate::Error::InvalidInput {
|
||||
message: format!("invalid SQL predicate: {err}"),
|
||||
})?;
|
||||
|
||||
Ok(tokens
|
||||
.into_iter()
|
||||
.map(|token| match token {
|
||||
Token::Word(word) if word.quote_style == Some('"') => {
|
||||
// with_unescape(false) retains doubled double quotes. Decode
|
||||
// those before escaping any backticks for Lance's delimiter.
|
||||
let identifier = word.value.replace("\"\"", "\"").replace('`', "``");
|
||||
format!("`{identifier}`")
|
||||
}
|
||||
other => other.to_string(),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Prefix for placeholder strings inserted in place of binary literals. Chosen
|
||||
/// to be extremely unlikely to occur in user data.
|
||||
const BINARY_PLACEHOLDER_PREFIX: &str = "__lancedb_binary_placeholder_";
|
||||
@@ -113,3 +174,51 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result<String> {
|
||||
}
|
||||
Ok(sql)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::canonicalize_sql_predicate;
|
||||
|
||||
#[test]
|
||||
fn normalizes_double_quoted_identifiers() {
|
||||
assert_eq!(
|
||||
canonicalize_sql_predicate(r#""PartyAbbrev" = 'D'"#).unwrap(),
|
||||
"`PartyAbbrev` = 'D'"
|
||||
);
|
||||
assert_eq!(
|
||||
canonicalize_sql_predicate(r#""MetaData"."userId" = 5"#).unwrap(),
|
||||
"`MetaData`.`userId` = 5"
|
||||
);
|
||||
assert_eq!(
|
||||
canonicalize_sql_predicate(r#""a""b" = 1"#).unwrap(),
|
||||
"`a\"b` = 1"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preserves_quotes_inside_literals_and_backticks() {
|
||||
let filter = r#"name = 'Alice "Ace"' AND `quoted"field` = 1"#;
|
||||
assert_eq!(canonicalize_sql_predicate(filter).unwrap(), filter);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preserves_literals_and_comments_using_lance_dialect_rules() {
|
||||
let predicate = r#"path = '\' AND "PartyAbbrev" = 'D' -- unmatched " in comment"#;
|
||||
assert_eq!(
|
||||
canonicalize_sql_predicate(predicate).unwrap(),
|
||||
r#"path = '\' AND `PartyAbbrev` = 'D' -- unmatched " in comment"#
|
||||
);
|
||||
|
||||
let predicate = r#"id = 1 /* unmatched " in block comment */"#;
|
||||
assert_eq!(canonicalize_sql_predicate(predicate).unwrap(), predicate);
|
||||
|
||||
let predicate = r#"id = 1 /*! OR "PartyAbbrev" = 'D' */"#;
|
||||
assert_eq!(canonicalize_sql_predicate(predicate).unwrap(), predicate);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_unterminated_double_quoted_identifier() {
|
||||
let error = canonicalize_sql_predicate(r#""PartyAbbrev = 'D'"#).unwrap_err();
|
||||
assert!(matches!(error, crate::Error::InvalidInput { .. }));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -170,6 +170,15 @@ pub(crate) fn plan(
|
||||
filter: Option<&str>,
|
||||
limit: Option<u64>,
|
||||
) -> Result<(MaterializedViewDefinition, Vec<ArrowField>, Lineage)> {
|
||||
let filter = filter
|
||||
.map(crate::expr::canonicalize_sql_predicate)
|
||||
.transpose()
|
||||
.map_err(|err| match err {
|
||||
Error::InvalidInput { message } => Error::InvalidInput {
|
||||
message: format!("invalid view filter: {message}"),
|
||||
},
|
||||
err => err,
|
||||
})?;
|
||||
let projections: Vec<(String, String)> = if projections.is_empty() {
|
||||
source_schema
|
||||
.fields()
|
||||
@@ -274,7 +283,7 @@ pub(crate) fn plan(
|
||||
declared.push(output);
|
||||
}
|
||||
|
||||
if let Some(filter) = filter {
|
||||
if let Some(filter) = filter.as_deref() {
|
||||
let expr = planner
|
||||
.parse_filter(filter)
|
||||
.map_err(|e| Error::InvalidInput {
|
||||
@@ -314,7 +323,7 @@ pub(crate) fn plan(
|
||||
.into_iter()
|
||||
.map(|(output, expression)| ViewProjection { output, expression })
|
||||
.collect(),
|
||||
filter: filter.map(String::from),
|
||||
filter,
|
||||
limit,
|
||||
inputs,
|
||||
};
|
||||
|
||||
@@ -46,8 +46,9 @@ use lance_table::format::Fragment;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::{
|
||||
INCARNATION_META_KEY, MaterializedViewDefinition, REFRESHED_AT_MS_META_KEY,
|
||||
SOURCE_ROW_ID_COLUMN, SOURCE_VERSION_META_KEY,
|
||||
DEFINITION_META_KEY, INCARNATION_META_KEY, MaterializedViewDefinition,
|
||||
REFRESHED_AT_MS_META_KEY, SOURCE_ROW_ID_COLUMN, SOURCE_VERSION_META_KEY,
|
||||
definition_to_metadata,
|
||||
};
|
||||
use crate::database::OpenTableRequest;
|
||||
use crate::table::{NativeTable, NativeTableExt, Table};
|
||||
@@ -197,8 +198,28 @@ pub(crate) async fn execute_refresh(
|
||||
),
|
||||
});
|
||||
}
|
||||
let definition_changed =
|
||||
definition.filter != replanned.filter || definition.inputs != replanned.inputs;
|
||||
let definition = &replanned;
|
||||
|
||||
// A watermark written for a legacy raw filter certifies the rows that
|
||||
// filter produced, not the canonical predicate above. Rebuild instead of
|
||||
// accepting or advancing it, and persist the migrated definition in the
|
||||
// same metadata commit that certifies the replacement rows.
|
||||
if definition_changed {
|
||||
return rebuild(
|
||||
view_native,
|
||||
&view_ds,
|
||||
&source_ds,
|
||||
source_version,
|
||||
source_ts,
|
||||
definition,
|
||||
true,
|
||||
expected_incarnation,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let metadata = &view_ds.schema().metadata;
|
||||
let watermark: Option<u64> = metadata
|
||||
.get(SOURCE_VERSION_META_KEY)
|
||||
@@ -257,6 +278,7 @@ pub(crate) async fn execute_refresh(
|
||||
source_version,
|
||||
source_ts,
|
||||
definition,
|
||||
false,
|
||||
expected_incarnation,
|
||||
)
|
||||
.await
|
||||
@@ -271,6 +293,7 @@ pub(crate) async fn execute_refresh(
|
||||
source_version,
|
||||
source_ts,
|
||||
definition,
|
||||
false,
|
||||
expected_incarnation,
|
||||
)
|
||||
.await
|
||||
@@ -683,6 +706,7 @@ async fn incremental(
|
||||
view_ds.clone(),
|
||||
source_version,
|
||||
source_ts,
|
||||
None,
|
||||
expected_incarnation,
|
||||
)
|
||||
.await?;
|
||||
@@ -704,6 +728,7 @@ async fn incremental(
|
||||
published,
|
||||
source_version,
|
||||
source_ts,
|
||||
None,
|
||||
expected_incarnation,
|
||||
)
|
||||
.await?;
|
||||
@@ -775,6 +800,7 @@ async fn incremental(
|
||||
published,
|
||||
source_version,
|
||||
source_ts,
|
||||
None,
|
||||
expected_incarnation,
|
||||
)
|
||||
.await?;
|
||||
@@ -824,12 +850,14 @@ async fn incremental(
|
||||
appended,
|
||||
source_version,
|
||||
source_ts,
|
||||
None,
|
||||
expected_incarnation,
|
||||
)
|
||||
.await?;
|
||||
Ok(Some(result))
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn rebuild(
|
||||
view_native: &NativeTable,
|
||||
view_ds: &Dataset,
|
||||
@@ -837,6 +865,7 @@ async fn rebuild(
|
||||
source_version: u64,
|
||||
source_ts: u128,
|
||||
definition: &MaterializedViewDefinition,
|
||||
persist_definition: bool,
|
||||
expected_incarnation: Option<&str>,
|
||||
) -> Result<RefreshMaterializedViewResult> {
|
||||
let rows_written = Arc::new(AtomicU64::new(0));
|
||||
@@ -867,6 +896,7 @@ async fn rebuild(
|
||||
replaced,
|
||||
source_version,
|
||||
source_ts,
|
||||
persist_definition.then_some(definition),
|
||||
expected_incarnation,
|
||||
)
|
||||
.await?;
|
||||
@@ -981,6 +1011,7 @@ async fn stamp_watermark(
|
||||
mut dataset: Dataset,
|
||||
source_version: u64,
|
||||
source_ts: u128,
|
||||
definition: Option<&MaterializedViewDefinition>,
|
||||
expected_incarnation: Option<&str>,
|
||||
) -> Result<u64> {
|
||||
ensure_incarnation(&dataset, expected_incarnation, dataset.uri()).await?;
|
||||
@@ -993,27 +1024,32 @@ async fn stamp_watermark(
|
||||
.get(INCARNATION_META_KEY)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
|
||||
dataset
|
||||
.update_schema_metadata([
|
||||
(INCARNATION_META_KEY.to_string(), Some(incarnation)),
|
||||
(
|
||||
SOURCE_VERSION_META_KEY.to_string(),
|
||||
Some(source_version.to_string()),
|
||||
),
|
||||
(
|
||||
SOURCE_VERSION_TS_META_KEY.to_string(),
|
||||
Some(source_ts.to_string()),
|
||||
),
|
||||
(
|
||||
REFRESHED_AT_MS_META_KEY.to_string(),
|
||||
Some(now_ms().to_string()),
|
||||
),
|
||||
(
|
||||
VIEW_VERSION_META_KEY.to_string(),
|
||||
Some(predicted.to_string()),
|
||||
),
|
||||
])
|
||||
.await?;
|
||||
let mut metadata = vec![(INCARNATION_META_KEY.to_string(), Some(incarnation))];
|
||||
if let Some(definition) = definition {
|
||||
metadata.push((
|
||||
DEFINITION_META_KEY.to_string(),
|
||||
Some(definition_to_metadata(definition)?),
|
||||
));
|
||||
}
|
||||
metadata.extend([
|
||||
(
|
||||
SOURCE_VERSION_META_KEY.to_string(),
|
||||
Some(source_version.to_string()),
|
||||
),
|
||||
(
|
||||
SOURCE_VERSION_TS_META_KEY.to_string(),
|
||||
Some(source_ts.to_string()),
|
||||
),
|
||||
(
|
||||
REFRESHED_AT_MS_META_KEY.to_string(),
|
||||
Some(now_ms().to_string()),
|
||||
),
|
||||
(
|
||||
VIEW_VERSION_META_KEY.to_string(),
|
||||
Some(predicted.to_string()),
|
||||
),
|
||||
]);
|
||||
dataset.update_schema_metadata(metadata).await?;
|
||||
let actual = dataset.version().version;
|
||||
if actual != predicted {
|
||||
return Err(Error::Runtime {
|
||||
@@ -1585,6 +1621,106 @@ mod tests {
|
||||
assert_eq!(read(view.table(), "x").await, vec![20, 40]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_mixed_case_filter_is_canonicalized_for_lineage_and_refresh() {
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
let batch = record_batch!(
|
||||
("id", Int32, [1, 2, 3]),
|
||||
("PartyAbbrev", Utf8, ["D", "R", "D"])
|
||||
)
|
||||
.unwrap();
|
||||
conn.create_table("src", batch)
|
||||
.write_options(crate::materialized_view::tests::stable_row_ids())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
conn.create_materialized_view("democrats", "src")
|
||||
.select([("id", "id")])
|
||||
.only_if(r#""PartyAbbrev" = 'D'"#)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Reopen from schema metadata so these assertions cover the stored
|
||||
// predicate and lineage, not only the declaration-time handle.
|
||||
let view = conn.open_materialized_view("democrats").await.unwrap();
|
||||
assert_eq!(
|
||||
view.definition().filter.as_deref(),
|
||||
Some("`PartyAbbrev` = 'D'")
|
||||
);
|
||||
assert_eq!(view.definition().inputs, ["PartyAbbrev", "id"]);
|
||||
|
||||
let result = view.refresh().execute().await.unwrap();
|
||||
assert_eq!(result.rows_written, 2);
|
||||
assert_eq!(read(view.table(), "id").await, vec![1, 3]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_legacy_raw_filter_rebuilds_and_persists_canonical_definition() {
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
let batch = record_batch!(
|
||||
("id", Int32, [1, 2, 3]),
|
||||
("PartyAbbrev", Utf8, ["D", "R", "D"])
|
||||
)
|
||||
.unwrap();
|
||||
conn.create_table("legacy_src", batch)
|
||||
.write_options(crate::materialized_view::tests::stable_row_ids())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let view = conn
|
||||
.create_materialized_view("legacy_view", "legacy_src")
|
||||
.select([("id", "id")])
|
||||
.only_if(r#""PartyAbbrev" = 'X'"#)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(view.refresh().execute().await.unwrap().rows_written, 0);
|
||||
|
||||
// Model a definition and up-to-date watermark written before filter
|
||||
// canonicalization was applied to materialized views.
|
||||
let mut legacy = view.definition().clone();
|
||||
legacy.filter = Some(r#""PartyAbbrev" = 'D'"#.into());
|
||||
legacy.inputs = vec!["id".into()];
|
||||
let native = view.table().as_native().unwrap();
|
||||
let mut dataset = native.dataset.get().await.unwrap().as_ref().clone();
|
||||
let predicted = dataset.version().version + 1;
|
||||
dataset
|
||||
.update_schema_metadata([
|
||||
(
|
||||
DEFINITION_META_KEY.to_string(),
|
||||
Some(definition_to_metadata(&legacy).unwrap()),
|
||||
),
|
||||
(
|
||||
VIEW_VERSION_META_KEY.to_string(),
|
||||
Some(predicted.to_string()),
|
||||
),
|
||||
])
|
||||
.await
|
||||
.unwrap();
|
||||
native.dataset.update(dataset);
|
||||
|
||||
let reopened = conn.open_materialized_view("legacy_view").await.unwrap();
|
||||
let result = reopened.refresh().execute().await.unwrap();
|
||||
assert_eq!(result.mode, RefreshMode::Rebuild);
|
||||
assert_eq!(result.rows_written, 2);
|
||||
assert_eq!(read(reopened.table(), "id").await, vec![1, 3]);
|
||||
|
||||
// A fresh handle proves the migration was stored alongside the new
|
||||
// watermark and therefore happens only once.
|
||||
let migrated = conn.open_materialized_view("legacy_view").await.unwrap();
|
||||
assert_eq!(
|
||||
migrated.definition().filter.as_deref(),
|
||||
Some("`PartyAbbrev` = 'D'")
|
||||
);
|
||||
assert_eq!(migrated.definition().inputs, ["PartyAbbrev", "id"]);
|
||||
assert_eq!(
|
||||
migrated.refresh().execute().await.unwrap().mode,
|
||||
RefreshMode::NoOp
|
||||
);
|
||||
assert_eq!(read(migrated.table(), "id").await, vec![1, 3]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_append_refreshes_incrementally() {
|
||||
let (_conn, source, view) = refreshed_doubled(vec![1, 2]).await;
|
||||
@@ -2767,7 +2903,7 @@ mod tests {
|
||||
let stale = view_native.dataset.get().await.unwrap().as_ref().clone();
|
||||
view.table().delete("x = 1").await.unwrap();
|
||||
|
||||
let err = stamp_watermark(view_native, stale, 99, 99, None).await;
|
||||
let err = stamp_watermark(view_native, stale, 99, 99, None, None).await;
|
||||
assert!(err.is_err());
|
||||
|
||||
let result = view.refresh().execute().await.unwrap();
|
||||
|
||||
+273
-9
@@ -415,6 +415,9 @@ pub trait QueryBase {
|
||||
/// x > 5 OR y = 'test'
|
||||
/// ```
|
||||
///
|
||||
/// Identifiers may be delimited with SQL-standard double quotes or
|
||||
/// backticks. String literals must use single quotes.
|
||||
///
|
||||
/// Filtering performance can often be improved by creating a scalar index
|
||||
/// on the filter column(s).
|
||||
///
|
||||
@@ -938,6 +941,17 @@ impl QueryRequest {
|
||||
/// use different representations) the error is recorded and surfaced later
|
||||
/// by [`Self::check_filter`].
|
||||
pub(crate) fn add_filter(&mut self, new: QueryFilter) {
|
||||
let new = match new {
|
||||
QueryFilter::Sql(filter) => match crate::expr::canonicalize_sql_predicate(&filter) {
|
||||
Ok(filter) => QueryFilter::Sql(filter),
|
||||
Err(err) => {
|
||||
self.filter_error = Some(err.to_string());
|
||||
return;
|
||||
}
|
||||
},
|
||||
other => other,
|
||||
};
|
||||
|
||||
self.filter = Some(match self.filter.take() {
|
||||
None => new,
|
||||
Some(existing) => match and_filters(existing, new) {
|
||||
@@ -1199,12 +1213,12 @@ impl VectorQuery {
|
||||
|
||||
/// Add another query vector to the search.
|
||||
///
|
||||
/// Multiple searches will be dispatched as part of the query.
|
||||
/// This is a convenience method for adding multiple query vectors
|
||||
/// to the search. It is not expected to be faster than issuing
|
||||
/// multiple queries concurrently.
|
||||
/// Multiple searches will be dispatched as a batch. Flat searches share
|
||||
/// one table scan across the query vectors, avoiding the scan and memory
|
||||
/// amplification of issuing the searches concurrently. Indexed searches
|
||||
/// may still perform per-vector index work.
|
||||
///
|
||||
/// The output data will contain an additional columns `query_index` which
|
||||
/// The output data will contain an additional column `query_index` which
|
||||
/// will contain the index of the query vector that was used to generate the
|
||||
/// result.
|
||||
pub fn add_query_vector(mut self, vector: impl IntoQueryVector) -> Result<Self> {
|
||||
@@ -2263,10 +2277,14 @@ mod tests {
|
||||
use std::{collections::HashSet, sync::Arc};
|
||||
|
||||
use super::*;
|
||||
use arrow::{array::downcast_array, compute::concat_batches, datatypes::Int32Type};
|
||||
use arrow::{
|
||||
array::downcast_array,
|
||||
compute::concat_batches,
|
||||
datatypes::{Int32Type, UInt8Type},
|
||||
};
|
||||
use arrow_array::{
|
||||
FixedSizeListArray, Float32Array, Int32Array, RecordBatch, StringArray, cast::AsArray,
|
||||
types::Float32Type,
|
||||
FixedSizeListArray, Float32Array, Int32Array, RecordBatch, RecordBatchIterator,
|
||||
StringArray, cast::AsArray, types::Float32Type,
|
||||
};
|
||||
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
|
||||
use datafusion_physical_plan::display::DisplayableExecutionPlan;
|
||||
@@ -2496,6 +2514,157 @@ mod tests {
|
||||
query.execute().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_double_quoted_predicates_across_table_operations() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let dataset_path = tmp_dir.path().join("test.lance");
|
||||
let uri = dataset_path.to_str().unwrap();
|
||||
let schema = Arc::new(ArrowSchema::new(vec![
|
||||
ArrowField::new("id", DataType::Int32, false),
|
||||
ArrowField::new("PartyAbbrev", DataType::Utf8, false),
|
||||
ArrowField::new("path", DataType::Utf8, false),
|
||||
]));
|
||||
let batch = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(Int32Array::from(vec![1, 2, 3, 4])),
|
||||
Arc::new(StringArray::from(vec!["D", "R", "R", "D"])),
|
||||
Arc::new(StringArray::from(vec!["\\", "\\", "x", "x"])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let conn = connect(uri).execute().await.unwrap();
|
||||
let table = conn.create_table("parties", batch).execute().await.unwrap();
|
||||
let batches = table
|
||||
.query()
|
||||
.only_if(r#""PartyAbbrev" = 'D'"#)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
|
||||
assert_eq!(
|
||||
table
|
||||
.count_rows(Some(r#""PartyAbbrev" = 'D'"#.to_string()))
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
|
||||
// Public BaseTable dispatch cannot bypass canonicalization.
|
||||
let query = AnyQuery::Query(QueryRequest {
|
||||
filter: Some(QueryFilter::Sql(r#""PartyAbbrev" = 'D'"#.to_string())),
|
||||
..Default::default()
|
||||
});
|
||||
let batches = table
|
||||
.base_table()
|
||||
.query(&query, Default::default())
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
|
||||
assert_eq!(
|
||||
table
|
||||
.base_table()
|
||||
.count_rows(Some(crate::table::Filter::Sql(
|
||||
r#""PartyAbbrev" = 'D'"#.to_string(),
|
||||
)))
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
|
||||
for predicate in [
|
||||
r#"id = 1 -- unmatched " in a valid SQL comment"#,
|
||||
r#"id = 1 /* unmatched " in a valid SQL comment */"#,
|
||||
r#"id = 1 /*! OR "PartyAbbrev" = 'D' */"#,
|
||||
r#"path = '\' AND "PartyAbbrev" = 'D'"#,
|
||||
] {
|
||||
let batches = table
|
||||
.query()
|
||||
.only_if(predicate)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
|
||||
}
|
||||
|
||||
// The same canonical predicate contract applies to both merge filters.
|
||||
let source = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(Int32Array::from(vec![1, 2, 3])),
|
||||
Arc::new(StringArray::from(vec!["D", "R", "R"])),
|
||||
Arc::new(StringArray::from(vec!["\\", "\\", "x"])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let mut merge = table.merge_insert(&["id"]);
|
||||
merge.when_not_matched_by_source_delete(Some(r#""PartyAbbrev" = 'D'"#.to_string()));
|
||||
let result = table
|
||||
.base_table()
|
||||
.merge_insert(
|
||||
merge,
|
||||
Box::new(RecordBatchIterator::new(vec![Ok(source)], schema.clone())),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(result.num_deleted_rows, 1);
|
||||
|
||||
let source = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(Int32Array::from(vec![1, 2, 3])),
|
||||
Arc::new(StringArray::from(vec!["U", "U", "U"])),
|
||||
Arc::new(StringArray::from(vec!["\\", "\\", "x"])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let mut merge = table.merge_insert(&["id"]);
|
||||
merge.when_matched_update_all(Some(r#"target."PartyAbbrev" = 'D'"#.to_string()));
|
||||
merge
|
||||
.execute(Box::new(RecordBatchIterator::new(vec![Ok(source)], schema)))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
table
|
||||
.count_rows(Some(r#""PartyAbbrev" = 'U'"#.to_string()))
|
||||
.await
|
||||
.unwrap(),
|
||||
1
|
||||
);
|
||||
|
||||
let update = table
|
||||
.update()
|
||||
.only_if(r#""PartyAbbrev" = 'R'"#)
|
||||
.column("PartyAbbrev", "'X'");
|
||||
table.base_table().update(update).await.unwrap();
|
||||
assert_eq!(
|
||||
table
|
||||
.count_rows(Some(r#""PartyAbbrev" = 'X'"#.to_string()))
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
|
||||
let result = table
|
||||
.base_table()
|
||||
.delete(crate::table::Predicate::String(r#""PartyAbbrev" = 'X'"#))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(result.num_deleted_rows, 2);
|
||||
assert_eq!(table.count_rows(None).await.unwrap(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_select_with_transform() {
|
||||
let batches = make_non_empty_batches();
|
||||
@@ -2952,7 +3121,8 @@ mod tests {
|
||||
.limit(1);
|
||||
|
||||
let plan = query.explain_plan(true).await.unwrap();
|
||||
assert!(plan.contains("UnionExec"));
|
||||
assert!(plan.contains("KNNVectorDistance: queries=2"));
|
||||
assert!(!plan.contains("UnionExec"));
|
||||
|
||||
let results = query
|
||||
.execute()
|
||||
@@ -2967,6 +3137,100 @@ mod tests {
|
||||
// We don't guarantee order.
|
||||
assert!(query_index.values().contains(&0));
|
||||
assert!(query_index.values().contains(&1));
|
||||
|
||||
// Batch KNN does not support a per-query offset, so offset queries keep
|
||||
// the legacy per-vector plan to preserve their result semantics.
|
||||
let offset_query = table
|
||||
.query()
|
||||
.nearest_to(&[0.1, 0.2, 0.3, 0.4])
|
||||
.unwrap()
|
||||
.add_query_vector(&[0.5, 0.6, 0.7, 0.8])
|
||||
.unwrap()
|
||||
.limit(1)
|
||||
.offset(1);
|
||||
assert!(
|
||||
offset_query
|
||||
.explain_plan(true)
|
||||
.await
|
||||
.unwrap()
|
||||
.contains("UnionExec")
|
||||
);
|
||||
let offset_results = offset_query
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
offset_results
|
||||
.iter()
|
||||
.map(RecordBatch::num_rows)
|
||||
.sum::<usize>(),
|
||||
2
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multiple_binary_query_vectors() {
|
||||
let vectors = FixedSizeListArray::from_iter_primitive::<UInt8Type, _, _>(
|
||||
vec![
|
||||
Some(vec![Some(0), Some(0)]),
|
||||
Some(vec![Some(255), Some(255)]),
|
||||
],
|
||||
2,
|
||||
);
|
||||
let schema = Arc::new(ArrowSchema::new(vec![
|
||||
ArrowField::new("id", DataType::Int32, false),
|
||||
ArrowField::new("vector", vectors.data_type().clone(), false),
|
||||
]));
|
||||
let batch = RecordBatch::try_new(
|
||||
schema,
|
||||
vec![Arc::new(Int32Array::from(vec![0, 1])), Arc::new(vectors)],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
let table = conn
|
||||
.create_table("binary_batch", batch)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let query = table
|
||||
.query()
|
||||
.nearest_to(&[0.0, 0.0])
|
||||
.unwrap()
|
||||
.add_query_vector(&[255.0, 255.0])
|
||||
.unwrap()
|
||||
.distance_type(DistanceType::Hamming)
|
||||
.limit(1);
|
||||
|
||||
// Binary queries retain the per-vector plan because Lance's binary
|
||||
// nearest path requires primitive UInt8 query arrays.
|
||||
assert!(
|
||||
query
|
||||
.explain_plan(true)
|
||||
.await
|
||||
.unwrap()
|
||||
.contains("UnionExec")
|
||||
);
|
||||
|
||||
let results = query
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
let results = concat_batches(&results[0].schema(), &results).unwrap();
|
||||
assert_eq!(results.num_rows(), 2);
|
||||
|
||||
let ids = results["id"].as_primitive::<Int32Type>();
|
||||
assert!(ids.values().contains(&0));
|
||||
assert!(ids.values().contains(&1));
|
||||
let query_index = results["query_index"].as_primitive::<Int32Type>();
|
||||
assert!(query_index.values().contains(&0));
|
||||
assert!(query_index.values().contains(&1));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -1379,10 +1379,11 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
query: &AnyQuery,
|
||||
version: Option<u64>,
|
||||
) -> Result<Vec<serde_json::Value>> {
|
||||
let query = query.canonicalized()?;
|
||||
let mut base_body = serde_json::json!({ "version": version });
|
||||
self.apply_branch_body(&mut base_body);
|
||||
|
||||
match query {
|
||||
match &query {
|
||||
AnyQuery::Query(query) => {
|
||||
let mut body = base_body.clone();
|
||||
self.apply_query_params(&mut body, query)?;
|
||||
@@ -2494,7 +2495,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
|
||||
let mut body = if let Some(filter) = filter {
|
||||
let filter_sql = match filter {
|
||||
Filter::Sql(sql) => sql.clone(),
|
||||
Filter::Sql(sql) => crate::expr::canonicalize_sql_predicate(&sql)?,
|
||||
Filter::Datafusion(expr) => expr_to_sql_string(&expr)?,
|
||||
};
|
||||
serde_json::json!({ "predicate": filter_sql, "version": read_snapshot.version })
|
||||
@@ -2795,7 +2796,8 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
Ok(final_analyze)
|
||||
}
|
||||
|
||||
async fn update(&self, update: UpdateBuilder) -> Result<UpdateResult> {
|
||||
async fn update(&self, mut update: UpdateBuilder) -> Result<UpdateResult> {
|
||||
update.canonicalize_filter()?;
|
||||
self.check_mutable().await?;
|
||||
let request = self
|
||||
.client
|
||||
@@ -2842,7 +2844,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
async fn delete(&self, predicate: Predicate<'_>) -> Result<DeleteResult> {
|
||||
self.check_mutable().await?;
|
||||
let predicate_sql = match predicate {
|
||||
Predicate::String(s) => s.to_string(),
|
||||
Predicate::String(s) => crate::expr::canonicalize_sql_predicate(s)?,
|
||||
Predicate::Expr(expr) => expr_to_sql_string(expr)?,
|
||||
};
|
||||
let mut body = serde_json::json!({ "predicate": predicate_sql });
|
||||
@@ -2899,9 +2901,10 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
|
||||
async fn merge_insert(
|
||||
&self,
|
||||
params: MergeInsertBuilder,
|
||||
mut params: MergeInsertBuilder,
|
||||
new_data: Box<dyn RecordBatchReader + Send>,
|
||||
) -> Result<MergeResult> {
|
||||
params.canonicalize_filters()?;
|
||||
self.check_mutable().await?;
|
||||
|
||||
let timeout = params.timeout;
|
||||
@@ -3912,13 +3915,17 @@ mod tests {
|
||||
);
|
||||
assert_eq!(
|
||||
request.body().unwrap().as_bytes().unwrap(),
|
||||
br#"{"predicate":"a > 10","version":null}"#
|
||||
br#"{"predicate":"`A` > 10","version":null}"#
|
||||
);
|
||||
|
||||
http::Response::builder().status(200).body("42").unwrap()
|
||||
});
|
||||
|
||||
let count = table.count_rows(Some("a > 10".into())).await.unwrap();
|
||||
let count = table
|
||||
.base_table()
|
||||
.count_rows(Some(Filter::Sql(r#""A" > 10"#.into())))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(count, 42);
|
||||
}
|
||||
|
||||
@@ -4401,7 +4408,7 @@ mod tests {
|
||||
assert_eq!(expression, "b - 1");
|
||||
|
||||
let only_if = value.get("predicate").unwrap().as_str().unwrap();
|
||||
assert_eq!(only_if, "b > 10");
|
||||
assert_eq!(only_if, "`B` > 10");
|
||||
}
|
||||
|
||||
if old_server {
|
||||
@@ -4417,14 +4424,12 @@ mod tests {
|
||||
}
|
||||
});
|
||||
|
||||
let result = table
|
||||
let update = table
|
||||
.update()
|
||||
.column("a", "a + 1")
|
||||
.column("b", "b - 1")
|
||||
.only_if("b > 10")
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
.only_if(r#""B" > 10"#);
|
||||
let result = table.base_table().update(update).await.unwrap();
|
||||
|
||||
assert_eq!(result.version, if old_server { 0 } else { 43 });
|
||||
assert_eq!(result.rows_updated, if old_server { 0 } else { 5 });
|
||||
@@ -4511,10 +4516,10 @@ mod tests {
|
||||
|
||||
let params = request.url().query_pairs().collect::<HashMap<_, _>>();
|
||||
assert_eq!(params["on"], "some_col");
|
||||
assert_eq!(params["when_matched_update_all"], "false");
|
||||
assert_eq!(params["when_matched_update_all"], "true");
|
||||
assert_eq!(params["when_not_matched_insert_all"], "false");
|
||||
assert_eq!(params["when_not_matched_by_source_delete"], "false");
|
||||
assert!(!params.contains_key("when_matched_update_all_filt"));
|
||||
assert_eq!(params["when_matched_update_all_filt"], "target.`A` > 0");
|
||||
assert!(!params.contains_key("when_not_matched_by_source_delete_filt"));
|
||||
assert!(!params.contains_key("use_index"));
|
||||
|
||||
@@ -4531,11 +4536,9 @@ mod tests {
|
||||
}
|
||||
});
|
||||
|
||||
let result = table
|
||||
.merge_insert(&["some_col"])
|
||||
.execute(data)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut merge = table.merge_insert(&["some_col"]);
|
||||
merge.when_matched_update_all(Some(r#"target."A" > 0"#.into()));
|
||||
let result = table.base_table().merge_insert(merge, data).await.unwrap();
|
||||
|
||||
assert_eq!(result.version, if old_server { 0 } else { 43 });
|
||||
if !old_server {
|
||||
@@ -4597,7 +4600,7 @@ mod tests {
|
||||
let body = request.body().unwrap().as_bytes().unwrap();
|
||||
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
|
||||
let predicate = body.get("predicate").unwrap().as_str().unwrap();
|
||||
assert_eq!(predicate, "id in (1, 2, 3)");
|
||||
assert_eq!(predicate, "`ID` in (1, 2, 3)");
|
||||
|
||||
if old_server {
|
||||
http::Response::builder()
|
||||
@@ -4615,7 +4618,11 @@ mod tests {
|
||||
}
|
||||
});
|
||||
|
||||
let result = table.delete("id in (1, 2, 3)").await.unwrap();
|
||||
let result = table
|
||||
.base_table()
|
||||
.delete(Predicate::String(r#""ID" in (1, 2, 3)"#))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(result.version, if old_server { 0 } else { 43 });
|
||||
}
|
||||
|
||||
@@ -4707,6 +4714,7 @@ mod tests {
|
||||
let body = request.body().unwrap().as_bytes().unwrap();
|
||||
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
|
||||
let expected_body = serde_json::json!({
|
||||
"filter": "`A` > 0",
|
||||
"k": isize::MAX as usize,
|
||||
"prefilter": true,
|
||||
"vector": [], // Empty vector means no vector query.
|
||||
@@ -4722,9 +4730,13 @@ mod tests {
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let query = AnyQuery::Query(QueryRequest {
|
||||
filter: Some(QueryFilter::Sql(r#""A" > 0"#.into())),
|
||||
..Default::default()
|
||||
});
|
||||
let data = table
|
||||
.query()
|
||||
.execute()
|
||||
.base_table()
|
||||
.query(&query, Default::default())
|
||||
.await
|
||||
.unwrap()
|
||||
.collect::<Vec<_>>()
|
||||
|
||||
@@ -1172,7 +1172,10 @@ impl Table {
|
||||
///
|
||||
/// * `filter` if present, only count rows matching the filter
|
||||
pub async fn count_rows(&self, filter: Option<String>) -> Result<usize> {
|
||||
self.inner.count_rows(filter.map(Filter::Sql)).await
|
||||
let filter = filter
|
||||
.map(|predicate| crate::expr::canonicalize_sql_predicate(&predicate).map(Filter::Sql))
|
||||
.transpose()?;
|
||||
self.inner.count_rows(filter).await
|
||||
}
|
||||
|
||||
/// Names of the blob v2 columns in this table, in declaration order.
|
||||
@@ -1372,7 +1375,13 @@ impl Table {
|
||||
/// # });
|
||||
/// ```
|
||||
pub async fn delete(&self, predicate: impl Into<Predicate<'_>>) -> Result<DeleteResult> {
|
||||
self.inner.delete(predicate.into()).await
|
||||
match predicate.into() {
|
||||
Predicate::String(predicate) => {
|
||||
let predicate = crate::expr::canonicalize_sql_predicate(predicate)?;
|
||||
self.inner.delete(Predicate::String(&predicate)).await
|
||||
}
|
||||
predicate @ Predicate::Expr(_) => self.inner.delete(predicate).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an index on the provided column(s).
|
||||
@@ -1785,7 +1794,23 @@ impl Table {
|
||||
self.inner.alter_columns(alterations).await
|
||||
}
|
||||
|
||||
/// Update per-field metadata (merges by default).
|
||||
/// Update per-field (column) metadata.
|
||||
///
|
||||
/// Each [`FieldMetadataUpdate`] is merged into the field's existing metadata
|
||||
/// by default; use [`FieldMetadataUpdate::remove`] to delete a key, or
|
||||
/// [`FieldMetadataUpdate::replace`] to swap the field's entire metadata map.
|
||||
///
|
||||
/// The following keys are treated specially, by convention, and should be
|
||||
/// used when appropriate:
|
||||
///
|
||||
/// - `lancedb:description`: for a human-readable description of a field.
|
||||
/// - `lancedb:tag:<name>`: for a user-defined key-value tag, where the suffix
|
||||
/// names the tag category; e.g. `lancedb:tag:model: "clip"`.
|
||||
/// - `lancedb:logical-column`: for a column grouping; e.g. `feature_v1` and
|
||||
/// `feature_v2` might be in the same logical column.
|
||||
/// - `lancedb:status`: for status options (`production`, `candidate`,
|
||||
/// `deprecated`, `archived`) to designate the current life cycle state of
|
||||
/// this column.
|
||||
pub async fn update_field_metadata(
|
||||
&self,
|
||||
updates: &[FieldMetadataUpdate],
|
||||
@@ -3231,7 +3256,10 @@ impl BaseTable for NativeTable {
|
||||
let dataset = self.dataset.get().await?;
|
||||
match filter {
|
||||
None => Ok(dataset.count_rows(None).await?),
|
||||
Some(Filter::Sql(sql)) => Ok(dataset.count_rows(Some(sql)).await?),
|
||||
Some(Filter::Sql(sql)) => {
|
||||
let sql = crate::expr::canonicalize_sql_predicate(&sql)?;
|
||||
Ok(dataset.count_rows(Some(sql)).await?)
|
||||
}
|
||||
Some(Filter::Datafusion(_)) => Err(Error::NotSupported {
|
||||
message: "Datafusion filters are not yet supported".to_string(),
|
||||
}),
|
||||
|
||||
@@ -31,8 +31,9 @@ pub(crate) async fn execute_delete(
|
||||
table.dataset.ensure_mutable()?;
|
||||
match predicate {
|
||||
Predicate::String(s) => {
|
||||
let predicate = crate::expr::canonicalize_sql_predicate(s)?;
|
||||
let mut dataset = (*table.dataset.get().await?).clone();
|
||||
let delete_result = dataset.delete(s).boxed().await?;
|
||||
let delete_result = dataset.delete(&predicate).boxed().await?;
|
||||
let num_deleted_rows = delete_result.num_deleted_rows;
|
||||
let version = dataset.version().version;
|
||||
table.dataset.update(dataset);
|
||||
|
||||
@@ -220,9 +220,32 @@ impl MergeInsertBuilder {
|
||||
///
|
||||
/// Returns version and statistics about the merge operation including the number of rows
|
||||
/// inserted, updated, and deleted.
|
||||
pub async fn execute(self, new_data: Box<dyn RecordBatchReader + Send>) -> Result<MergeResult> {
|
||||
pub async fn execute(
|
||||
mut self,
|
||||
new_data: Box<dyn RecordBatchReader + Send>,
|
||||
) -> Result<MergeResult> {
|
||||
self.canonicalize_filters()?;
|
||||
self.table.clone().merge_insert(self, new_data).await
|
||||
}
|
||||
|
||||
pub(crate) fn canonicalize_filters(&mut self) -> Result<()> {
|
||||
self.when_matched_update_all_filt =
|
||||
canonicalize_merge_filter(self.when_matched_update_all_filt.take())?;
|
||||
self.when_not_matched_by_source_delete_filt =
|
||||
canonicalize_merge_filter(self.when_not_matched_by_source_delete_filt.take())?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn canonicalize_merge_filter(filter: Option<MergeFilter>) -> Result<Option<MergeFilter>> {
|
||||
filter
|
||||
.map(|filter| match filter {
|
||||
MergeFilter::Sql(predicate) => {
|
||||
crate::expr::canonicalize_sql_predicate(&predicate).map(MergeFilter::Sql)
|
||||
}
|
||||
filter @ MergeFilter::Expr(_) => Ok(filter),
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
/// Internal implementation of the merge insert logic
|
||||
@@ -230,9 +253,10 @@ impl MergeInsertBuilder {
|
||||
/// This logic was moved from NativeTable::merge_insert to keep table.rs clean.
|
||||
pub(crate) async fn execute_merge_insert(
|
||||
table: &NativeTable,
|
||||
params: MergeInsertBuilder,
|
||||
mut params: MergeInsertBuilder,
|
||||
new_data: Box<dyn RecordBatchReader + Send>,
|
||||
) -> Result<MergeResult> {
|
||||
params.canonicalize_filters()?;
|
||||
super::computed_columns::ensure_no_function_bindings_for_mutation(
|
||||
table.schema().await?.as_ref(),
|
||||
"merge_insert",
|
||||
|
||||
@@ -21,7 +21,6 @@ use datafusion_physical_plan::ExecutionPlan;
|
||||
use datafusion_physical_plan::projection::ProjectionExec;
|
||||
use datafusion_physical_plan::repartition::RepartitionExec;
|
||||
use datafusion_physical_plan::union::UnionExec;
|
||||
use futures::future::try_join_all;
|
||||
use lance::dataset::mem_wal::DatasetMemWalExt;
|
||||
use lance::dataset::scanner::DatasetRecordBatchStream;
|
||||
use lance::dataset::scanner::Scanner;
|
||||
@@ -45,6 +44,22 @@ impl AnyQuery {
|
||||
Self::VectorQuery(query) => &query.base,
|
||||
}
|
||||
}
|
||||
|
||||
fn base_mut(&mut self) -> &mut QueryRequest {
|
||||
match self {
|
||||
Self::Query(query) => query,
|
||||
Self::VectorQuery(query) => &mut query.base,
|
||||
}
|
||||
}
|
||||
|
||||
/// Canonicalize any raw SQL filter immediately before backend dispatch.
|
||||
pub(crate) fn canonicalized(&self) -> Result<Self> {
|
||||
let mut query = self.clone();
|
||||
if let Some(QueryFilter::Sql(predicate)) = &mut query.base_mut().filter {
|
||||
*predicate = crate::expr::canonicalize_sql_predicate(predicate)?;
|
||||
}
|
||||
Ok(query)
|
||||
}
|
||||
}
|
||||
|
||||
//Decide between namespace or local
|
||||
@@ -53,15 +68,16 @@ pub async fn execute_query(
|
||||
query: &AnyQuery,
|
||||
options: QueryExecutionOptions,
|
||||
) -> Result<DatasetRecordBatchStream> {
|
||||
let query = query.canonicalized()?;
|
||||
// QueryTable pushdown runs the query server-side, but only on the main
|
||||
// branch: the namespace request carries no branch yet, so a branch handle
|
||||
// must fall through to local execution.
|
||||
if can_execute_namespace_query(table, query).await?
|
||||
if can_execute_namespace_query(table, &query).await?
|
||||
&& let Some(ref namespace_client) = table.namespace_client
|
||||
{
|
||||
return execute_namespace_query(table, namespace_client.clone(), query, options).await;
|
||||
return execute_namespace_query(table, namespace_client.clone(), &query, options).await;
|
||||
}
|
||||
execute_generic_query(table, query, options).await
|
||||
execute_generic_query(table, &query, options).await
|
||||
}
|
||||
|
||||
async fn can_execute_namespace_query(table: &NativeTable, query: &AnyQuery) -> Result<bool> {
|
||||
@@ -136,7 +152,8 @@ pub async fn create_plan(
|
||||
query: &AnyQuery,
|
||||
options: QueryExecutionOptions,
|
||||
) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
if let AnyQuery::Query(request) = query
|
||||
let query = query.canonicalized()?;
|
||||
if let AnyQuery::Query(request) = &query
|
||||
&& let Some(offsets) = &request.take_offsets
|
||||
{
|
||||
return crate::query::create_take_offsets_plan(table, request, offsets, options, false)
|
||||
@@ -144,8 +161,8 @@ pub async fn create_plan(
|
||||
}
|
||||
|
||||
let query = match query {
|
||||
AnyQuery::VectorQuery(query) => query.clone(),
|
||||
AnyQuery::Query(query) => VectorQueryRequest::from_plain_query(query.clone()),
|
||||
AnyQuery::VectorQuery(query) => query,
|
||||
AnyQuery::Query(query) => VectorQueryRequest::from_plain_query(query),
|
||||
};
|
||||
query.base.check_filter()?;
|
||||
|
||||
@@ -177,6 +194,7 @@ pub async fn create_plan(
|
||||
let mut column = query.column.clone();
|
||||
|
||||
let mut query_vector = query.query_vector.first().cloned();
|
||||
let mut is_batch_query = false;
|
||||
if query.query_vector.len() > 1 {
|
||||
if column.is_none() {
|
||||
// Infer a vector column with the same dimension of the query vector.
|
||||
@@ -187,16 +205,37 @@ pub async fn create_plan(
|
||||
)?);
|
||||
}
|
||||
let vector_field = schema.field(column.as_ref().unwrap()).unwrap();
|
||||
if let DataType::List(_) = vector_field.data_type() {
|
||||
// Multivector handling: concatenate into FixedSizeList<FixedSizeList<_>>
|
||||
let (_, element_type) =
|
||||
lance::index::vector::utils::get_vector_type(schema, column.as_ref().unwrap())?;
|
||||
let is_binary = matches!(element_type, DataType::UInt8);
|
||||
if matches!(vector_field.data_type(), DataType::List(_))
|
||||
|| (query.base.offset.unwrap_or(0) == 0 && !is_binary)
|
||||
{
|
||||
// Lance distinguishes these cases from the vector column type: a
|
||||
// list-like query against a List column is one multivector query,
|
||||
// while the same query against a FixedSizeList column is a batch of
|
||||
// independent queries. The batch path shares a single flat scan and
|
||||
// bounds retained candidate data instead of running one scan per
|
||||
// query vector.
|
||||
let vectors = query
|
||||
.query_vector
|
||||
.iter()
|
||||
.map(|arr| arr.as_ref())
|
||||
.collect::<Vec<_>>();
|
||||
let dim = vectors[0].len();
|
||||
if let Some((query_index, actual_dim)) = vectors
|
||||
.iter()
|
||||
.enumerate()
|
||||
.find_map(|(index, vector)| (vector.len() != dim).then_some((index, vector.len())))
|
||||
{
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!(
|
||||
"query vector at index {query_index} has dimension {actual_dim}, expected {dim}"
|
||||
),
|
||||
});
|
||||
}
|
||||
let mut fsl_builder = FixedSizeListBuilder::with_capacity(
|
||||
Float32Builder::with_capacity(dim),
|
||||
Float32Builder::with_capacity(dim * vectors.len()),
|
||||
dim as i32,
|
||||
vectors.len(),
|
||||
);
|
||||
@@ -207,8 +246,12 @@ pub async fn create_plan(
|
||||
fsl_builder.append(true);
|
||||
}
|
||||
query_vector = Some(Arc::new(fsl_builder.finish()));
|
||||
is_batch_query = !matches!(vector_field.data_type(), DataType::List(_));
|
||||
} else {
|
||||
// Multiple query vectors: create a plan for each and union them
|
||||
// Lance's batch path has no per-query offset, and its binary path
|
||||
// requires primitive UInt8 queries rather than a fixed-size list.
|
||||
// Keep the prior plan shape for these cases so offsets are applied
|
||||
// per query and binary query vectors retain their primitive shape.
|
||||
let query_vecs = query.query_vector.clone();
|
||||
let plan_futures = query_vecs
|
||||
.into_iter()
|
||||
@@ -221,7 +264,7 @@ pub async fn create_plan(
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let plans = try_join_all(plan_futures).await?;
|
||||
let plans = futures::future::try_join_all(plan_futures).await?;
|
||||
return create_multi_vector_plan(plans);
|
||||
}
|
||||
}
|
||||
@@ -258,10 +301,14 @@ pub async fn create_plan(
|
||||
}
|
||||
}
|
||||
|
||||
scanner.limit(
|
||||
query.base.limit.map(|limit| limit as i64),
|
||||
query.base.offset.map(|offset| offset as i64),
|
||||
)?;
|
||||
// For a batch query, `nearest` already applies k to each query vector.
|
||||
// Adding Scanner's global limit would truncate the combined result to k rows.
|
||||
if !is_batch_query {
|
||||
scanner.limit(
|
||||
query.base.limit.map(|limit| limit as i64),
|
||||
query.base.offset.map(|offset| offset as i64),
|
||||
)?;
|
||||
}
|
||||
|
||||
if let Some(ef) = query.ef {
|
||||
scanner.ef(ef);
|
||||
@@ -1095,7 +1142,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_plan_multivector_structure() {
|
||||
async fn test_create_plan_batch_vector_uses_shared_scan() {
|
||||
use arrow_array::{Float32Array, RecordBatch};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use datafusion_physical_plan::display::DisplayableExecutionPlan;
|
||||
@@ -1122,11 +1169,18 @@ mod tests {
|
||||
.unwrap();
|
||||
let native_table = table.as_native().unwrap();
|
||||
|
||||
// This triggers the "create_multi_vector_plan" logic branch
|
||||
// A batch of vectors against a fixed-size vector column should use
|
||||
// Lance's native batch KNN path instead of independent scan plans.
|
||||
let q1 = Arc::new(Float32Array::from(vec![1.0, 2.0]));
|
||||
let q2 = Arc::new(Float32Array::from(vec![3.0, 4.0]));
|
||||
|
||||
let req = VectorQueryRequest {
|
||||
base: QueryRequest {
|
||||
filter: Some(QueryFilter::Sql("id >= 0".to_string())),
|
||||
limit: Some(1),
|
||||
select: Select::Columns(vec!["id".to_string()]),
|
||||
..Default::default()
|
||||
},
|
||||
column: Some("vector".to_string()),
|
||||
query_vector: vec![q1, q2],
|
||||
..Default::default()
|
||||
@@ -1143,19 +1197,17 @@ mod tests {
|
||||
.indent(true)
|
||||
.to_string();
|
||||
|
||||
// We expect a RepartitionExec wrapping a UnionExec
|
||||
assert!(
|
||||
display.contains("RepartitionExec"),
|
||||
"Plan should include Repartitioning"
|
||||
display.contains("KNNVectorDistance: queries=2"),
|
||||
"plan should use native batch KNN, got:\n{display}"
|
||||
);
|
||||
assert!(
|
||||
display.contains("UnionExec"),
|
||||
"Plan should include a Union of multiple searches"
|
||||
!display.contains("UnionExec"),
|
||||
"flat batch KNN should share one scan, got:\n{display}"
|
||||
);
|
||||
// We expect the projection to add the 'query_index' column (logic inside multi_vector_plan)
|
||||
assert!(
|
||||
display.contains("query_index"),
|
||||
"Plan should add query_index column"
|
||||
"plan should add query_index column, got:\n{display}"
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -55,7 +55,9 @@ pub struct DropColumnsResult {
|
||||
pub struct FieldMetadataUpdate {
|
||||
/// Dot-separated path to the field (e.g. `"embedding"` or `"address.zip"`).
|
||||
pub path: String,
|
||||
/// Keys to set (`Some`) or delete (`None`).
|
||||
/// Keys to set (`Some`) or delete (`None`). See
|
||||
/// [`Table::update_field_metadata`](crate::Table::update_field_metadata) for
|
||||
/// the conventional `lancedb:*` keys.
|
||||
pub metadata: HashMap<String, Option<String>>,
|
||||
/// If `true`, replace the field's entire metadata map instead of merging.
|
||||
pub replace: bool,
|
||||
|
||||
@@ -62,22 +62,33 @@ impl UpdateBuilder {
|
||||
}
|
||||
|
||||
/// Executes the update operation.
|
||||
pub async fn execute(self) -> Result<UpdateResult> {
|
||||
pub async fn execute(mut self) -> Result<UpdateResult> {
|
||||
if self.columns.is_empty() {
|
||||
Err(Error::InvalidInput {
|
||||
message: "at least one column must be specified in an update operation".to_string(),
|
||||
})
|
||||
} else {
|
||||
self.canonicalize_filter()?;
|
||||
self.parent.clone().update(self).await
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn canonicalize_filter(&mut self) -> Result<()> {
|
||||
self.filter = self
|
||||
.filter
|
||||
.take()
|
||||
.map(|predicate| crate::expr::canonicalize_sql_predicate(&predicate))
|
||||
.transpose()?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Internal implementation of the update logic
|
||||
pub(crate) async fn execute_update(
|
||||
table: &NativeTable,
|
||||
update: UpdateBuilder,
|
||||
mut update: UpdateBuilder,
|
||||
) -> Result<UpdateResult> {
|
||||
update.canonicalize_filter()?;
|
||||
table.dataset.ensure_mutable()?;
|
||||
|
||||
// 1. Snapshot the current dataset
|
||||
|
||||
Reference in New Issue
Block a user