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, };