diff --git a/rust/lancedb/src/materialized_view.rs b/rust/lancedb/src/materialized_view.rs index c4456dbe1..8530b7a37 100644 --- a/rust/lancedb/src/materialized_view.rs +++ b/rust/lancedb/src/materialized_view.rs @@ -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) -> Result> { + metadata + .get(ROWS_PER_FRAGMENT_META_KEY) + .map(|raw| { + raw.parse::() + .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, } /// 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> { + /// 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 { + 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, limit: Option, with_no_data: bool, + rows_per_fragment: Option, } 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?; diff --git a/rust/lancedb/src/materialized_view/refresh.rs b/rust/lancedb/src/materialized_view/refresh.rs index 6f19c3c91..e6e1cfa91 100644 --- a/rust/lancedb/src/materialized_view/refresh.rs +++ b/rust/lancedb/src/materialized_view/refresh.rs @@ -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) -> 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, + write: &WriteParams, expected_incarnation: Option<&str>, ) -> Result> { 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 { 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>, + write: &WriteParams, expected_incarnation: Option<&str>, ) -> Result { let ds = Arc::new(view_ds); @@ -1063,10 +1085,7 @@ async fn replace_retaining_indices( let removed_fragment_ids: Vec = 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 { + 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() + ); + } } diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 7061ac979..2481f164e 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -809,10 +809,18 @@ impl Database for RemoteDatabase { 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();