From b8af68728617a0d50f5ea67ce32cf6febf2cbbdf Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Tue, 29 Sep 2026 13:24:21 +0800 Subject: [PATCH] feat(sql): expose composable native vector relations (#4340) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Native vector dedup currently materializes through the native MV refresh protocol. Add composable DataFusion pair and retained-row relations so a host can use its existing SQL job and MV writer instead. `VectorQueryTable` binds a retained snapshot, selects all/first/random nonempty IVF scopes before dispatch, and exposes `DuplicatePairsExec` as its physical child. Pair queries can include segment, partition and snapshot columns and fetch a scalar application key for each endpoint without reading source vectors for scoring. The default primitive schema remains unchanged. `VectorDedupExec` consumes the pair stream, reuses the spillable greedy direct-representative selection, and scans retained original rows from that snapshot. A–B–C without A–C retains A and C. This query path has no durable partition receipts or writes; publication and failure belong to the enclosing SQL sink. Existing native MV APIs remain available. `BaseTableAdapter::vector_query_snapshot(None)` captures latest without changing a shared table handle's pin; `Some(version)` selects a retained version. A host can retain unbound SQL and resolve latest again for each refresh. Single-segment random sampling matches Sophon's `ivf_partition` SplitMix64 policy; multi-segment sampling ranks canonical nonempty scopes. Host integration and SQL syntax: https://github.com/lancedb/sophon/pull/8135. The host authorizes sources, binds named arguments, distributes the pair child, and uses its existing SQL Job Registry and atomic MV replacement. ## Validation All manual validation ran on an AWS r8i.8xlarge VM, not the developer laptop. - Focused Rust integration suite: 2 passed, including multiple segments/partitions, empty results, threshold boundaries, 8,193 rows per segment, deterministic sampling, mapped application IDs, and native-oracle pair equality. - Additional compacted stable/nonstable row-ID and blob test: 1 passed, including application-key mapping, null vectors, snapshot deletions and explicit/latest version capture. - Rust formatting and focused Clippy with warnings denied passed. - Dependent Sophon SQL-job E2E passes with two worker PEs: pairs, sampling, direct representatives, latest/fixed snapshot refreshes, and blob payloads. Killing a SQL worker after partial staging preserves the old MV; full rerun commits the complete replacement. Performance measurement is in progress in the dependent PR; no performance claim is made here. --- rust/lancedb/src/materialized_view.rs | 2 +- .../src/materialized_view/vector_dedup.rs | 2 +- rust/lancedb/src/table/datafusion/udtf.rs | 2 + .../table/datafusion/udtf/duplicate_pairs.rs | 124 +++++- .../src/table/datafusion/udtf/vector_query.rs | 352 ++++++++++++++++++ rust/lancedb/tests/duplicate_pairs.rs | 191 +++++++++- 6 files changed, 664 insertions(+), 9 deletions(-) create mode 100644 rust/lancedb/src/table/datafusion/udtf/vector_query.rs diff --git a/rust/lancedb/src/materialized_view.rs b/rust/lancedb/src/materialized_view.rs index 616bd42d4..7dec300a2 100644 --- a/rust/lancedb/src/materialized_view.rs +++ b/rust/lancedb/src/materialized_view.rs @@ -13,7 +13,7 @@ mod duplicate_pairs; mod grouped; mod grouped_units; mod partitioned; -mod vector_dedup; +pub(crate) mod vector_dedup; pub use duplicate_pairs::{VectorSource, VectorSourceKind}; pub use grouped::IVF_PARTITION; /// Refreshing a view grouped by [`IVF_PARTITION`] in units, one per index diff --git a/rust/lancedb/src/materialized_view/vector_dedup.rs b/rust/lancedb/src/materialized_view/vector_dedup.rs index 38c007076..dc2759433 100644 --- a/rust/lancedb/src/materialized_view/vector_dedup.rs +++ b/rust/lancedb/src/materialized_view/vector_dedup.rs @@ -156,7 +156,7 @@ fn edge_schema() -> SchemaRef { /// Sorting canonical edges makes the first unremoved endpoint the smallest /// unassigned representative. The sort spills; only vertex masks stay in RAM. -async fn select_representatives( +pub async fn select_representatives( stream: SendableRecordBatchStream, memory: usize, ) -> Result<(RoaringTreemap, RoaringTreemap)> { diff --git a/rust/lancedb/src/table/datafusion/udtf.rs b/rust/lancedb/src/table/datafusion/udtf.rs index a7a72f800..53eb55fcc 100644 --- a/rust/lancedb/src/table/datafusion/udtf.rs +++ b/rust/lancedb/src/table/datafusion/udtf.rs @@ -6,3 +6,5 @@ pub mod fts; pub mod duplicate_pairs; + +pub mod vector_query; diff --git a/rust/lancedb/src/table/datafusion/udtf/duplicate_pairs.rs b/rust/lancedb/src/table/datafusion/udtf/duplicate_pairs.rs index 47a35cbd2..20de93148 100644 --- a/rust/lancedb/src/table/datafusion/udtf/duplicate_pairs.rs +++ b/rust/lancedb/src/table/datafusion/udtf/duplicate_pairs.rs @@ -330,6 +330,96 @@ pub async fn plan_duplicate_pairs( Ok(tasks) } +/// Optional query columns. The default preserves the primitive's three-column +/// schema; SQL convenience calls include physical scope identity and may map +/// endpoints to an application key without fetching source vectors. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct DuplicatePairsOutput { + pub scope: bool, + pub id_column: Option, +} + +impl DuplicatePairsOutput { + pub fn schema(&self, dataset: &Dataset) -> Result { + let mut fields = duplicate_pair_schema().fields().to_vec(); + if self.scope { + fields.extend([ + Arc::new(Field::new("segment_id", DataType::Utf8, false)), + Arc::new(Field::new("partition_id", DataType::UInt64, false)), + Arc::new(Field::new("dataset_version", DataType::UInt64, false)), + ]); + } + if let Some(column) = &self.id_column { + let schema = Schema::from(dataset.schema()); + let field = schema.field_with_name(column)?; + if column.contains('.') + || matches!( + field.data_type(), + DataType::FixedSizeList(..) + | DataType::List(..) + | DataType::LargeList(..) + | DataType::Struct(..) + | DataType::Map(..) + ) + { + return plan_err!("id_column must be a top-level scalar column"); + } + fields.push(Arc::new(Field::new( + "id_a", + field.data_type().clone(), + field.is_nullable(), + ))); + fields.push(Arc::new(Field::new( + "id_b", + field.data_type().clone(), + field.is_nullable(), + ))); + } + Ok(Arc::new(Schema::new(fields))) + } + + async fn enrich( + &self, + batch: arrow_array::RecordBatch, + dataset: &Dataset, + task: &DuplicatePairTask, + ) -> Result { + use arrow_array::{StringArray, UInt64Array}; + let mut columns = batch.columns().to_vec(); + let n = batch.num_rows(); + if self.scope { + columns.extend([ + Arc::new(StringArray::from(vec![task.segment_id.to_string(); n])) + as arrow_array::ArrayRef, + Arc::new(UInt64Array::from(vec![task.partition_id as u64; n])), + Arc::new(UInt64Array::from(vec![dataset.version().version; n])), + ]); + } + if let Some(column) = &self.id_column { + for endpoint in 0..2 { + let ids = batch + .column(endpoint) + .as_any() + .downcast_ref::() + .expect("native row IDs"); + let projection = lance::dataset::ProjectionRequest::from_columns( + [column.as_str()], + dataset.schema(), + ); + let values = dataset.take_rows(ids.values(), projection).await?; + if values.num_rows() != n { + return plan_err!("pair row ID is missing from the pinned snapshot"); + } + columns.push(values.column(0).clone()); + } + } + Ok(arrow_array::RecordBatch::try_new( + self.schema(dataset)?, + columns, + )?) + } +} + /// 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. @@ -341,6 +431,7 @@ pub struct DuplicatePairsExec { tasks: Vec, properties: Arc, metrics: ExecutionPlanMetricsSet, + output: DuplicatePairsOutput, } impl DuplicatePairsExec { @@ -374,8 +465,23 @@ impl DuplicatePairsExec { tasks, properties, metrics: ExecutionPlanMetricsSet::new(), + output: DuplicatePairsOutput::default(), }) } + /// Configure optional output fields before dispatching this plan. + pub fn with_output(mut self, output: DuplicatePairsOutput) -> Result { + self.properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(output.schema(&self.dataset)?), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + self.output = output; + Ok(self) + } + pub fn output(&self) -> &DuplicatePairsOutput { + &self.output + } pub fn dataset(&self) -> &Arc { &self.dataset } @@ -430,6 +536,7 @@ impl ExecutionPlan for DuplicatePairsExec { if partition != 0 { return plan_err!("duplicate-pair output partition out of bounds"); } + let output = self.output.clone(); let tasks = self.tasks.clone(); let dataset = self.dataset.clone(); let config = self.config.clone(); @@ -467,6 +574,7 @@ impl ExecutionPlan for DuplicatePairsExec { // into one-task PE requests before execution. let stream = stream::iter(tasks) .then(move |task| { + let output = output.clone(); let dataset = dataset.clone(); let config = config.clone(); let reservation = reservation.clone(); @@ -474,7 +582,7 @@ impl ExecutionPlan for DuplicatePairsExec { let completed = completed.clone(); async move { let reader = lance::index::vector::dedup::find_duplicate_pairs_in_partition_with_options( - dataset, + dataset.clone(), &config.column, task.segment_id, task.partition_id, @@ -506,15 +614,19 @@ impl ExecutionPlan for DuplicatePairsExec { } }, ); - Ok::<_, DataFusionError>(Box::pin(RecordBatchStreamAdapter::new( - duplicate_pair_schema(), - measured, - )) as SendableRecordBatchStream) + let schema = output.schema(&dataset)?; + let measured = measured.and_then(move |batch| { + let output = output.clone(); + let dataset = dataset.clone(); + let task = task.clone(); + async move { output.enrich(batch, &dataset, &task).await } + }); + Ok::<_, DataFusionError>(Box::pin(RecordBatchStreamAdapter::new(schema, measured)) as SendableRecordBatchStream) } }) .try_flatten(); Ok(Box::pin(RecordBatchStreamAdapter::new( - duplicate_pair_schema(), + self.schema(), stream, ))) } diff --git a/rust/lancedb/src/table/datafusion/udtf/vector_query.rs b/rust/lancedb/src/table/datafusion/udtf/vector_query.rs new file mode 100644 index 000000000..0d4f23d92 --- /dev/null +++ b/rust/lancedb/src/table/datafusion/udtf/vector_query.rs @@ -0,0 +1,352 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Composable index-backed pair and dedup relations. Hosts bind a snapshot once +//! during planning, distribute the pair child, and use their ordinary SQL sink. + +use std::sync::Arc; + +use arrow_schema::SchemaRef; +use async_trait::async_trait; +use datafusion_catalog::{Session, TableProvider}; +use datafusion_common::{Result, plan_err}; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_expr::{Expr, TableType}; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + coalesce_partitions::CoalescePartitionsExec, + empty::EmptyExec, + execution_plan::{Boundedness, EmissionType}, + projection::ProjectionExec, + stream::RecordBatchStreamAdapter, +}; +use futures::{TryStreamExt, stream}; +use lance::dataset::scanner::{RowAddrMask, RowAddrTreeMap}; +use lance::{Dataset, index::DatasetIndexInternalExt}; +use lance_core::datatypes::BlobHandling; +use lance_index::metrics::NoOpMetricsCollector; + +use super::duplicate_pairs::{ + DuplicatePairTask, DuplicatePairsConfig, DuplicatePairsExec, DuplicatePairsOutput, + plan_duplicate_pairs, +}; +use crate::table::{NativeTableExt, datafusion::BaseTableAdapter}; + +const REPRESENTATIVE_SORT_MEMORY_BYTES: usize = 256 * 1024 * 1024; +const MATERIALIZE_BATCH_ROWS: usize = 1024; + +/// Select nonempty physical IVF scopes before dispatch, never by filtering the +/// resulting pairs. A count spans all segments; partition IDs are segment-local. +#[derive(Debug, Clone, Default)] +pub enum PartitionSelection { + #[default] + All, + First { + count: usize, + }, + Random { + count: usize, + seed: u64, + }, +} + +/// Query-only options. Snapshot binding is separate so multiple calls in one +/// statement can share the same immutable dataset. +#[derive(Debug, Clone)] +pub struct VectorQueryOptions { + pub column: String, + pub distance_threshold: f32, + pub selection: PartitionSelection, + pub output: DuplicatePairsOutput, + pub dedup: bool, +} + +impl BaseTableAdapter { + /// Open a retained version, or refresh and capture the latest manifest once. + /// No worker is allowed to independently resolve "latest". + pub async fn vector_query_snapshot(&self, version: Option) -> Result> { + let native = self.table.as_native().ok_or_else(|| { + datafusion_common::DataFusionError::Plan( + "vector queries require a native indexed table".into(), + ) + })?; + if version == Some(0) { + return plan_err!("dataset_version must be positive"); + } + let current = native + .dataset + .get() + .await + .map_err(|e| datafusion_common::DataFusionError::External(Box::new(e)))?; + Ok(match version { + Some(version) => Arc::new(current.checkout_version(version).await?), + None => { + // Do not change a shared table handle's pinned version. + let mut latest = (*current).clone(); + latest.checkout_latest().await?; + Arc::new(latest) + } + }) + } +} + +/// A table provider bound to one source snapshot. The host retains the user's +/// unbound SQL for subsequent refreshes; this provider lives only for this run. +#[derive(Debug)] +pub struct VectorQueryTable { + dataset: Arc, + config: DuplicatePairsConfig, + tasks: Vec, + options: VectorQueryOptions, + schema: SchemaRef, +} + +impl VectorQueryTable { + pub async fn try_new(dataset: Arc, options: VectorQueryOptions) -> Result { + if options.dedup + && (!matches!(options.selection, PartitionSelection::All) + || options.output.id_column.is_some()) + { + return plan_err!( + "vector_dedup requires all partitions and returns original rows; sampling and id_column are pair-query options" + ); + } + let config = DuplicatePairsConfig { + dataset_version: dataset.version().version, + column: options.column.clone(), + distance_threshold: options.distance_threshold, + }; + let mut tasks = plan_duplicate_pairs(dataset.clone(), &config).await?; + // Canonical scope order is independent of metadata enumeration order. + tasks.sort_by_key(|task| (task.segment_id, task.partition_id)); + let mut nonempty = Vec::new(); + for group in tasks.chunk_by(|a, b| a.segment_id == b.segment_id) { + let index = dataset + .open_vector_index(&config.column, &group[0].segment_id, &NoOpMetricsCollector) + .await?; + nonempty.extend( + group + .iter() + .filter(|task| index.partition_size(task.partition_id) > 0) + .cloned(), + ); + } + match options.selection { + PartitionSelection::All => {} + PartitionSelection::First { count } | PartitionSelection::Random { count, .. } + if count == 0 => + { + return plan_err!("partition count must be positive"); + } + PartitionSelection::First { count } => nonempty.truncate(count), + PartitionSelection::Random { count, seed } => { + // A single segment matches ivf_partition's SplitMix64 sampling. + // With multiple segments the canonical scope ordinal extends the + // same policy without conflating equal partition IDs. + let segments = tasks + .iter() + .map(|t| t.segment_id) + .collect::>() + .len(); + let mut ranked = nonempty + .into_iter() + .enumerate() + .map(|(ordinal, task)| { + let key = if segments == 1 { + task.partition_id as u64 + } else { + ordinal as u64 + }; + (sample_rank(seed, key), task) + }) + .collect::>(); + ranked.sort_by_key(|(rank, task)| (*rank, task.segment_id, task.partition_id)); + ranked.truncate(count); + nonempty = ranked.into_iter().map(|(_, task)| task).collect(); + nonempty.sort_by_key(|task| (task.segment_id, task.partition_id)); + } + } + let schema = if options.dedup { + let mut scanner = dataset.scan(); + scanner.blob_handling(BlobHandling::AllBinary); + scanner.schema().await? + } else { + options.output.schema(&dataset)? + }; + Ok(Self { + dataset, + config, + tasks: nonempty, + options, + schema, + }) + } +} + +fn sample_rank(seed: u64, ordinal: u64) -> u64 { + let mut value = seed.wrapping_add(ordinal).wrapping_add(0x9E3779B97F4A7C15); + value = (value ^ (value >> 30)).wrapping_mul(0xBF58476D1CE4E5B9); + value = (value ^ (value >> 27)).wrapping_mul(0x94D049BB133111EB); + value ^ (value >> 31) +} + +#[async_trait] +impl TableProvider for VectorQueryTable { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + fn table_type(&self) -> TableType { + TableType::View + } + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + let output = if self.options.dedup { + DuplicatePairsOutput::default() + } else { + self.options.output.clone() + }; + let pairs: Arc = if self.tasks.is_empty() { + Arc::new(EmptyExec::new(output.schema(&self.dataset)?)) + } else { + Arc::new( + DuplicatePairsExec::try_new( + self.dataset.clone(), + self.config.clone(), + self.tasks.clone(), + )? + .with_output(output)?, + ) + }; + let plan: Arc = if self.options.dedup { + Arc::new(VectorDedupExec::new( + self.dataset.clone(), + pairs, + self.schema.clone(), + )) + } else { + pairs + }; + if let Some(indices) = projection { + let expressions: Vec<_> = indices + .iter() + .map(|&i| { + ( + Arc::new(datafusion_physical_expr::expressions::Column::new( + self.schema.field(i).name(), + i, + )) + as Arc, + self.schema.field(i).name().clone(), + ) + }) + .collect(); + Ok(Arc::new(ProjectionExec::try_new(expressions, plan)?)) + } else { + Ok(plan) + } + } +} + +/// Greedy direct representatives over the complete pair relation. Sorting spills +/// to disk; only vertex masks are retained. This has no durable checkpoint and +/// publishes nothing: failure is handled by the enclosing query/MV transaction. +#[derive(Debug)] +pub struct VectorDedupExec { + dataset: Arc, + pairs: Arc, + properties: Arc, +} + +impl VectorDedupExec { + fn new(dataset: Arc, pairs: Arc, schema: SchemaRef) -> Self { + Self { + dataset, + pairs, + properties: Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )), + } + } +} + +impl DisplayAs for VectorDedupExec { + fn fmt_as(&self, _: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "VectorDedupExec: source_version={}, policy=direct_representative", + self.dataset.version().version + ) + } +} + +impl ExecutionPlan for VectorDedupExec { + fn name(&self) -> &str { + "VectorDedupExec" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![&self.pairs] + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return plan_err!("VectorDedupExec requires one pair input"); + } + Ok(Arc::new(Self::new( + self.dataset.clone(), + children[0].clone(), + self.schema(), + ))) + } + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + if partition != 0 { + return plan_err!("invalid dedup output partition"); + } + let dataset = self.dataset.clone(); + let pairs = CoalescePartitionsExec::new(self.pairs.clone()).execute(0, context)?; + let stream = stream::once(async move { + let (removed, _) = crate::materialized_view::vector_dedup::select_representatives( + pairs, + REPRESENTATIVE_SORT_MEMORY_BYTES, + ) + .await + .map_err(|e| datafusion_common::DataFusionError::External(Box::new(e)))?; + let mut scanner = dataset.scan(); + scanner + .with_row_addr_prefilter(RowAddrMask::from_block(RowAddrTreeMap::from_iter( + removed, + ))) + .blob_handling(BlobHandling::AllBinary) + .batch_size(MATERIALIZE_BATCH_ROWS) + .strict_batch_size(true); + Ok::<_, datafusion_common::DataFusionError>( + scanner + .try_into_stream() + .await? + .map_err(datafusion_common::DataFusionError::from), + ) + }) + .try_flatten(); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } +} diff --git a/rust/lancedb/tests/duplicate_pairs.rs b/rust/lancedb/tests/duplicate_pairs.rs index 910088bc8..f7d7ed3a2 100644 --- a/rust/lancedb/tests/duplicate_pairs.rs +++ b/rust/lancedb/tests/duplicate_pairs.rs @@ -76,7 +76,7 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result .to_str() .unwrap() .to_string(); - let n = 1025; + let n = 8193; let vectors = (0..n * 2) .flat_map(|i| { let i = i % n; @@ -196,6 +196,89 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result ); let actual = ctx.sql(&sql).await?.collect().await?; assert_eq!(pairs(&actual), pairs(&expected)); + use lancedb::table::datafusion::udtf::{ + duplicate_pairs::DuplicatePairsOutput, + vector_query::{PartitionSelection, VectorQueryOptions, VectorQueryTable}, + }; + let query_options = VectorQueryOptions { + column: "vector".into(), + distance_threshold: threshold, + selection: PartitionSelection::All, + output: DuplicatePairsOutput { + scope: true, + id_column: Some("id".into()), + }, + dedup: false, + }; + let provider = VectorQueryTable::try_new(ds.clone(), query_options.clone()).await?; + let extended = ctx.read_table(Arc::new(provider))?.collect().await?; + assert_eq!(pairs(&extended), pairs(&expected)); + for batch in &extended { + let rowids = batch + .column_by_name("row_id_a") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let ids = batch + .column_by_name("id_a") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + let expected_id = (rowids.value(i) >> 32) as i64 * n as i64 + + (rowids.value(i) & 0xffffffff) as i64; + assert_eq!(ids.value(i), expected_id); + } + } + let mut sampled_options = query_options.clone(); + sampled_options.selection = PartitionSelection::Random { count: 1, seed: 42 }; + let sample = ctx + .read_table(Arc::new( + VectorQueryTable::try_new(ds.clone(), sampled_options.clone()).await?, + ))? + .collect() + .await?; + let replay = ctx + .read_table(Arc::new( + VectorQueryTable::try_new(ds.clone(), sampled_options).await?, + ))? + .collect() + .await?; + assert_eq!(sample, replay); + let scopes = sample + .iter() + .flat_map(|batch| { + let segments = batch + .column_by_name("segment_id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let partitions = batch + .column_by_name("partition_id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()) + .map(|i| (segments.value(i).to_string(), partitions.value(i))) + .collect::>() + }) + .collect::>(); + assert!(scopes.len() <= 1); + let mut dedup_options = query_options; + dedup_options.dedup = true; + dedup_options.output = DuplicatePairsOutput::default(); + let clean_query = ctx + .read_table(Arc::new( + VectorQueryTable::try_new(ds.clone(), dedup_options).await?, + ))? + .select(vec![datafusion::prelude::col("id")])? + .collect() + .await?; + assert_eq!( pairs(&actual).len(), match threshold { @@ -388,6 +471,21 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result .collect::>(); actual_ids.sort_unstable(); expected_ids.sort_unstable(); + let mut query_ids = clean_query + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + query_ids.sort_unstable(); + assert_eq!(query_ids, expected_ids); + assert_eq!(actual_ids, expected_ids); if threshold == 1.0 { assert!(actual_ids.contains(&0) && actual_ids.contains(&2)); @@ -510,6 +608,90 @@ async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions() ) .await?; let version = ds.version().version; + use lancedb::table::datafusion::udtf::{ + duplicate_pairs::DuplicatePairsOutput, + vector_query::{PartitionSelection, VectorQueryOptions, VectorQueryTable}, + }; + let options = VectorQueryOptions { + column: "vector".into(), + distance_threshold: 1.0, + selection: PartitionSelection::All, + output: DuplicatePairsOutput::default(), + dedup: true, + }; + let mapped = VectorQueryTable::try_new( + Arc::new(ds.clone()), + VectorQueryOptions { + dedup: false, + output: DuplicatePairsOutput { + scope: true, + id_column: Some("id".into()), + }, + ..options.clone() + }, + ) + .await?; + let mapped = SessionContext::new() + .read_table(Arc::new(mapped))? + .collect() + .await?; + let mut mapped_edges = Vec::new(); + for batch in mapped { + let a = batch + .column_by_name("id_a") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let b = batch + .column_by_name("id_b") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + mapped_edges.push((a.value(i).min(b.value(i)), a.value(i).max(b.value(i)))); + } + } + mapped_edges.sort_unstable(); + assert_eq!(mapped_edges, [(0, 1), (1, 2)]); + let provider = VectorQueryTable::try_new(Arc::new(ds.clone()), options).await?; + let query_rows = SessionContext::new() + .read_table(Arc::new(provider))? + .collect() + .await?; + let mut query_ids = query_rows + .iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + query_ids.sort_unstable(); + assert_eq!(query_ids, [0, 2, 3, 4]); + for batch in &query_rows { + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let images = batch + .column_by_name("image") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + assert_eq!(images.value(i), vec![ids.value(i) as u8; 4096]); + } + } let definition = MaterializedViewDefinition::from_sql(&format!( "SELECT * FROM vector_dedup('source', {version}, 'vector', 1)" ))?; @@ -519,6 +701,13 @@ async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions() .await?; // A later source delete does not change the declared result snapshot. table.delete("id = 0").await?; + let adapter = BaseTableAdapter::try_new(table.base_table().clone()).await?; + let explicit = adapter.vector_query_snapshot(Some(version)).await?; + let latest = adapter.vector_query_snapshot(None).await?; + assert_eq!(explicit.version().version, version); + assert!(latest.version().version > version); + assert_eq!(explicit.count_rows(None).await?, 5); + assert_eq!(latest.count_rows(None).await?, 4); clean.refresh().execute().await?; assert_eq!(clean.table().count_rows(None).await?, 4); let output = Dataset::open(&clean.table().uri().await?).await?;