fix(remote): preserve schema for empty merge insert sources (#4425)

Remote `merge_insert` buffered its reader into a `Vec<RecordBatch>`,
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<RecordBatch>`; 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

<!-- lance-gatekeeper-fix:v1 agent=a48c9fc83772479990d76124f3a2b3be
generation=1 -->

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
This commit is contained in:
lancedb-gatefixer[bot]
2026-10-03 05:53:00 -05:00
committed by GitHub
parent 3d8c9e2ddc
commit cc6b4aa4b4
+105 -1
View File
@@ -3101,7 +3101,13 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
// 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::<std::result::Result<Vec<_>, _>>()?;
let schema = RecordBatchReader::schema(new_data.as_ref());
let mut batches = new_data.collect::<std::result::Result<Vec<_>, _>>()?;
// 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<dyn Scannable> = Box::new(batches);
let rescannable = source.rescannable();
let input: Arc<dyn ExecutionPlan> =
@@ -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<dyn RecordBatchReader + Send> =
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::<HashMap<_, _>>();
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(