mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-13 07:12:27 +00:00
feat(sql): expose composable native vector relations (#4340)
## 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.
This commit is contained in:
Yang Cen
@@ -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
|
||||
|
||||
@@ -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)> {
|
||||
|
||||
@@ -6,3 +6,5 @@
|
||||
pub mod fts;
|
||||
|
||||
pub mod duplicate_pairs;
|
||||
|
||||
pub mod vector_query;
|
||||
|
||||
@@ -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<String>,
|
||||
}
|
||||
|
||||
impl DuplicatePairsOutput {
|
||||
pub fn schema(&self, dataset: &Dataset) -> Result<SchemaRef> {
|
||||
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<arrow_array::RecordBatch> {
|
||||
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::<UInt64Array>()
|
||||
.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<DuplicatePairTask>,
|
||||
properties: Arc<PlanProperties>,
|
||||
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> {
|
||||
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<Dataset> {
|
||||
&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,
|
||||
)))
|
||||
}
|
||||
|
||||
@@ -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<u64>) -> Result<Arc<Dataset>> {
|
||||
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<Dataset>,
|
||||
config: DuplicatePairsConfig,
|
||||
tasks: Vec<DuplicatePairTask>,
|
||||
options: VectorQueryOptions,
|
||||
schema: SchemaRef,
|
||||
}
|
||||
|
||||
impl VectorQueryTable {
|
||||
pub async fn try_new(dataset: Arc<Dataset>, options: VectorQueryOptions) -> Result<Self> {
|
||||
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::<std::collections::BTreeSet<_>>()
|
||||
.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::<Vec<_>>();
|
||||
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<usize>>,
|
||||
_filters: &[Expr],
|
||||
_limit: Option<usize>,
|
||||
) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
let output = if self.options.dedup {
|
||||
DuplicatePairsOutput::default()
|
||||
} else {
|
||||
self.options.output.clone()
|
||||
};
|
||||
let pairs: Arc<dyn ExecutionPlan> = 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<dyn ExecutionPlan> = 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<dyn datafusion_physical_expr::PhysicalExpr>,
|
||||
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<Dataset>,
|
||||
pairs: Arc<dyn ExecutionPlan>,
|
||||
properties: Arc<PlanProperties>,
|
||||
}
|
||||
|
||||
impl VectorDedupExec {
|
||||
fn new(dataset: Arc<Dataset>, pairs: Arc<dyn ExecutionPlan>, 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<PlanProperties> {
|
||||
&self.properties
|
||||
}
|
||||
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
|
||||
vec![&self.pairs]
|
||||
}
|
||||
fn with_new_children(
|
||||
self: Arc<Self>,
|
||||
children: Vec<Arc<dyn ExecutionPlan>>,
|
||||
) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
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<TaskContext>,
|
||||
) -> Result<SendableRecordBatchStream> {
|
||||
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,
|
||||
)))
|
||||
}
|
||||
}
|
||||
@@ -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::<UInt64Array>()
|
||||
.unwrap();
|
||||
let ids = batch
|
||||
.column_by_name("id_a")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<Int64Array>()
|
||||
.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::<arrow_array::StringArray>()
|
||||
.unwrap();
|
||||
let partitions = batch
|
||||
.column_by_name("partition_id")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<UInt64Array>()
|
||||
.unwrap();
|
||||
(0..batch.num_rows())
|
||||
.map(|i| (segments.value(i).to_string(), partitions.value(i)))
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.collect::<std::collections::BTreeSet<_>>();
|
||||
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::<Vec<_>>();
|
||||
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::<Int64Array>()
|
||||
.unwrap()
|
||||
.values()
|
||||
.to_vec()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
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::<Int64Array>()
|
||||
.unwrap();
|
||||
let b = batch
|
||||
.column_by_name("id_b")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<Int64Array>()
|
||||
.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::<Int64Array>()
|
||||
.unwrap()
|
||||
.values()
|
||||
.to_vec()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
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::<Int64Array>()
|
||||
.unwrap();
|
||||
let images = batch
|
||||
.column_by_name("image")
|
||||
.unwrap()
|
||||
.as_any()
|
||||
.downcast_ref::<LargeBinaryArray>()
|
||||
.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?;
|
||||
|
||||
Reference in New Issue
Block a user