From cc6b4aa4b47607f92c5a49893342a44a366340cc Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Sat, 3 Oct 2026 05:53:00 -0500 Subject: [PATCH] fix(remote): preserve schema for empty merge insert sources (#4425) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remote `merge_insert` buffered its reader into a `Vec`, losing the schema when the reader yielded no batches. Scanning that vector failed before the request could reach the server. Preserve the reader's schema in an empty record batch when buffering produces no batches. The source remains replayable for retries, and empty-source merges still reach the server so `when_not_matched_by_source_delete` can apply deletes. The regression tests decode the outgoing Arrow IPC and verify the source schema, zero rows, merge options, and result statistics on an initial request and a retry after HTTP 409. They cover readers with no batches and readers with an empty batch, in insert and delete modes, including an ID-only delete source. Validation: - Before the fix, both no-batch cases reproduced `Cannot scan an empty Vec`; the empty-batch cases passed. - `cargo test --quiet --profile ci --features remote -p lancedb --lib remote::table::tests::test_merge_insert` — 9 passed. - `cargo fmt --all` and `cargo fmt --all -- --check`. - `cargo clippy --quiet --profile ci --features remote -p lancedb --tests --examples`. Fixes #4380 Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- rust/lancedb/src/remote/table.rs | 106 ++++++++++++++++++++++++++++++- 1 file changed, 105 insertions(+), 1 deletion(-) diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 16c546147..d43bab9db 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3101,7 +3101,13 @@ impl BaseTable for RemoteTable { // loop can re-execute the plan (and re-stream the body) on each retry. // This mirrors the old `send_streaming(with_retry=true)` path, which // likewise buffered the reader to support retries. - let batches = new_data.collect::, _>>()?; + let schema = RecordBatchReader::schema(new_data.as_ref()); + let mut batches = new_data.collect::, _>>()?; + // An empty reader still carries a schema. Keep it in an empty batch so + // the buffered source remains scannable and can be replayed on retries. + if batches.is_empty() { + batches.push(RecordBatch::new_empty(schema)); + } let source: Box = Box::new(batches); let rescannable = source.rescannable(); let input: Arc = @@ -4801,6 +4807,104 @@ mod tests { } } + #[rstest] + #[case::no_batches_insert(false, false)] + #[case::empty_batch_insert(true, false)] + #[case::no_batches_delete(false, true)] + #[case::empty_batch_delete(true, true)] + #[tokio::test] + async fn test_merge_insert_empty_source( + #[case] has_batch: bool, + #[case] delete_unmatched: bool, + ) { + let mut fields = vec![Field::new("id", DataType::Int64, false)]; + if !delete_unmatched { + fields.extend([ + Field::new("k", DataType::Int64, false), + Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + 2, + ), + true, + ), + Field::new("s", DataType::Utf8, true), + ]); + } + let schema = Arc::new(Schema::new_with_metadata( + fields, + HashMap::from([("source".to_string(), "empty".to_string())]), + )); + let batches = if has_batch { + vec![Ok(RecordBatch::new_empty(schema.clone()))] + } else { + vec![] + }; + let data: Box = + Box::new(RecordBatchIterator::new(batches, schema.clone())); + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_ref = attempts.clone(); + let num_deleted_rows = if delete_unmatched { 3 } else { 0 }; + + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/"); + assert_eq!(request.headers()[CONTENT_TYPE], ARROW_STREAM_CONTENT_TYPE); + let params = request.url().query_pairs().collect::>(); + assert_eq!(params["on"], "id"); + assert_eq!( + params["when_not_matched_insert_all"], + (!delete_unmatched).to_string() + ); + assert_eq!( + params["when_not_matched_by_source_delete"], + delete_unmatched.to_string() + ); + + let body = request.body().unwrap().as_bytes().unwrap(); + let reader = StreamReader::try_new(Cursor::new(body), None).unwrap(); + assert_eq!(reader.schema(), schema); + for batch in reader { + assert_eq!(batch.unwrap().num_rows(), 0); + } + + // The empty source must retain its schema when replayed after a conflict. + if attempts_ref.fetch_add(1, Ordering::SeqCst) == 0 { + http::Response::builder() + .status(409) + .body(String::new()) + .unwrap() + } else { + http::Response::builder() + .status(200) + .body( + json!({ + "version": 43, + "num_deleted_rows": num_deleted_rows, + "num_inserted_rows": 0, + "num_updated_rows": 0, + }) + .to_string(), + ) + .unwrap() + } + }); + + let mut merge = table.merge_insert(&["id"]); + if delete_unmatched { + merge.when_not_matched_by_source_delete(None); + } else { + merge.when_not_matched_insert_all(); + } + let result = merge.execute(data).await.unwrap(); + assert_eq!(attempts.load(Ordering::SeqCst), 2); + assert_eq!(result.version, 43); + assert_eq!(result.num_deleted_rows, num_deleted_rows); + assert_eq!(result.num_inserted_rows, 0); + assert_eq!(result.num_updated_rows, 0); + } + #[tokio::test] async fn test_merge_insert_composite_key() { let batch = RecordBatch::try_new(