This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 174ea898 [core] Support column-subset append writes and upserts (#1054)
174ea898 is described below
commit 174ea898bac1652fec02a3ae638c55e6d52152f2
Author: Jingsong Lee <[email protected]>
AuthorDate: Wed Oct 7 12:49:01 2026 +0800
[core] Support column-subset append writes and upserts (#1054)
---
.github/workflows/ci.yml | 3 +-
.../python/python/pypaimon_rust/datafusion.pyi | 2 +
bindings/python/src/write.rs | 36 +++
bindings/python/tests/test_write.py | 25 ++
.../src/table/dedicated_format_file_writer.rs | 144 ++++++---
crates/paimon/src/table/table_update.rs | 9 +-
crates/paimon/src/table/table_upsert.rs | 65 +++-
crates/paimon/src/table/table_write.rs | 126 +++++++-
crates/paimon/src/table/upsert_key_matcher.rs | 2 +
crates/paimon/tests/table_partial_write_test.rs | 344 +++++++++++++++++++++
crates/paimon/tests/table_update_test.rs | 278 +++++++++++++++++
11 files changed, 970 insertions(+), 64 deletions(-)
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index d0cca364..d568322f 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -377,4 +377,5 @@ jobs:
run: |
cargo test --locked -p paimon --features
storage-oss-cpp,storage-jindo --lib \
io::storage_oss_cpp -- --include-ignored
- cargo clippy --locked -p paimon --features storage-oss-cpp --lib
--tests -- -D warnings
+ rustup toolchain install 1.98.0 --profile minimal --component clippy
+ cargo +1.98.0 clippy --locked -p paimon --features storage-oss-cpp
--lib --tests -- -D warnings
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index fcfe3c58..517222e3 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -235,6 +235,7 @@ class CommitMessage:
...
class BatchTableWrite:
+ def with_write_type(self, columns: List[str]) -> "BatchTableWrite": ...
def close(self) -> None: ...
def write_arrow(self, batch: pyarrow.RecordBatch) -> None: ...
def prepare_commit(self) -> List[CommitMessage]:
@@ -287,6 +288,7 @@ class StreamTableUpdate:
...
class StreamTableWrite:
+ def with_write_type(self, columns: List[str]) -> "StreamTableWrite": ...
def close(self) -> None: ...
def write_arrow(self, batch: pyarrow.RecordBatch) -> None: ...
def prepare_commit(self, wait_compaction: bool, commit_identifier: int) ->
List[CommitMessage]:
diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs
index 23d0cd62..3057a440 100644
--- a/bindings/python/src/write.rs
+++ b/bindings/python/src/write.rs
@@ -280,6 +280,18 @@ enum WriteTarget {
}
impl WriteTarget {
+ fn with_write_type(&mut self, columns: Vec<String>) -> paimon::Result<()> {
+ match self {
+ Self::Table(writer) => {
+ writer.with_write_type(columns)?;
+ Ok(())
+ }
+ Self::PostponeFixed(_) => Err(paimon::Error::Unsupported {
+ message: "with_write_type requires a Paimon append
table".into(),
+ }),
+ }
+ }
+
async fn write_arrow_batch(&mut self, batch: &RecordBatch) ->
paimon::Result<()> {
match self {
Self::Table(writer) => writer.write_arrow_batch(batch).await,
@@ -590,6 +602,14 @@ impl UpdateContext {
}
impl WriteState {
+ fn with_write_type(&mut self, columns: Vec<String>) -> PyResult<()> {
+ self.inner
+ .as_mut()
+ .ok_or_else(|| PyRuntimeError::new_err("TableWrite is closed"))?
+ .with_write_type(columns)
+ .map_err(to_py_err)
+ }
+
fn write_arrow(&mut self, py: Python<'_>, batch: &Bound<'_, PyAny>) ->
PyResult<()> {
let batch = RecordBatch::from_pyarrow_bound(batch)?;
let inner = self
@@ -814,6 +834,14 @@ impl PyBatchTableUpdate {
#[pymethods]
impl PyBatchTableWrite {
+ fn with_write_type(
+ mut slf: PyRefMut<'_, Self>,
+ columns: Vec<String>,
+ ) -> PyResult<PyRefMut<'_, Self>> {
+ slf.state.with_write_type(columns)?;
+ Ok(slf)
+ }
+
fn close(&mut self, py: Python<'_>) {
if let Some(mut writer) = self.state.inner.take() {
py.detach(|| runtime().block_on(writer.close()));
@@ -846,6 +874,14 @@ pub struct PyStreamTableWrite {
#[pymethods]
impl PyStreamTableWrite {
+ fn with_write_type(
+ mut slf: PyRefMut<'_, Self>,
+ columns: Vec<String>,
+ ) -> PyResult<PyRefMut<'_, Self>> {
+ slf.state.with_write_type(columns)?;
+ Ok(slf)
+ }
+
fn close(&mut self, py: Python<'_>) {
if let Some(mut writer) = self.state.inner.take() {
py.detach(|| runtime().block_on(writer.close()));
diff --git a/bindings/python/tests/test_write.py
b/bindings/python/tests/test_write.py
index 59e256a3..900a1774 100644
--- a/bindings/python/tests/test_write.py
+++ b/bindings/python/tests/test_write.py
@@ -64,6 +64,31 @@ def test_write_commit_read_roundtrip():
assert result == {"id": [1, 2, 3], "name": ["a", "b", "c"]}
[email protected]("streaming", [False, True])
+def test_partial_write_type_roundtrip_and_lifecycle(tmp_path, streaming):
+ ctx = _make_empty_table(str(tmp_path))
+ table = _get_table(str(tmp_path))
+ builder = table.new_stream_write_builder() if streaming else
table.new_batch_write_builder()
+ writer = builder.new_write()
+ assert writer.with_write_type(["id"]) is writer
+ writer.write_arrow(pa.record_batch([pa.array([], type=pa.int32())],
names=["id"]))
+ assert writer.with_write_type(["id"]) is writer
+ for columns in [[], ["missing"], ["id", "id"]]:
+ with pytest.raises(ValueError):
+ writer.with_write_type(columns)
+ writer.write_arrow(pa.record_batch([pa.array([1, 2], type=pa.int32())],
names=["id"]))
+ with pytest.raises(ValueError, match="before any data"):
+ writer.with_write_type(["id", "name"])
+ messages = writer.prepare_commit(True, 7) if streaming else
writer.prepare_commit()
+ commit = builder.new_commit()
+ commit.commit(7, messages) if streaming else commit.commit(messages)
+ writer.close()
+ with pytest.raises(RuntimeError, match="closed"):
+ writer.with_write_type(["id"])
+ assert pa.Table.from_batches(ctx.sql("SELECT id, name FROM paimon.wdb.t
ORDER BY id")).to_pydict() == {
+ "id": [1, 2], "name": [None, None]}
+
+
@pytest.mark.parametrize("primary_key", [False, True])
def test_custom_data_file_prefix_matches_table_option(tmp_path, primary_key):
diff --git a/crates/paimon/src/table/dedicated_format_file_writer.rs
b/crates/paimon/src/table/dedicated_format_file_writer.rs
index d02c4618..975aee76 100644
--- a/crates/paimon/src/table/dedicated_format_file_writer.rs
+++ b/crates/paimon/src/table/dedicated_format_file_writer.rs
@@ -53,13 +53,15 @@ struct VectorFieldWriter {
/// If a blob value is already a serialized `BlobDescriptor`, the actual data
is
/// resolved from the referenced URI and written to the `.blob` file.
pub(crate) struct AppendDedicatedFormatFileWriter {
- normal_writer: DataFileWriter,
+ normal_writer: Option<DataFileWriter>,
blob_writers: Vec<BlobFieldWriter>,
vector_writer: Option<VectorFieldWriter>,
normal_column_indices: Vec<usize>,
normal_schema: Arc<arrow_schema::Schema>,
// Completed normal + dedicated file groups, still owned until
prepare_commit.
written_files: Vec<DataFileMeta>,
+ target_file_row_num: i64,
+ current_group_row_count: i64,
}
impl AppendDedicatedFormatFileWriter {
@@ -80,6 +82,7 @@ impl AppendDedicatedFormatFileWriter {
vector_target_file_size: i64,
vector_file_format: Option<&str>,
input_schema: &arrow_schema::Schema,
+ write_fields: &[DataField],
table_fields: &[DataField],
format_options: &HashMap<String, String>,
blob_inline_fields: &HashSet<String>,
@@ -99,7 +102,7 @@ impl AppendDedicatedFormatFileWriter {
let mut vector_table_fields = Vec::new();
let mut vector_field_names = Vec::new();
- for (idx, field) in table_fields.iter().enumerate() {
+ for (idx, field) in write_fields.iter().enumerate() {
let is_blob = field.data_type().is_blob_file_field();
let is_inline = blob_inline_fields.contains(field.name());
let is_dedicated_vector =
@@ -191,32 +194,38 @@ impl AppendDedicatedFormatFileWriter {
super::data_file_index_writer::FileIndexOptions::parse(format_options,
table_fields)?
.and_then(|options|
options.project_to_fields(&normal_table_fields))
.map(Arc::new);
- let normal_writer = DataFileWriter::new(
- file_io.clone(),
- table_location,
- partition_path,
- bucket,
- schema_id,
- target_file_size,
- file_compression,
- file_compression_zstd_level,
- write_buffer_size,
- file_format,
- normal_table_fields,
- format_options.clone(),
- Some(0),
- None,
- normal_write_cols,
- )?
- .with_path_factory(paths.clone())
- .with_file_index(normal_index)
- .with_target_file_row_num(target_file_row_num);
- // A view-only schema also uses this adapter to translate references,
- // but Java treats its physical writer as an ordinary inline writer.
- let normal_writer = if !blob_writers.is_empty() ||
vector_writer.is_some() {
- normal_writer.without_row_sidecar()
+ let normal_writer = if normal_table_fields.is_empty() {
+ // Java has no normal writer when the write type is entirely
+ // dedicated columns. Do not create an empty Parquet file.
+ None
} else {
- normal_writer
+ let normal_writer = DataFileWriter::new(
+ file_io.clone(),
+ table_location,
+ partition_path,
+ bucket,
+ schema_id,
+ target_file_size,
+ file_compression,
+ file_compression_zstd_level,
+ write_buffer_size,
+ file_format,
+ normal_table_fields,
+ format_options.clone(),
+ Some(0),
+ None,
+ normal_write_cols,
+ )?
+ .with_path_factory(paths.clone())
+ .with_file_index(normal_index)
+ .with_target_file_row_num(target_file_row_num);
+ // A view-only schema also uses this adapter to translate
references,
+ // but Java treats its physical writer as an ordinary inline
writer.
+ Some(if !blob_writers.is_empty() || vector_writer.is_some() {
+ normal_writer.without_row_sidecar()
+ } else {
+ normal_writer
+ })
};
Ok(Self {
@@ -226,11 +235,15 @@ impl AppendDedicatedFormatFileWriter {
vector_writer,
normal_column_indices,
normal_schema,
+ target_file_row_num,
+ current_group_row_count: 0,
})
}
pub(crate) fn with_resources(mut self, resources: Option<ResourceContext>)
-> Self {
- self.normal_writer.set_resources(resources.clone());
+ if let Some(normal) = &mut self.normal_writer {
+ normal.set_resources(resources.clone());
+ }
for blob in &mut self.blob_writers {
blob.writer.set_resources(resources.clone());
}
@@ -246,17 +259,19 @@ impl AppendDedicatedFormatFileWriter {
}
// Write normal columns
- let normal_columns: Vec<Arc<dyn arrow_array::Array>> = self
- .normal_column_indices
- .iter()
- .map(|&idx| batch.column(idx).clone())
- .collect();
- let normal_batch = RecordBatch::try_new(self.normal_schema.clone(),
normal_columns)
- .map_err(|e| crate::Error::DataInvalid {
- message: format!("Failed to project normal columns: {e}"),
- source: None,
- })?;
- self.normal_writer.write(&normal_batch).await?;
+ if let Some(normal) = &mut self.normal_writer {
+ let normal_columns: Vec<Arc<dyn arrow_array::Array>> = self
+ .normal_column_indices
+ .iter()
+ .map(|&idx| batch.column(idx).clone())
+ .collect();
+ let normal_batch =
RecordBatch::try_new(self.normal_schema.clone(), normal_columns)
+ .map_err(|e| crate::Error::DataInvalid {
+ message: format!("Failed to project normal columns: {e}"),
+ source: None,
+ })?;
+ normal.write(&normal_batch).await?;
+ }
// Write each blob column directly — BlobFormatWriter resolves
descriptors inline
for blob_writer in &mut self.blob_writers {
@@ -302,16 +317,34 @@ impl AppendDedicatedFormatFileWriter {
vector_writer.writer.write(&vector_batch).await?;
}
- if !self.normal_writer.has_open_file() {
+ self.current_group_row_count = self
+ .current_group_row_count
+ .saturating_add(batch.num_rows() as i64);
+ if self.current_group_row_count >= self.target_file_row_num
+ || self
+ .normal_writer
+ .as_ref()
+ .is_some_and(|writer| !writer.has_open_file())
+ {
self.close_group().await?;
}
Ok(())
}
pub(crate) async fn abort(&mut self) {
- self.normal_writer.delete_files(&self.written_files).await;
+ let cleanup = self
+ .normal_writer
+ .as_mut()
+ .or_else(|| self.blob_writers.first_mut().map(|blob| &mut
blob.writer))
+ .or_else(|| self.vector_writer.as_mut().map(|vector| &mut
vector.writer));
+ if let Some(writer) = cleanup {
+ writer.delete_files(&self.written_files).await;
+ }
self.written_files.clear();
- self.normal_writer.abort().await;
+ self.current_group_row_count = 0;
+ if let Some(normal) = &mut self.normal_writer {
+ normal.abort().await;
+ }
for writer in &mut self.blob_writers {
writer.writer.abort().await;
}
@@ -326,7 +359,9 @@ impl AppendDedicatedFormatFileWriter {
}
async fn close_group(&mut self) -> Result<()> {
- let writers = std::iter::once(&mut self.normal_writer)
+ let writers = self
+ .normal_writer
+ .iter_mut()
.chain(
self.blob_writers
.iter_mut()
@@ -344,6 +379,7 @@ impl AppendDedicatedFormatFileWriter {
return Err(error);
}
}
+ self.current_group_row_count = 0;
Ok(())
}
}
@@ -357,7 +393,16 @@ mod tests {
#[tokio::test]
async fn failed_blob_close_cleans_every_physical_column() {
- for (external, rolled) in [(false, false), (true, false), (false,
true), (true, true)] {
+ for (external, rolled, normal) in [
+ (false, false, true),
+ (true, false, true),
+ (false, true, true),
+ (true, true, true),
+ (false, false, false),
+ (true, false, false),
+ (false, true, false),
+ (true, true, false),
+ ] {
let io = FileIOBuilder::new("memory").build().unwrap();
let batch = RecordBatch::try_from_iter([
("id", Arc::new(Int32Array::from(vec![1])) as ArrayRef),
@@ -376,6 +421,16 @@ mod tests {
DataField::new(1, "a".into(), DataType::Blob(BlobType::new())),
DataField::new(2, "b".into(), DataType::Blob(BlobType::new())),
];
+ let write_fields = if normal {
+ fields.clone()
+ } else {
+ fields[1..].to_vec()
+ };
+ let batch = if normal {
+ batch
+ } else {
+ batch.project(&[1, 2]).unwrap()
+ };
let mut options = if external {
HashMap::from([
(
@@ -411,6 +466,7 @@ mod tests {
i64::MAX,
None,
batch.schema().as_ref(),
+ &write_fields,
&fields,
&options,
&HashSet::new(),
diff --git a/crates/paimon/src/table/table_update.rs
b/crates/paimon/src/table/table_update.rs
index e84c0053..49765fed 100644
--- a/crates/paimon/src/table/table_update.rs
+++ b/crates/paimon/src/table/table_update.rs
@@ -43,8 +43,8 @@ fn invalid(message: impl Into<String>) -> crate::Error {
/// [`TableCommit`](super::TableCommit).
///
/// Row-ID updates infer columns from the input unless configured with
-/// [`with_update_type`](Self::with_update_type). Key-based upserts currently
-/// require complete Arrow rows and match keys within each partition.
+/// [`with_update_type`](Self::with_update_type). Key-based upserts accept a
+/// column subset including the keys and match keys within each partition.
#[derive(Clone)]
pub struct TableUpdate {
table: Table,
@@ -190,10 +190,11 @@ impl TableUpdate {
.await
}
- /// Upsert complete Arrow rows by composite key through the core upsert
+ /// Upsert Arrow rows by composite key through the core upsert
/// writer. Keys match within each partition, including when partition
/// columns are omitted from `upsert_keys`. Existing keys update every
- /// matching row ID; new keys append.
+ /// matching row ID; new keys append all input columns. Omitted columns on
+ /// matched rows retain their values; on new rows they are absent from the
file.
pub async fn upsert_by_arrow_with_key(
&self,
batches: Vec<RecordBatch>,
diff --git a/crates/paimon/src/table/table_upsert.rs
b/crates/paimon/src/table/table_upsert.rs
index 7860825a..5ff0534f 100644
--- a/crates/paimon/src/table/table_upsert.rs
+++ b/crates/paimon/src/table/table_upsert.rs
@@ -26,6 +26,7 @@ use arrow_schema::{DataType, Field, Schema};
use arrow_select::{concat::concat_batches, take::take};
use futures::TryStreamExt;
+use super::update_input::unique_column_index;
use super::upsert_key_matcher::UpsertKeyMatcher;
use super::write_batch_normalize::normalize_write_array;
use crate::spec::{
@@ -74,7 +75,7 @@ fn partition_key(row: &BinaryRow, fields: &[DataField]) ->
crate::Result<Vec<u8>
Ok(builder.build().to_serialized_bytes())
}
-/// Upsert full Arrow rows into a data-evolution table without primary keys.
+/// Upsert Arrow rows into a data-evolution table without primary keys.
/// Internal executor for `TableUpdate::upsert_by_arrow_with_key`. Existing
/// keys are updated by row ID; new keys are appended.
#[must_use = "upsert must be used to call prepare_commit()"]
@@ -137,27 +138,48 @@ impl TableUpsert {
})
}
- /// Add full rows. Column order may differ from the table schema; names and
- /// Arrow layouts follow the same normalization as ordinary writes.
Multiple
- /// batches form one logical upsert input.
+ /// Add rows containing the keys and columns to write. Column order may
+ /// differ from the table schema; names and Arrow layouts follow the same
+ /// normalization as ordinary writes. Batches form one logical upsert
input.
pub(super) fn add_batch(&mut self, batch: RecordBatch) ->
crate::Result<()> {
let target =
crate::arrow::build_target_arrow_schema(self.table.schema().fields())?;
- if batch.num_columns() != target.fields().len() {
- return Err(invalid("native upsert requires all table columns"));
+ let input_schema = batch.schema();
+ for field in input_schema.fields() {
+ unique_column_index(&input_schema, field.name())?;
+ target.field_with_name(field.name()).map_err(|_| {
+ invalid(format!(
+ "upsert column '{}' is not in table schema",
+ field.name()
+ ))
+ })?;
+ }
+ for key in &self.keys {
+ unique_column_index(&input_schema, key)?;
}
let mut columns = Vec::with_capacity(target.fields().len());
+ let mut fields = Vec::new();
for field in target.fields() {
- let column = batch
- .column_by_name(field.name())
- .ok_or_else(|| invalid(format!("missing upsert column '{}'",
field.name())))?;
+ let Some(column) = batch.column_by_name(field.name()) else {
+ continue;
+ };
columns.push(
normalize_write_array(column,
field.data_type()).map_err(|error| {
invalid(format!("Invalid upsert column '{}': {error}",
field.name()))
})?,
);
+ fields.push(field.clone());
}
- let ordered = RecordBatch::try_new(target, columns)
+ let ordered = RecordBatch::try_new(Arc::new(Schema::new(fields)),
columns)
.map_err(|error| invalid(format!("cannot order upsert columns:
{error}")))?;
+ if self
+ .source
+ .first()
+ .is_some_and(|first| first.schema() != ordered.schema())
+ {
+ return Err(invalid(
+ "Arrow batches in one upsert input must have the same columns",
+ ));
+ }
self.source.push(ordered);
Ok(())
}
@@ -183,7 +205,20 @@ impl TableUpsert {
.map_err(|error| invalid(error.to_string()))
})
.collect::<crate::Result<Vec<_>>>()?;
- let partitions: HashSet<_> = batch_to_serialized_bytes(source,
&indices, schema.fields())?
+ let input_fields = source
+ .schema()
+ .fields()
+ .iter()
+ .map(|field| {
+ schema
+ .fields()
+ .iter()
+ .find(|table_field| table_field.name() == field.name())
+ .unwrap()
+ .clone()
+ })
+ .collect::<Vec<_>>();
+ let partitions: HashSet<_> = batch_to_serialized_bytes(source,
&indices, &input_fields)?
.into_iter()
.collect();
let fields = schema.partition_fields();
@@ -251,6 +286,14 @@ impl TableUpsert {
.new_write_builder()
.with_commit_user(self.commit_user.clone())?
.new_write()?;
+ append.with_write_type(
+ new_rows
+ .schema()
+ .fields()
+ .iter()
+ .map(|field| field.name().to_string())
+ .collect(),
+ )?;
if let Err(error) = append.write_arrow_batch(&new_rows).await {
append.close().await;
return Err(error);
diff --git a/crates/paimon/src/table/table_write.rs
b/crates/paimon/src/table/table_write.rs
index 25f8e1df..d5d1a7e2 100644
--- a/crates/paimon/src/table/table_write.rs
+++ b/crates/paimon/src/table/table_write.rs
@@ -133,6 +133,8 @@ pub struct TableWrite {
table: Table,
branch_schema_checked: bool,
write_schema: Arc<arrow_schema::Schema>,
+ write_fields: Vec<DataField>,
+ written: bool,
partition_writers: HashMap<PartitionBucketKey, FileWriter>,
partition_computer: PartitionComputer,
partition_keys: Vec<String>,
@@ -198,6 +200,8 @@ impl TableWrite {
table: table.clone(),
branch_schema_checked: table.is_main_branch() ||
table.rest_env.is_some(),
write_schema: build_target_arrow_schema(schema.fields())?,
+ write_fields: schema.fields().to_vec(),
+ written: false,
partition_writers: HashMap::new(),
partition_computer: PartitionComputer::new(
schema.partition_keys(),
@@ -482,6 +486,8 @@ impl TableWrite {
table: table.clone(),
branch_schema_checked: table.is_main_branch() ||
table.rest_env.is_some(),
write_schema,
+ write_fields: fields.to_vec(),
+ written: false,
partition_writers: HashMap::new(),
partition_computer,
partition_keys,
@@ -602,6 +608,97 @@ impl TableWrite {
self
}
+ /// Select the top-level columns written by an append table, as Java
+ /// `BaseAppendFileStoreWrite.withWriteType` does. Input follows this
order;
+ /// omitted columns are recorded in file metadata rather than filled with
NULL.
+ /// Configure this before writing any nonempty batch.
+ pub fn with_write_type(&mut self, write_columns: Vec<String>) ->
Result<&mut Self> {
+ self.ensure_active()?;
+ if self.written {
+ return Err(crate::Error::DataInvalid {
+ message: "with_write_type must be called before any data is
written".into(),
+ source: None,
+ });
+ }
+ if self.format_writer.is_some() ||
!self.table.schema().primary_keys().is_empty() {
+ return Err(crate::Error::Unsupported {
+ message: "with_write_type requires a Paimon append
table".into(),
+ });
+ }
+ if write_columns.is_empty() {
+ return Err(crate::Error::DataInvalid {
+ message: "write columns must not be empty".into(),
+ source: None,
+ });
+ }
+ let schema = self.table.schema();
+ let mut seen = HashSet::new();
+ let fields = write_columns
+ .iter()
+ .map(|name| {
+ if !seen.insert(name) {
+ return Err(crate::Error::DataInvalid {
+ message: format!("Duplicate write column '{name}'"),
+ source: None,
+ });
+ }
+ schema
+ .fields()
+ .iter()
+ .find(|field| field.name() == name)
+ .cloned()
+ .ok_or_else(|| crate::Error::DataInvalid {
+ message: format!("Column '{name}' is not in table
schema"),
+ source: None,
+ })
+ })
+ .collect::<Result<Vec<_>>>()?;
+ let indices = |names: &[String]| {
+ names
+ .iter()
+ .map(|name| {
+ fields
+ .iter()
+ .position(|field| field.name() == name)
+ .ok_or_else(|| crate::Error::DataInvalid {
+ message: format!("Write type must include routing
column '{name}'"),
+ source: None,
+ })
+ })
+ .collect::<Result<Vec<_>>>()
+ };
+ let partition_indices = indices(&self.partition_keys)?;
+ let options = CoreOptions::new(schema.options());
+ let total_buckets = options.bucket();
+ let bucket_indices = if total_buckets > 1 {
+ indices(&schema.bucket_keys())?
+ } else {
+ Vec::new()
+ };
+ let assigner = if total_buckets <= 1 || bucket_indices.is_empty() {
+ BucketAssignerEnum::Constant(ConstantBucketAssigner::new(
+ partition_indices,
+ if total_buckets == POSTPONE_BUCKET {
+ POSTPONE_BUCKET
+ } else {
+ 0
+ },
+ ))
+ } else {
+ BucketAssignerEnum::Fixed(FixedBucketAssigner::new(
+ partition_indices,
+ bucket_indices,
+ options.bucket_function_type()?,
+ total_buckets,
+ ))
+ };
+ let write_schema = build_target_arrow_schema(&fields)?;
+ self.bucket_assigner = assigner;
+ self.write_schema = write_schema;
+ self.write_fields = fields;
+ Ok(self)
+ }
+
/// Write an Arrow RecordBatch. Rows are routed to the correct partition
and bucket.
pub async fn write_arrow_batch(&mut self, batch: &RecordBatch) ->
Result<()> {
if let Some(writer) = self.format_writer.as_mut() {
@@ -611,6 +708,7 @@ impl TableWrite {
let Some(batch) = self.normalize_write_batch(batch)? else {
return Ok(());
};
+ self.written = true;
let grouped = match self.divide_by_partition_bucket(&batch).await {
Ok(grouped) => grouped,
@@ -788,7 +886,7 @@ impl TableWrite {
}
}
- let fields = self.table.schema().fields().to_vec();
+ let fields = self.write_fields.clone();
let output = self.bucket_assigner.assign_batch(batch, &fields).await?;
if matches!(self.bucket_assigner,
BucketAssignerEnum::CrossPartition(_)) {
return self.divide_cross_partition_batch(batch, output);
@@ -1222,6 +1320,7 @@ impl TableWrite {
self.vector_target_file_size,
self.vector_file_format.as_deref(),
&self.write_schema,
+ &self.write_fields,
fields,
self.table.schema().options(),
&self.blob_inline_fields,
@@ -1241,13 +1340,32 @@ impl TableWrite {
self.file_compression_zstd_level,
self.write_buffer_size,
self.file_format.clone(),
- self.table.schema().fields().to_vec(),
+ self.write_fields.clone(),
self.table.schema().options().clone(),
Some(0),
None,
- None,
+ (!super::data_evolution_fields::can_omit_normal_write_cols(
+ self.table.schema().fields(),
+ &self
+ .write_fields
+ .iter()
+ .map(|field| field.name().to_string())
+ .collect::<Vec<_>>(),
+ &CoreOptions::new(self.table.schema().options()),
+ ))
+ .then(|| {
+ self.write_fields
+ .iter()
+ .map(|field| field.name().to_string())
+ .collect()
+ }),
)?
- .with_file_index(self.file_index_options.clone())
+ .with_file_index(
+ self.file_index_options
+ .as_ref()
+ .and_then(|options|
options.project_to_fields(&self.write_fields))
+ .map(Arc::new),
+ )
.with_target_file_row_num(self.target_file_row_num)
.with_resources(self.resources.clone()),
))
diff --git a/crates/paimon/src/table/upsert_key_matcher.rs
b/crates/paimon/src/table/upsert_key_matcher.rs
index e5257dc4..38e8e41f 100644
--- a/crates/paimon/src/table/upsert_key_matcher.rs
+++ b/crates/paimon/src/table/upsert_key_matcher.rs
@@ -52,6 +52,8 @@ pub(super) fn supported_key_type(data_type: &DataType) ->
bool {
| DataType::Int16
| DataType::Int32
| DataType::Int64
+ | DataType::Float32
+ | DataType::Float64
| DataType::UInt8
| DataType::UInt16
| DataType::UInt32
diff --git a/crates/paimon/tests/table_partial_write_test.rs
b/crates/paimon/tests/table_partial_write_test.rs
new file mode 100644
index 00000000..248b2735
--- /dev/null
+++ b/crates/paimon/tests/table_partial_write_test.rs
@@ -0,0 +1,344 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+mod common;
+
+use std::sync::Arc;
+
+use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch};
+use futures::TryStreamExt;
+use paimon::spec::{BlobType, DataType, IntType, Schema, TableSchema};
+use paimon::table::Table;
+
+use common::incremental_helpers::{memory_table, persist_table_schema,
setup_dirs};
+
+async fn table(evolution: bool, fixed_bucket: bool, blob: bool) -> Table {
+ let mut schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column(
+ "omitted",
+ if blob {
+ DataType::Blob(BlobType::new())
+ } else {
+ DataType::Int(IntType::new())
+ },
+ )
+ .column("pt", DataType::Int(IntType::new()))
+ .partition_keys(["pt"])
+ .option("row-tracking.enabled", evolution.to_string())
+ .option("data-evolution.enabled", evolution.to_string());
+ if fixed_bucket {
+ schema = schema.option("bucket", "4").option("bucket-key", "id");
+ }
+ let (io, table) = memory_table(
+ "memory:/partial_write",
+ TableSchema::new(0, &schema.build().unwrap()),
+ );
+ setup_dirs(&io, table.location()).await;
+ persist_table_schema(&io, table.location(), table.schema()).await;
+ table
+}
+
+fn partial(ids: Vec<i32>, partitions: Vec<Option<i32>>) -> RecordBatch {
+ RecordBatch::try_from_iter([
+ ("pt", Arc::new(Int32Array::from(partitions)) as ArrayRef),
+ ("id", Arc::new(Int32Array::from(ids)) as ArrayRef),
+ ])
+ .unwrap()
+}
+
+async fn rows(table: &Table) -> Vec<Vec<Option<i32>>> {
+ let reader = table.new_read_builder();
+ let plan = reader.new_scan().plan().await.unwrap();
+ let batches: Vec<RecordBatch> = reader
+ .new_read()
+ .unwrap()
+ .to_arrow(plan.splits())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ let mut rows = Vec::new();
+ for batch in batches {
+ for row in 0..batch.num_rows() {
+ rows.push(
+ batch
+ .columns()
+ .iter()
+ .map(|column| {
+ let column =
column.as_any().downcast_ref::<Int32Array>().unwrap();
+ column.is_valid(row).then(|| column.value(row))
+ })
+ .collect(),
+ );
+ }
+ }
+ rows.sort();
+ rows
+}
+
+#[tokio::test]
+async fn partial_append_routes_projected_partition_and_bucket_indices() {
+ for evolution in [false, true] {
+ // Java disallows fixed buckets for data evolution tables.
+ for fixed_bucket in [false, true]
+ .into_iter()
+ .filter(|fixed| !evolution || !fixed)
+ {
+ let table = table(evolution, fixed_bucket, false).await;
+ let builder = table.new_write_builder();
+ let mut writer = builder.new_write().unwrap();
+ writer
+ .with_write_type(vec!["pt".into(), "id".into()])
+ .unwrap();
+ writer
+ .write_arrow_batch(&partial(vec![1, 2, 3], vec![Some(7),
Some(8), None]))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert_eq!(
+ messages
+ .iter()
+ .flat_map(|m| &m.new_files)
+ .map(|f| f.row_count)
+ .sum::<i64>(),
+ 3
+ );
+ for file in messages.iter().flat_map(|m| &m.new_files) {
+ assert_eq!(file.write_cols, Some(vec!["pt".into(),
"id".into()]));
+ }
+ builder.new_commit().commit(messages).await.unwrap();
+ // Stream checkpoints must retain the selected type.
+ writer
+ .write_arrow_batch(&partial(vec![4], vec![Some(7)]))
+ .await
+ .unwrap();
+ builder
+ .new_commit()
+ .commit(writer.prepare_commit().await.unwrap())
+ .await
+ .unwrap();
+ assert_eq!(
+ rows(&table).await,
+ vec![
+ vec![Some(1), None, Some(7)],
+ vec![Some(2), None, Some(8)],
+ vec![Some(3), None, None],
+ vec![Some(4), None, Some(7)],
+ ]
+ );
+ }
+ }
+}
+
+#[tokio::test]
+async fn invalid_write_type_does_not_replace_previous_configuration() {
+ let table = table(false, true, false).await;
+ let builder = table.new_write_builder();
+ let mut writer = builder.new_write().unwrap();
+ writer
+ .with_write_type(vec!["pt".into(), "id".into()])
+ .unwrap();
+ for columns in [
+ vec![],
+ vec!["pt", "id", "id"],
+ vec!["pt", "missing"],
+ vec!["id"],
+ vec!["pt"],
+ ] {
+ assert!(writer
+ .with_write_type(columns.into_iter().map(str::to_string).collect())
+ .is_err());
+ }
+ // Empty batches do not lock the configuration; nonempty ones do, even
+ // after their files have been prepared.
+ writer
+ .write_arrow_batch(&partial(vec![], vec![]))
+ .await
+ .unwrap();
+ writer
+ .with_write_type(vec!["pt".into(), "id".into()])
+ .unwrap();
+ writer
+ .write_arrow_batch(&partial(vec![1], vec![Some(2)]))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert!(writer
+ .with_write_type(vec!["id".into(), "omitted".into(), "pt".into()])
+ .is_err());
+ builder.new_commit().commit(messages).await.unwrap();
+ assert_eq!(rows(&table).await, vec![vec![Some(1), None, Some(2)]]);
+}
+
+#[tokio::test]
+async fn omitted_managed_blob_does_not_create_dedicated_files() {
+ let table = table(true, false, true).await;
+ let builder = table.new_write_builder();
+ let mut writer = builder.new_write().unwrap();
+ writer
+ .with_write_type(vec!["pt".into(), "id".into()])
+ .unwrap();
+ writer
+ .write_arrow_batch(&partial(vec![1], vec![Some(2)]))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ let files = messages
+ .iter()
+ .flat_map(|m| &m.new_files)
+ .collect::<Vec<_>>();
+ assert_eq!(files.len(), 1);
+ assert_eq!(files[0].write_cols, Some(vec!["pt".into(), "id".into()]));
+ assert!(files[0].file_name.ends_with(".parquet"));
+ builder.new_commit().commit(messages).await.unwrap();
+ let reader = table.new_read_builder();
+ let plan = reader.new_scan().plan().await.unwrap();
+ let batches: Vec<RecordBatch> = reader
+ .new_read()
+ .unwrap()
+ .to_arrow(plan.splits())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ assert_eq!(batches[0].column(1).null_count(), 1);
+}
+
+#[tokio::test]
+async fn
blob_only_partial_write_has_no_empty_normal_file_and_keeps_row_count() {
+ use arrow_array::LargeBinaryArray;
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("payload", DataType::Blob(BlobType::new()))
+ .option("row-tracking.enabled", "true")
+ .option("data-evolution.enabled", "true")
+ .option("target-file-row-num", "2")
+ .build()
+ .unwrap();
+ let (io, table) = memory_table("memory:/blob_only_partial",
TableSchema::new(0, &schema));
+ setup_dirs(&io, table.location()).await;
+ persist_table_schema(&io, table.location(), table.schema()).await;
+ let builder = table.new_write_builder();
+ let mut seed = builder.new_write().unwrap();
+ seed.with_write_type(vec!["id".into()]).unwrap();
+ let seed_input =
+ RecordBatch::try_from_iter([("id", Arc::new(Int32Array::from(vec![1,
2, 3])) as ArrayRef)])
+ .unwrap();
+ seed.write_arrow_batch(&seed_input.slice(0, 2))
+ .await
+ .unwrap();
+ seed.write_arrow_batch(&seed_input.slice(2, 1))
+ .await
+ .unwrap();
+ builder
+ .new_commit()
+ .commit(seed.prepare_commit().await.unwrap())
+ .await
+ .unwrap();
+ let mut writer = builder.new_write().unwrap();
+ writer.with_write_type(vec!["payload".into()]).unwrap();
+ let values = vec![Some(b"first".as_slice()), None,
Some(b"third".as_slice())];
+ let input = RecordBatch::try_from_iter([(
+ "payload",
+ Arc::new(LargeBinaryArray::from(values.clone())) as ArrayRef,
+ )])
+ .unwrap();
+ writer.write_arrow_batch(&input).await.unwrap();
+ let mut messages = writer.prepare_commit().await.unwrap();
+ let files = messages
+ .iter()
+ .flat_map(|m| &m.new_files)
+ .collect::<Vec<_>>();
+ assert_eq!(
+ files.iter().map(|f| f.row_count).collect::<Vec<_>>(),
+ vec![2, 1]
+ );
+ assert!(files
+ .iter()
+ .all(|file| file.file_name.ends_with(".blob")
+ && file.write_cols == Some(vec!["payload".into()])));
+ // Dedicated-only files update an existing row range. As in Java, they
+ // require explicit row IDs because no normal file allocates new IDs.
+ let mut first_row_id = 0;
+ for file in messages
+ .iter_mut()
+ .flat_map(|message| &mut message.new_files)
+ {
+ file.first_row_id = Some(first_row_id);
+ first_row_id += file.row_count;
+ }
+ builder.new_commit().commit(messages).await.unwrap();
+ let reader = table.new_read_builder();
+ let plan = reader.new_scan().plan().await.unwrap();
+ let batches: Vec<RecordBatch> = reader
+ .new_read()
+ .unwrap()
+ .to_arrow(plan.splits())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ let mut actual = Vec::new();
+ for batch in batches {
+ assert_eq!(batch.column(0).null_count(), 0);
+ let payload = batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<LargeBinaryArray>()
+ .unwrap();
+ actual.extend(
+ (0..batch.num_rows())
+ .map(|row| payload.is_valid(row).then(||
payload.value(row).to_vec())),
+ );
+ }
+ let mut expected = values
+ .into_iter()
+ .map(|value| value.map(|bytes| bytes.to_vec()))
+ .collect::<Vec<_>>();
+ actual.sort();
+ expected.sort();
+ assert_eq!(actual, expected);
+}
+
+#[tokio::test]
+async fn selected_columns_validate_input_before_staging() {
+ let table = table(false, false, false).await;
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ writer
+ .with_write_type(vec!["pt".into(), "id".into()])
+ .unwrap();
+ let input = partial(vec![1], vec![Some(2)]);
+ for bad in [
+ input.project(&[1, 0]).unwrap(),
+ input.project(&[0]).unwrap(),
+ ] {
+ assert!(writer.write_arrow_batch(&bad).await.is_err());
+ }
+ assert!(table
+ .file_io()
+ .list_status_recursive(table.location())
+ .await
+ .unwrap()
+ .iter()
+ .all(|status| !status.path.ends_with(".parquet")));
+ writer.write_arrow_batch(&input).await.unwrap();
+ assert_eq!(
+ writer.prepare_commit().await.unwrap()[0].new_files[0].row_count,
+ 1
+ );
+}
diff --git a/crates/paimon/tests/table_update_test.rs
b/crates/paimon/tests/table_update_test.rs
index 3908a3dc..220f7471 100644
--- a/crates/paimon/tests/table_update_test.rs
+++ b/crates/paimon/tests/table_update_test.rs
@@ -754,6 +754,284 @@ async fn
grouped_updates_on_an_empty_table_preserve_empty_input_semantics() {
assert!(parquet_files(&table).await.is_empty());
}
+#[tokio::test]
+async fn partial_upsert_updates_matches_and_appends_only_input_columns() {
+ use arrow_array::Array;
+ let table = evolution_table().await;
+ write_batch(
+ &table,
+ &batch(&[
+ ("id", vec![1, 1, 2]),
+ ("value", vec![10, 11, 20]),
+ ("score", vec![100, 101, 200]),
+ ]),
+ )
+ .await;
+ let builder = table.new_write_builder();
+ let mut update = builder.new_update().unwrap();
+ update.with_update_type(vec!["value".into()]).unwrap();
+ let messages = update
+ .upsert_by_arrow_with_key(
+ vec![batch(&[("value", vec![12, 30, 13]), ("id", vec![1, 3, 1])])],
+ vec!["id".into()],
+ )
+ .await
+ .unwrap();
+ for file in messages.iter().flat_map(|message| &message.new_files) {
+ assert!(file
+ .write_cols
+ .as_ref()
+ .unwrap()
+ .iter()
+ .all(|column| column != "score"));
+ }
+ builder.new_commit().commit(messages).await.unwrap();
+ let reader = table.new_read_builder();
+ let plan = reader.new_scan().plan().await.unwrap();
+ let batches: Vec<RecordBatch> = reader
+ .new_read()
+ .unwrap()
+ .to_arrow(plan.splits())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ let mut rows = Vec::new();
+ for batch in batches {
+ let columns = batch
+ .columns()
+ .iter()
+ .map(|array| array.as_any().downcast_ref::<Int32Array>().unwrap())
+ .collect::<Vec<_>>();
+ for row in 0..batch.num_rows() {
+ rows.push(
+ columns
+ .iter()
+ .map(|column| column.is_valid(row).then(||
column.value(row)))
+ .collect::<Vec<_>>(),
+ );
+ }
+ }
+ rows.sort();
+ assert_eq!(
+ rows,
+ vec![
+ vec![Some(1), Some(13), Some(100)],
+ vec![Some(1), Some(13), Some(101)],
+ vec![Some(2), Some(20), Some(200)],
+ vec![Some(3), Some(30), None],
+ ]
+ );
+}
+
+#[tokio::test]
+async fn partial_upsert_requires_update_columns_only_for_matches() {
+ let table = evolution_table().await;
+ seed(&table).await;
+ let mut update = table.new_write_builder().new_update().unwrap();
+ update.with_update_type(vec!["score".into()]).unwrap();
+ let before = parquet_files(&table).await;
+ let error = update
+ .upsert_by_arrow_with_key(
+ vec![batch(&[("id", vec![1]), ("value", vec![99])])],
+ vec!["id".into()],
+ )
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("score"));
+ assert_eq!(parquet_files(&table).await, before);
+ // The update selection does not restrict the columns of unmatched rows.
+ let messages = update
+ .upsert_by_arrow_with_key(
+ vec![batch(&[("id", vec![4]), ("value", vec![40])])],
+ vec!["id".into()],
+ )
+ .await
+ .unwrap();
+ assert_eq!(
+ messages[0].new_files[0].write_cols,
+ Some(vec!["id".into(), "value".into()])
+ );
+ commit(&table, messages).await;
+}
+
+#[tokio::test]
+async fn partial_upsert_rejects_ambiguous_unknown_and_inconsistent_input() {
+ let table = evolution_table().await;
+ seed(&table).await;
+ let before = parquet_files(&table).await;
+ let mut update = table.new_write_builder().new_update().unwrap();
+ update.with_update_type(vec!["value".into()]).unwrap();
+ for batches in [
+ vec![batch(&[
+ ("id", vec![1]),
+ ("id", vec![2]),
+ ("value", vec![99]),
+ ])],
+ vec![batch(&[("id", vec![1]), ("missing", vec![99])])],
+ vec![batch(&[("value", vec![99])])],
+ vec![
+ batch(&[("id", vec![1]), ("value", vec![99])]),
+ batch(&[("id", vec![1]), ("score", vec![99])]),
+ ],
+ ] {
+ assert!(update
+ .upsert_by_arrow_with_key(batches, vec!["id".into()])
+ .await
+ .is_err());
+ assert_eq!(parquet_files(&table).await, before);
+ }
+ let messages = update
+ .upsert_by_arrow_with_key(
+ vec![
+ batch(&[("value", vec![41]), ("id", vec![1])]),
+ batch(&[("id", vec![1]), ("value", vec![42])]),
+ ],
+ vec!["id".into()],
+ )
+ .await
+ .unwrap();
+ commit(&table, messages).await;
+ assert_eq!(
+ read_rows(&table).await,
+ vec![vec![1, 42, 100], vec![2, 20, 200], vec![3, 30, 300]]
+ );
+}
+
+#[tokio::test]
+async fn floating_upsert_keys_preserve_fractions_nulls_and_duplicate_targets()
{
+ use arrow_array::{Array, Float32Array, Float64Array};
+ use paimon::spec::{DoubleType, FloatType};
+ for wide in [false, true] {
+ let schema = Schema::builder()
+ .column(
+ "key",
+ if wide {
+ DataType::Double(DoubleType::new())
+ } else {
+ DataType::Float(FloatType::new())
+ },
+ )
+ .column("value", DataType::Int(IntType::new()))
+ .column("keep", DataType::Int(IntType::new()))
+ .option("row-tracking.enabled", "true")
+ .option("data-evolution.enabled", "true")
+ .build()
+ .unwrap();
+ let (io, table) = memory_table("memory:/floating_upsert",
TableSchema::new(0, &schema));
+ setup_dirs(&io, table.location()).await;
+ persist_table_schema(&io, table.location(), table.schema()).await;
+ let keys = |values: Vec<Option<f64>>| -> ArrayRef {
+ if wide {
+ Arc::new(Float64Array::from(values))
+ } else {
+ Arc::new(Float32Array::from(
+ values
+ .into_iter()
+ .map(|value| value.map(|v| v as f32))
+ .collect::<Vec<_>>(),
+ ))
+ }
+ };
+ let original = RecordBatch::try_from_iter([
+ ("key", keys(vec![Some(-1.25), Some(-1.25), None, Some(1.5)])),
+ (
+ "value",
+ Arc::new(Int32Array::from(vec![10, 11, 20, 30])) as ArrayRef,
+ ),
+ (
+ "keep",
+ Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as ArrayRef,
+ ),
+ ])
+ .unwrap();
+ write_batch(&table, &original).await;
+ let input = RecordBatch::try_from_iter([
+ (
+ "key",
+ keys(vec![None, Some(-1.25), Some(2.75), Some(-1.25)]),
+ ),
+ (
+ "value",
+ Arc::new(Int32Array::from(vec![21, 12, 40, 13])) as ArrayRef,
+ ),
+ ])
+ .unwrap();
+ let mut update = table.new_write_builder().new_update().unwrap();
+ update.with_update_type(vec!["value".into()]).unwrap();
+ commit(
+ &table,
+ update
+ .upsert_by_arrow_with_key(vec![input], vec!["key".into()])
+ .await
+ .unwrap(),
+ )
+ .await;
+ let reader = table.new_read_builder();
+ let plan = reader.new_scan().plan().await.unwrap();
+ let batches: Vec<RecordBatch> = reader
+ .new_read()
+ .unwrap()
+ .to_arrow(plan.splits())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ let mut actual = Vec::new();
+ for batch in batches {
+ let value = batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ let keep = batch
+ .column(2)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ for row in 0..batch.num_rows() {
+ let key = if batch.column(0).is_null(row) {
+ None
+ } else if wide {
+ Some(
+ batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Float64Array>()
+ .unwrap()
+ .value(row),
+ )
+ } else {
+ Some(f64::from(
+ batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Float32Array>()
+ .unwrap()
+ .value(row),
+ ))
+ };
+ actual.push((
+ key,
+ value.value(row),
+ keep.is_valid(row).then(|| keep.value(row)),
+ ));
+ }
+ }
+ actual.sort_by_key(|row| row.1);
+ assert_eq!(
+ actual,
+ vec![
+ (Some(-1.25), 13, Some(1)),
+ (Some(-1.25), 13, Some(2)),
+ (None, 21, Some(3)),
+ (Some(1.5), 30, Some(4)),
+ (Some(2.75), 40, None),
+ ]
+ );
+ }
+}
+
#[tokio::test]
async fn configured_upsert_and_delete_are_core_operations() {
let table = evolution_table().await;
