rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-29 14:56:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(mv): refresh native duplicate pairs in partition units

Browse Source
This commit is contained in:
Yang Cen
2026-09-26 15:44:44 +08:00
parent 131157ecb0
commit 7f5779fa9d
9 changed files with 822 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/lancedb/src/materialized_view.rs
+66 -15
View File
@@ -9,8 +9,11 @@
//! query this version cannot maintain reads back as unrefreshable, not as a
//! plain table. Queries, indexes and search work on the view unchanged.
mod duplicate_pairs;
mod grouped;
mod grouped_units;
mod partitioned;
pub use duplicate_pairs::DuplicatePairsView;
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
@@ -74,6 +77,10 @@ pub use grouped_units::{
GroupedRefreshPlan, WrittenUnit, commit_grouped_refresh, plan_grouped_refresh,
write_grouped_unit,
};
pub use partitioned::{
PartitionedRefreshPlan, WrittenPartition, commit_partitioned_refresh, plan_partitioned_refresh,
write_refresh_partition,
};
mod query;
pub mod refresh;
@@ -146,18 +153,23 @@ const COLUMN_DEFINITIONS_META_KEY: &str = "lancedb::column_definitions";
/// The newest layout this version reads under [`DEFINITION_META_KEY`]:
/// `{"format": N, "query": "<SQL>"}`, the query as
/// [`MaterializedViewDefinition::to_sql`] renders it. Format 2 is a query
/// with `GROUP BY`; every other query is written as format 1, so a reader
/// that predates grouping reports a grouped view as unrefreshable rather
/// than failing to parse it. A reader refuses a newer format rather than
/// [`MaterializedViewDefinition::to_sql`] renders it. Format 3 identifies a
/// native duplicate-pair 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 = 2;
pub const DEFINITION_FORMAT: u64 = 3;
/// The format `definition` is written in; see [`DEFINITION_FORMAT`].
fn format_of(definition: &MaterializedViewDefinition) -> u64 {
if definition.is_grouped() { 2 } else { 1 }
if definition.duplicate_pairs.is_some() {
3
} else if definition.is_grouped() {
2
} else {
1
}
}
/// Legacy `kind` tag of the structured layout written before
@@ -300,6 +312,8 @@ pub fn read_staging(metadata: &HashMap<String, String>) -> Result<Option<Staging
/// for the shape.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MaterializedViewDefinition {
/// A table-bound native pair source. It is never evaluated per source row.
pub duplicate_pairs: Option<DuplicatePairsView>,
/// 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.
@@ -498,6 +512,7 @@ pub fn read_definition(metadata: &HashMap<String, String>) -> Result<Option<Stor
)));
}
Ok(Some(StoredDefinition::Query(MaterializedViewDefinition {
duplicate_pairs: None,
source_table: legacy.source_table,
source_namespace: legacy.source_namespace,
lateral: None,
@@ -554,6 +569,9 @@ pub(crate) fn plan(
definition: &MaterializedViewDefinition,
staging: Option<&StagingBinding>,
) -> Result<Planned> {
if definition.duplicate_pairs.is_some() {
return duplicate_pairs::plan(source_schema, definition);
}
let filter = definition
.filter
.as_deref()
@@ -712,6 +730,7 @@ pub(crate) fn plan(
}
let definition = MaterializedViewDefinition {
duplicate_pairs: None,
source_table: definition.source_table.clone(),
source_namespace: definition.source_namespace.clone(),
lateral: definition.lateral.clone(),
@@ -1145,6 +1164,11 @@ 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<String> {
if self.definition.duplicate_pairs.is_some() {
return Err(Error::InvalidInput {
message: "native pair views have no per-source-row inputs".into(),
});
}
// A grouped view's rows are groups; no source row carries a value into one.
if self.definition.is_grouped() {
return Err(Error::InvalidInput {
@@ -1245,6 +1269,11 @@ impl PreparedDeclaration {
columns: Vec<(usize, ArrowField)>,
bindings: &[FunctionBinding],
) -> Result<Self> {
if self.definition.duplicate_pairs.is_some() {
return Err(Error::InvalidInput {
message: "computed columns on native pair views are not supported".into(),
});
}
let invalid = |message: String| Error::InvalidInput { message };
if columns.is_empty() {
return Err(invalid("at least one computed column is needed".into()));
@@ -1451,6 +1480,7 @@ pub async fn prepare_declaration(
limit: Option<u64>,
) -> Result<PreparedDeclaration> {
let definition = MaterializedViewDefinition {
duplicate_pairs: None,
source_table: source.name().to_string(),
source_namespace: source.namespace().to_vec(),
lateral: None,
@@ -1627,7 +1657,9 @@ async fn prepare_with(
),
});
}
if !native.dataset.get().await?.manifest.uses_stable_row_ids() {
if definition.duplicate_pairs.is_none()
&& !native.dataset.get().await?.manifest.uses_stable_row_ids()
{
return Err(Error::InvalidInput {
message: format!(
"materialized views require stable row ids on the source table; \
@@ -1653,7 +1685,22 @@ async fn prepare_with(
message: format!("view column name '{}' is reserved", reserved.output),
});
}
let source_schema = resolved.schema().await?;
let source_schema = if let Some(pairs) = &definition.duplicate_pairs {
let ds = native
.dataset
.get()
.await?
.checkout_version(pairs.dataset_version)
.await?;
crate::table::datafusion::udtf::duplicate_pairs::plan_duplicate_pairs(
Arc::new(ds.clone()),
&pairs.config()?,
)
.await?;
Arc::new(ArrowSchema::from(ds.schema()))
} else {
resolved.schema().await?
};
let source_metadata = source_schema.metadata().clone();
let Planned {
definition,
@@ -1670,11 +1717,13 @@ async fn prepare_with(
None => source_schema.clone(),
Some(unnest) => flattened_schema(&source_schema, &unnest)?,
};
fields.push(ArrowField::new(
SOURCE_ROW_ID_COLUMN,
DataType::UInt64,
false,
));
if definition.duplicate_pairs.is_none() {
fields.push(ArrowField::new(
SOURCE_ROW_ID_COLUMN,
DataType::UInt64,
false,
));
}
// Only column-describing metadata comes along: structural declarations
// describe how a table is written, and a view is written by refresh alone.
let mut metadata: HashMap<String, String> = HashMap::new();
@@ -2186,6 +2235,7 @@ mod tests {
assert_eq!(
view.definition(),
&MaterializedViewDefinition {
duplicate_pairs: None,
source_table: "people".into(),
source_namespace: Vec::new(),
projections: vec![
@@ -3235,6 +3285,7 @@ mod tests {
fn definition(source_namespace: Vec<String>) -> MaterializedViewDefinition {
MaterializedViewDefinition {
duplicate_pairs: None,
source_table: "people".to_string(),
source_namespace,
lateral: None,
@@ -3321,8 +3372,8 @@ mod tests {
#[test]
fn a_newer_format_is_reported_not_guessed() {
assert_eq!(
read(r#"{"format":3,"query":"SELECT name FROM people"}"#).unwrap(),
Some(StoredDefinition::Newer { format: "3".into() })
read(r#"{"format":4,"query":"SELECT name FROM people"}"#).unwrap(),
Some(StoredDefinition::Newer { format: "4".into() })
);
assert_eq!(
read(r#"{"kind":"join"}"#).unwrap(),
rust/lancedb/src/materialized_view/duplicate_pairs.rs
+422
View File
@@ -0,0 +1,422 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Native pair generation as partitioned materialized-view refresh work.
use std::sync::Arc;
use arrow_schema::{Schema, SchemaRef};
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_physical_plan::ExecutionPlan;
use datafusion_sql::sqlparser::ast::{
Expr, FunctionArg, FunctionArgExpr, TableFunctionArgs, Value,
};
use futures::TryStreamExt;
use lance::Dataset;
use lance::dataset::transaction::Operation;
use lance::dataset::{InsertBuilder, WriteDestination, WriteMode, WriteParams};
use lance_table::format::Fragment;
use serde::{Deserialize, Serialize};
use super::grouped_units::GroupedRefreshPlan;
use super::refresh::{
RefreshMaterializedViewResult, ensure_incarnation, ensure_no_mem_wal, open_source,
};
use super::{MaterializedViewDefinition, Planned, ViewProjection};
use crate::table::Table;
use crate::table::datafusion::udtf::duplicate_pairs::{
DUPLICATE_PAIRS_OPERATOR_VERSION, DuplicatePairTask, DuplicatePairsConfig, DuplicatePairsExec,
duplicate_pair_schema, plan_duplicate_pairs,
};
use crate::{Error, Result};
pub(crate) const FUNCTION_NAME: &str = "vector_duplicate_pairs";
/// 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 dataset_version: u64,
pub column: String,
/// Validated numeric SQL spelling; execution and task identity use f32 bits.
pub distance_threshold: String,
}
impl DuplicatePairsView {
pub(crate) fn config(&self) -> Result<DuplicatePairsConfig> {
let threshold: f32 = self
.distance_threshold
.parse()
.map_err(|_| invalid("invalid pair threshold"))?;
if self.dataset_version == 0 || self.column.is_empty() || !threshold.is_finite() {
return Err(invalid(
"duplicate pairs require a positive snapshot, a column and a finite threshold",
));
}
Ok(DuplicatePairsConfig {
dataset_version: self.dataset_version,
column: self.column.clone(),
distance_threshold: threshold,
})
}
}
fn invalid(message: impl Into<String>) -> Error {
Error::InvalidInput {
message: message.into(),
}
}
pub(super) fn source(
args: &TableFunctionArgs,
) -> Result<(Vec<String>, String, DuplicatePairsView)> {
if args.settings.is_some() {
return Err(invalid("duplicate pairs do not accept function settings"));
}
let exprs = args
.args
.iter()
.map(|arg| match arg {
FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => Ok(expr),
_ => Err(invalid(
"duplicate pairs require four positional literal arguments",
)),
})
.collect::<Result<Vec<_>>>()?;
let [table, version, column, threshold] = exprs.as_slice() else {
return Err(invalid(
"vector_duplicate_pairs requires (table, dataset_version, column, distance_threshold)",
));
};
let string = |expr: &Expr| match expr {
Expr::Value(value) => match &value.value {
Value::SingleQuotedString(value) => Ok(value.clone()),
_ => Err(invalid("table and column must be string literals")),
},
_ => Err(invalid("table and column must be string literals")),
};
let table = string(table)?;
let plain = MaterializedViewDefinition::from_sql(&format!("SELECT * FROM {table}"))?;
if plain.duplicate_pairs.is_some()
|| plain.lateral.is_some()
|| plain.filter.is_some()
|| plain.is_grouped()
|| plain.limit.is_some()
{
return Err(invalid("duplicate pairs must name one source table"));
}
let dataset_version = match version {
Expr::Value(value) => match &value.value {
Value::Number(number, _) => number
.parse()
.map_err(|_| invalid("invalid source version"))?,
_ => return Err(invalid("source version must be an integer literal")),
},
_ => return Err(invalid("source version must be an integer literal")),
};
let config = DuplicatePairsView {
dataset_version,
column: string(column)?,
distance_threshold: threshold.to_string(),
};
config.config()?;
Ok((plain.source_namespace, plain.source_table, config))
}
pub(super) fn check_shape(definition: &MaterializedViewDefinition) -> Result<()> {
let schema = duplicate_pair_schema();
let projections = schema
.fields()
.iter()
.map(|field| ViewProjection {
output: field.name().clone(),
expression: field.name().clone(),
})
.collect::<Vec<_>>();
if (!definition.selects_star() && definition.projections != projections)
|| definition.filter.is_some()
|| definition.limit.is_some()
|| definition.lateral.is_some()
|| definition.is_grouped()
{
return Err(invalid(
"a native pair view selects *, or row_id_a, row_id_b, distance, without other clauses",
));
}
Ok(())
}
pub(super) fn plan(
source_schema: SchemaRef,
definition: &MaterializedViewDefinition,
) -> Result<Planned> {
check_shape(definition)?;
let config = definition.duplicate_pairs.as_ref().expect("native source");
config.config()?;
source_schema.field_with_name(&config.column)?;
let schema = duplicate_pair_schema();
let projections = schema
.fields()
.iter()
.map(|field| ViewProjection {
output: field.name().clone(),
expression: field.name().clone(),
})
.collect();
let mut definition = definition.clone();
let inputs = vec![config.column.clone()];
definition.projections = projections;
Ok(Planned {
definition,
fields: schema.fields().iter().map(|f| f.as_ref().clone()).collect(),
lineage: Default::default(),
inputs,
})
}
/// A saved native refresh plan, including every physical partition identity.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct DuplicatePairsRefreshPlan {
pub snapshot: GroupedRefreshPlan,
pub source_uri: String,
pub definition_sql: String,
pub operator_version: u32,
pub threshold_bits: u32,
pub tasks: Vec<DuplicatePairTask>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WrittenPairsPartition {
pub(super) unit: u32,
pub(super) plan_id: String,
pub(super) task: DuplicatePairTask,
pub(super) fragments: Vec<Fragment>,
pub(super) rows: u64,
}
impl DuplicatePairsRefreshPlan {
fn identity(&self) -> Result<String> {
use sha2::{Digest, Sha256};
let encoded = serde_json::to_vec(&(
&self.snapshot,
&self.source_uri,
&self.definition_sql,
self.operator_version,
self.threshold_bits,
))
.map_err(|e| invalid(e.to_string()))?;
Ok(format!("{:x}", Sha256::digest(encoded)))
}
}
async fn open(view: &Table) -> Result<(Dataset, MaterializedViewDefinition, Arc<Dataset>)> {
let native = view
.as_native()
.ok_or_else(|| invalid("native pair refresh requires a local view"))?;
native.dataset.ensure_mutable()?;
native.dataset.reload().await?;
let view_ds = native.dataset.get().await?.as_ref().clone();
let Some(super::StoredDefinition::Query(definition)) =
super::read_definition(&view_ds.schema().metadata)?
else {
return Err(invalid("view does not carry a supported definition"));
};
let config = definition
.duplicate_pairs
.as_ref()
.ok_or_else(|| invalid("not a native pair view"))?
.config()?;
let source = Arc::new(
open_source(view, &definition, None)
.await?
.checkout_version(config.dataset_version)
.await?,
);
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 physical = Schema::from(view_ds.schema());
if physical.fields().len() != expected.fields().len()
|| physical
.fields()
.iter()
.zip(expected.fields())
.any(|(a, b)| {
a.name() != b.name()
|| a.data_type() != b.data_type()
|| a.is_nullable() != b.is_nullable()
})
{
return Err(invalid("native pair view schema changed since creation"));
}
Ok((view_ds, planned.definition, source))
}
pub(super) async fn plan_refresh(
view: &Table,
pinned: Option<u64>,
) -> Result<DuplicatePairsRefreshPlan> {
let (view_ds, definition, source) = open(view).await?;
let config = definition
.duplicate_pairs
.as_ref()
.expect("native")
.config()?;
if pinned.is_some_and(|v| v != config.dataset_version) {
return Err(invalid(
"refresh source version differs from the native pair view's fixed snapshot",
));
}
let tasks = plan_duplicate_pairs(source.clone(), &config).await?;
let units = u32::try_from(tasks.len()).map_err(|_| invalid("too many index partitions"))?;
let incarnation = view_ds
.schema()
.metadata
.get(super::INCARNATION_META_KEY)
.cloned()
.ok_or_else(|| invalid("native pair view has no incarnation"))?;
Ok(DuplicatePairsRefreshPlan {
snapshot: GroupedRefreshPlan {
units,
source_version: config.dataset_version,
view_version: view_ds.version().version,
incarnation,
},
source_uri: source.uri().to_string(),
definition_sql: definition.to_sql(),
operator_version: DUPLICATE_PAIRS_OPERATOR_VERSION,
threshold_bits: config.distance_threshold.to_bits(),
tasks,
})
}
async fn open_planned(
view: &Table,
plan: &DuplicatePairsRefreshPlan,
) -> Result<(Dataset, MaterializedViewDefinition, Arc<Dataset>)> {
let (view_ds, definition, source) = open(view).await?;
ensure_incarnation(&view_ds, Some(&plan.snapshot.incarnation), view.name()).await?;
let config = definition
.duplicate_pairs
.as_ref()
.expect("native")
.config()?;
if source.uri() != plan.source_uri
|| source.version().version != plan.snapshot.source_version
|| definition.to_sql() != plan.definition_sql
|| config.distance_threshold.to_bits() != plan.threshold_bits
|| plan.operator_version != DUPLICATE_PAIRS_OPERATOR_VERSION
|| plan.tasks.len() != plan.snapshot.units as usize
{
return Err(invalid(
"native pair task belongs to a different snapshot, configuration or operator",
));
}
Ok((view_ds, definition, source))
}
pub(super) async fn write_unit(
view: &Table,
unit: u32,
plan: &DuplicatePairsRefreshPlan,
) -> Result<WrittenPairsPartition> {
let (view_ds, definition, source) = open_planned(view, plan).await?;
let task = plan
.tasks
.get(unit as usize)
.ok_or_else(|| invalid("pair partition is outside the plan"))?
.clone();
let config = definition
.duplicate_pairs
.as_ref()
.expect("native")
.config()?;
let exec = DuplicatePairsExec::try_new(source, config, vec![task.clone()])?;
let stream = exec.execute(0, Arc::new(TaskContext::default()))?;
let rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
let count = rows.clone();
let stream = stream.inspect_ok(move |batch| {
count.fetch_add(
batch.num_rows() as u64,
std::sync::atomic::Ordering::Relaxed,
);
});
let stream: SendableRecordBatchStream = Box::pin(
datafusion_physical_plan::stream::RecordBatchStreamAdapter::new(
duplicate_pair_schema(),
stream,
),
);
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(
"native pair partition did not stage append fragments",
));
};
Ok(WrittenPairsPartition {
unit,
plan_id: plan.identity()?,
task,
fragments,
rows: rows.load(std::sync::atomic::Ordering::Relaxed),
})
}
pub(super) async fn commit(
view: &Table,
plan: &DuplicatePairsRefreshPlan,
units: Vec<WrittenPairsPartition>,
expected: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
let (_, definition, source) = open_planned(view, plan).await?;
let config = definition
.duplicate_pairs
.as_ref()
.expect("native")
.config()?;
if plan_duplicate_pairs(source.clone(), &config).await? != plan.tasks {
return Err(invalid(
"pair plan does not cover the pinned index partitions exactly",
));
}
let identity = plan.identity()?;
let mut ordered = vec![None; plan.snapshot.units as usize];
for written in units {
if written.plan_id != identity
|| plan.tasks.get(written.unit as usize) != Some(&written.task)
{
return Err(invalid("pair partition receipt belongs to another plan"));
}
let slot = ordered
.get_mut(written.unit as usize)
.ok_or_else(|| invalid("unknown pair partition receipt"))?;
if slot.is_some() {
return Err(invalid("duplicate pair partition receipt"));
}
*slot = Some(written);
}
let units = ordered
.into_iter()
.collect::<Option<Vec<_>>>()
.ok_or_else(|| invalid("missing pair partition receipt"))?;
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
}
rust/lancedb/src/materialized_view/grouped_units.rs
+42 -8
View File
@@ -53,6 +53,7 @@ const PROBE_ROWS: usize = 256;
/// How a grouped view can be refreshed in units.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct GroupedRefreshPlan {
/// Units `0..units`: unit `p` before the last holds index partition
/// `p`'s groups; the last holds the rows the index cannot place (null or
@@ -177,6 +178,7 @@ pub async fn plan_grouped_refresh(
/// accounted for has to be refusable. Travels to whoever commits, so it
/// serializes, but its fields are not a caller's to write.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WrittenUnit {
unit: u32,
plan: GroupedRefreshPlan,
@@ -491,10 +493,38 @@ pub async fn commit_grouped_refresh(
plan: &GroupedRefreshPlan,
units: Vec<WrittenUnit>,
expected_incarnation: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
let grouped = open_planned(view, plan).await?;
let units = complete_set(view.name(), plan, units)?;
let rows_written = units.iter().map(|unit| unit.rows).sum();
let new_fragments = units.into_iter().flat_map(|unit| unit.fragments).collect();
commit_fragments(
view,
plan,
&grouped.definition,
grouped.source.manifest.timestamp_nanos,
new_fragments,
rows_written,
expected_incarnation,
)
.await
}
/// The common all-or-nothing unit sink, shared by grouped and native views.
#[allow(clippy::too_many_arguments)]
pub(super) async fn commit_fragments(
view: &Table,
plan: &GroupedRefreshPlan,
definition: &MaterializedViewDefinition,
source_ts: u128,
new_fragments: Vec<Fragment>,
rows_written: u64,
expected_incarnation: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
let view_native = view.as_native().ok_or_else(|| Error::NotSupported {
message: "materialized views are supported only on local tables".into(),
})?;
view_native.dataset.ensure_mutable()?;
let lock = refresh_lock(view_native.dataset.get().await?.uri());
let _guard = lock.lock().await;
if let Some(expected) = expected_incarnation
@@ -508,10 +538,17 @@ pub async fn commit_grouped_refresh(
),
});
}
let grouped = open_planned(view, plan).await?;
let source_ts = grouped.source.manifest.timestamp_nanos;
let ds = Arc::new(grouped.view_ds);
view_native.dataset.reload().await?;
let ds = view_native.dataset.get().await?.as_ref().clone();
ensure_incarnation(&ds, Some(&plan.incarnation), view.name()).await?;
if super::read_definition(&ds.schema().metadata)?
!= Some(super::StoredDefinition::Query(definition.clone()))
{
return Err(Error::InvalidInput {
message: "view definition changed while its units were written".into(),
});
}
let ds = Arc::new(ds);
if ds.version().version != plan.view_version {
return Err(Error::Runtime {
message: format!(
@@ -523,11 +560,8 @@ pub async fn commit_grouped_refresh(
),
});
}
let units = complete_set(view.name(), plan, units)?;
let removed_fragment_ids: Vec<u64> = ds.get_fragments().iter().map(|f| f.id() as u64).collect();
// Fragment ids are assigned at commit; the units' fragments all carry 0.
let rows_written = units.iter().map(|unit| unit.rows).sum();
let new_fragments = units.into_iter().flat_map(|unit| unit.fragments).collect();
let transaction = Transaction::new(
plan.view_version,
Operation::Update {
@@ -566,7 +600,7 @@ pub async fn commit_grouped_refresh(
committed,
plan.source_version,
source_ts,
Some(&grouped.definition),
Some(definition),
expected_incarnation,
)
.await?;
rust/lancedb/src/materialized_view/partitioned.rs
+132
View File
@@ -0,0 +1,132 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Shared plan/write/commit surface for materialized views with independent units.
use serde::{Deserialize, Serialize};
use super::duplicate_pairs::{self, DuplicatePairsRefreshPlan, WrittenPairsPartition};
use super::{GroupedRefreshPlan, RefreshMaterializedViewResult, StoredDefinition, WrittenUnit};
use crate::table::Table;
use crate::{Error, Result};
/// Serializable work description consumed by a host's existing job workers.
/// The grouped variant preserves the previous saved-plan representation.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum PartitionedRefreshPlan {
DuplicatePairs(DuplicatePairsRefreshPlan),
Grouped(GroupedRefreshPlan),
}
impl PartitionedRefreshPlan {
/// The shared source snapshot and destination generation fence.
pub fn snapshot(&self) -> &GroupedRefreshPlan {
match self {
Self::DuplicatePairs(plan) => &plan.snapshot,
Self::Grouped(plan) => plan,
}
}
/// Number of independent units in this immutable plan.
pub fn units(&self) -> u32 {
self.snapshot().units
}
}
/// One completed unit's uncommitted fragments. Failed attempts produce no receipt.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum WrittenPartition {
DuplicatePairs(WrittenPairsPartition),
Grouped(WrittenUnit),
}
impl WrittenPartition {
/// Number of pair or grouped rows staged by this unit.
pub fn rows(&self) -> u64 {
match self {
Self::DuplicatePairs(unit) => unit.rows,
Self::Grouped(unit) => unit.rows(),
}
}
/// Uncommitted output files, retained only after task completion.
pub fn fragments(&self) -> &[lance_table::format::Fragment] {
match self {
Self::DuplicatePairs(unit) => &unit.fragments,
Self::Grouped(unit) => unit.fragments(),
}
}
}
/// Plan native pair or indexed grouped work; return `None` for ordinary views.
/// `pinned`, when supplied, must match a native view's declared source version.
pub async fn plan_partitioned_refresh(
view: &Table,
pinned: Option<u64>,
) -> Result<Option<PartitionedRefreshPlan>> {
if let Some(StoredDefinition::Query(definition)) =
super::read_definition(view.schema().await?.metadata())?
&& definition.duplicate_pairs.is_some()
{
return duplicate_pairs::plan_refresh(view, pinned)
.await
.map(|p| Some(PartitionedRefreshPlan::DuplicatePairs(p)));
}
super::plan_grouped_refresh(view, pinned)
.await
.map(|plan| plan.map(PartitionedRefreshPlan::Grouped))
}
/// Stream one unit into uncommitted fragments. Persist the returned receipt only
/// on success. A failed or interrupted call must be retried from the unit's start.
pub async fn write_refresh_partition(
view: &Table,
unit: u32,
plan: &PartitionedRefreshPlan,
) -> Result<WrittenPartition> {
match plan {
PartitionedRefreshPlan::DuplicatePairs(plan) => {
duplicate_pairs::write_unit(view, unit, plan)
.await
.map(WrittenPartition::DuplicatePairs)
}
PartitionedRefreshPlan::Grouped(plan) => super::write_grouped_unit(view, unit, plan)
.await
.map(WrittenPartition::Grouped),
}
}
/// Publish a complete set of unit receipts through the common MV commit sink.
/// `expected` fences the destination incarnation. No partial set is publishable.
pub async fn commit_partitioned_refresh(
view: &Table,
plan: &PartitionedRefreshPlan,
units: Vec<WrittenPartition>,
expected: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
let foreign = || Error::InvalidInput {
message: "refresh receipt does not match the plan kind".into(),
};
match plan {
PartitionedRefreshPlan::DuplicatePairs(plan) => {
let units = units
.into_iter()
.map(|unit| match unit {
WrittenPartition::DuplicatePairs(unit) => Ok(unit),
_ => Err(foreign()),
})
.collect::<Result<Vec<_>>>()?;
duplicate_pairs::commit(view, plan, units, expected).await
}
PartitionedRefreshPlan::Grouped(plan) => {
let units = units
.into_iter()
.map(|unit| match unit {
WrittenPartition::Grouped(unit) => Ok(unit),
_ => Err(foreign()),
})
.collect::<Result<Vec<_>>>()?;
super::commit_grouped_refresh(view, plan, units, expected).await
}
}
}
rust/lancedb/src/materialized_view/query.rs
+77 -1
View File
@@ -264,7 +264,29 @@ fn extract(query: &Query) -> Result<MaterializedViewDefinition> {
};
let mut from = select.from.iter();
let mut duplicate_pairs = None;
let (source_namespace, source_table) = match from.next().map(|f| &f.relation) {
Some(TableFactor::Table {
name,
args: Some(args),
alias: None,
with_hints,
version: None,
with_ordinality: false,
partitions,
json_path: None,
sample: None,
index_hints,
..
}) if name.to_string() == super::duplicate_pairs::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(config);
(namespace, table)
}
Some(TableFactor::Table { args: Some(_), .. }) => {
return Err(invalid(
"a view reads a table; a Function in FROM position follows it: \
@@ -416,6 +438,7 @@ fn extract(query: &Query) -> Result<MaterializedViewDefinition> {
};
let definition = MaterializedViewDefinition {
duplicate_pairs,
source_table,
source_namespace,
lateral,
@@ -424,6 +447,9 @@ fn extract(query: &Query) -> Result<MaterializedViewDefinition> {
group_by,
limit,
};
if definition.duplicate_pairs.is_some() {
super::duplicate_pairs::check_shape(&definition)?;
}
check_grouping(&definition)?;
Ok(definition)
}
@@ -483,7 +509,18 @@ pub fn render(definition: &MaterializedViewDefinition) -> String {
.chain(std::iter::once(&definition.source_table))
.map(|part| ident_sql(part))
.collect();
sql.push_str(&table.join("."));
if let Some(native) = &definition.duplicate_pairs {
sql.push_str(&format!(
"{}('{}', {}, '{}', {})",
super::duplicate_pairs::FUNCTION_NAME,
table.join(".").replace('\'', "''"),
native.dataset_version,
native.column.replace('\'', "''"),
native.distance_threshold
));
} else {
sql.push_str(&table.join("."));
}
if let Some(lateral) = &definition.lateral {
match &lateral.source {
LateralSource::Unnest { column } => {
@@ -511,6 +548,45 @@ pub fn render(definition: &MaterializedViewDefinition) -> String {
mod tests {
use super::*;
#[test]
fn native_pairs_are_a_table_source_with_literal_snapshot_arguments() {
for threshold in ["4", "4.0", "-1", "0.05"] {
let sql =
format!("SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', {threshold})");
let definition = parse(&sql).unwrap();
assert_eq!(definition.source_table, "images");
assert!(definition.lateral.is_none());
assert_eq!(
definition.duplicate_pairs.as_ref().unwrap().dataset_version,
3
);
assert_eq!(parse(&render(&definition)).unwrap(), definition);
let metadata = std::collections::HashMap::from([(
super::super::DEFINITION_META_KEY.to_string(),
definition.to_json().unwrap(),
)]);
assert_eq!(
super::super::read_definition(&metadata).unwrap(),
Some(super::super::StoredDefinition::Query(definition))
);
}
for sql in [
"SELECT row_id_a FROM vector_duplicate_pairs('images', 3, 'phash', 4)",
"SELECT * FROM vector_duplicate_pairs('images', 0, 'phash', 4)",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 1e100)",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', cutoff)",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4) LIMIT 1",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4) WHERE distance < 2",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4) AS p",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4), f(phash) AS p",
] {
assert!(
parse(sql).is_err(),
"accepted unsupported native view: {sql}"
);
}
}
#[test]
fn a_query_round_trips_through_its_canonical_form() {
for (sql, canonical) in [
rust/lancedb/src/materialized_view/refresh.rs
+20 -1
View File
@@ -117,6 +117,23 @@ pub(crate) async fn execute_refresh(
pinned: Option<u64>,
expected_incarnation: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
// Native table functions use the same partition writer and commit sink in
// 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()
{
let plan = super::plan_partitioned_refresh(view, pinned)
.await?
.ok_or_else(|| Error::Runtime {
message: "native pair view has no partition plan".into(),
})?;
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?);
}
return super::commit_partitioned_refresh(view, &plan, units, expected_incarnation).await;
}
let view_native = view.as_native().ok_or_else(|| Error::NotSupported {
message: "materialized views are supported only on local tables".into(),
})?;
@@ -728,7 +745,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 !dataset.manifest.uses_stable_row_ids() {
if definition.duplicate_pairs.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 \
@@ -3674,6 +3691,7 @@ pub(crate) mod tests {
let (_conn, _, view) = refreshed_doubled(vec![1, 2]).await;
let replacement = crate::materialized_view::MaterializedViewDefinition {
duplicate_pairs: None,
source_table: "src".into(),
source_namespace: Vec::new(),
projections: vec![
@@ -3715,6 +3733,7 @@ pub(crate) mod tests {
let (_conn, _, view) = refreshed_doubled(vec![1, 2]).await;
let narrower = crate::materialized_view::MaterializedViewDefinition {
duplicate_pairs: None,
source_table: "src".into(),
source_namespace: Vec::new(),
projections: vec![crate::materialized_view::ViewProjection {
rust/lancedb/src/remote/table.rs
+1
View File
@@ -2140,6 +2140,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
let definition = match response.query {
Some(query) => MaterializedViewDefinition::from_sql(&query)?,
None => MaterializedViewDefinition {
duplicate_pairs: None,
source_table: response.source_table,
source_namespace: response.source_namespace,
lateral: None,
rust/lancedb/src/table/datafusion/udtf/duplicate_pairs.rs
+1 -1
View File
@@ -60,7 +60,7 @@ pub struct DuplicatePairsConfig {
/// One worker input, never a row of source vectors. The complete task identity
/// also includes the dataset URI/version, column, threshold and operator version.
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct DuplicatePairTask {
pub segment_id: Uuid,
pub partition_id: usize,
rust/lancedb/tests/duplicate_pairs.rs
+61
View File
@@ -206,6 +206,67 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result
);
let retry = ctx.sql(&sql).await?.collect().await?;
assert_eq!(pairs(&retry), pairs(&actual));
// The same table function can be maintained through the MV unit
// lifecycle. Staging attempts do not make any rows visible.
use lancedb::materialized_view::{
MaterializedViewDefinition, commit_partitioned_refresh, plan_partitioned_refresh,
prepare_definition, write_refresh_partition,
};
let view_name = format!("pairs_{}", threshold.to_bits());
let view = prepare_definition(&table, MaterializedViewDefinition::from_sql(&sql)?)
.await?
.create(&view_name)
.await?;
assert_eq!(view.table().schema().await?.fields().len(), 3);
let plan = plan_partitioned_refresh(view.table(), Some(version))
.await?
.unwrap();
assert_eq!(plan.units() as usize, tasks.len());
let plan = serde_json::from_slice(&serde_json::to_vec(&plan)?)?;
let first = write_refresh_partition(view.table(), 0, &plan).await?;
assert_eq!(view.table().count_rows(None).await?, 0);
assert!(
commit_partitioned_refresh(
view.table(),
&plan,
vec![first.clone(), first.clone()],
view.incarnation()
)
.await
.is_err()
);
// Lost receipt: its files remain unreferenced; the retry writes
// another attempt, of which only one receipt is ever committed.
drop(write_refresh_partition(view.table(), 0, &plan).await?);
let mut units = vec![first];
for unit in 1..plan.units() {
units.push(write_refresh_partition(view.table(), unit, &plan).await?);
}
assert_eq!(view.table().count_rows(None).await?, 0);
let outcome =
commit_partitioned_refresh(view.table(), &plan, units.clone(), view.incarnation())
.await?;
assert_eq!(outcome.rows_written as usize, pairs(&expected).len());
assert!(
commit_partitioned_refresh(view.table(), &plan, units, view.incarnation())
.await
.is_err()
);
use lancedb::query::ExecutableQuery;
let written = view
.table()
.query()
.execute()
.await?
.try_collect::<Vec<_>>()
.await?;
assert_eq!(pairs(&written), pairs(&expected));
// Repeating a refresh replaces the view rather than appending.
view.refresh().execute().await?;
assert_eq!(
view.table().count_rows(None).await? as usize,
pairs(&expected).len()
);
assert!(
DuplicatePairsExec::try_new(
ds.clone(),