rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-13 15:22:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: cap the rows per fragment a materialized view refresh writes

Browse Source 
A Function column on a view is filled one fragment per worker, and a
refresh writes the view's rows from one stream under lance's default
cap of a million rows per file. Every view under that size is one
fragment, so its Function fill runs on one worker however many the
refresh was given, and a grouped view whose rows are each a whole
Function call runs them one after another.

A declaration can now cap the rows per fragment; the cap is stored in the
view's metadata and every refresh, rebuild or increment, writes under it.
A view of Function calls declares a small cap to spread them, one row per
fragment when each row is a group.
This commit is contained in:
Wyatt Alt
2026-09-19 00:03:00 +00:00
parent 99adb1e621
commit e298f68b40
3 changed files with 169 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/lancedb/src/materialized_view.rs
+76 -1
View File
@@ -179,6 +179,27 @@ pub(crate) struct ViewUnnest {
/// whose query calls a Function in `FROM` position.
pub const STAGING_META_KEY: &str = "mv.staging";
/// Schema metadata key holding the view's rows-per-fragment cap, when one was
/// declared: refresh writes the view in fragments of at most that many rows.
/// A Function column is filled one fragment per worker, so this is what lets
/// a fill fan out over a view a refresh would otherwise write as one fragment.
pub const ROWS_PER_FRAGMENT_META_KEY: &str = "mv.rows_per_fragment";
/// Read the rows-per-fragment cap off a view's schema metadata, if declared.
pub fn read_rows_per_fragment(metadata: &HashMap<String, String>) -> Result<Option<u64>> {
metadata
.get(ROWS_PER_FRAGMENT_META_KEY)
.map(|raw| {
raw.parse::<u64>()
.ok()
.filter(|rows| *rows >= 1)
.ok_or_else(|| Error::Runtime {
message: format!("unreadable materialized view rows-per-fragment '{raw}'"),
})
})
.transpose()
}
/// Where a Function in `FROM` position has its output staged: a hidden
/// table carrying every source column plus `column`, the Function's list
/// output. Refresh scans this table in place of the query's source.
@@ -345,6 +366,9 @@ pub struct CreateMaterializedViewRequest {
pub query: String,
/// Whether to skip the initial population job.
pub with_no_data: bool,
/// Cap on rows per fragment the refresh writes; see
/// [`ROWS_PER_FRAGMENT_META_KEY`].
pub rows_per_fragment: Option<u64>,
}
/// Prefix of the internal columns holding source columns a computed column
@@ -1071,6 +1095,44 @@ impl PreparedDeclaration {
&self.definition
}
/// Cap the fragments refresh writes at `rows` rows each. A Function
/// column is filled one fragment per worker, so a view of expensive
/// Function calls declares a small cap to spread them.
///
/// ```
/// # #![recursion_limit = "256"]
/// use lancedb::materialized_view::{MaterializedViewDefinition, prepare_definition};
///
/// # async fn declare(images: &lancedb::Table) -> Result<(), Box<dyn std::error::Error>> {
/// let definition = MaterializedViewDefinition::from_sql(
/// "SELECT ivf_partition(vec) AS bucket, array_agg(id) AS ids FROM images \
/// GROUP BY ivf_partition(vec)",
/// )?;
/// // One group per fragment: every bucket becomes its own unit of work.
/// let view = prepare_definition(images, definition)
/// .await?
/// .with_rows_per_fragment(1)?
/// .create("buckets")
/// .await?;
/// view.refresh().execute().await?;
/// # Ok(())
/// # }
/// ```
pub fn with_rows_per_fragment(mut self, rows: u64) -> Result<Self> {
if rows == 0 {
return Err(Error::InvalidInput {
message: "rows_per_fragment must be at least 1".to_string(),
});
}
let mut metadata = self.schema.metadata().clone();
metadata.insert(ROWS_PER_FRAGMENT_META_KEY.to_string(), rows.to_string());
self.schema = Arc::new(ArrowSchema::new_with_metadata(
self.schema.fields().clone(),
metadata,
));
Ok(self)
}
/// The schema the view will have: the declared columns in order, any
/// internal projections added by [`PreparedDeclaration::input_column`],
/// then [`SOURCE_ROW_ID_COLUMN`].
@@ -1659,6 +1721,7 @@ pub struct CreateMaterializedViewBuilder {
filter: Option<String>,
limit: Option<u64>,
with_no_data: bool,
rows_per_fragment: Option<u64>,
}
impl CreateMaterializedViewBuilder {
@@ -1673,6 +1736,7 @@ impl CreateMaterializedViewBuilder {
filter: None,
limit: None,
with_no_data: false,
rows_per_fragment: None,
}
}
@@ -1721,6 +1785,13 @@ impl CreateMaterializedViewBuilder {
self
}
/// Cap the fragments refresh writes at `rows` rows each; see
/// [`PreparedDeclaration::with_rows_per_fragment`].
pub fn rows_per_fragment(mut self, rows: u64) -> Self {
self.rows_per_fragment = Some(rows);
self
}
fn query(&self) -> String {
fn quote(name: &str) -> String {
format!("\"{}\"", name.replace('"', "\"\""))
@@ -1766,6 +1837,7 @@ impl CreateMaterializedViewBuilder {
namespace_path: self.namespace.clone(),
query: self.query(),
with_no_data: self.with_no_data,
rows_per_fragment: self.rows_per_fragment,
})
.await;
}
@@ -1798,13 +1870,16 @@ impl CreateMaterializedViewBuilder {
.namespace(self.source_namespace.clone())
.execute()
.await?;
let prepared = prepare_declaration(
let mut prepared = prepare_declaration(
&source,
(!self.projections.is_empty()).then_some(self.projections.as_slice()),
self.filter.as_deref(),
self.limit,
)
.await?;
if let Some(rows) = self.rows_per_fragment {
prepared = prepared.with_rows_per_fragment(rows)?;
}
let view = prepared.create_in(&self.namespace, &self.name).await?;
if !self.with_no_data {
view.refresh().execute().await?;
rust/lancedb/src/materialized_view/refresh.rs
+80 -9
View File
@@ -153,6 +153,7 @@ pub(crate) async fn execute_refresh(
};
let definition = &definition;
let staging = super::read_staging(&view_ds.schema().metadata)?;
let write = write_params(super::read_rows_per_fragment(&view_ds.schema().metadata)?);
ensure_no_mem_wal(&view_ds, "materialized view", view.name()).await?;
let source_ds = open_source(view, definition, staging.as_ref()).await?;
@@ -262,6 +263,7 @@ pub(crate) async fn execute_refresh(
&inputs,
unnest.as_ref(),
true,
&write,
expected_incarnation,
)
.await;
@@ -309,6 +311,7 @@ pub(crate) async fn execute_refresh(
&inputs,
None,
persist.is_some(),
&write,
expected_incarnation,
)
.await;
@@ -338,6 +341,7 @@ pub(crate) async fn execute_refresh(
unnest.as_ref(),
persist,
watermark,
&write,
expected_incarnation,
)
.await?;
@@ -355,6 +359,7 @@ pub(crate) async fn execute_refresh(
&inputs,
unnest.as_ref(),
persist.is_some(),
&write,
expected_incarnation,
)
.await
@@ -372,6 +377,7 @@ pub(crate) async fn execute_refresh(
&inputs,
unnest.as_ref(),
persist.is_some(),
&write,
expected_incarnation,
)
.await
@@ -379,6 +385,21 @@ pub(crate) async fn execute_refresh(
}
}
/// How refresh writes the view's rows: appended, in fragments capped at the
/// declared rows-per-fragment when the view has one.
fn write_params(rows_per_fragment: Option<u64>) -> WriteParams {
let mut params = WriteParams {
mode: WriteMode::Append,
..Default::default()
};
if let Some(rows) = rows_per_fragment {
let rows = usize::try_from(rows).unwrap_or(usize::MAX);
params.max_rows_per_file = rows;
params.max_rows_per_group = params.max_rows_per_group.min(rows);
}
params
}
/// The source fragments whose rows are new since the watermark, or `None`
/// where the view has to rebuild. Two tiers: the transaction walk is exact
/// where it applies; the fragment-signature check is the fallback for deltas
@@ -753,6 +774,7 @@ async fn incremental(
unnest: Option<&super::ViewUnnest>,
persist: Option<&MaterializedViewDefinition>,
watermark: Option<u64>,
write: &WriteParams,
expected_incarnation: Option<&str>,
) -> Result<Option<RefreshMaterializedViewResult>> {
let new_fragments = increment.appended;
@@ -954,10 +976,7 @@ async fn incremental(
let ds = Arc::new(view_ds.clone());
let write_txn = InsertBuilder::new(WriteDestination::Dataset(ds.clone()))
.with_params(&WriteParams {
mode: WriteMode::Append,
..Default::default()
})
.with_params(write)
.execute_uncommitted_stream(stream)
.await?;
let Operation::Append {
@@ -1001,6 +1020,7 @@ async fn rebuild(
inputs: &[String],
unnest: Option<&super::ViewUnnest>,
persist_definition: bool,
write: &WriteParams,
expected_incarnation: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
let rows_written = Arc::new(AtomicU64::new(0));
@@ -1027,7 +1047,8 @@ async fn rebuild(
// that raced in the way an overwrite (which adopts its stream's schema)
// durably would -- and it must land on the planned generation or abort.
let replaced =
replace_retaining_indices(view_ds.clone(), stream, keys, expected_incarnation).await?;
replace_retaining_indices(view_ds.clone(), stream, keys, write, expected_incarnation)
.await?;
let version = stamp_watermark(
view_native,
replaced,
@@ -1053,6 +1074,7 @@ async fn replace_retaining_indices(
view_ds: Dataset,
stream: SendableRecordBatchStream,
keys: Arc<StdMutex<KeyExistenceFilterBuilder>>,
write: &WriteParams,
expected_incarnation: Option<&str>,
) -> Result<Dataset> {
let ds = Arc::new(view_ds);
@@ -1063,10 +1085,7 @@ async fn replace_retaining_indices(
let removed_fragment_ids: Vec<u64> = ds.get_fragments().iter().map(|f| f.id() as u64).collect();
let write_txn = InsertBuilder::new(WriteDestination::Dataset(ds.clone()))
.with_params(&WriteParams {
mode: WriteMode::Append,
..Default::default()
})
.with_params(write)
.execute_uncommitted_stream(stream)
.await?;
let Operation::Append {
@@ -4787,4 +4806,56 @@ mod tests {
["0 10 0 9", "1 10 10 19"]
);
}
async fn fragment_sizes(table: &Table) -> Vec<usize> {
let dataset = table.as_native().unwrap().dataset.get().await.unwrap();
let mut sizes = Vec::new();
for fragment in dataset.get_fragments() {
sizes.push(fragment.count_rows(None).await.unwrap());
}
sizes
}
/// A declared rows-per-fragment caps every fragment refresh writes, on
/// the rebuild and on each increment; without one the default applies
/// and a small view is one fragment.
#[tokio::test]
async fn rows_per_fragment_caps_what_refresh_writes() {
let (conn, source) = db_with_source(vec![1, 2, 3, 4, 5]).await;
let capped = crate::materialized_view::prepare_definition(
&source,
MaterializedViewDefinition::from_sql("SELECT x FROM src").unwrap(),
)
.await
.unwrap()
.with_rows_per_fragment(2)
.unwrap()
.create("capped")
.await
.unwrap();
capped.refresh().execute().await.unwrap();
assert_eq!(fragment_sizes(capped.table()).await, [2, 2, 1]);
append(&source, vec![6, 7, 8]).await;
let result = capped.refresh().execute().await.unwrap();
assert_eq!(result.mode, RefreshMode::Incremental);
assert_eq!(fragment_sizes(capped.table()).await, [2, 2, 1, 2, 1]);
assert_eq!(
read(capped.table(), "x").await,
vec![1, 2, 3, 4, 5, 6, 7, 8]
);
let plain = declare(&source, "plain", "SELECT x FROM src")
.await
.unwrap();
plain.refresh().execute().await.unwrap();
assert_eq!(fragment_sizes(plain.table()).await, [8]);
assert!(
conn.create_materialized_view("zero", "src")
.rows_per_fragment(0)
.execute()
.await
.is_err()
);
}
}
rust/lancedb/src/remote/db.rs
+13 -4
View File
@@ -809,10 +809,18 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
let req = self
.client
.post(&format!("/v1/materialized_view/{identifier}/create"))
.json(&serde_json::json!({
"query": request.query,
"with_no_data": request.with_no_data,
}));
.json(&{
let mut body = serde_json::json!({
"query": request.query,
"with_no_data": request.with_no_data,
});
// Sent only when set: a server that predates the option
// rejects unknown fields.
if let Some(rows) = request.rows_per_fragment {
body["rows_per_fragment"] = serde_json::json!(rows);
}
body
});
let (request_id, response) = self.client.send(req).await?;
let response = self.client.check_response(&request_id, response).await?;
let status = response.status();
@@ -1756,6 +1764,7 @@ mod tests {
query: "SELECT age AS \"age\" FROM \"raw\".\"people\" WHERE age >= 18 LIMIT 10"
.into(),
with_no_data: false,
rows_per_fragment: None,
})
.await
.unwrap();