mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-29 23:06:51 +00:00
feat(sql): add vector duplicate pair table functions
This commit is contained in:
Yang Cen
Generated
+46
-44
@@ -3588,8 +3588,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
|
||||
|
||||
[[package]]
|
||||
name = "fsst"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"rand 0.9.5",
|
||||
@@ -5143,8 +5143,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a"
|
||||
|
||||
[[package]]
|
||||
name = "lance"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arrow",
|
||||
@@ -5216,8 +5216,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-arrow"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5239,7 +5239,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "lance-arrow-scalar"
|
||||
version = "58.0.0"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5253,7 +5253,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "lance-arrow-stats"
|
||||
version = "58.0.0"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5262,8 +5262,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-bitpacking"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrayref",
|
||||
"crunchy",
|
||||
@@ -5273,8 +5273,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-core"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5311,8 +5311,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-datafusion"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5343,8 +5343,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-datagen"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5361,8 +5361,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-derive"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -5371,8 +5371,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-encoding"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
@@ -5405,8 +5405,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-file"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
@@ -5438,8 +5438,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-geo"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"datafusion",
|
||||
"geo-traits",
|
||||
@@ -5453,8 +5453,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-index"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arrow",
|
||||
@@ -5481,6 +5481,7 @@ dependencies = [
|
||||
"geoarrow-array",
|
||||
"geoarrow-schema",
|
||||
"half",
|
||||
"hex",
|
||||
"itertools 0.14.0",
|
||||
"jieba-rs",
|
||||
"jsonb",
|
||||
@@ -5518,12 +5519,13 @@ dependencies = [
|
||||
"tempfile",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"twox-hash",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lance-index-core"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5545,8 +5547,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-io"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5588,8 +5590,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-linalg"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5603,8 +5605,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -5618,8 +5620,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace-impls"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-ipc",
|
||||
@@ -5672,8 +5674,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-select"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5687,8 +5689,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-table"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5728,8 +5730,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-testing"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5742,8 +5744,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-tokenizer"
|
||||
version = "13.0.0-beta.8"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d"
|
||||
version = "13.0.0-beta.11"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.11#84c8a40720d278f847fed5eaccae96d806dc375e"
|
||||
dependencies = [
|
||||
"frostem",
|
||||
"icu_segmenter",
|
||||
|
||||
+14
-14
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
lance = { "version" = "=13.0.0-beta.8", default-features = false, "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=13.0.0-beta.8", default-features = false, "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=13.0.0-beta.8", default-features = false, "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance = { "version" = "=13.0.0-beta.11", default-features = false, "tag" = "v13.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=13.0.0-beta.11", "tag" = "v13.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=13.0.0-beta.11", "tag" = "v13.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=13.0.0-beta.11", "tag" = "v13.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=13.0.0-beta.11", default-features = false, "tag" = "v13.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=13.0.0-beta.11", "tag" = "v13.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=13.0.0-beta.11", "tag" = "v13.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=13.0.0-beta.11", "tag" = "v13.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=13.0.0-beta.11", default-features = false, "tag" = "v13.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=13.0.0-beta.11", "tag" = "v13.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=13.0.0-beta.11", "tag" = "v13.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=13.0.0-beta.11", "tag" = "v13.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=13.0.0-beta.11", "tag" = "v13.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=13.0.0-beta.11", "tag" = "v13.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lancedb = { path = "rust/lancedb", default-features = false }
|
||||
ahash = "0.8"
|
||||
# Note that this one does not include pyarrow
|
||||
|
||||
@@ -4,3 +4,5 @@
|
||||
//! User-Defined Table Functions (UDTFs) for DataFusion integration
|
||||
|
||||
pub mod fts;
|
||||
|
||||
pub mod duplicate_pairs;
|
||||
|
||||
@@ -0,0 +1,571 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! Index-partition duplicate pairs. No source-vector scan or shuffle is needed.
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_schema::{DataType, Field, Schema, SchemaRef};
|
||||
use async_trait::async_trait;
|
||||
use datafusion::catalog::TableFunctionImpl;
|
||||
use datafusion_catalog::{Session, TableProvider};
|
||||
use datafusion_common::{DataFusionError, Result, ScalarValue, plan_err};
|
||||
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
|
||||
use datafusion_expr::{Expr, TableType};
|
||||
use datafusion_physical_expr::EquivalenceProperties;
|
||||
use datafusion_physical_plan::{
|
||||
DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
|
||||
execution_plan::{Boundedness, EmissionType},
|
||||
metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet},
|
||||
projection::ProjectionExec,
|
||||
stream::RecordBatchStreamAdapter,
|
||||
};
|
||||
use futures::{StreamExt, TryStreamExt, stream};
|
||||
use lance::{
|
||||
Dataset,
|
||||
index::{DatasetIndexExt, DatasetIndexInternalExt},
|
||||
};
|
||||
use lance_index::metrics::NoOpMetricsCollector;
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::super::BaseTableAdapter;
|
||||
use crate::table::NativeTableExt;
|
||||
|
||||
// Admission estimates for Lance 13's bounded pairwise kernel. Keep these with
|
||||
// the minimum dependency when the kernel's buffering contract changes.
|
||||
/// Version of the native task descriptor and its execution semantics.
|
||||
pub const DUPLICATE_PAIRS_OPERATOR_VERSION: u32 = 1;
|
||||
|
||||
const CODE_STAGING_BYTES: usize = 16 * 1024 * 1024;
|
||||
const INDEX_READ_ROWS: usize = 8192;
|
||||
const VECTOR_BATCH_ROWS: usize = 1024;
|
||||
const MAX_COORDINATE_BYTES: usize = 8;
|
||||
const PAIR_ROW_BYTES: usize = 20;
|
||||
|
||||
/// Configuration belongs to one immutable source snapshot. Threshold uses
|
||||
/// Lance's float32 index distance, including symmetric reconstructed distances
|
||||
/// for quantized indexes; it is not an exact source-embedding distance.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DuplicatePairsConfig {
|
||||
pub dataset_version: u64,
|
||||
pub column: String,
|
||||
pub distance_threshold: f32,
|
||||
}
|
||||
|
||||
/// One worker input, never a row of source vectors. The complete task identity
|
||||
/// also includes the dataset URI/version, column, threshold and operator version.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct DuplicatePairTask {
|
||||
pub segment_id: Uuid,
|
||||
pub partition_id: usize,
|
||||
}
|
||||
|
||||
/// Resolve an authorized source table and preserve its catalog identity.
|
||||
pub trait DuplicatePairsResolver: std::fmt::Debug + Send + Sync {
|
||||
fn resolve(
|
||||
&self,
|
||||
table: &str,
|
||||
config: DuplicatePairsConfig,
|
||||
tasks_only: bool,
|
||||
) -> Result<Arc<dyn TableProvider>>;
|
||||
}
|
||||
|
||||
/// `vector_duplicate_pairs('table', version, 'column', threshold)`.
|
||||
/// Registration is explicit so the host controls catalog access and execution.
|
||||
#[derive(Debug)]
|
||||
pub struct DuplicatePairsTableFunction {
|
||||
resolver: Arc<dyn DuplicatePairsResolver>,
|
||||
tasks_only: bool,
|
||||
}
|
||||
|
||||
impl DuplicatePairsTableFunction {
|
||||
pub fn new(resolver: Arc<dyn DuplicatePairsResolver>) -> Self {
|
||||
Self {
|
||||
resolver,
|
||||
tasks_only: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl DuplicatePairsTableFunction {
|
||||
/// Register as `vector_duplicate_pair_tasks` to inspect or persist exactly
|
||||
/// the descriptors used by the pair operator, without producing any pairs.
|
||||
pub fn task_manifest(resolver: Arc<dyn DuplicatePairsResolver>) -> Self {
|
||||
Self {
|
||||
resolver,
|
||||
tasks_only: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TableFunctionImpl for DuplicatePairsTableFunction {
|
||||
fn call(&self, args: &[Expr]) -> Result<Arc<dyn TableProvider>> {
|
||||
if args.len() != 4 {
|
||||
return plan_err!(
|
||||
"vector_duplicate_pairs requires (table, dataset_version, column, distance_threshold)"
|
||||
);
|
||||
}
|
||||
let string = |expr: &Expr| match expr {
|
||||
Expr::Literal(ScalarValue::Utf8(Some(s)) | ScalarValue::LargeUtf8(Some(s)), _) => {
|
||||
Ok(s.clone())
|
||||
}
|
||||
_ => plan_err!("table and column must be string literals"),
|
||||
};
|
||||
let dataset_version = match &args[1] {
|
||||
Expr::Literal(ScalarValue::Int64(Some(v)), _) if *v > 0 => *v as u64,
|
||||
Expr::Literal(ScalarValue::UInt64(Some(v)), _) if *v > 0 => *v,
|
||||
_ => return plan_err!("dataset_version must be a positive integer literal"),
|
||||
};
|
||||
let distance_threshold = match &args[3] {
|
||||
Expr::Literal(value, _) if value.data_type().is_numeric() => {
|
||||
value.cast_to(&DataType::Float32)?
|
||||
}
|
||||
_ => return plan_err!("distance_threshold must be a finite numeric literal"),
|
||||
};
|
||||
let ScalarValue::Float32(Some(distance_threshold)) = distance_threshold else {
|
||||
return plan_err!("distance_threshold cannot be null");
|
||||
};
|
||||
if !distance_threshold.is_finite() {
|
||||
return plan_err!("distance_threshold must be finite");
|
||||
}
|
||||
self.resolver.resolve(
|
||||
&string(&args[0])?,
|
||||
DuplicatePairsConfig {
|
||||
dataset_version,
|
||||
column: string(&args[2])?,
|
||||
distance_threshold,
|
||||
},
|
||||
self.tasks_only,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Three non-null columns; pair direction is storage traversal order, not row-ID order.
|
||||
pub fn duplicate_pair_schema() -> SchemaRef {
|
||||
Arc::new(Schema::new(vec![
|
||||
Field::new("row_id_a", DataType::UInt64, false),
|
||||
Field::new("row_id_b", DataType::UInt64, false),
|
||||
Field::new("distance", DataType::Float32, false),
|
||||
]))
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct DuplicatePairsTable {
|
||||
source: Arc<dyn crate::table::BaseTable>,
|
||||
config: DuplicatePairsConfig,
|
||||
tasks_only: bool,
|
||||
}
|
||||
|
||||
impl BaseTableAdapter {
|
||||
/// Bind the builtin to this authorized table; only metadata is read while
|
||||
/// planning. The source snapshot must remain retained until execution ends.
|
||||
pub fn duplicate_pairs(
|
||||
&self,
|
||||
config: DuplicatePairsConfig,
|
||||
tasks_only: bool,
|
||||
) -> DuplicatePairsTable {
|
||||
DuplicatePairsTable {
|
||||
source: self.table.clone(),
|
||||
config,
|
||||
tasks_only,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TableProvider for DuplicatePairsTable {
|
||||
fn schema(&self) -> SchemaRef {
|
||||
if self.tasks_only {
|
||||
duplicate_pair_task_schema()
|
||||
} else {
|
||||
duplicate_pair_schema()
|
||||
}
|
||||
}
|
||||
fn table_type(&self) -> TableType {
|
||||
TableType::View
|
||||
}
|
||||
async fn scan(
|
||||
&self,
|
||||
_state: &dyn Session,
|
||||
projection: Option<&Vec<usize>>,
|
||||
_filters: &[Expr],
|
||||
_limit: Option<usize>,
|
||||
) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
let native = self.source.as_native().ok_or_else(|| {
|
||||
DataFusionError::Plan("duplicate pairs requires a native indexed table".into())
|
||||
})?;
|
||||
let current = native
|
||||
.dataset
|
||||
.get()
|
||||
.await
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?;
|
||||
let dataset = Arc::new(
|
||||
current
|
||||
.checkout_version(self.config.dataset_version)
|
||||
.await?,
|
||||
);
|
||||
let tasks = plan_duplicate_pairs(dataset.clone(), &self.config).await?;
|
||||
if self.tasks_only {
|
||||
let batch = duplicate_pair_task_batch(&dataset, &self.config, &tasks)?;
|
||||
let table =
|
||||
datafusion::datasource::MemTable::try_new(batch.schema(), vec![vec![batch]])?;
|
||||
return table.scan(_state, projection, _filters, _limit).await;
|
||||
}
|
||||
let plan: Arc<dyn ExecutionPlan> = Arc::new(DuplicatePairsExec::try_new(
|
||||
dataset,
|
||||
self.config.clone(),
|
||||
tasks,
|
||||
)?);
|
||||
match projection {
|
||||
Some(indices) => {
|
||||
let schema = plan.schema();
|
||||
let exprs = indices
|
||||
.iter()
|
||||
.map(|&i| {
|
||||
(
|
||||
Arc::new(datafusion_physical_expr::expressions::Column::new(
|
||||
schema.field(i).name(),
|
||||
i,
|
||||
)) as _,
|
||||
schema.field(i).name().clone(),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
Ok(Arc::new(ProjectionExec::try_new(exprs, plan)?))
|
||||
}
|
||||
None => Ok(plan),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Enumerate segment-local IVF partitions using metadata only. Live-row masks
|
||||
/// and index codes are opened by the scoped worker, never by a source-vector
|
||||
/// scan on the coordinator. The high-level pair API is only a test oracle.
|
||||
pub async fn plan_duplicate_pairs(
|
||||
dataset: Arc<Dataset>,
|
||||
config: &DuplicatePairsConfig,
|
||||
) -> Result<Vec<DuplicatePairTask>> {
|
||||
if dataset.version().version != config.dataset_version || !config.distance_threshold.is_finite()
|
||||
{
|
||||
return plan_err!("duplicate pairs source snapshot or threshold is invalid");
|
||||
}
|
||||
let field = dataset.schema().field_id(&config.column)?;
|
||||
let indices = dataset.load_indices().await?;
|
||||
let segments = indices
|
||||
.iter()
|
||||
.filter(|meta| {
|
||||
meta.keyed_fields() == [field]
|
||||
&& meta.index_details.as_ref().map_or_else(
|
||||
|| {
|
||||
meta.files.as_ref().is_some_and(|files| {
|
||||
files
|
||||
.iter()
|
||||
.any(|file| file.path == lance_index::INDEX_FILE_NAME)
|
||||
})
|
||||
},
|
||||
|details| details.type_url.ends_with("VectorIndexDetails"),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let names = segments
|
||||
.iter()
|
||||
.map(|meta| meta.name.as_str())
|
||||
.collect::<BTreeSet<_>>();
|
||||
if names.len() != 1 {
|
||||
return plan_err!(
|
||||
"column '{}' requires exactly one current-format vector index, found {}",
|
||||
config.column,
|
||||
names.len()
|
||||
);
|
||||
}
|
||||
let name = names.into_iter().next().expect("one index name");
|
||||
if !dataset.unindexed_fragments(name).await?.is_empty() {
|
||||
return plan_err!(
|
||||
"vector index on column '{}' does not cover all fragments; optimize the index first",
|
||||
config.column
|
||||
);
|
||||
}
|
||||
let mut tasks = Vec::new();
|
||||
for meta in segments {
|
||||
if dataset.fragments().iter().any(|fragment| {
|
||||
meta.fragment_bitmap
|
||||
.as_ref()
|
||||
.is_none_or(|ids| ids.contains(fragment.id as u32))
|
||||
&& fragment.overlays.iter().any(|overlay| {
|
||||
overlay.committed_version > meta.dataset_version
|
||||
&& overlay.data_file.fields.contains(&field)
|
||||
})
|
||||
}) {
|
||||
return plan_err!(
|
||||
"column '{}' has stale index values in segment {}; rebuild the index",
|
||||
config.column,
|
||||
meta.uuid
|
||||
);
|
||||
}
|
||||
let index = dataset
|
||||
.open_vector_index(&config.column, &meta.uuid, &NoOpMetricsCollector)
|
||||
.await?;
|
||||
if !index.supports_pairwise_vectors() {
|
||||
return plan_err!(
|
||||
"segment {} requires a current-format pairwise vector index; rebuild the index",
|
||||
meta.uuid
|
||||
);
|
||||
}
|
||||
for partition_id in 0..index.ivf_model().num_partitions() {
|
||||
tasks.push(DuplicatePairTask {
|
||||
segment_id: meta.uuid,
|
||||
partition_id,
|
||||
});
|
||||
}
|
||||
}
|
||||
if tasks.is_empty() {
|
||||
return plan_err!("no vector index partitions found");
|
||||
}
|
||||
Ok(tasks)
|
||||
}
|
||||
|
||||
/// Streaming native operator over partition descriptors, never source-vector
|
||||
/// rows. Assigned descriptors are consumed sequentially; distributed hosts send
|
||||
/// one descriptor per worker request. No Python or per-vector ANN calls.
|
||||
/// Distributed merging deliberately advertises no global row ordering.
|
||||
#[derive(Debug)]
|
||||
pub struct DuplicatePairsExec {
|
||||
dataset: Arc<Dataset>,
|
||||
config: DuplicatePairsConfig,
|
||||
tasks: Vec<DuplicatePairTask>,
|
||||
properties: Arc<PlanProperties>,
|
||||
metrics: ExecutionPlanMetricsSet,
|
||||
}
|
||||
|
||||
impl DuplicatePairsExec {
|
||||
pub fn try_new(
|
||||
dataset: Arc<Dataset>,
|
||||
config: DuplicatePairsConfig,
|
||||
tasks: Vec<DuplicatePairTask>,
|
||||
) -> Result<Self> {
|
||||
if dataset.version().version != config.dataset_version
|
||||
|| !config.distance_threshold.is_finite()
|
||||
|| tasks.is_empty()
|
||||
{
|
||||
return plan_err!("invalid duplicate-pair snapshot, threshold, or empty task list");
|
||||
}
|
||||
let unique = tasks
|
||||
.iter()
|
||||
.map(|t| (t.segment_id, t.partition_id))
|
||||
.collect::<BTreeSet<_>>();
|
||||
if unique.len() != tasks.len() {
|
||||
return plan_err!("duplicate partition tasks");
|
||||
}
|
||||
let properties = Arc::new(PlanProperties::new(
|
||||
EquivalenceProperties::new(duplicate_pair_schema()),
|
||||
Partitioning::UnknownPartitioning(1),
|
||||
EmissionType::Incremental,
|
||||
Boundedness::Bounded,
|
||||
));
|
||||
Ok(Self {
|
||||
dataset,
|
||||
config,
|
||||
tasks,
|
||||
properties,
|
||||
metrics: ExecutionPlanMetricsSet::new(),
|
||||
})
|
||||
}
|
||||
pub fn dataset(&self) -> &Arc<Dataset> {
|
||||
&self.dataset
|
||||
}
|
||||
pub fn config(&self) -> &DuplicatePairsConfig {
|
||||
&self.config
|
||||
}
|
||||
pub fn tasks(&self) -> &[DuplicatePairTask] {
|
||||
&self.tasks
|
||||
}
|
||||
}
|
||||
|
||||
impl DisplayAs for DuplicatePairsExec {
|
||||
fn fmt_as(&self, _: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"DuplicatePairsExec: source_version={}, column={}, threshold={}, tasks={}",
|
||||
self.config.dataset_version,
|
||||
self.config.column,
|
||||
self.config.distance_threshold,
|
||||
self.tasks.len()
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl ExecutionPlan for DuplicatePairsExec {
|
||||
fn name(&self) -> &str {
|
||||
"DuplicatePairsExec"
|
||||
}
|
||||
fn properties(&self) -> &Arc<PlanProperties> {
|
||||
&self.properties
|
||||
}
|
||||
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
|
||||
vec![]
|
||||
}
|
||||
fn with_new_children(
|
||||
self: Arc<Self>,
|
||||
children: Vec<Arc<dyn ExecutionPlan>>,
|
||||
) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
if !children.is_empty() {
|
||||
return plan_err!("DuplicatePairsExec has no children");
|
||||
}
|
||||
Ok(self)
|
||||
}
|
||||
fn metrics(&self) -> Option<MetricsSet> {
|
||||
Some(self.metrics.clone_inner())
|
||||
}
|
||||
fn execute(
|
||||
&self,
|
||||
partition: usize,
|
||||
context: Arc<TaskContext>,
|
||||
) -> Result<SendableRecordBatchStream> {
|
||||
if partition != 0 {
|
||||
return plan_err!("duplicate-pair output partition out of bounds");
|
||||
}
|
||||
let tasks = self.tasks.clone();
|
||||
let dataset = self.dataset.clone();
|
||||
let config = self.config.clone();
|
||||
let output_rows = MetricBuilder::new(&self.metrics).counter("output_rows", partition);
|
||||
let completed =
|
||||
MetricBuilder::new(&self.metrics).counter("completed_partitions", partition);
|
||||
// Admission estimate for the released Lance kernel: compact-code staging,
|
||||
// one 8192-row read and two decoded 1024-row batches. This reservation
|
||||
// limits concurrent kernels; Lance still owns its bounded buffers/spill.
|
||||
let dimension = match dataset
|
||||
.schema()
|
||||
.field(&config.column)
|
||||
.map(|f| f.data_type())
|
||||
{
|
||||
Some(DataType::FixedSizeList(_, dimension)) if dimension > 0 => dimension as usize,
|
||||
_ => return plan_err!("duplicate pairs requires a fixed-size vector column"),
|
||||
};
|
||||
let budget = CODE_STAGING_BYTES
|
||||
+ (INDEX_READ_ROWS + 2 * VECTOR_BATCH_ROWS) * dimension * MAX_COORDINATE_BYTES
|
||||
+ VECTOR_BATCH_ROWS * PAIR_ROW_BYTES;
|
||||
let reservation =
|
||||
datafusion_execution::memory_pool::MemoryConsumer::new("DuplicatePairsExec")
|
||||
.register(context.memory_pool());
|
||||
reservation.try_grow(budget)?;
|
||||
let reservation = Arc::new(reservation);
|
||||
// Sequential local work prevents one query from reserving a kernel for
|
||||
// every IVF partition at once. Distributed hosts split these descriptors
|
||||
// into one-task PE requests before execution.
|
||||
let stream = stream::iter(tasks)
|
||||
.then(move |task| {
|
||||
let dataset = dataset.clone();
|
||||
let config = config.clone();
|
||||
let reservation = reservation.clone();
|
||||
let output_rows = output_rows.clone();
|
||||
let completed = completed.clone();
|
||||
async move {
|
||||
let reader = lance::index::vector::dedup::find_duplicate_pairs_in_partition(
|
||||
dataset,
|
||||
&config.column,
|
||||
task.segment_id,
|
||||
task.partition_id,
|
||||
config.distance_threshold,
|
||||
)
|
||||
.await?;
|
||||
let measured = stream::try_unfold(
|
||||
(reader, reservation),
|
||||
move |(mut reader, reservation)| {
|
||||
let output_rows = output_rows.clone();
|
||||
let completed = completed.clone();
|
||||
async move {
|
||||
match reader.try_next().await? {
|
||||
Some(batch) => {
|
||||
output_rows.add(batch.num_rows());
|
||||
Ok::<_, DataFusionError>(Some((
|
||||
batch,
|
||||
(reader, reservation),
|
||||
)))
|
||||
}
|
||||
None => {
|
||||
completed.add(1);
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
);
|
||||
Ok::<_, DataFusionError>(Box::pin(RecordBatchStreamAdapter::new(
|
||||
duplicate_pair_schema(),
|
||||
measured,
|
||||
)) as SendableRecordBatchStream)
|
||||
}
|
||||
})
|
||||
.try_flatten();
|
||||
Ok(Box::pin(RecordBatchStreamAdapter::new(
|
||||
duplicate_pair_schema(),
|
||||
stream,
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
/// Inspectable immutable task manifest. Credentials are never included.
|
||||
pub fn duplicate_pair_task_schema() -> SchemaRef {
|
||||
Arc::new(Schema::new(vec![
|
||||
Field::new("task_id", DataType::Utf8, false),
|
||||
Field::new("dataset_uri", DataType::Utf8, false),
|
||||
Field::new("dataset_version", DataType::UInt64, false),
|
||||
Field::new("column", DataType::Utf8, false),
|
||||
Field::new("segment_id", DataType::Utf8, false),
|
||||
Field::new("partition_id", DataType::UInt64, false),
|
||||
Field::new("distance_threshold", DataType::Float32, false),
|
||||
Field::new("operator_version", DataType::UInt32, false),
|
||||
]))
|
||||
}
|
||||
|
||||
pub fn duplicate_pair_task_batch(
|
||||
dataset: &Dataset,
|
||||
config: &DuplicatePairsConfig,
|
||||
tasks: &[DuplicatePairTask],
|
||||
) -> Result<arrow_array::RecordBatch> {
|
||||
use arrow_array::{Float32Array, StringArray, UInt32Array, UInt64Array};
|
||||
use sha2::{Digest, Sha256};
|
||||
let ids = tasks
|
||||
.iter()
|
||||
.map(|task| {
|
||||
// Length-delimited strings prevent ambiguous concatenation; threshold
|
||||
// bits preserve the exact float32 value sent to the native primitive.
|
||||
let bytes = serde_json::to_vec(&(
|
||||
DUPLICATE_PAIRS_OPERATOR_VERSION,
|
||||
dataset.uri(),
|
||||
config.dataset_version,
|
||||
&config.column,
|
||||
task.segment_id.to_string(),
|
||||
task.partition_id,
|
||||
config.distance_threshold.to_bits(),
|
||||
))
|
||||
.expect("task fields are serializable");
|
||||
format!("{:x}", Sha256::digest(bytes))
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let n = tasks.len();
|
||||
Ok(arrow_array::RecordBatch::try_new(
|
||||
duplicate_pair_task_schema(),
|
||||
vec![
|
||||
Arc::new(StringArray::from(ids)),
|
||||
Arc::new(StringArray::from(vec![dataset.uri(); n])),
|
||||
Arc::new(UInt64Array::from(vec![config.dataset_version; n])),
|
||||
Arc::new(StringArray::from(vec![config.column.as_str(); n])),
|
||||
Arc::new(StringArray::from(
|
||||
tasks
|
||||
.iter()
|
||||
.map(|t| t.segment_id.to_string())
|
||||
.collect::<Vec<_>>(),
|
||||
)),
|
||||
Arc::new(UInt64Array::from(
|
||||
tasks
|
||||
.iter()
|
||||
.map(|t| t.partition_id as u64)
|
||||
.collect::<Vec<_>>(),
|
||||
)),
|
||||
Arc::new(Float32Array::from(vec![config.distance_threshold; n])),
|
||||
Arc::new(UInt32Array::from(vec![DUPLICATE_PAIRS_OPERATOR_VERSION; n])),
|
||||
],
|
||||
)?)
|
||||
}
|
||||
@@ -0,0 +1,239 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use arrow_array::{
|
||||
Array, FixedSizeListArray, Float32Array, Int64Array, RecordBatch, RecordBatchIterator,
|
||||
UInt64Array,
|
||||
};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use datafusion::prelude::SessionContext;
|
||||
use datafusion_catalog::TableProvider;
|
||||
use datafusion_common::Result;
|
||||
use futures::TryStreamExt;
|
||||
use lance::{
|
||||
Dataset,
|
||||
dataset::WriteParams,
|
||||
index::{DatasetIndexExt, vector::VectorIndexParams},
|
||||
};
|
||||
use lance_index::{IndexType, vector::ivf::IvfBuildParams};
|
||||
use lance_linalg::distance::MetricType;
|
||||
use lancedb::table::datafusion::{
|
||||
BaseTableAdapter,
|
||||
udtf::duplicate_pairs::{
|
||||
DuplicatePairsConfig, DuplicatePairsExec, DuplicatePairsResolver,
|
||||
DuplicatePairsTableFunction, duplicate_pair_task_batch, plan_duplicate_pairs,
|
||||
},
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Resolver(Arc<BaseTableAdapter>);
|
||||
impl DuplicatePairsResolver for Resolver {
|
||||
fn resolve(
|
||||
&self,
|
||||
table: &str,
|
||||
config: DuplicatePairsConfig,
|
||||
tasks_only: bool,
|
||||
) -> Result<Arc<dyn TableProvider>> {
|
||||
assert_eq!(table, "source");
|
||||
Ok(Arc::new(self.0.duplicate_pairs(config, tasks_only)))
|
||||
}
|
||||
}
|
||||
|
||||
fn pairs(batches: &[RecordBatch]) -> Vec<(u64, u64, u32)> {
|
||||
let mut result = Vec::new();
|
||||
for batch in batches {
|
||||
let a = batch
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<UInt64Array>()
|
||||
.unwrap();
|
||||
let b = batch
|
||||
.column(1)
|
||||
.as_any()
|
||||
.downcast_ref::<UInt64Array>()
|
||||
.unwrap();
|
||||
let d = batch
|
||||
.column(2)
|
||||
.as_any()
|
||||
.downcast_ref::<Float32Array>()
|
||||
.unwrap();
|
||||
for i in 0..batch.num_rows() {
|
||||
result.push((a.value(i), b.value(i), d.value(i).to_bits()));
|
||||
}
|
||||
}
|
||||
result.sort_unstable();
|
||||
result
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result<()> {
|
||||
for partitions in [1, 8] {
|
||||
let dir = tempfile::tempdir()?;
|
||||
let uri = dir
|
||||
.path()
|
||||
.join("source.lance")
|
||||
.to_str()
|
||||
.unwrap()
|
||||
.to_string();
|
||||
let n = 1025;
|
||||
let vectors = (0..n * 2)
|
||||
.flat_map(|i| {
|
||||
let i = i % n;
|
||||
let x = if i == n - 1 {
|
||||
0.0
|
||||
} else if i < 3 {
|
||||
i as f32
|
||||
} else {
|
||||
i as f32 * 10.0
|
||||
};
|
||||
[x, 0.0]
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let vectors = FixedSizeListArray::try_new(
|
||||
Arc::new(Field::new("item", DataType::Float32, true)),
|
||||
2,
|
||||
Arc::new(Float32Array::from(vectors)),
|
||||
None,
|
||||
)?;
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int64, false),
|
||||
Field::new("vector", vectors.data_type().clone(), false),
|
||||
]));
|
||||
let batch = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(Int64Array::from_iter_values(0..(n * 2) as i64)),
|
||||
Arc::new(vectors),
|
||||
],
|
||||
)?;
|
||||
let mut ds = Dataset::write(
|
||||
RecordBatchIterator::new(vec![Ok(batch)], schema),
|
||||
&uri,
|
||||
Some(WriteParams {
|
||||
max_rows_per_file: n,
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
let fragments = ds
|
||||
.fragments()
|
||||
.iter()
|
||||
.map(|f| f.id as u32)
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(fragments.len(), 2);
|
||||
let mut segments = Vec::new();
|
||||
for (i, fragment) in fragments.into_iter().enumerate() {
|
||||
let centroids = FixedSizeListArray::try_new(
|
||||
Arc::new(Field::new("item", DataType::Float32, true)),
|
||||
2,
|
||||
Arc::new(Float32Array::from(
|
||||
(0..partitions)
|
||||
.flat_map(|p| [p as f32 * 1400.0 + i as f32 * 0.125, 0.0])
|
||||
.collect::<Vec<_>>(),
|
||||
)),
|
||||
None,
|
||||
)?;
|
||||
let ivf = IvfBuildParams {
|
||||
centroids: Some(Arc::new(centroids)),
|
||||
..IvfBuildParams::new(partitions)
|
||||
};
|
||||
let params = VectorIndexParams::with_ivf_flat_params(MetricType::L2, ivf);
|
||||
segments.push(
|
||||
ds.create_index_builder(&["vector"], IndexType::Vector, ¶ms)
|
||||
.name("vector_idx".into())
|
||||
.fragments(vec![fragment])
|
||||
.execute_uncommitted()
|
||||
.await?,
|
||||
);
|
||||
}
|
||||
ds.commit_existing_index_segments("vector_idx", "vector", segments)
|
||||
.await?;
|
||||
let version = ds.version().version;
|
||||
let ds = Arc::new(ds);
|
||||
let conn = lancedb::connect(dir.path().to_str().unwrap())
|
||||
.execute()
|
||||
.await?;
|
||||
let table = conn.open_table("source").execute().await?;
|
||||
let resolver = Arc::new(Resolver(Arc::new(
|
||||
BaseTableAdapter::try_new(table.base_table().clone()).await?,
|
||||
)));
|
||||
// More partitions than this pool could hold concurrently: local
|
||||
// execution must consume scoped readers sequentially.
|
||||
let runtime = Arc::new(
|
||||
datafusion_execution::runtime_env::RuntimeEnvBuilder::new()
|
||||
.with_memory_limit(32 * 1024 * 1024, 1.0)
|
||||
.build()?,
|
||||
);
|
||||
let ctx =
|
||||
SessionContext::new_with_config_rt(datafusion::prelude::SessionConfig::new(), runtime);
|
||||
ctx.register_udtf(
|
||||
"vector_duplicate_pairs",
|
||||
Arc::new(DuplicatePairsTableFunction::new(resolver.clone())),
|
||||
);
|
||||
ctx.register_udtf(
|
||||
"vector_duplicate_pair_tasks",
|
||||
Arc::new(DuplicatePairsTableFunction::task_manifest(resolver)),
|
||||
);
|
||||
for threshold in [-1.0, 0.0, 1.0] {
|
||||
let config = DuplicatePairsConfig {
|
||||
dataset_version: version,
|
||||
column: "vector".into(),
|
||||
distance_threshold: threshold,
|
||||
};
|
||||
let tasks = plan_duplicate_pairs(ds.clone(), &config).await?;
|
||||
assert_eq!(tasks.len(), partitions * 2);
|
||||
let manifest = duplicate_pair_task_batch(&ds, &config, &tasks)?;
|
||||
let sql_manifest = ctx.sql(&format!("SELECT * FROM vector_duplicate_pair_tasks('source', {version}, 'vector', {threshold})")).await?.collect().await?;
|
||||
assert_eq!(sql_manifest, vec![manifest]);
|
||||
let expected =
|
||||
lance::index::vector::dedup::find_duplicate_pairs(ds.clone(), "vector", threshold)
|
||||
.await?
|
||||
.try_collect::<Vec<_>>()
|
||||
.await?;
|
||||
let sql = format!(
|
||||
"SELECT * FROM vector_duplicate_pairs('source', {version}, 'vector', {threshold})"
|
||||
);
|
||||
let actual = ctx.sql(&sql).await?.collect().await?;
|
||||
assert_eq!(pairs(&actual), pairs(&expected));
|
||||
assert_eq!(
|
||||
pairs(&actual).len(),
|
||||
match threshold {
|
||||
-1.0 => 0,
|
||||
0.0 => 2,
|
||||
_ => 8,
|
||||
}
|
||||
);
|
||||
let retry = ctx.sql(&sql).await?.collect().await?;
|
||||
assert_eq!(pairs(&retry), pairs(&actual));
|
||||
assert!(
|
||||
DuplicatePairsExec::try_new(
|
||||
ds.clone(),
|
||||
config.clone(),
|
||||
vec![tasks[0].clone(), tasks[0].clone()]
|
||||
)
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
// Latest no longer contains B, but pinned queries must retain the same pairs.
|
||||
table.delete("id = 1").await?;
|
||||
let actual = ctx
|
||||
.sql(&format!(
|
||||
"SELECT * FROM vector_duplicate_pairs('source', {version}, 'vector', 1.0)"
|
||||
))
|
||||
.await?
|
||||
.collect()
|
||||
.await?;
|
||||
assert_eq!(pairs(&actual).len(), 8);
|
||||
assert!(
|
||||
ctx.sql(&format!(
|
||||
"SELECT * FROM vector_duplicate_pairs('source', {version}, 'missing', 1.0)"
|
||||
))
|
||||
.await?
|
||||
.collect()
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user