From 50666cdcc8eb475e11a8b60531b046c4d93ee297 Mon Sep 17 00:00:00 2001 From: Yang Cen <159225399+BubbleCal@users.noreply.github.com> Date: Sun, 27 Sep 2026 03:15:54 +0800 Subject: [PATCH] feat(mv): materialize dedup results with direct representatives --- nodejs/__test__/materialized_view.test.ts | 12 +- nodejs/lancedb/materialized_view.ts | 5 +- python/python/lancedb/materialized_view.py | 5 +- .../python/tests/test_materialized_views.py | 6 +- rust/lancedb/src/materialized_view.rs | 70 +- .../src/materialized_view/duplicate_pairs.rs | 86 ++- .../src/materialized_view/partitioned.rs | 64 +- rust/lancedb/src/materialized_view/query.rs | 27 +- rust/lancedb/src/materialized_view/refresh.rs | 13 +- .../src/materialized_view/vector_dedup.rs | 636 ++++++++++++++++++ rust/lancedb/src/remote/table.rs | 2 +- rust/lancedb/tests/duplicate_pairs.rs | 261 +++++++ 12 files changed, 1125 insertions(+), 62 deletions(-) create mode 100644 rust/lancedb/src/materialized_view/vector_dedup.rs diff --git a/nodejs/__test__/materialized_view.test.ts b/nodejs/__test__/materialized_view.test.ts index 5e5c08745..6739320aa 100644 --- a/nodejs/__test__/materialized_view.test.ts +++ b/nodejs/__test__/materialized_view.test.ts @@ -52,9 +52,19 @@ describe("materialized views", () => { "SELECT * FROM people", ); + for (const [format, fn] of [ + [3, "vector_duplicate_pairs"], + [4, "vector_dedup"], + ]) { + const native = `SELECT * FROM ${fn}('images', 3, 'phash', 4)`; + expect(read(JSON.stringify({ format, query: native })).query).toBe( + native, + ); + } + // A newer writer's layout is reported, never guessed at. for (const newer of [ - `{"format":3,"query":${JSON.stringify(query)}}`, + `{"format":5,"query":${JSON.stringify(query)}}`, '{"kind":"select_v3","source_table":"people"}', ]) { expect(() => read(newer)).toThrow(/cannot refresh/); diff --git a/nodejs/lancedb/materialized_view.ts b/nodejs/lancedb/materialized_view.ts index 51f8cfead..a59a3a5f6 100644 --- a/nodejs/lancedb/materialized_view.ts +++ b/nodejs/lancedb/materialized_view.ts @@ -9,9 +9,10 @@ export const DEFINITION_META_KEY = "mv.definition"; /** * The newest stored layout this version reads: `{"format": N, "query": ""}`, - * format 2 being a query with `GROUP BY`. + * format 2 being a query with `GROUP BY`, format 3 native pairs and format 4 + * native dedup results. */ -export const DEFINITION_FORMAT = 2; +export const DEFINITION_FORMAT = 4; /** * The query that defines a materialized view, as stored: diff --git a/python/python/lancedb/materialized_view.py b/python/python/lancedb/materialized_view.py index a836210cb..d888ae365 100644 --- a/python/python/lancedb/materialized_view.py +++ b/python/python/lancedb/materialized_view.py @@ -29,9 +29,10 @@ SelectArg = Union[ ] -DEFINITION_FORMAT = 2 +DEFINITION_FORMAT = 4 """The newest stored layout this version reads: ``{"format": N, "query": ""}``, -format 2 being a query with ``GROUP BY``. A ``kind`` key beside it is for +format 2 being a query with ``GROUP BY``, format 3 native pairs and format 4 +native dedup results. A ``kind`` key beside it is for readers older than the format number.""" diff --git a/python/python/tests/test_materialized_views.py b/python/python/tests/test_materialized_views.py index c030ac1df..01f45aac4 100644 --- a/python/python/tests/test_materialized_views.py +++ b/python/python/tests/test_materialized_views.py @@ -520,9 +520,13 @@ def test_stored_queries_and_legacy_layouts_are_read(): "SELECT * FROM people" ) + for fmt, function in [(3, "vector_duplicate_pairs"), (4, "vector_dedup")]: + native = f"SELECT * FROM {function}('images', 3, 'phash', 4)" + assert read({"format": fmt, "query": native}).query == native + # A newer writer's layout is reported, never guessed at. for newer in ( - {"format": 3, "query": query}, + {"format": 5, "query": query}, {"kind": "select_v3", "source_table": "people"}, ): with pytest.raises(NotImplementedError, match="cannot refresh"): diff --git a/rust/lancedb/src/materialized_view.rs b/rust/lancedb/src/materialized_view.rs index 2316e8f9f..535c2ace0 100644 --- a/rust/lancedb/src/materialized_view.rs +++ b/rust/lancedb/src/materialized_view.rs @@ -13,7 +13,8 @@ mod duplicate_pairs; mod grouped; mod grouped_units; mod partitioned; -pub use duplicate_pairs::DuplicatePairsView; +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 /// partition, so the work can be spread over several processes. The caller @@ -79,7 +80,7 @@ pub use grouped_units::{ }; pub use partitioned::{ PartitionedRefreshPlan, WrittenPartition, commit_partitioned_refresh, plan_partitioned_refresh, - write_refresh_partition, + write_refresh_partition, write_refresh_partition_with_inputs, }; mod query; pub mod refresh; @@ -154,17 +155,20 @@ const COLUMN_DEFINITIONS_META_KEY: &str = "lancedb::column_definitions"; /// The newest layout this version reads under [`DEFINITION_META_KEY`]: /// `{"format": N, "query": ""}`, the query as /// [`MaterializedViewDefinition::to_sql`] renders it. Format 3 identifies a -/// native duplicate-pair source, format 2 a query with `GROUP BY`, and +/// native duplicate-pair source, format 4 a dedup result source, format 2 a query with `GROUP BY`, and /// format 1 an ordinary query. A reader refuses a newer format rather than /// guess at it. The layout also carries `"kind": "query"`, which readers /// older than the format number report as an unrefreshable view instead of /// failing to read the metadata. -pub const DEFINITION_FORMAT: u64 = 3; +pub const DEFINITION_FORMAT: u64 = 4; /// The format `definition` is written in; see [`DEFINITION_FORMAT`]. fn format_of(definition: &MaterializedViewDefinition) -> u64 { - if definition.duplicate_pairs.is_some() { - 3 + if let Some(source) = &definition.vector_source { + match source.kind { + VectorSourceKind::Pairs => 3, + VectorSourceKind::Dedup => 4, + } } else if definition.is_grouped() { 2 } else { @@ -312,8 +316,8 @@ pub fn read_staging(metadata: &HashMap) -> Result>, + /// A table-bound native vector source. It is never evaluated per source row. + pub vector_source: Option>, /// Name of the source table, in the same database as the view. pub source_table: String, /// Namespace path holding the source table; empty is the root namespace. @@ -512,7 +516,7 @@ pub fn read_definition(metadata: &HashMap) -> Result, ) -> Result { - if definition.duplicate_pairs.is_some() { + if definition.vector_source.is_some() { return duplicate_pairs::plan(source_schema, definition); } let filter = definition @@ -730,7 +734,7 @@ pub(crate) fn plan( } let definition = MaterializedViewDefinition { - duplicate_pairs: None, + vector_source: None, source_table: definition.source_table.clone(), source_namespace: definition.source_namespace.clone(), lateral: definition.lateral.clone(), @@ -1164,9 +1168,9 @@ impl PreparedDeclaration { /// read: the column the view projects it to, if any, otherwise an /// internal projection added here, named by [`input_column_name`]. pub fn input_column(&mut self, source_column: &str) -> Result { - if self.definition.duplicate_pairs.is_some() { + if self.definition.vector_source.is_some() { return Err(Error::InvalidInput { - message: "native pair views have no per-source-row inputs".into(), + message: "native vector views have no per-source-row inputs".into(), }); } // A grouped view's rows are groups; no source row carries a value into one. @@ -1269,9 +1273,9 @@ impl PreparedDeclaration { columns: Vec<(usize, ArrowField)>, bindings: &[FunctionBinding], ) -> Result { - if self.definition.duplicate_pairs.is_some() { + if self.definition.vector_source.is_some() { return Err(Error::InvalidInput { - message: "computed columns on native pair views are not supported".into(), + message: "computed columns on native vector views are not supported".into(), }); } let invalid = |message: String| Error::InvalidInput { message }; @@ -1480,7 +1484,7 @@ pub async fn prepare_declaration( limit: Option, ) -> Result { let definition = MaterializedViewDefinition { - duplicate_pairs: None, + vector_source: None, source_table: source.name().to_string(), source_namespace: source.namespace().to_vec(), lateral: None, @@ -1542,6 +1546,26 @@ pub async fn prepare_declaration( /// # Ok(()) /// # } /// ``` +/// +/// An indexed, pinned source can materialize retained original rows using the +/// default greedy direct-representative policy. A row is removed only when it +/// directly matches a retained row; an A-B-C chain keeps A and C. The source is +/// unchanged. Hosts distribute this through [`plan_partitioned_refresh`] and +/// [`write_refresh_partition_with_inputs`], with a barrier between selection +/// tasks and source-fragment materialization tasks. +/// +/// ``` +/// use lancedb::materialized_view::{MaterializedViewDefinition, prepare_definition}; +/// # async fn dedup(images: &lancedb::Table) -> Result<(), Box> { +/// let version = images.version().await?; +/// let definition = MaterializedViewDefinition::from_sql(&format!( +/// "SELECT * FROM vector_dedup('images', {version}, 'phash', 4)" +/// ))?; +/// let clean = prepare_definition(images, definition).await?.create("images_clean").await?; +/// clean.refresh().execute().await?; +/// # Ok(()) +/// # } +/// ``` pub async fn prepare_definition( source: &Table, definition: MaterializedViewDefinition, @@ -1657,7 +1681,7 @@ async fn prepare_with( ), }); } - if definition.duplicate_pairs.is_none() + if definition.vector_source.is_none() && !native.dataset.get().await?.manifest.uses_stable_row_ids() { return Err(Error::InvalidInput { @@ -1685,7 +1709,7 @@ async fn prepare_with( message: format!("view column name '{}' is reserved", reserved.output), }); } - let source_schema = if let Some(pairs) = &definition.duplicate_pairs { + let source_schema = if let Some(pairs) = &definition.vector_source { let ds = native .dataset .get() @@ -1717,7 +1741,7 @@ async fn prepare_with( None => source_schema.clone(), Some(unnest) => flattened_schema(&source_schema, &unnest)?, }; - if definition.duplicate_pairs.is_none() { + if definition.vector_source.is_none() { fields.push(ArrowField::new( SOURCE_ROW_ID_COLUMN, DataType::UInt64, @@ -2235,7 +2259,7 @@ mod tests { assert_eq!( view.definition(), &MaterializedViewDefinition { - duplicate_pairs: None, + vector_source: None, source_table: "people".into(), source_namespace: Vec::new(), projections: vec![ @@ -3285,7 +3309,7 @@ mod tests { fn definition(source_namespace: Vec) -> MaterializedViewDefinition { MaterializedViewDefinition { - duplicate_pairs: None, + vector_source: None, source_table: "people".to_string(), source_namespace, lateral: None, @@ -3372,8 +3396,8 @@ mod tests { #[test] fn a_newer_format_is_reported_not_guessed() { assert_eq!( - read(r#"{"format":4,"query":"SELECT name FROM people"}"#).unwrap(), - Some(StoredDefinition::Newer { format: "4".into() }) + read(r#"{"format":5,"query":"SELECT name FROM people"}"#).unwrap(), + Some(StoredDefinition::Newer { format: "5".into() }) ); assert_eq!( read(r#"{"kind":"join"}"#).unwrap(), diff --git a/rust/lancedb/src/materialized_view/duplicate_pairs.rs b/rust/lancedb/src/materialized_view/duplicate_pairs.rs index 765f01852..d0dc219ad 100644 --- a/rust/lancedb/src/materialized_view/duplicate_pairs.rs +++ b/rust/lancedb/src/materialized_view/duplicate_pairs.rs @@ -31,18 +31,35 @@ use crate::table::datafusion::udtf::duplicate_pairs::{ use crate::{Error, Result}; pub(super) const FUNCTION_NAME: &str = "vector_duplicate_pairs"; +pub(super) const DEDUP_FUNCTION_NAME: &str = "vector_dedup"; + +/// The result of a pinned native vector source. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum VectorSourceKind { + /// The qualifying pair relation. + Pairs, + /// Original source rows retained by greedy direct representatives. + Dedup, +} /// A native table-function source, pinned for the lifetime of the view. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] -pub struct DuplicatePairsView { +pub struct VectorSource { + pub kind: VectorSourceKind, pub dataset_version: u64, pub column: String, /// Validated numeric SQL spelling; execution and task identity use f32 bits. pub distance_threshold: String, } -impl DuplicatePairsView { +impl VectorSource { + pub(super) fn function_name(&self) -> &'static str { + match self.kind { + VectorSourceKind::Pairs => FUNCTION_NAME, + VectorSourceKind::Dedup => DEDUP_FUNCTION_NAME, + } + } pub(crate) fn config(&self) -> Result { let threshold: f32 = self .distance_threshold @@ -69,7 +86,8 @@ fn invalid(message: impl Into) -> Error { pub(super) fn source( args: &TableFunctionArgs, -) -> Result<(Vec, String, DuplicatePairsView)> { + kind: VectorSourceKind, +) -> Result<(Vec, String, VectorSource)> { if args.settings.is_some() { return Err(invalid("duplicate pairs do not accept function settings")); } @@ -85,7 +103,7 @@ pub(super) fn source( .collect::>>()?; let [table, version, column, threshold] = exprs.as_slice() else { return Err(invalid( - "vector_duplicate_pairs requires (table, dataset_version, column, distance_threshold)", + "native vector sources require (table, dataset_version, column, distance_threshold)", )); }; let string = |expr: &Expr| match expr { @@ -97,7 +115,7 @@ pub(super) fn source( }; let table = string(table)?; let plain = MaterializedViewDefinition::from_sql(&format!("SELECT * FROM {table}"))?; - if plain.duplicate_pairs.is_some() + if plain.vector_source.is_some() || plain.lateral.is_some() || plain.filter.is_some() || plain.is_grouped() @@ -114,7 +132,8 @@ pub(super) fn source( }, _ => return Err(invalid("source version must be an integer literal")), }; - let config = DuplicatePairsView { + let config = VectorSource { + kind, dataset_version, column: string(column)?, distance_threshold: threshold.to_string(), @@ -124,6 +143,23 @@ pub(super) fn source( } pub(super) fn check_shape(definition: &MaterializedViewDefinition) -> Result<()> { + if definition + .vector_source + .as_ref() + .is_some_and(|s| s.kind == VectorSourceKind::Dedup) + { + if !definition.selects_star() + || definition.filter.is_some() + || definition.limit.is_some() + || definition.lateral.is_some() + || definition.is_grouped() + { + return Err(invalid( + "a vector_dedup view must select * without other clauses", + )); + } + return Ok(()); + } let schema = duplicate_pair_schema(); let projections = schema .fields() @@ -151,9 +187,29 @@ pub(super) fn plan( definition: &MaterializedViewDefinition, ) -> Result { check_shape(definition)?; - let config = definition.duplicate_pairs.as_ref().expect("native source"); + let config = definition.vector_source.as_ref().expect("native source"); config.config()?; source_schema.field_with_name(&config.column)?; + if config.kind == VectorSourceKind::Dedup { + return Ok(Planned { + definition: definition.clone(), + fields: source_schema + .fields() + .iter() + .map(|f| super::without_declarations(f)) + .collect(), + lineage: source_schema + .fields() + .iter() + .map(|f| (f.name().clone(), vec![f.name().clone()])) + .collect(), + inputs: source_schema + .fields() + .iter() + .map(|f| f.name().clone()) + .collect(), + }); + } let schema = duplicate_pair_schema(); let projections = schema .fields() @@ -197,7 +253,7 @@ pub struct WrittenPairsPartition { } impl DuplicatePairsRefreshPlan { - fn identity(&self) -> Result { + pub(super) fn identity(&self) -> Result { use sha2::{Digest, Sha256}; let encoded = serde_json::to_vec(&( &self.snapshot, @@ -224,7 +280,7 @@ async fn open(view: &Table) -> Result<(Dataset, MaterializedViewDefinition, Arc< return Err(invalid("view does not carry a supported definition")); }; let config = definition - .duplicate_pairs + .vector_source .as_ref() .ok_or_else(|| invalid("not a native pair view"))? .config()?; @@ -236,7 +292,7 @@ async fn open(view: &Table) -> Result<(Dataset, MaterializedViewDefinition, Arc< ); ensure_no_mem_wal(&source, "source table", &definition.source_table).await?; let planned = plan(Arc::new(Schema::from(source.schema())), &definition)?; - let expected = duplicate_pair_schema(); + let expected = Schema::new(planned.fields); let physical = Schema::from(view_ds.schema()); if physical.fields().len() != expected.fields().len() || physical @@ -260,7 +316,7 @@ pub(super) async fn plan_refresh( ) -> Result { let (view_ds, definition, source) = open(view).await?; let config = definition - .duplicate_pairs + .vector_source .as_ref() .expect("native") .config()?; @@ -292,14 +348,14 @@ pub(super) async fn plan_refresh( }) } -async fn open_planned( +pub(super) async fn open_planned( view: &Table, plan: &DuplicatePairsRefreshPlan, ) -> Result<(Dataset, MaterializedViewDefinition, Arc)> { let (view_ds, definition, source) = open(view).await?; ensure_incarnation(&view_ds, Some(&plan.snapshot.incarnation), view.name()).await?; let config = definition - .duplicate_pairs + .vector_source .as_ref() .expect("native") .config()?; @@ -329,7 +385,7 @@ pub(super) async fn write_unit( .ok_or_else(|| invalid("pair partition is outside the plan"))? .clone(); let config = definition - .duplicate_pairs + .vector_source .as_ref() .expect("native") .config()?; @@ -378,7 +434,7 @@ pub(super) async fn commit( ) -> Result { let (_, definition, source) = open_planned(view, plan).await?; let config = definition - .duplicate_pairs + .vector_source .as_ref() .expect("native") .config()?; diff --git a/rust/lancedb/src/materialized_view/partitioned.rs b/rust/lancedb/src/materialized_view/partitioned.rs index 5bdd88f0d..58192089b 100644 --- a/rust/lancedb/src/materialized_view/partitioned.rs +++ b/rust/lancedb/src/materialized_view/partitioned.rs @@ -5,7 +5,10 @@ use serde::{Deserialize, Serialize}; -use super::duplicate_pairs::{self, DuplicatePairsRefreshPlan, WrittenPairsPartition}; +use super::duplicate_pairs::{ + self, DuplicatePairsRefreshPlan, VectorSourceKind, WrittenPairsPartition, +}; +use super::vector_dedup::{self, VectorDedupPlan, WrittenDedupUnit}; use super::{GroupedRefreshPlan, RefreshMaterializedViewResult, StoredDefinition, WrittenUnit}; use crate::table::Table; use crate::{Error, Result}; @@ -15,6 +18,7 @@ use crate::{Error, Result}; #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(untagged)] pub enum PartitionedRefreshPlan { + VectorDedup(VectorDedupPlan), DuplicatePairs(DuplicatePairsRefreshPlan), Grouped(GroupedRefreshPlan), } @@ -23,10 +27,18 @@ impl PartitionedRefreshPlan { /// The shared source snapshot and destination generation fence. pub fn snapshot(&self) -> &GroupedRefreshPlan { match self { + Self::VectorDedup(plan) => &plan.snapshot, Self::DuplicatePairs(plan) => &plan.snapshot, Self::Grouped(plan) => plan, } } + /// Selection tasks that must complete before remaining tasks can start. + pub fn dependency_units(&self) -> u32 { + match self { + Self::VectorDedup(plan) => plan.dependencies(), + _ => 0, + } + } /// Number of independent units in this immutable plan. pub fn units(&self) -> u32 { self.snapshot().units @@ -37,6 +49,7 @@ impl PartitionedRefreshPlan { #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(untagged)] pub enum WrittenPartition { + VectorDedup(WrittenDedupUnit), DuplicatePairs(WrittenPairsPartition), Grouped(WrittenUnit), } @@ -45,6 +58,7 @@ impl WrittenPartition { /// Number of pair or grouped rows staged by this unit. pub fn rows(&self) -> u64 { match self { + Self::VectorDedup(unit) => unit.rows, Self::DuplicatePairs(unit) => unit.rows, Self::Grouped(unit) => unit.rows(), } @@ -52,6 +66,7 @@ impl WrittenPartition { /// Uncommitted output files, retained only after task completion. pub fn fragments(&self) -> &[lance_table::format::Fragment] { match self { + Self::VectorDedup(unit) => &unit.fragments, Self::DuplicatePairs(unit) => &unit.fragments, Self::Grouped(unit) => unit.fragments(), } @@ -66,8 +81,17 @@ pub async fn plan_partitioned_refresh( ) -> Result> { if let Some(StoredDefinition::Query(definition)) = super::read_definition(view.schema().await?.metadata())? - && definition.duplicate_pairs.is_some() + && definition.vector_source.is_some() { + if definition + .vector_source + .as_ref() + .is_some_and(|s| s.kind == VectorSourceKind::Dedup) + { + return vector_dedup::plan_refresh(view, pinned) + .await + .map(|p| Some(PartitionedRefreshPlan::VectorDedup(p))); + } return duplicate_pairs::plan_refresh(view, pinned) .await .map(|p| Some(PartitionedRefreshPlan::DuplicatePairs(p))); @@ -83,8 +107,34 @@ pub async fn write_refresh_partition( view: &Table, unit: u32, plan: &PartitionedRefreshPlan, +) -> Result { + write_refresh_partition_with_inputs(view, unit, plan, &[]).await +} + +/// Write a unit after its durable selection dependencies have completed. +/// Supply exactly the first `plan.dependency_units()` receipts for a dependent +/// unit. A missing, duplicate or foreign receipt fails before output is staged. +pub async fn write_refresh_partition_with_inputs( + view: &Table, + unit: u32, + plan: &PartitionedRefreshPlan, + dependencies: &[WrittenPartition], ) -> Result { match plan { + PartitionedRefreshPlan::VectorDedup(plan) => { + let dependencies = dependencies + .iter() + .map(|unit| match unit { + WrittenPartition::VectorDedup(unit) => Ok(unit.clone()), + _ => Err(Error::InvalidInput { + message: "foreign dedup dependency".into(), + }), + }) + .collect::>>()?; + vector_dedup::write_unit(view, unit, plan, &dependencies) + .await + .map(WrittenPartition::VectorDedup) + } PartitionedRefreshPlan::DuplicatePairs(plan) => { duplicate_pairs::write_unit(view, unit, plan) .await @@ -108,6 +158,16 @@ pub async fn commit_partitioned_refresh( message: "refresh receipt does not match the plan kind".into(), }; match plan { + PartitionedRefreshPlan::VectorDedup(plan) => { + let units = units + .into_iter() + .map(|unit| match unit { + WrittenPartition::VectorDedup(unit) => Ok(unit), + _ => Err(foreign()), + }) + .collect::>>()?; + vector_dedup::commit(view, plan, units, expected).await + } PartitionedRefreshPlan::DuplicatePairs(plan) => { let units = units .into_iter() diff --git a/rust/lancedb/src/materialized_view/query.rs b/rust/lancedb/src/materialized_view/query.rs index 11bd05fda..4bfd5542e 100644 --- a/rust/lancedb/src/materialized_view/query.rs +++ b/rust/lancedb/src/materialized_view/query.rs @@ -264,7 +264,7 @@ fn extract(query: &Query) -> Result { }; let mut from = select.from.iter(); - let mut duplicate_pairs = None; + let mut vector_source = None; let (source_namespace, source_table) = match from.next().map(|f| &f.relation) { Some(TableFactor::Table { name, @@ -278,13 +278,20 @@ fn extract(query: &Query) -> Result { sample: None, index_hints, .. - }) if name.to_string() == super::duplicate_pairs::FUNCTION_NAME - && with_hints.is_empty() + }) if matches!( + name.to_string().as_str(), + super::duplicate_pairs::FUNCTION_NAME | super::duplicate_pairs::DEDUP_FUNCTION_NAME + ) && with_hints.is_empty() && partitions.is_empty() && index_hints.is_empty() => { - let (namespace, table, config) = super::duplicate_pairs::source(args)?; - duplicate_pairs = Some(Box::new(config)); + let kind = if name.to_string() == super::duplicate_pairs::FUNCTION_NAME { + super::VectorSourceKind::Pairs + } else { + super::VectorSourceKind::Dedup + }; + let (namespace, table, config) = super::duplicate_pairs::source(args, kind)?; + vector_source = Some(Box::new(config)); (namespace, table) } Some(TableFactor::Table { args: Some(_), .. }) => { @@ -438,7 +445,7 @@ fn extract(query: &Query) -> Result { }; let definition = MaterializedViewDefinition { - duplicate_pairs, + vector_source, source_table, source_namespace, lateral, @@ -447,7 +454,7 @@ fn extract(query: &Query) -> Result { group_by, limit, }; - if definition.duplicate_pairs.is_some() { + if definition.vector_source.is_some() { super::duplicate_pairs::check_shape(&definition)?; } check_grouping(&definition)?; @@ -509,10 +516,10 @@ pub fn render(definition: &MaterializedViewDefinition) -> String { .chain(std::iter::once(&definition.source_table)) .map(|part| ident_sql(part)) .collect(); - if let Some(native) = &definition.duplicate_pairs { + if let Some(native) = &definition.vector_source { sql.push_str(&format!( "{}('{}', {}, '{}', {})", - super::duplicate_pairs::FUNCTION_NAME, + native.function_name(), table.join(".").replace('\'', "''"), native.dataset_version, native.column.replace('\'', "''"), @@ -557,7 +564,7 @@ mod tests { assert_eq!(definition.source_table, "images"); assert!(definition.lateral.is_none()); assert_eq!( - definition.duplicate_pairs.as_ref().unwrap().dataset_version, + definition.vector_source.as_ref().unwrap().dataset_version, 3 ); assert_eq!(parse(&render(&definition)).unwrap(), definition); diff --git a/rust/lancedb/src/materialized_view/refresh.rs b/rust/lancedb/src/materialized_view/refresh.rs index 07754af76..9f0ea2646 100644 --- a/rust/lancedb/src/materialized_view/refresh.rs +++ b/rust/lancedb/src/materialized_view/refresh.rs @@ -121,7 +121,7 @@ pub(crate) async fn execute_refresh( // local execution; Sophon drives these units through its job registry. if let Some(super::StoredDefinition::Query(definition)) = super::read_definition(view.schema().await?.metadata())? - && definition.duplicate_pairs.is_some() + && definition.vector_source.is_some() { let plan = super::plan_partitioned_refresh(view, pinned) .await? @@ -130,7 +130,10 @@ pub(crate) async fn execute_refresh( })?; let mut units = Vec::with_capacity(plan.units() as usize); for unit in 0..plan.units() { - units.push(super::write_refresh_partition(view, unit, &plan).await?); + let dependencies = &units[..units.len().min(plan.dependency_units() as usize)]; + units.push( + super::write_refresh_partition_with_inputs(view, unit, &plan, dependencies).await?, + ); } return super::commit_partitioned_refresh(view, &plan, units, expected_incarnation).await; } @@ -745,7 +748,7 @@ pub(super) async fn open_source( message: "materialized views are supported only on local tables".into(), })?; let dataset = native.dataset.get().await?.as_ref().clone(); - if definition.duplicate_pairs.is_none() && !dataset.manifest.uses_stable_row_ids() { + if definition.vector_source.is_none() && !dataset.manifest.uses_stable_row_ids() { return Err(Error::InvalidInput { message: format!( "source table '{}' does not have stable row ids; it is not the \ @@ -3691,7 +3694,7 @@ pub(crate) mod tests { let (_conn, _, view) = refreshed_doubled(vec![1, 2]).await; let replacement = crate::materialized_view::MaterializedViewDefinition { - duplicate_pairs: None, + vector_source: None, source_table: "src".into(), source_namespace: Vec::new(), projections: vec![ @@ -3733,7 +3736,7 @@ pub(crate) mod tests { let (_conn, _, view) = refreshed_doubled(vec![1, 2]).await; let narrower = crate::materialized_view::MaterializedViewDefinition { - duplicate_pairs: None, + vector_source: None, source_table: "src".into(), source_namespace: Vec::new(), projections: vec![crate::materialized_view::ViewProjection { diff --git a/rust/lancedb/src/materialized_view/vector_dedup.rs b/rust/lancedb/src/materialized_view/vector_dedup.rs new file mode 100644 index 000000000..38c007076 --- /dev/null +++ b/rust/lancedb/src/materialized_view/vector_dedup.rs @@ -0,0 +1,636 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Two dependent waves in the ordinary refresh fleet: index partition selection, +//! then source fragment materialization. Only the latter publishes data files. + +use std::collections::BTreeMap; +use std::io::Cursor; +use std::sync::Arc; + +use arrow_array::{Array, RecordBatch, UInt64Array}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use datafusion::prelude::{SessionConfig, SessionContext, col}; +use datafusion_common::DataFusionError; +use datafusion_execution::{ + SendableRecordBatchStream, TaskContext, runtime_env::RuntimeEnvBuilder, +}; +use datafusion_physical_plan::{ExecutionPlan, stream::RecordBatchStreamAdapter}; +use futures::TryStreamExt; +use lance::Dataset; +use lance::dataset::rowids::{get_row_id_index, load_row_id_sequence}; +use lance::dataset::scanner::{RowAddrMask, RowAddrTreeMap}; +use lance::dataset::transaction::Operation; +use lance::dataset::{InsertBuilder, WriteDestination, WriteMode, WriteParams}; +use lance_core::datatypes::BlobHandling; +use lance_core::utils::address::RowAddress; +use lance_datafusion::exec::{OneShotExec, SessionContextExt}; +use lance_table::format::Fragment; +use roaring::{RoaringBitmap, RoaringTreemap}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use uuid::Uuid; + +use super::duplicate_pairs::{self, DuplicatePairsRefreshPlan, VectorSourceKind}; +use super::{GroupedRefreshPlan, RefreshMaterializedViewResult}; +use crate::table::Table; +use crate::table::datafusion::udtf::duplicate_pairs::{DuplicatePairsExec, plan_duplicate_pairs}; +use crate::{Error, Result}; + +const POLICY_VERSION: u32 = 1; +const SORT_MEMORY_BYTES: usize = 256 * 1024 * 1024; +const MATERIALIZE_BATCH_ROWS: usize = 1024; +const MASK_FILE: &str = "dedup.bitmap"; +const MASK_MAGIC: &[u8] = b"LANCEDB_DEDUP_1"; + +fn invalid(message: impl Into) -> Error { + Error::InvalidInput { + message: message.into(), + } +} +fn io_error(error: std::io::Error) -> Error { + invalid(format!("invalid dedup bitmap: {error}")) +} + +/// Immutable snapshot, partition tasks and source fragments for both waves. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct VectorDedupPlan { + pub snapshot: GroupedRefreshPlan, + native: DuplicatePairsRefreshPlan, + source_fragments: Vec, + policy_version: u32, +} +impl VectorDedupPlan { + pub(super) fn dependencies(&self) -> u32 { + self.native.snapshot.units + } + fn identity(&self) -> Result { + Ok(format!( + "{:x}", + Sha256::digest(serde_json::to_vec(self).map_err(|e| invalid(e.to_string()))?) + )) + } +} + +/// Small receipt for a durable selection mask. Row IDs and pairs are not stored +/// in the host's task registry. UUID directories use ordinary index-file GC. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct MaskFile { + source_fragment: u32, + object_id: Uuid, + bytes: usize, + digest: String, +} + +/// A selection task produces masks; a materialization task produces data files. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct WrittenDedupUnit { + unit: u32, + plan_id: String, + masks: Vec, + pub(super) fragments: Vec, + pub(super) rows: u64, +} + +pub(super) async fn plan_refresh(view: &Table, pinned: Option) -> Result { + let native = duplicate_pairs::plan_refresh(view, pinned).await?; + let (_, _, source) = duplicate_pairs::open_planned(view, &native).await?; + let source_fragments = source + .get_fragments() + .iter() + .map(|f| f.id() as u32) + .collect::>(); + let mut snapshot = native.snapshot.clone(); + snapshot.units = snapshot + .units + .checked_add( + u32::try_from(source_fragments.len()) + .map_err(|_| invalid("too many source fragments"))?, + ) + .ok_or_else(|| invalid("too many dedup refresh units"))?; + Ok(VectorDedupPlan { + snapshot, + native, + source_fragments, + policy_version: POLICY_VERSION, + }) +} + +async fn open( + view: &Table, + plan: &VectorDedupPlan, +) -> Result<(Dataset, super::MaterializedViewDefinition, Arc)> { + let (view_ds, definition, source) = duplicate_pairs::open_planned(view, &plan.native).await?; + let mut snapshot = plan.native.snapshot.clone(); + snapshot.units = snapshot + .units + .checked_add( + u32::try_from(plan.source_fragments.len()) + .map_err(|_| invalid("too many dedup fragments"))?, + ) + .ok_or_else(|| invalid("too many dedup units"))?; + if plan.policy_version != POLICY_VERSION + || snapshot != plan.snapshot + || definition.vector_source.as_ref().map(|s| s.kind) != Some(VectorSourceKind::Dedup) + || source + .get_fragments() + .iter() + .map(|f| f.id() as u32) + .collect::>() + != plan.source_fragments + { + return Err(invalid("dedup task belongs to another snapshot or policy")); + } + Ok((view_ds, definition, source)) +} + +fn edge_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt64, false), + Field::new("b", DataType::UInt64, false), + ])) +} + +/// 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( + stream: SendableRecordBatchStream, + memory: usize, +) -> Result<(RoaringTreemap, RoaringTreemap)> { + let schema = edge_schema(); + let batch_schema = schema.clone(); + let canonical = stream + .map_ok(move |batch| { + let a = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("native row id"); + let b = batch + .column(1) + .as_any() + .downcast_ref::() + .expect("native row id"); + RecordBatch::try_new( + batch_schema.clone(), + vec![ + Arc::new(UInt64Array::from_iter_values( + a.values().iter().zip(b.values()).map(|(&a, &b)| a.min(b)), + )), + Arc::new(UInt64Array::from_iter_values( + a.values().iter().zip(b.values()).map(|(&a, &b)| a.max(b)), + )), + ], + ) + .map_err(DataFusionError::from) + }) + .and_then(futures::future::ready); + let canonical = Box::pin(RecordBatchStreamAdapter::new(schema, canonical)); + let mut config = SessionConfig::default().with_batch_size(MATERIALIZE_BATCH_ROWS); + // Leave enough room to merge spilled runs while the final in-memory run + // is still resident. Small batches also bound each merge input buffer. + config.options_mut().execution.sort_spill_reservation_bytes = memory / 2; + let ctx = SessionContext::new_with_config_rt( + config, + RuntimeEnvBuilder::new() + .with_memory_limit(memory, 1.0) + .build_arc()?, + ); + let mut sorted = ctx + .read_one_shot(canonical)? + .sort_by(vec![col("a"), col("b")])? + .execute_stream() + .await?; + let mut removed = RoaringTreemap::new(); + let mut representatives = RoaringTreemap::new(); + while let Some(batch) = sorted.try_next().await? { + let a = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("sorted id"); + let b = batch + .column(1) + .as_any() + .downcast_ref::() + .expect("sorted id"); + for (&a, &b) in a.values().iter().zip(b.values()) { + if a != b && !removed.contains(a) && !removed.contains(b) { + representatives.insert(a); + removed.insert(b); + } + } + } + Ok((removed, representatives)) +} + +#[derive(Default)] +struct Masks { + removed: RoaringBitmap, + representatives: RoaringBitmap, +} + +async fn persist_masks( + view: &Dataset, + source: &Dataset, + removed: RoaringTreemap, + representatives: RoaringTreemap, +) -> Result> { + let row_index = get_row_id_index(source).await?; + let mut masks: BTreeMap = BTreeMap::new(); + for (ids, is_removed) in [(removed, true), (representatives, false)] { + for id in ids { + let address = match &row_index { + Some(index) => index + .get(id)? + .ok_or_else(|| invalid("pair row id missing from pinned source"))?, + None => RowAddress::new_from_u64(id), + }; + let mask = masks.entry(address.fragment_id()).or_default(); + if is_removed { + mask.removed.insert(address.row_offset()); + } else { + mask.representatives.insert(address.row_offset()); + } + } + } + let store = view.object_store(None).await?; + let mut files = Vec::new(); + for (source_fragment, mask) in masks { + if source.get_fragment(source_fragment as usize).is_none() { + return Err(invalid("pair row address is outside the pinned snapshot")); + } + let mut bytes = MASK_MAGIC.to_vec(); + mask.removed.serialize_into(&mut bytes).map_err(io_error)?; + mask.representatives + .serialize_into(&mut bytes) + .map_err(io_error)?; + let object_id = Uuid::new_v4(); + store + .put( + &view + .indices_dir() + .join(object_id.to_string()) + .join(MASK_FILE), + &bytes, + ) + .await?; + files.push(MaskFile { + source_fragment, + object_id, + bytes: bytes.len(), + digest: format!("{:x}", Sha256::digest(&bytes)), + }); + } + Ok(files) +} + +fn validate_receipts<'a>( + plan: &VectorDedupPlan, + receipts: &'a [WrittenDedupUnit], + count: u32, +) -> Result> { + let identity = plan.identity()?; + let mut ordered = vec![None; count as usize]; + for receipt in receipts { + if receipt.plan_id != identity { + return Err(invalid("dedup receipt belongs to another plan")); + } + let slot = ordered + .get_mut(receipt.unit as usize) + .ok_or_else(|| invalid("unexpected dedup receipt"))?; + if slot.is_some() { + return Err(invalid("duplicate dedup receipt")); + } + if receipt.unit < plan.dependencies() { + if !receipt.fragments.is_empty() + || receipt.rows != 0 + || receipt + .masks + .iter() + .any(|m| !plan.source_fragments.contains(&m.source_fragment)) + { + return Err(invalid("invalid selection receipt")); + } + } else if !receipt.masks.is_empty() { + return Err(invalid("invalid materialization receipt")); + } + *slot = Some(receipt); + } + ordered + .into_iter() + .collect::>>() + .ok_or_else(|| invalid("missing dedup dependency receipt")) +} + +async fn load_masks( + view: &Dataset, + source_fragment: u32, + dependencies: &[&WrittenDedupUnit], +) -> Result { + let store = view.object_store(None).await?; + let mut masks = Masks::default(); + for file in dependencies + .iter() + .flat_map(|u| &u.masks) + .filter(|m| m.source_fragment == source_fragment) + { + let reader = store + .open( + &view + .indices_dir() + .join(file.object_id.to_string()) + .join(MASK_FILE), + ) + .await?; + if reader.size().await? != file.bytes { + return Err(invalid("dedup mask length mismatch")); + } + let bytes = reader.get_range(0..file.bytes).await?; + if !bytes.starts_with(MASK_MAGIC) || format!("{:x}", Sha256::digest(&bytes)) != file.digest + { + return Err(invalid("dedup mask checksum mismatch")); + } + let mut cursor = Cursor::new(&bytes[MASK_MAGIC.len()..]); + masks.removed |= RoaringBitmap::deserialize_from(&mut cursor).map_err(io_error)?; + masks.representatives |= RoaringBitmap::deserialize_from(&mut cursor).map_err(io_error)?; + if cursor.position() as usize != bytes.len() - MASK_MAGIC.len() { + return Err(invalid("dedup mask has trailing data")); + } + } + // Native index scopes normally have disjoint rows. Fail closed if an + // overlapping scope would delete a representative used by another scope. + if !masks.removed.is_disjoint(&masks.representatives) { + return Err(invalid( + "overlapping index scopes disagree on retained representatives", + )); + } + Ok(masks.removed) +} + +pub(super) async fn write_unit( + view: &Table, + unit: u32, + plan: &VectorDedupPlan, + dependencies: &[WrittenDedupUnit], +) -> Result { + let (view_ds, definition, source) = open(view, plan).await?; + let mut written = WrittenDedupUnit { + unit, + plan_id: plan.identity()?, + masks: Vec::new(), + fragments: Vec::new(), + rows: 0, + }; + if unit < plan.dependencies() { + let config = definition + .vector_source + .as_ref() + .expect("dedup source") + .config()?; + let exec = DuplicatePairsExec::try_new( + source.clone(), + config, + vec![plan.native.tasks[unit as usize].clone()], + )?; + let stream = exec.execute(0, Arc::new(TaskContext::default()))?; + let (removed, representatives) = select_representatives(stream, SORT_MEMORY_BYTES).await?; + written.masks = persist_masks(&view_ds, &source, removed, representatives).await?; + return Ok(written); + } + let source_fragment = *plan + .source_fragments + .get((unit - plan.dependencies()) as usize) + .ok_or_else(|| invalid("dedup unit is outside the plan"))?; + let dependencies = validate_receipts(plan, dependencies, plan.dependencies())?; + let removed = load_masks(&view_ds, source_fragment, &dependencies).await?; + let fragment = source + .get_fragment(source_fragment as usize) + .ok_or_else(|| invalid("source fragment missing"))?; + // The scanner prefilter uses snapshot row IDs. Our persisted masks use + // physical offsets so each writer only loads its own fragment's artifacts. + let blocked = if source.manifest.uses_stable_row_ids() { + let sequence = load_row_id_sequence(&source, fragment.metadata()).await?; + if removed + .max() + .is_some_and(|offset| offset as u64 >= sequence.len()) + { + return Err(invalid("dedup mask offset is outside the source fragment")); + } + RowAddrTreeMap::from_iter(sequence.select(removed.iter().map(|offset| offset as usize))) + } else { + RowAddrTreeMap::from_iter( + removed + .iter() + .map(|offset| u64::from(RowAddress::new_from_parts(source_fragment, offset))), + ) + }; + let mut scanner = source.scan(); + scanner + .with_fragments(vec![fragment.metadata().clone()]) + .with_row_addr_prefilter(RowAddrMask::from_block(blocked)) + .blob_handling(BlobHandling::AllBinary) + .batch_size(MATERIALIZE_BATCH_ROWS) + .strict_batch_size(true); + let stream = scanner.try_into_stream().await?; + let schema = lance_io::stream::RecordBatchStream::schema(&stream); + let rows = Arc::new(std::sync::atomic::AtomicU64::new(0)); + let count = rows.clone(); + let stream = stream + .map_err(|e| DataFusionError::External(Box::new(e))) + .inspect_ok(move |batch| { + count.fetch_add( + batch.num_rows() as u64, + std::sync::atomic::Ordering::Relaxed, + ); + }); + let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new(schema, stream)); + // Materialized binary payloads must become the destination's logical blob + // structs before the Lance writer encodes them. Reuse the normal add path's + // coercions, including nulls and nested fields. + let coerced = crate::table::datafusion::cast::cast_to_table_schema( + Arc::new(OneShotExec::new(stream)), + &Schema::from(view_ds.schema()), + )?; + let stream = coerced.execute(0, Arc::new(TaskContext::default()))?; + let txn = InsertBuilder::new(WriteDestination::Dataset(Arc::new(view_ds))) + .with_params(&WriteParams { + mode: WriteMode::Append, + ..Default::default() + }) + .execute_uncommitted_stream(stream) + .await?; + let Operation::Append { fragments } = txn.operation else { + return Err(invalid( + "dedup materialization did not stage append fragments", + )); + }; + written.fragments = fragments; + written.rows = rows.load(std::sync::atomic::Ordering::Relaxed); + Ok(written) +} + +pub(super) async fn commit( + view: &Table, + plan: &VectorDedupPlan, + units: Vec, + expected: Option<&str>, +) -> Result { + let (_, definition, source) = open(view, plan).await?; + let config = definition + .vector_source + .as_ref() + .expect("dedup source") + .config()?; + if plan_duplicate_pairs(source.clone(), &config).await? != plan.native.tasks { + return Err(invalid("dedup plan does not cover the pinned index")); + } + validate_receipts(plan, &units, plan.snapshot.units)?; + let rows = units.iter().map(|u| u.rows).sum(); + let fragments = units.into_iter().flat_map(|u| u.fragments).collect(); + super::grouped_units::commit_fragments( + view, + &plan.snapshot, + &definition, + source.manifest.timestamp_nanos, + fragments, + rows, + expected, + ) + .await +} + +#[cfg(test)] +mod tests { + use super::*; + + fn edges(values: &[(u64, u64)], batch_size: usize) -> SendableRecordBatchStream { + let schema = edge_schema(); + let batches = values + .chunks(batch_size) + .map(|edges| { + Ok(RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(UInt64Array::from_iter_values(edges.iter().map(|e| e.0))), + Arc::new(UInt64Array::from_iter_values(edges.iter().map(|e| e.1))), + ], + ) + .unwrap()) + }) + .collect::>(); + Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(batches), + )) + } + + #[tokio::test] + async fn direct_representatives_handle_chains_orientation_and_uint64() { + let high = u64::MAX - 2; + let pairs = [ + (2, 1), + (1, 0), + (high + 2, high + 1), + (high, high + 1), + (1, 0), + ]; + let (removed, representatives) = + select_representatives(edges(&pairs, 1), SORT_MEMORY_BYTES) + .await + .unwrap(); + assert_eq!(removed.iter().collect::>(), [1, high + 1]); + assert_eq!(representatives.iter().collect::>(), [0, high]); + } + + #[tokio::test] + async fn masks_reject_overlapping_representatives_and_corrupt_artifacts() { + let dir = tempfile::tempdir().unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::UInt64, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(UInt64Array::from(vec![0, 1, 2]))], + ) + .unwrap(); + let ds = Dataset::write( + arrow_array::RecordBatchIterator::new(vec![Ok(batch)], schema), + dir.path().join("source.lance").to_str().unwrap(), + None, + ) + .await + .unwrap(); + let first = persist_masks( + &ds, + &ds, + [1].into_iter().collect(), + [0].into_iter().collect(), + ) + .await + .unwrap(); + let second = persist_masks( + &ds, + &ds, + [2].into_iter().collect(), + [1].into_iter().collect(), + ) + .await + .unwrap(); + let receipt = |masks| WrittenDedupUnit { + unit: 0, + plan_id: "test".into(), + masks, + fragments: vec![], + rows: 0, + }; + let first = receipt(first); + let second = receipt(second); + assert_eq!( + load_masks(&ds, 0, &[&first]).await.unwrap(), + [1].into_iter().collect() + ); + assert!( + load_masks(&ds, 0, &[&first, &second]) + .await + .unwrap_err() + .to_string() + .contains("overlapping") + ); + let file = &first.masks[0]; + ds.object_store(None) + .await + .unwrap() + .put( + &ds.indices_dir() + .join(file.object_id.to_string()) + .join(MASK_FILE), + &vec![0; file.bytes], + ) + .await + .unwrap(); + assert!( + load_masks(&ds, 0, &[&first]) + .await + .unwrap_err() + .to_string() + .contains("checksum") + ); + } + + #[tokio::test] + async fn external_sort_handles_more_edges_than_its_budget() { + let pairs = (0..262144) + .rev() + .map(|i| ((i + 1) as u64, i as u64)) + .collect::>(); + let (removed, reps) = select_representatives(edges(&pairs, 4096), 2 * 1024 * 1024) + .await + .unwrap(); + assert_eq!(removed.len(), 131072); + assert_eq!(reps.len(), 131072); + assert!(removed.iter().all(|id| id % 2 == 1)); + assert!(reps.iter().all(|id| id % 2 == 0)); + } +} diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 624340ffb..cabddbe7b 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -2140,7 +2140,7 @@ impl BaseTable for RemoteTable { let definition = match response.query { Some(query) => MaterializedViewDefinition::from_sql(&query)?, None => MaterializedViewDefinition { - duplicate_pairs: None, + vector_source: None, source_table: response.source_table, source_namespace: response.source_namespace, lateral: None, diff --git a/rust/lancedb/tests/duplicate_pairs.rs b/rust/lancedb/tests/duplicate_pairs.rs index e63b20ecf..910088bc8 100644 --- a/rust/lancedb/tests/duplicate_pairs.rs +++ b/rust/lancedb/tests/duplicate_pairs.rs @@ -267,6 +267,135 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result view.table().count_rows(None).await? as usize, pairs(&expected).len() ); + // The complete declarative pipeline keeps direct representatives + // and materializes original rows, including isolated source rows. + use lancedb::materialized_view::write_refresh_partition_with_inputs; + use lancedb::query::QueryBase; + let definition = MaterializedViewDefinition::from_sql(&format!( + "SELECT * FROM vector_dedup('source', {version}, 'vector', {threshold})" + ))?; + let clean = prepare_definition(&table, definition) + .await? + .create(&format!("clean_{}", threshold.to_bits())) + .await?; + let clean_plan = plan_partitioned_refresh(clean.table(), Some(version)) + .await? + .unwrap(); + assert_eq!(clean_plan.dependency_units() as usize, tasks.len()); + assert_eq!(clean_plan.units() as usize, tasks.len() + 2); + let clean_plan: lancedb::materialized_view::PartitionedRefreshPlan = + serde_json::from_slice(&serde_json::to_vec(&clean_plan)?)?; + assert!( + write_refresh_partition(clean.table(), tasks.len() as u32, &clean_plan) + .await + .is_err() + ); + let mut receipts = Vec::new(); + for unit in 0..clean_plan.units() { + let dependencies = &receipts[..receipts.len().min(tasks.len())]; + if unit == 0 || unit == tasks.len() as u32 { + // Both selection and materialization can lose a receipt. + drop( + write_refresh_partition_with_inputs( + clean.table(), + unit, + &clean_plan, + dependencies, + ) + .await?, + ); + } + let receipt = write_refresh_partition_with_inputs( + clean.table(), + unit, + &clean_plan, + dependencies, + ) + .await?; + receipts.push(serde_json::from_slice(&serde_json::to_vec(&receipt)?)?); + } + assert_eq!(clean.table().count_rows(None).await?, 0); + assert!( + commit_partitioned_refresh( + clean.table(), + &clean_plan, + receipts[..receipts.len() - 1].to_vec(), + clean.incarnation() + ) + .await + .is_err() + ); + commit_partitioned_refresh(clean.table(), &clean_plan, receipts, clean.incarnation()) + .await?; + let mut edges = pairs(&expected) + .into_iter() + .map(|(a, b, _)| (a.min(b), a.max(b))) + .collect::>(); + edges.sort_unstable(); + let mut removed = std::collections::BTreeSet::new(); + for (a, b) in edges { + if !removed.contains(&a) { + removed.insert(b); + } + } + let source_ids = ds + .scan() + .project(&["id"])? + .with_row_id() + .try_into_stream() + .await? + .try_collect::>() + .await?; + let mut expected_ids = Vec::new(); + for batch in source_ids { + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let rowids = batch + .column_by_name("_rowid") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + if !removed.contains(&rowids.value(i)) { + expected_ids.push(ids.value(i)); + } + } + } + let kept = clean + .table() + .query() + .select(lancedb::query::Select::columns(&["id"])) + .execute() + .await? + .try_collect::>() + .await?; + let mut actual_ids = kept + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + actual_ids.sort_unstable(); + expected_ids.sort_unstable(); + assert_eq!(actual_ids, expected_ids); + if threshold == 1.0 { + assert!(actual_ids.contains(&0) && actual_ids.contains(&2)); + assert!(!actual_ids.contains(&1), "A-B-C retains A and C"); + } + clean.refresh().execute().await?; + assert_eq!(clean.table().count_rows(None).await?, expected_ids.len()); + assert_eq!(table.count_rows(None).await?, n * 2); assert!( DuplicatePairsExec::try_new( ds.clone(), @@ -298,3 +427,135 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result } Ok(()) } + +#[tokio::test] +async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions() -> anyhow::Result<()> +{ + use arrow_array::{LargeBinaryArray, StringArray, StructArray, types::Float32Type}; + use lancedb::materialized_view::{MaterializedViewDefinition, prepare_definition}; + for stable in [false, true] { + let dir = tempfile::tempdir()?; + let conn = lancedb::connect(dir.path().to_str().unwrap()) + .execute() + .await?; + let vectors = FixedSizeListArray::from_iter_primitive::( + [ + Some(vec![Some(0.0), Some(0.0)]), + Some(vec![Some(1.0), Some(0.0)]), + Some(vec![Some(2.0), Some(0.0)]), + None, + Some(vec![Some(30.0), Some(0.0)]), + Some(vec![Some(40.0), Some(0.0)]), + ], + 2, + ); + let blob = lancedb::blob::blob("image", true); + let DataType::Struct(fields) = blob.data_type().clone() else { + unreachable!() + }; + let images = StructArray::new( + fields, + vec![ + Arc::new(LargeBinaryArray::from_iter_values( + (0..6).map(|id| vec![id; 4096]), + )), + Arc::new(StringArray::from(vec![None::<&str>; 6])), + ], + None, + ); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("vector", vectors.data_type().clone(), true), + blob, + ])), + vec![ + Arc::new(Int64Array::from_iter_values(0..6)), + Arc::new(vectors), + Arc::new(images), + ], + )?; + let table = conn + .create_table("source", batch) + .storage_option("new_table_enable_stable_row_ids", stable.to_string()) + .execute() + .await?; + table.delete("id = 5").await?; + table + .optimize(lancedb::table::OptimizeAction::Compact { + options: lancedb::table::CompactionOptions::default(), + remap_options: None, + }) + .await?; + let uri = table.uri().await?; + let mut ds = Dataset::open(&uri).await?; + assert!( + ds.fragments().iter().all(|f| f.id > 0), + "compaction must change physical row addresses" + ); + let centroids = FixedSizeListArray::from_iter_primitive::( + [Some(vec![Some(0.0), Some(0.0)])], + 2, + ); + let ivf = IvfBuildParams { + centroids: Some(Arc::new(centroids)), + ..IvfBuildParams::new(1) + }; + ds.create_index( + &["vector"], + IndexType::Vector, + Some("vector_idx".into()), + &VectorIndexParams::with_ivf_flat_params(MetricType::L2, ivf), + true, + ) + .await?; + let version = ds.version().version; + let definition = MaterializedViewDefinition::from_sql(&format!( + "SELECT * FROM vector_dedup('source', {version}, 'vector', 1)" + ))?; + let clean = prepare_definition(&table, definition) + .await? + .create("clean") + .await?; + // A later source delete does not change the declared result snapshot. + table.delete("id = 0").await?; + clean.refresh().execute().await?; + assert_eq!(clean.table().count_rows(None).await?, 4); + let output = Dataset::open(&clean.table().uri().await?).await?; + let batches = output + .scan() + .blob_handling(lance_core::datatypes::BlobHandling::AllBinary) + .try_into_stream() + .await? + .try_collect::>() + .await?; + let mut ids = Vec::new(); + for batch in batches { + let row_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() { + let id = row_ids.value(i); + assert_eq!(images.value(i), vec![id as u8; 4096]); + ids.push(id); + } + } + ids.sort_unstable(); + assert_eq!(ids, [0, 2, 3, 4]); + assert_eq!( + table.count_rows(None).await?, + 4, + "refresh must not mutate the source" + ); + } + Ok(()) +}