From 1e49757d799a0745c2ec142984c52e37841cb980 Mon Sep 17 00:00:00 2001 From: ytyky Date: Sun, 5 Jul 2026 02:13:03 +0000 Subject: [PATCH] perf(merge_insert): stream missing partial-schema columns instead of collecting them in the join A partial-schema merge_insert filled the dataset columns missing from the source by copying them out of the target side of the join (with_column(col("target.x"))). The target scan feeds the CollectLeft build side of the hash join, so every payload column of every target row was collected into memory before the source streamed through, and the plan executes with an unbounded memory pool. On wide tables (rows carrying large binary columns, datasets in the hundreds of GB) this reliably OOM-kills the process even when the merge only touches a small fraction of rows. The join now carries only the join keys, _rowid/_rowaddr, the source columns, and the action column. The write exec fetches the missing columns per output batch with a take by row address (TargetColumnFiller): updated rows read the matched target row's values, inserted rows fill NULL. Peak memory is bounded by the batch size, and only the rows actually rewritten are read, instead of the full column for every target row. --- rust/lance/src/dataset/write/merge_insert.rs | 257 +++++++++--- .../dataset/write/merge_insert/exec/write.rs | 394 ++++++++++++------ .../write/merge_insert/logical_plan.rs | 13 +- 3 files changed, 480 insertions(+), 184 deletions(-) diff --git a/rust/lance/src/dataset/write/merge_insert.rs b/rust/lance/src/dataset/write/merge_insert.rs index ca0e51ed0f7..a8c463246ba 100644 --- a/rust/lance/src/dataset/write/merge_insert.rs +++ b/rust/lance/src/dataset/write/merge_insert.rs @@ -1646,15 +1646,6 @@ impl MergeInsertJob { .collect::>(); let on_cols_refs = on_cols.iter().map(|s| s.as_str()).collect::>(); let source_df = session_ctx.read_one_shot(source)?; - // Capture the source field names *before* aliasing / joining so we - // can tell which dataset columns are missing from the source and - // need to be filled from the target side of the join below. - let source_field_names: std::collections::HashSet = source_df - .schema() - .fields() - .iter() - .map(|f| f.name().clone()) - .collect(); // Inject a sentinel literal column so we can reliably determine, after the join, // whether the source side contributed a row. This is NULL-safe: even when every // ON column is NULL the sentinel lets us distinguish a source-only row from a @@ -1666,7 +1657,15 @@ impl MergeInsertJob { let scan_aliased = scan.alias("target")?; let join_type = self.create_plan_join_type(); let dataset_schema: Schema = self.dataset.schema().into(); - let mut df = scan_aliased + // Partial-schema upsert note: dataset columns missing from the source + // are deliberately NOT copied out of the target side of this join. + // The join's target side is collected into memory (CollectLeft build + // side), so copying payload columns through it materializes every + // target row's payload at once and OOMs on wide tables. Instead the + // write exec fetches those columns per output batch, by row address + // (see `TargetColumnFiller`), which keeps memory bounded and only + // reads the rows that are actually written. + let df = scan_aliased .join( source_df_aliased, join_type, @@ -1679,26 +1678,6 @@ impl MergeInsertJob { merge_insert_action(&self.params, Some(&dataset_schema))?, )?; - // Partial-schema upsert: for every dataset column missing from the - // source, add a synthetic unqualified column that copies the target - // side's value for that column. For matched rows this carries the - // existing target value (preserving non-source columns on update); - // for unmatched source rows (inserts) the outer join leaves the - // target side NULL, so inserts get NULL for missing columns. The - // unqualified name matches the dataset field and becomes a normal - // data column from the write exec's perspective. - // - // We iterate the dataset schema in order so that the resulting - // physical plan is deterministic and easy to inspect in tests. - for field in dataset_schema.fields() { - if !source_field_names.contains(field.name()) { - df = df.with_column( - field.name(), - logical_expr::col(format!("target.\"{}\"", field.name())), - )?; - } - } - let (session_state, logical_plan) = df.into_parts(); let write_node = logical_plan::MergeInsertWriteNode::new( @@ -1829,8 +1808,8 @@ impl MergeInsertJob { ); // Partial-schema upsert: every source field must exist in the target - // and have a compatible data type. Missing target columns will be - // filled from the target side of the join in `create_plan`. + // and have a compatible data type. Missing target columns are + // fetched from the target by row address in the write exec. let is_subset_schema = !is_full_schema && lance_schema.fields.iter().all(|sf| { full_schema @@ -5231,7 +5210,7 @@ mod tests { new_data: RecordBatch, } - async fn setup(scalar_index: bool) -> Fixtures { + async fn setup(scalar_index: bool, stable_row_ids: bool) -> Fixtures { let data = lance_datagen::gen_batch() .with_seed(Seed::from(1)) .col("other", array::rand_utf8(4.into(), false)) @@ -5250,6 +5229,7 @@ mod tests { let write_params = WriteParams { max_rows_per_file: 256, max_rows_per_group: 32, // Non-standard group size to hit edge cases + enable_stable_row_ids: stable_row_ids, ..Default::default() }; let mut ds = Dataset::write(reader, "memory://", Some(write_params.clone())) @@ -5299,7 +5279,7 @@ mod tests { // it uniformly through the action column, so the previously- // rejected configuration now succeeds. This test asserts the // successful path to keep the negative-test history explicit. - let Fixtures { ds, new_data } = Box::pin(setup(false)).await; + let Fixtures { ds, new_data } = Box::pin(setup(false, false)).await; let rows_before = ds.count_rows(None).await.unwrap() as u64; @@ -5338,7 +5318,7 @@ mod tests { #[tokio::test] async fn test_errors_on_bad_schema() { - let Fixtures { ds, new_data } = Box::pin(setup(false)).await; + let Fixtures { ds, new_data } = Box::pin(setup(false, false)).await; // Schema with different names, which should be rejected. let bad_schema = Arc::new(Schema::new(vec![ @@ -5375,7 +5355,7 @@ mod tests { #[values(false, true)] scalar_index: bool, #[values(false, true)] insert: bool, ) { - let Fixtures { ds, new_data } = Box::pin(setup(scalar_index)).await; + let Fixtures { ds, new_data } = Box::pin(setup(scalar_index, false)).await; let reader = Box::new(RecordBatchIterator::new( [Ok(new_data.clone())], new_data.schema(), @@ -5385,6 +5365,31 @@ mod tests { .iter() .map(|f| f.metadata().clone()) .collect::>(); + + // Capture `other` per key before the merge so we can verify the + // exact value is preserved for updated rows (the write path must + // fetch the matched target row's value, not just any value). + let other_before: std::collections::HashMap = { + let data = ds.scan().try_into_batch().await.unwrap(); + let keys = data + .column_by_name("key") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + let others = data + .column_by_name("other") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + (0..data.num_rows()) + .map(|i| (keys.value(i).to_string(), others.value(i).to_string())) + .collect() + }; + let job = MergeInsertBuilder::try_new(ds.clone(), vec!["key".to_string()]) .unwrap() .when_matched(WhenMatched::UpdateAll) @@ -5517,12 +5522,10 @@ mod tests { .as_any() .downcast_ref::() .unwrap(); - let mut row_by_key: HashMap = HashMap::new(); + let mut row_by_key: HashMap)> = HashMap::new(); for i in 0..data.num_rows() { - row_by_key.insert( - key_col.value(i).to_string(), - (value_col.value(i), other_col.value(i).to_string()), - ); + let other = (!other_col.is_null(i)).then(|| other_col.value(i).to_string()); + row_by_key.insert(key_col.value(i).to_string(), (value_col.value(i), other)); } // Pull original column data for reference lookups. @@ -5540,15 +5543,17 @@ mod tests { .downcast_ref::() .unwrap(); // Every updated source row (270 of them) should be present - // with its new value and a preserved `other` string. + // with its new value and the exact `other` string it had + // before the merge. for i in 0..(new_data.num_rows() - 2) { let key = new_keys.value(i).to_string(); let (value, other) = row_by_key .get(&key) .unwrap_or_else(|| panic!("updated key {} missing from result", key)); assert_eq!(*value, new_values.value(i)); - assert!( - !other.is_empty(), + assert_eq!( + other.as_deref(), + other_before.get(&key).map(|s| s.as_str()), "updated row for key {} should retain its original `other` value", key ); @@ -5558,9 +5563,14 @@ mod tests { let key = new_keys.value(i).to_string(); let found = row_by_key.get(&key); if insert { - let (value, _) = + let (value, other) = found.unwrap_or_else(|| panic!("inserted key {} missing from result", key)); assert_eq!(*value, new_values.value(i)); + assert!( + other.is_none(), + "inserted row for key {} has no matched target row, so `other` must be NULL", + key + ); } else { assert!( found.is_none(), @@ -5571,6 +5581,142 @@ mod tests { } } + /// Partial-schema batches that contain no matched row at all: + /// an insert-only source must fill the missing column with NULL + /// without issuing a target take (every `_rowaddr` is NULL), and a + /// fully-unmatched source with inserts disabled must write nothing. + #[tokio::test] + async fn test_merge_insert_subcols_insert_only_and_unmatched() { + let Fixtures { ds, new_data } = Box::pin(setup(false, false)).await; + let rows_before = ds.count_rows(None).await.unwrap(); + + // The last two rows of new_data are keys that do not exist in + // the target dataset. + let unmatched = new_data.slice(new_data.num_rows() - 2, 2); + + // Inserts disabled: every source row resolves to a no-op, so the + // write stream produces only empty batches and commits nothing. + let reader = Box::new(RecordBatchIterator::new( + [Ok(unmatched.clone())], + unmatched.schema(), + )); + let job = MergeInsertBuilder::try_new(ds.clone(), vec!["key".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::DoNothing) + .try_build() + .unwrap(); + let (ds, stats) = job.execute_reader(reader).await.unwrap(); + assert_eq!(stats.num_updated_rows, 0); + assert_eq!(stats.num_inserted_rows, 0); + assert_eq!(stats.num_deleted_rows, 0); + assert_eq!(ds.count_rows(None).await.unwrap(), rows_before); + + // Insert-only merge: the batch reaching the write exec has no + // matched row, so the missing `other` column is NULL-filled + // without reading the target. + let reader = Box::new(RecordBatchIterator::new( + [Ok(unmatched.clone())], + unmatched.schema(), + )); + let job = MergeInsertBuilder::try_new(ds.clone(), vec!["key".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap(); + let (ds, stats) = job.execute_reader(reader).await.unwrap(); + assert_eq!(stats.num_updated_rows, 0); + assert_eq!(stats.num_inserted_rows, 2); + assert_eq!(ds.count_rows(None).await.unwrap(), rows_before + 2); + + let inserted_keys = unmatched + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let filter = format!( + "key IN ('{}', '{}')", + inserted_keys.value(0), + inserted_keys.value(1) + ); + let data = ds + .scan() + .filter(&filter) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(data.num_rows(), 2); + let other_col = data.column_by_name("other").unwrap(); + assert_eq!( + other_col.null_count(), + 2, + "inserted rows must have NULL for the column missing from the source" + ); + } + + /// Same partial-schema upsert as `test_merge_insert_subcols`, but on + /// a dataset with stable row ids. Stable row ids route the write + /// through the ordered update/insert split stream, so this covers + /// the target-column fill on that path as well. + #[tokio::test] + async fn test_merge_insert_subcols_stable_row_ids() { + let Fixtures { ds, new_data } = Box::pin(setup(false, true)).await; + let reader = Box::new(RecordBatchIterator::new( + [Ok(new_data.clone())], + new_data.schema(), + )); + let job = MergeInsertBuilder::try_new(ds.clone(), vec!["key".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap(); + let (ds, stats) = job.execute_reader(reader).await.unwrap(); + + assert_eq!(stats.num_updated_rows, (new_data.num_rows() - 2) as u64); + assert_eq!(stats.num_inserted_rows, 2); + assert_eq!(stats.num_deleted_rows, 0); + + let data = ds.scan().try_into_batch().await.unwrap(); + assert_eq!(data.num_rows(), 1024 + 2); + + // Updated rows keep a non-null `other` (fetched from the matched + // target row); the 2 inserted rows get NULL. + let key_col = data + .column_by_name("key") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let other_col = data.column_by_name("other").unwrap(); + let new_keys = new_data + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let updated: std::collections::HashSet<&str> = (0..new_data.num_rows() - 2) + .map(|i| new_keys.value(i)) + .collect(); + let inserted: std::collections::HashSet<&str> = (new_data.num_rows() - 2 + ..new_data.num_rows()) + .map(|i| new_keys.value(i)) + .collect(); + for i in 0..data.num_rows() { + let key = key_col.value(i); + if inserted.contains(key) { + assert!( + other_col.is_null(i), + "inserted key {} must have NULL other", + key + ); + } else if updated.contains(key) { + assert!(!other_col.is_null(i), "updated key {} must keep other", key); + } + } + } + /// Verifies that `explain_plan` succeeds for a partial-schema upsert /// and emits a plan that uses the v2 `FullSchemaMergeInsertExec` /// path. This is the explicit acceptance criterion for #6442: the @@ -5578,7 +5724,7 @@ mod tests { /// full-schema upserts instead of falling back to v1. #[tokio::test] async fn test_merge_insert_subcols_v2_explain_plan() { - let Fixtures { ds, new_data } = Box::pin(setup(false)).await; + let Fixtures { ds, new_data } = Box::pin(setup(false, false)).await; let job = MergeInsertBuilder::try_new(ds.clone(), vec!["key".to_string()]) .unwrap() @@ -5608,18 +5754,15 @@ mod tests { "expected HashJoinExec in plan, got: {}", plan ); - // Evidence that the partial-schema fix is active: the target - // side of the join reads the `other` column (which is missing - // from the source) and an explicit projection carries it - // through to the write exec alongside source columns. - assert!( - plan.contains("LanceRead") && plan.contains("projection=[other"), - "target-side scan should include the filled `other` column: {}", - plan - ); + // The `other` column (missing from the source) must NOT be + // routed through the join: the target scan feeds the collected + // build side of the hash join, so pulling payload columns + // through it materializes the whole target in memory. The write + // exec fetches `other` per batch by row address instead, so the + // target-side scan stays narrow (join key + row id + row addr). assert!( - plan.contains("other@0 as other"), - "expected post-join projection to carry `other` from the target side: {}", + plan.contains("LanceRead") && !plan.contains("other"), + "target-side scan must not read the `other` payload column: {}", plan ); } diff --git a/rust/lance/src/dataset/write/merge_insert/exec/write.rs b/rust/lance/src/dataset/write/merge_insert/exec/write.rs index d5b51b3d97f..f436abf26ca 100644 --- a/rust/lance/src/dataset/write/merge_insert/exec/write.rs +++ b/rust/lance/src/dataset/write/merge_insert/exec/write.rs @@ -1,10 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::sync::{Arc, Mutex}; -use arrow_array::{Array, RecordBatch, UInt8Array, UInt64Array}; +use arrow_array::{Array, ArrayRef, RecordBatch, UInt8Array, UInt32Array, UInt64Array}; use arrow_schema::Schema; use arrow_select; use datafusion::common::{DataFusionError, Result as DFResult}; @@ -20,6 +20,7 @@ use datafusion::{ use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use futures::{StreamExt, stream}; use lance_core::{Error, ROW_ADDR, ROW_ID}; +use lance_datafusion::projection::ProjectionPlan; use lance_table::format::RowIdMeta; use roaring::RoaringTreemap; @@ -32,6 +33,7 @@ use crate::dataset::write::merge_insert::{ MERGE_SOURCE_SENTINEL, SourceDedupeBehavior, create_duplicate_row_error, format_key_values_on_columns, resolve_target_bases, }; +use crate::dataset::{ProjectionRequest, TakeBuilder}; use crate::{ Dataset, dataset::{ @@ -193,6 +195,146 @@ impl MergeState { } } +/// Where each output column of the write stream comes from, in dataset +/// schema order. +#[derive(Debug, Clone, Copy)] +enum WriteColumnSource { + /// The column is present in the exec's input (provided by the source). + Input(usize), + /// The column is missing from the source (partial-schema upsert) and is + /// fetched from the target by row address; index into + /// [`TargetColumnFiller::fields`]. + FillFromTarget(usize), +} + +/// Fetches target-side values for dataset columns that a partial-schema +/// source does not provide. +/// +/// The values are fetched one batch at a time, by the `_rowaddr` of the +/// matched target rows, so peak memory is bounded by the batch size. This +/// deliberately replaces the earlier approach of copying these columns +/// through the join in the logical plan: there the target scan fed the +/// build side of a `CollectLeft` hash join, so every payload column of +/// every target row was collected into memory at once, which caused OOM on +/// wide tables. +#[derive(Debug)] +struct TargetColumnFiller { + dataset: Arc, + projection: Arc, + /// The filled fields, matching the projection. Nullable because + /// inserted rows (no matched target row) receive NULL. + fields: Vec, +} + +impl TargetColumnFiller { + fn try_new(dataset: Arc, fields: Vec) -> DFResult { + let names = fields.iter().map(|f| f.name().as_str()).collect::>(); + let projection = ProjectionRequest::from_columns(names, dataset.schema()) + .into_projection_plan(dataset.clone()) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + Ok(Self { + dataset, + projection: Arc::new(projection), + fields, + }) + } + + /// Returns one array per filled field, aligned with `row_addrs`: rows + /// with a valid row address (updates) get the target's current value, + /// rows with a NULL row address (inserts) get NULL. + async fn fetch(&self, row_addrs: &UInt64Array) -> DFResult> { + let valid_addrs: Vec = (0..row_addrs.len()) + .filter(|&i| row_addrs.is_valid(i)) + .map(|i| row_addrs.value(i)) + .collect(); + + if valid_addrs.is_empty() { + return Ok(self + .fields + .iter() + .map(|field| arrow_array::new_null_array(field.data_type(), row_addrs.len())) + .collect()); + } + + let mut sorted_addrs = valid_addrs; + sorted_addrs.sort_unstable(); + sorted_addrs.dedup(); + + let taken = TakeBuilder::try_new_from_addresses( + self.dataset.clone(), + sorted_addrs, + self.projection.clone(), + ) + .map_err(|e| DataFusionError::External(Box::new(e)))? + .with_row_address(true) + .execute() + .await + .map_err(|e| DataFusionError::External(Box::new(e)))?; + + // Map each returned row address back to its position in the taken + // batch so we do not depend on the take preserving request order. + let taken_addrs = taken + .column_by_name(ROW_ADDR) + .ok_or_else(|| { + DataFusionError::Internal("take result is missing the _rowaddr column".to_string()) + })? + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Internal("Expected UInt64Array for _rowaddr column".to_string()) + })?; + let addr_to_taken_idx: HashMap = taken_addrs + .values() + .iter() + .enumerate() + .map(|(idx, addr)| (*addr, idx as u32)) + .collect(); + + let indices = + (0..row_addrs.len()) + .map(|i| { + if row_addrs.is_valid(i) { + let addr = row_addrs.value(i); + addr_to_taken_idx.get(&addr).copied().map(Some).ok_or_else(|| { + DataFusionError::Internal(format!( + "Row address {} matched by the merge join was not returned by take", + addr + )) + }) + } else { + Ok(None) + } + }) + .collect::>()?; + + self.fields + .iter() + .map(|field| { + let column = taken.column_by_name(field.name()).ok_or_else(|| { + DataFusionError::Internal(format!( + "take result is missing the {:?} column", + field.name() + )) + })?; + arrow_select::take::take(column, &indices, None).map_err(DataFusionError::from) + }) + .collect() + } +} + +/// Precomputed layout of the write stream: where the control columns live +/// in the input, how to assemble output columns in dataset schema order, +/// and (for partial-schema sources) how to fill the missing columns. +#[derive(Debug)] +struct WriteStreamContext { + rowaddr_idx: usize, + rowid_idx: usize, + action_idx: usize, + column_sources: Vec, + filler: Option, + output_schema: Arc, +} + /// Inserts new rows and updates existing rows in the target table. /// /// This does the actual write. @@ -328,48 +470,58 @@ impl FullSchemaMergeInsertExec { input_stream: SendableRecordBatchStream, merge_state: Arc>, ) -> DFResult { - let (_, rowaddr_idx, rowid_idx, action_idx, data_column_indices, output_schema) = - self.prepare_stream_schema(input_stream.schema())?; - - let output_schema_clone = output_schema.clone(); - let stream = input_stream.map(move |batch_result| -> DFResult { - let batch = batch_result?; - let (row_addr_array, row_id_array, action_array) = - Self::extract_control_arrays(&batch, rowaddr_idx, rowid_idx, action_idx)?; - - // Process each row using the shared state - let mut keep_rows: Vec = Vec::with_capacity(batch.num_rows()); - - let mut merge_state = merge_state.lock().map_err(|e| { - datafusion::error::DataFusionError::Internal(format!( - "Failed to lock merge state: {}", - e - )) - })?; - - for row_idx in 0..batch.num_rows() { - let action_code = action_array.value(row_idx); - let action = Action::try_from(action_code).map_err(|e| { - datafusion::error::DataFusionError::Internal(format!( - "Invalid action code {}: {}", - action_code, e - )) - })?; + let ctx = self.prepare_stream_schema(input_stream.schema())?; + + let output_schema = ctx.output_schema.clone(); + let stream = input_stream.then(move |batch_result| { + let ctx = ctx.clone(); + let merge_state = merge_state.clone(); + async move { + let batch = batch_result?; + let (row_addr_array, row_id_array, action_array) = Self::extract_control_arrays( + &batch, + ctx.rowaddr_idx, + ctx.rowid_idx, + ctx.action_idx, + )?; + + // Process each row using the shared state + let mut keep_rows: Vec = Vec::with_capacity(batch.num_rows()); - if merge_state - .process_row_action(action, row_idx, row_addr_array, row_id_array, &batch)? - .is_some() { - keep_rows.push(row_idx as u32); + let mut merge_state = merge_state.lock().map_err(|e| { + datafusion::error::DataFusionError::Internal(format!( + "Failed to lock merge state: {}", + e + )) + })?; + + for row_idx in 0..batch.num_rows() { + let action_code = action_array.value(row_idx); + let action = Action::try_from(action_code).map_err(|e| { + datafusion::error::DataFusionError::Internal(format!( + "Invalid action code {}: {}", + action_code, e + )) + })?; + + if merge_state + .process_row_action( + action, + row_idx, + row_addr_array, + row_id_array, + &batch, + )? + .is_some() + { + keep_rows.push(row_idx as u32); + } + } } - } - Self::create_filtered_batch( - &batch, - keep_rows, - &data_column_indices, - output_schema_clone.clone(), - ) + Self::create_write_batch(&batch, keep_rows, &ctx).await + } }); Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -403,18 +555,10 @@ impl FullSchemaMergeInsertExec { } /// Common schema preparation logic - #[allow(clippy::type_complexity)] fn prepare_stream_schema( &self, input_schema: arrow_schema::SchemaRef, - ) -> DFResult<( - arrow_schema::SchemaRef, - usize, - usize, - usize, - Vec, - Arc, - )> { + ) -> DFResult> { // Find column indices let (rowaddr_idx, _) = input_schema.column_with_name(ROW_ADDR).ok_or_else(|| { datafusion::error::DataFusionError::Internal( @@ -437,17 +581,15 @@ impl FullSchemaMergeInsertExec { )) })?; - // Emit data columns in dataset-schema order, keyed by name. This - // matters for partial-schema upserts, where `create_plan` adds - // synthetic unqualified columns (filled from the target side) at - // the end of the logical projection. Without name-based reordering - // they would land at the end of the write stream and mismatch the - // intended writer schema (which is `dataset.schema()`). Using name - // lookup is also a strictly-safer choice for the full-schema path: - // it turns an implicit positional assumption into an explicit - // name-based invariant. - let mut name_to_idx: std::collections::HashMap<&str, usize> = - std::collections::HashMap::with_capacity(input_schema.fields().len()); + // Emit data columns in dataset-schema order, keyed by name. Columns + // present in the input come from the source; dataset columns that a + // partial-schema source omits are fetched from the target per batch + // (see [`TargetColumnFiller`]) rather than routed through the join. + // Name-based lookup is also a strictly-safer choice for the + // full-schema path: it turns an implicit positional assumption into + // an explicit name-based invariant. + let mut name_to_idx: HashMap<&str, usize> = + HashMap::with_capacity(input_schema.fields().len()); for (idx, field) in input_schema.fields().iter().enumerate() { let name = field.name(); // Skip special columns: _rowaddr, _rowid, __action, and the @@ -467,39 +609,55 @@ impl FullSchemaMergeInsertExec { let dataset_arrow_schema: arrow_schema::Schema = self.dataset.schema().into(); let dataset_fields = dataset_arrow_schema.fields(); - let mut data_column_indices: Vec = Vec::with_capacity(dataset_fields.len()); + let mut column_sources: Vec = Vec::with_capacity(dataset_fields.len()); + let mut fill_fields: Vec = Vec::new(); let mut output_fields: Vec> = Vec::with_capacity(dataset_fields.len()); + let mut num_input_columns = 0; for dataset_field in dataset_fields { - let idx = *name_to_idx - .get(dataset_field.name().as_str()) - .ok_or_else(|| { - datafusion::error::DataFusionError::Internal(format!( - "Dataset field {:?} missing from merge insert input schema \ - — this indicates a logical plan pruning bug", - dataset_field.name() - )) - })?; - data_column_indices.push(idx); - output_fields.push(Arc::new(input_schema.field(idx).clone())); + match name_to_idx.get(dataset_field.name().as_str()) { + Some(&idx) => { + column_sources.push(WriteColumnSource::Input(idx)); + output_fields.push(Arc::new(input_schema.field(idx).clone())); + num_input_columns += 1; + } + None => { + // Nullable regardless of the dataset field: inserted rows + // have no matched target row and receive NULL (the writer + // schema is what the commit validates against). + let field = Arc::new(dataset_field.as_ref().clone().with_nullable(true)); + column_sources.push(WriteColumnSource::FillFromTarget(fill_fields.len())); + fill_fields.push(field.clone()); + output_fields.push(field); + } + } } - if data_column_indices.is_empty() { + if num_input_columns == 0 { return Err(datafusion::error::DataFusionError::Internal( "No data columns found in merge insert input".to_string(), )); } + let filler = if fill_fields.is_empty() { + None + } else { + Some(TargetColumnFiller::try_new( + self.dataset.clone(), + fill_fields, + )?) + }; + let output_schema = Arc::new(Schema::new(output_fields)); - Ok(( - input_schema, + Ok(Arc::new(WriteStreamContext { rowaddr_idx, rowid_idx, action_idx, - data_column_indices, + column_sources, + filler, output_schema, - )) + })) } /// Extract control arrays from batch @@ -544,21 +702,24 @@ impl FullSchemaMergeInsertExec { Ok((row_addr_array, row_id_array, action_array)) } - /// Create filtered batch from selected rows - fn create_filtered_batch( + /// Create the batch to write from the selected rows, assembling columns + /// in dataset schema order. Columns a partial-schema source does not + /// provide are fetched from the target for this batch only, keeping + /// memory bounded by the batch size. + async fn create_write_batch( batch: &RecordBatch, keep_rows: Vec, - data_column_indices: &[usize], - output_schema: Arc, + ctx: &WriteStreamContext, ) -> DFResult { // If no rows to keep, return empty batch if keep_rows.is_empty() { - let empty_columns: Vec<_> = output_schema + let empty_columns: Vec<_> = ctx + .output_schema .fields() .iter() .map(|field| arrow_array::new_empty_array(field.data_type())) .collect(); - return RecordBatch::try_new(output_schema, empty_columns) + return RecordBatch::try_new(ctx.output_schema.clone(), empty_columns) .map_err(datafusion::error::DataFusionError::from); } @@ -568,13 +729,31 @@ impl FullSchemaMergeInsertExec { // Take only the rows we want to keep let filtered_batch = arrow_select::take::take_record_batch(batch, &indices)?; - // Project only the data columns - let output_columns: Vec<_> = data_column_indices + let fill_arrays = if let Some(filler) = &ctx.filler { + let row_addrs = filtered_batch + .column(ctx.rowaddr_idx) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + datafusion::error::DataFusionError::Internal( + "Expected UInt64Array for _rowaddr column".to_string(), + ) + })?; + filler.fetch(row_addrs).await? + } else { + Vec::new() + }; + + let output_columns: Vec<_> = ctx + .column_sources .iter() - .map(|&idx| filtered_batch.column(idx).clone()) + .map(|source| match source { + WriteColumnSource::Input(idx) => filtered_batch.column(*idx).clone(), + WriteColumnSource::FillFromTarget(idx) => fill_arrays[*idx].clone(), + }) .collect(); - RecordBatch::try_new(output_schema, output_columns) + RecordBatch::try_new(ctx.output_schema.clone(), output_columns) .map_err(datafusion::error::DataFusionError::from) } @@ -600,13 +779,12 @@ impl FullSchemaMergeInsertExec { input_stream: SendableRecordBatchStream, merge_state: Arc>, ) -> DFResult<(SendableRecordBatchStream, SendableRecordBatchStream)> { - let (_, rowaddr_idx, rowid_idx, action_idx, data_column_indices, output_schema) = - self.prepare_stream_schema(input_stream.schema())?; + let ctx = self.prepare_stream_schema(input_stream.schema())?; + let output_schema = ctx.output_schema.clone(); let (update_tx, update_rx) = tokio::sync::mpsc::unbounded_channel(); let (insert_tx, insert_rx) = tokio::sync::mpsc::unbounded_channel(); - let output_schema_clone = output_schema.clone(); let merge_state_clone = merge_state; tokio::spawn(async move { @@ -615,15 +793,9 @@ impl FullSchemaMergeInsertExec { while let Some(batch_result) = input_stream.next().await { match batch_result { Ok(batch) => { - match Self::process_and_split_batch( - &batch, - rowaddr_idx, - rowid_idx, - action_idx, - &data_column_indices, - output_schema_clone.clone(), - merge_state_clone.clone(), - ) { + match Self::process_and_split_batch(&batch, &ctx, merge_state_clone.clone()) + .await + { Ok((update_batch_opt, insert_batch_opt)) => { if let Some(update_batch) = update_batch_opt && update_tx.send(Ok(update_batch)).is_err() @@ -663,17 +835,13 @@ impl FullSchemaMergeInsertExec { Ok((update_stream, insert_stream)) } - fn process_and_split_batch( + async fn process_and_split_batch( batch: &RecordBatch, - rowaddr_idx: usize, - rowid_idx: usize, - action_idx: usize, - data_column_indices: &[usize], - output_schema: Arc, + ctx: &WriteStreamContext, merge_state: Arc>, ) -> DFResult<(Option, Option)> { let (row_addr_array, row_id_array, action_array) = - Self::extract_control_arrays(batch, rowaddr_idx, rowid_idx, action_idx)?; + Self::extract_control_arrays(batch, ctx.rowaddr_idx, ctx.rowid_idx, ctx.action_idx)?; let mut update_indices: Vec = Vec::new(); let mut insert_indices: Vec = Vec::new(); @@ -709,23 +877,13 @@ impl FullSchemaMergeInsertExec { } let update_batch = if !update_indices.is_empty() { - Some(Self::create_filtered_batch( - batch, - update_indices, - data_column_indices, - output_schema.clone(), - )?) + Some(Self::create_write_batch(batch, update_indices, ctx).await?) } else { None }; let insert_batch = if !insert_indices.is_empty() { - Some(Self::create_filtered_batch( - batch, - insert_indices, - data_column_indices, - output_schema, - )?) + Some(Self::create_write_batch(batch, insert_indices, ctx).await?) } else { None }; diff --git a/rust/lance/src/dataset/write/merge_insert/logical_plan.rs b/rust/lance/src/dataset/write/merge_insert/logical_plan.rs index 7f67972ff7e..d7ac803b54c 100644 --- a/rust/lance/src/dataset/write/merge_insert/logical_plan.rs +++ b/rust/lance/src/dataset/write/merge_insert/logical_plan.rs @@ -187,15 +187,10 @@ impl UserDefinedLogicalNodeCore for MergeInsertWriteNode { // Include unqualified columns like "__action" - tells us what operation to perform None if field.name() == MERGE_ACTION_COLUMN => true, - // Partial-schema upsert: the `create_plan` builder adds - // unqualified columns (named after dataset fields) for every - // column missing from the source, filled from the target - // side of the join. Those columns carry the values that - // should be written for non-source columns, so they must - // flow through to the write exec alongside `source.*`. - None if self.dataset.schema().field(field.name()).is_some() => true, - - // Skip other target columns (target.value, target.key, target._rowid) - not needed for write + // Skip other target columns (target.value, target.key, target._rowid) - not needed + // for write. In particular, dataset columns missing from a partial-schema source + // must NOT be pulled through the join: the write exec fetches them by row address + // instead, so that the join's collected build side stays narrow. _ => false, };