This is an automated email from the ASF dual-hosted git repository.
yihua pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/hudi-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 5ec4cc38 fix(core): accept list and fixed-size list columns on upsert
and update (#868)
5ec4cc38 is described below
commit 5ec4cc38b55959e86c44c01a420f9b49612f9206
Author: Y Ethan Guo <[email protected]>
AuthorDate: Thu Oct 29 22:00:25 2026 -0700
fix(core): accept list and fixed-size list columns on upsert and update
(#868)
An embedding column written by `append` could not be upserted or updated
afterwards: the table schema comes back from Avro with the list item named
`element`, while PyArrow and most Arrow producers name it `item`, and the write
path compared types exactly. `FixedSizeList` columns were rejected outright
because Avro has only `array`.
Every write (append, including the version 6 path, upsert, update,
overwrite and dynamic partition overwrite) now normalizes its batches first:
fixed-size lists become lists and every list item is named `element`, at any
depth. This needs no table schema read, and because the first write is
normalized too, every file of the table carries the same item name, so file
group reader version 1 can combine them. Struct children and anything else
about a column are left as they are, and the d [...]
---
crates/core/src/write/append.rs | 37 ++
crates/core/src/write/mod.rs | 189 ++++++++
crates/core/src/write/rewrite/mod.rs | 4 +
.../core/tests/table_write_list_columns_tests.rs | 537 +++++++++++++++++++++
4 files changed, 767 insertions(+)
diff --git a/crates/core/src/write/append.rs b/crates/core/src/write/append.rs
index 0fa25c79..c54a2e04 100644
--- a/crates/core/src/write/append.rs
+++ b/crates/core/src/write/append.rs
@@ -67,6 +67,7 @@ pub struct AppendResult {
/// Append one or more record batches as a single INSERT commit.
pub async fn append_batches(table: &mut Table, batches: &[RecordBatch]) ->
Result<AppendResult> {
+ let batches = &crate::write::normalize_for_write(batches)?;
if table_version(table) < 8 {
return append_version_6(table, batches).await;
}
@@ -1402,6 +1403,42 @@ mod version_6_tests {
}
}
+ #[tokio::test]
+ async fn test_fixed_size_list_appends_as_list() {
+ let dir = tempfile::tempdir().unwrap();
+ let uri = dir.path().to_str().unwrap();
+ let mut table = version_6_table(dir.path()).await;
+ let item = Arc::new(Field::new("item", DataType::Float32, true));
+ let batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new(
+ "embedding",
+ DataType::FixedSizeList(item.clone(), 2),
+ true,
+ )])),
+ vec![Arc::new(arrow::array::FixedSizeListArray::new(
+ item,
+ 2,
+ Arc::new(arrow::array::Float32Array::from(vec![1.0, 2.0, 3.0,
4.0])),
+ None,
+ ))],
+ )
+ .unwrap();
+ table.append([batch.clone()]).await.unwrap();
+ table.append([batch]).await.unwrap();
+ assert_eq!(row_count(uri).await, 4);
+ let element = DataType::List(Arc::new(Field::new("element",
DataType::Float32, true)));
+ // Reader version 1 hands back each file's own schema.
+ let options = ReadOptions::new().with_hudi_option(
+
crate::config::read::HudiReadConfig::FileGroupReaderVersion.as_ref(),
+ "1",
+ );
+ let batches = table.read(&options).await.unwrap();
+ assert_eq!(batches.len(), 2);
+ for batch in batches {
+ assert_eq!(batch.schema().field(0).data_type(), &element);
+ }
+ }
+
/// The commit records the write schema as Avro, so a batch whose schema
/// Avro cannot express is refused before any file is written.
#[tokio::test]
diff --git a/crates/core/src/write/mod.rs b/crates/core/src/write/mod.rs
index 76719655..bda91521 100644
--- a/crates/core/src/write/mod.rs
+++ b/crates/core/src/write/mod.rs
@@ -109,6 +109,87 @@ pub(crate) fn align_batch_to_schema(
arrow_array::RecordBatch::try_new(target.clone(),
batch.columns().to_vec()).ok()
}
+/// Bring a write's batches to the list shape the table schema reads back with.
+///
+/// Fixed-size lists become lists, and every list item is named `element`, at
+/// any depth: Avro has only `array`, which reads back with an `element` item,
+/// while PyArrow and most Arrow producers name it `item`. The first write is
+/// included so that every file of the table carries the same item name, which
+/// a reader taking each file's own schema needs in order to combine them.
+/// Anything else about a column is left for the schema checks downstream.
+pub(crate) fn normalize_for_write(
+ batches: &[arrow_array::RecordBatch],
+) -> Result<Vec<arrow_array::RecordBatch>> {
+ batches.iter().map(normalize_write_batch).collect()
+}
+
+pub(crate) fn normalize_write_batch(
+ batch: &arrow_array::RecordBatch,
+) -> Result<arrow_array::RecordBatch> {
+ use arrow_schema::Field;
+ let mut changed = false;
+ let mut fields = Vec::with_capacity(batch.num_columns());
+ let mut columns = Vec::with_capacity(batch.num_columns());
+ for (field, column) in batch.schema().fields().iter().zip(batch.columns())
{
+ let target = normalize_list_type(field.data_type());
+ if &target == field.data_type() {
+ fields.push(field.clone());
+ columns.push(column.clone());
+ continue;
+ }
+ changed = true;
+ let casted = arrow::compute::cast(column.as_ref(),
&target).map_err(|e| {
+ CoreError::Schema(format!(
+ "cannot convert column '{}' from {} to {target} for writing:
{e}",
+ field.name(),
+ field.data_type()
+ ))
+ })?;
+ fields.push(std::sync::Arc::new(
+ Field::new(field.name(), target, field.is_nullable())
+ .with_metadata(field.metadata().clone()),
+ ));
+ columns.push(casted);
+ }
+ if !changed {
+ return Ok(batch.clone());
+ }
+ let schema = std::sync::Arc::new(arrow_schema::Schema::new_with_metadata(
+ fields,
+ batch.schema().metadata().clone(),
+ ));
+ arrow_array::RecordBatch::try_new(schema, columns)
+ .map_err(|e| CoreError::Schema(format!("cannot rebuild batch for
writing: {e}")))
+}
+
+/// `dt` with every `FixedSizeList` replaced by a `List` of the same item, and
+/// every list item named `element`, at any depth.
+fn normalize_list_type(dt: &arrow_schema::DataType) -> arrow_schema::DataType {
+ use arrow_schema::{DataType, Field};
+ let normalize_field = |f: &std::sync::Arc<Field>, name: &str| {
+ std::sync::Arc::new(
+ Field::new(name, normalize_list_type(f.data_type()),
f.is_nullable())
+ .with_metadata(f.metadata().clone()),
+ )
+ };
+ match dt {
+ DataType::FixedSizeList(item, _) | DataType::List(item) => {
+ DataType::List(normalize_field(item, "element"))
+ }
+ DataType::LargeList(item) => DataType::LargeList(normalize_field(item,
"element")),
+ DataType::Struct(fields) => DataType::Struct(
+ fields
+ .iter()
+ .map(|f| normalize_field(f, f.name()))
+ .collect(),
+ ),
+ DataType::Map(entries, sorted) => {
+ DataType::Map(normalize_field(entries, entries.name()), *sorted)
+ }
+ other => other.clone(),
+ }
+}
+
pub(crate) fn new_file_id() -> String {
format!("{}-0", uuid::Uuid::new_v4())
}
@@ -668,4 +749,112 @@ mod tests {
let json = inflight_commit_metadata_bytes("UPSERT", false).unwrap();
assert!(std::str::from_utf8(&json).unwrap().contains("UPSERT"));
}
+
+ fn list_of(name: &str, item: arrow_schema::DataType) ->
arrow_schema::DataType {
+
arrow_schema::DataType::List(std::sync::Arc::new(arrow_schema::Field::new(
+ name, item, true,
+ )))
+ }
+
+ #[test]
+ fn
test_normalize_list_type_renames_items_and_widens_fixed_size_lists_at_any_depth()
{
+ use arrow_schema::{DataType, Field, Fields};
+ let item = std::sync::Arc::new(Field::new("item", DataType::Float32,
true));
+ let element = list_of("element", DataType::Float32);
+
+ assert_eq!(
+ normalize_list_type(&DataType::FixedSizeList(item.clone(), 2)),
+ element
+ );
+ assert_eq!(
+ normalize_list_type(&list_of("item", DataType::Float32)),
+ element
+ );
+ assert_eq!(
+ normalize_list_type(&DataType::LargeList(item.clone())),
+ DataType::LargeList(std::sync::Arc::new(Field::new(
+ "element",
+ DataType::Float32,
+ true
+ )))
+ );
+ assert_eq!(
+ normalize_list_type(&list_of("item",
DataType::FixedSizeList(item.clone(), 2))),
+ list_of("element", element.clone())
+ );
+
+ // Struct children and map entries keep their names; only list items
change.
+ let nested = |vec: DataType| {
+ let entries = Field::new(
+ "entries",
+ DataType::Struct(Fields::from(vec![
+ Field::new("key", DataType::Utf8, false),
+ Field::new("value", vec.clone(), true),
+ ])),
+ false,
+ );
+ DataType::Struct(Fields::from(vec![
+ Field::new("vec", vec, true),
+ Field::new(
+ "tags",
+ DataType::Map(std::sync::Arc::new(entries), false),
+ true,
+ ),
+ ]))
+ };
+ assert_eq!(
+ normalize_list_type(&nested(DataType::FixedSizeList(item, 2))),
+ nested(element.clone())
+ );
+
+ assert_eq!(normalize_list_type(&DataType::Int64), DataType::Int64);
+ assert_eq!(normalize_list_type(&element), element);
+ }
+
+ #[test]
+ fn test_normalize_write_batch_casts_list_columns_and_keeps_the_rest() {
+ use arrow_array::{Array, ArrayRef, FixedSizeListArray, Float32Array,
Int64Array};
+ use arrow_schema::{DataType, Field, Schema};
+ let item = std::sync::Arc::new(Field::new("item", DataType::Float32,
true));
+ let embedding = FixedSizeListArray::new(
+ item.clone(),
+ 2,
+ std::sync::Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0, 4.0])),
+ Some(vec![true, false].into()),
+ );
+ let field_metadata = HashMap::from([("k".to_string(),
"v".to_string())]);
+ let schema = std::sync::Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int64, false),
+ Field::new("embedding", DataType::FixedSizeList(item, 2), true)
+ .with_metadata(field_metadata.clone()),
+ ]));
+ let batch = arrow_array::RecordBatch::try_new(
+ schema,
+ vec![
+ std::sync::Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef,
+ std::sync::Arc::new(embedding),
+ ],
+ )
+ .unwrap();
+
+ let normalized = normalize_write_batch(&batch).unwrap();
+ let field = normalized.schema().field(1).clone();
+ assert_eq!(field.data_type(), &list_of("element", DataType::Float32));
+ assert_eq!(field.metadata(), &field_metadata);
+ assert_eq!(normalized.schema().field(0), batch.schema().field(0));
+ let list = normalized
+ .column(1)
+ .as_any()
+ .downcast_ref::<arrow_array::ListArray>()
+ .unwrap();
+ assert!(list.is_null(1));
+ let first = list.value(0);
+ let first = first.as_any().downcast_ref::<Float32Array>().unwrap();
+ assert_eq!(first.values(), &[1.0, 2.0]);
+
+ // A batch without lists, or already normalized, comes back unchanged.
+ assert_eq!(normalize_write_batch(&normalized).unwrap(), normalized);
+ let ids = batch.project(&[0]).unwrap();
+ assert_eq!(normalize_write_batch(&ids).unwrap(), ids);
+ }
}
diff --git a/crates/core/src/write/rewrite/mod.rs
b/crates/core/src/write/rewrite/mod.rs
index 577c1ce7..ea57cd25 100644
--- a/crates/core/src/write/rewrite/mod.rs
+++ b/crates/core/src/write/rewrite/mod.rs
@@ -155,6 +155,7 @@ pub async fn upsert_batches(
) -> Result<WriteResult> {
crate::write::ensure_supported_key_generator(&table.hudi_configs)?;
table.reload_timeline_for_write().await?;
+ let batches = &crate::write::normalize_for_write(batches)?;
ensure_configured_record_key(table)?;
if table.is_mor() {
return mor_upsert_batches(table, batches, options).await;
@@ -563,6 +564,7 @@ pub async fn update_filter(
updates: RecordBatch,
) -> Result<WriteResult> {
table.reload_timeline_for_write().await?;
+ let updates = crate::write::normalize_write_batch(&updates)?;
if updates.num_rows() != 1 {
return Err(CoreError::Write(
"update requires a single-row RecordBatch of SET
values".to_string(),
@@ -843,6 +845,7 @@ where
pub async fn overwrite_batches(table: &mut Table, batches: &[RecordBatch]) ->
Result<WriteResult> {
crate::write::ensure_supported_key_generator(&table.hudi_configs)?;
table.reload_timeline_for_write().await?;
+ let batches = &crate::write::normalize_for_write(batches)?;
// Insert-overwrite writes new BASE file groups + a replacecommit on both
// COW and MOR (Java uses the same executor for either table type).
ensure_supported_merge_configs(table)?;
@@ -893,6 +896,7 @@ pub async fn dynamic_partition_overwrite_batches(
) -> Result<WriteResult> {
crate::write::ensure_supported_key_generator(&table.hudi_configs)?;
table.reload_timeline_for_write().await?;
+ let batches = &crate::write::normalize_for_write(batches)?;
// Works on COW and MOR alike (see overwrite_batches).
ensure_supported_merge_configs(table)?;
if batches.is_empty() {
diff --git a/crates/core/tests/table_write_list_columns_tests.rs
b/crates/core/tests/table_write_list_columns_tests.rs
new file mode 100644
index 00000000..7327c7b7
--- /dev/null
+++ b/crates/core/tests/table_write_list_columns_tests.rs
@@ -0,0 +1,537 @@
+/*
+ * 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.
+ */
+//! Writes carrying list columns, the shape embedding vectors arrive in.
+
+use std::sync::Arc;
+
+use arrow::array::{
+ Array, ArrayRef, FixedSizeListArray, Float32Array, Int64Array, ListArray,
StringArray,
+};
+use arrow::datatypes::{DataType, Field, Schema};
+use arrow::record_batch::RecordBatch;
+use hudi_core::config::read::HudiReadConfig;
+use hudi_core::config::table::TableTypeValue;
+use hudi_core::error::CoreError;
+use hudi_core::table::{ReadOptions, Table};
+use tempfile::tempdir;
+
+fn embedding_values(rows: &[(&str, f32)]) -> Float32Array {
+ Float32Array::from(
+ rows.iter()
+ .flat_map(|(_, v)| [*v, *v + 1.0])
+ .collect::<Vec<_>>(),
+ )
+}
+
+/// `id`, `embedding: list<item: float32>`, `v`, with the item field named the
way
+/// PyArrow names it.
+fn list_batch(rows: &[(&str, f32)], version: i64) -> RecordBatch {
+ let item = Arc::new(Field::new("item", DataType::Float32, true));
+ let offsets =
arrow::buffer::OffsetBuffer::from_lengths(rows.iter().map(|_| 2));
+ let embedding = ListArray::new(
+ item.clone(),
+ offsets,
+ Arc::new(embedding_values(rows)),
+ None,
+ );
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Utf8, false),
+ Field::new("embedding", DataType::List(item), true),
+ Field::new("v", DataType::Int64, false),
+ ]));
+ RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(StringArray::from(
+ rows.iter().map(|(id, _)| *id).collect::<Vec<_>>(),
+ )),
+ Arc::new(embedding) as ArrayRef,
+ Arc::new(Int64Array::from(vec![version; rows.len()])),
+ ],
+ )
+ .unwrap()
+}
+
+/// The same rows with `embedding: fixed_size_list<item: float32>[2]`.
+fn fixed_batch(rows: &[(&str, f32)], version: i64) -> RecordBatch {
+ let item = Arc::new(Field::new("item", DataType::Float32, true));
+ let embedding =
+ FixedSizeListArray::new(item.clone(), 2,
Arc::new(embedding_values(rows)), None);
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Utf8, false),
+ Field::new("embedding", DataType::FixedSizeList(item, 2), true),
+ Field::new("v", DataType::Int64, false),
+ ]));
+ RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(StringArray::from(
+ rows.iter().map(|(id, _)| *id).collect::<Vec<_>>(),
+ )),
+ Arc::new(embedding) as ArrayRef,
+ Arc::new(Int64Array::from(vec![version; rows.len()])),
+ ],
+ )
+ .unwrap()
+}
+
+/// `(id, first embedding value)` for every row, sorted by id.
+///
+/// Read with both file group readers: version 1 takes each file's schema as it
+/// is, so it combines the files only when they all name the list item alike.
+async fn read_embeddings(table: &Table) -> Vec<(String, f32)> {
+ let v1 = read_embeddings_with(table, "1").await;
+ assert_eq!(v1, read_embeddings_with(table, "2").await);
+ v1
+}
+
+async fn read_embeddings_with(table: &Table, reader_version: &str) ->
Vec<(String, f32)> {
+ let options = ReadOptions::new().with_hudi_option(
+ HudiReadConfig::FileGroupReaderVersion.as_ref(),
+ reader_version,
+ );
+ let batches = table.read(&options).await.unwrap();
+ for batch in &batches {
+ assert_eq!(
+ batch.schema(),
+ batches[0].schema(),
+ "reader version {reader_version}"
+ );
+ }
+ let mut rows = Vec::new();
+ for batch in &batches {
+ let ids = batch
+ .column_by_name("id")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ let embeddings = batch
+ .column_by_name("embedding")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("embedding reads back as a variable-size list");
+ for i in 0..batch.num_rows() {
+ let values = embeddings.value(i);
+ let first = values
+ .as_any()
+ .downcast_ref::<Float32Array>()
+ .unwrap()
+ .value(0);
+ rows.push((ids.value(i).to_string(), first));
+ }
+ }
+ rows.sort_by(|a, b| a.0.cmp(&b.0));
+ rows
+}
+
+#[tokio::test]
+async fn
test_upsert_accepts_list_column_whose_item_name_differs_from_table_schema() {
+ let dir = tempdir().unwrap();
+ let mut table = Table::create(dir.path().to_str().unwrap())
+ .with_table_name("chunks")
+ .with_record_key_fields(["id"])
+ .with_ordering_fields(["v"])
+ .create()
+ .await
+ .unwrap();
+ table
+ .append([list_batch(&[("a", 1.0), ("b", 2.0)], 1)])
+ .await
+ .unwrap();
+ // The table schema now comes from Avro and names the item `element`; the
+ // incoming batch still names it `item`.
+ let result = table
+ .upsert([list_batch(&[("b", 20.0), ("c", 3.0)], 2)])
+ .await
+ .unwrap();
+ assert_eq!((result.num_updates, result.num_inserts), (1, 1));
+ assert_eq!(
+ read_embeddings(&table).await,
+ vec![("a".into(), 1.0), ("b".into(), 20.0), ("c".into(), 3.0)]
+ );
+
+ let result = table.delete("id = 'a'").await.unwrap();
+ assert_eq!(result.num_deletes, 1);
+ assert_eq!(
+ read_embeddings(&table).await,
+ vec![("b".into(), 20.0), ("c".into(), 3.0)]
+ );
+}
+
+#[tokio::test]
+async fn test_fixed_size_list_column_is_written_as_list() {
+ let dir = tempdir().unwrap();
+ let mut table = Table::create(dir.path().to_str().unwrap())
+ .with_table_name("chunks")
+ .with_record_key_fields(["id"])
+ .with_ordering_fields(["v"])
+ .create()
+ .await
+ .unwrap();
+ table
+ .append([fixed_batch(&[("a", 1.0), ("b", 2.0)], 1)])
+ .await
+ .unwrap();
+ let schema = table.get_schema().await.unwrap();
+ assert!(matches!(
+ schema.field_with_name("embedding").unwrap().data_type(),
+ DataType::List(_)
+ ));
+ let result = table
+ .upsert([fixed_batch(&[("b", 20.0)], 2)])
+ .await
+ .unwrap();
+ assert_eq!(result.num_updates, 1);
+ assert_eq!(
+ read_embeddings(&table).await,
+ vec![("a".into(), 1.0), ("b".into(), 20.0)]
+ );
+}
+
+async fn new_table(dir: &tempfile::TempDir, table_type: TableTypeValue) ->
Table {
+ Table::create(dir.path().to_str().unwrap())
+ .with_table_name("chunks")
+ .with_table_type(table_type)
+ .with_record_key_fields(["id"])
+ .with_ordering_fields(["v"])
+ .create()
+ .await
+ .unwrap()
+}
+
+#[tokio::test]
+async fn test_update_and_multi_batch_upsert_accept_list_item_name_mismatch() {
+ let dir = tempdir().unwrap();
+ let mut table = new_table(&dir, TableTypeValue::CopyOnWrite).await;
+ table
+ .append([list_batch(&[("a", 1.0), ("b", 2.0)], 1)])
+ .await
+ .unwrap();
+ // Two batches in one write, both named `item` against a table that says
`element`.
+ let result = table
+ .upsert([list_batch(&[("a", 10.0)], 2), fixed_batch(&[("c", 3.0)], 2)])
+ .await
+ .unwrap();
+ assert_eq!((result.num_updates, result.num_inserts), (1, 1));
+ // SET on the list column itself through a one-row batch named `item`.
+ let set = list_batch(&[("ignored", 7.0)], 3).project(&[1]).unwrap();
+ let result = table.update("id = 'b'", set).await.unwrap();
+ assert_eq!(result.num_updates, 1);
+ assert_eq!(
+ read_embeddings(&table).await,
+ vec![("a".into(), 10.0), ("b".into(), 7.0), ("c".into(), 3.0)]
+ );
+}
+
+#[tokio::test]
+async fn test_merge_on_read_upsert_accepts_list_column() {
+ let dir = tempdir().unwrap();
+ let mut table = new_table(&dir, TableTypeValue::MergeOnRead).await;
+ table
+ .append([list_batch(&[("a", 1.0), ("b", 2.0)], 1)])
+ .await
+ .unwrap();
+ let result = table
+ .upsert([fixed_batch(&[("b", 20.0), ("c", 3.0)], 2)])
+ .await
+ .unwrap();
+ assert_eq!((result.num_updates, result.num_inserts), (1, 1));
+ assert_eq!(
+ read_embeddings(&table).await,
+ vec![("a".into(), 1.0), ("b".into(), 20.0), ("c".into(), 3.0)]
+ );
+}
+
+#[tokio::test]
+async fn
test_nested_fixed_size_list_is_widened_and_struct_child_names_stay_strict() {
+ use arrow::array::StructArray;
+ let dir = tempdir().unwrap();
+ let mut table = new_table(&dir, TableTypeValue::CopyOnWrite).await;
+ // `payload: struct<vec: fixed_size_list<float32>[2]>`
+ let item = Arc::new(Field::new("item", DataType::Float32, true));
+ let vec_field = Arc::new(Field::new(
+ "vec",
+ DataType::FixedSizeList(item.clone(), 2),
+ true,
+ ));
+ let payload = StructArray::from(vec![(
+ vec_field.clone(),
+ Arc::new(FixedSizeListArray::new(
+ item,
+ 2,
+ Arc::new(Float32Array::from(vec![1.0, 2.0])),
+ None,
+ )) as ArrayRef,
+ )]);
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Utf8, false),
+ Field::new("payload", DataType::Struct(vec![vec_field].into()), true),
+ Field::new("v", DataType::Int64, false),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(StringArray::from(vec!["a"])),
+ Arc::new(payload) as ArrayRef,
+ Arc::new(Int64Array::from(vec![1])),
+ ],
+ )
+ .unwrap();
+ table.append([batch]).await.unwrap();
+ let stored = table.get_schema().await.unwrap();
+ let DataType::Struct(children) =
stored.field_with_name("payload").unwrap().data_type() else {
+ panic!("payload must stay a struct");
+ };
+ assert!(matches!(children[0].data_type(), DataType::List(_)));
+
+ // Same shape but a differently named struct child: not the same column.
+ let other = Arc::new(Field::new(
+ "other",
+ DataType::List(Arc::new(Field::new("item", DataType::Float32, true))),
+ true,
+ ));
+ let payload = StructArray::from(vec![(
+ other.clone(),
+ Arc::new(ListArray::new(
+ Arc::new(Field::new("item", DataType::Float32, true)),
+ arrow::buffer::OffsetBuffer::from_lengths([2]),
+ Arc::new(Float32Array::from(vec![3.0, 4.0])),
+ None,
+ )) as ArrayRef,
+ )]);
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Utf8, false),
+ Field::new("payload", DataType::Struct(vec![other].into()), true),
+ Field::new("v", DataType::Int64, false),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(StringArray::from(vec!["a"])),
+ Arc::new(payload) as ArrayRef,
+ Arc::new(Int64Array::from(vec![2])),
+ ],
+ )
+ .unwrap();
+ let err = table.upsert([batch]).await.unwrap_err();
+ assert!(
+ matches!(&err, CoreError::Schema(msg) if msg.contains("does not
match")),
+ "{err}"
+ );
+}
+
+#[tokio::test]
+async fn test_stale_handle_and_overwrite_accept_list_item_name_mismatch() {
+ let dir = tempdir().unwrap();
+ let mut stale = new_table(&dir, TableTypeValue::CopyOnWrite).await;
+ // Another handle writes first, so `stale` has not seen any commit.
+ let mut fresh = Table::new(dir.path().to_str().unwrap()).await.unwrap();
+ fresh
+ .append([list_batch(&[("a", 1.0), ("b", 2.0)], 1)])
+ .await
+ .unwrap();
+ let result = stale.upsert([list_batch(&[("b", 20.0)], 2)]).await.unwrap();
+ assert_eq!(result.num_updates, 1);
+ stale
+ .overwrite([fixed_batch(&[("c", 3.0)], 3)])
+ .await
+ .unwrap();
+ assert_eq!(read_embeddings(&stale).await, vec![("c".into(), 3.0)]);
+}
+
+#[tokio::test]
+async fn test_merge_on_read_update_sets_list_column() {
+ let dir = tempdir().unwrap();
+ let mut table = new_table(&dir, TableTypeValue::MergeOnRead).await;
+ table
+ .append([list_batch(&[("a", 1.0), ("b", 2.0)], 1)])
+ .await
+ .unwrap();
+ let set = fixed_batch(&[("ignored", 7.0)], 2).project(&[1]).unwrap();
+ let result = table.update("id = 'a'", set).await.unwrap();
+ assert_eq!(result.num_updates, 1);
+ assert_eq!(
+ read_embeddings(&table).await,
+ vec![("a".into(), 7.0), ("b".into(), 2.0)]
+ );
+}
+
+#[tokio::test]
+async fn test_appends_keep_one_list_item_name_across_files() {
+ let dir = tempdir().unwrap();
+ let mut table = new_table(&dir, TableTypeValue::CopyOnWrite).await;
+ table.append([list_batch(&[("a", 1.0)], 1)]).await.unwrap();
+ table.append([list_batch(&[("b", 2.0)], 1)]).await.unwrap();
+ assert_eq!(
+ read_embeddings(&table).await,
+ vec![("a".into(), 1.0), ("b".into(), 2.0)]
+ );
+}
+
+/// `batch` with a `part` column holding `part` on every row.
+fn with_part(batch: RecordBatch, part: &str) -> RecordBatch {
+ let mut fields = batch.schema().fields().to_vec();
+ fields.push(Arc::new(Field::new("part", DataType::Utf8, false)));
+ let mut columns = batch.columns().to_vec();
+ columns.push(Arc::new(StringArray::from(vec![part; batch.num_rows()])));
+ RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap()
+}
+
+#[tokio::test]
+async fn test_dynamic_partition_overwrite_accepts_list_columns() {
+ let dir = tempdir().unwrap();
+ let mut table = Table::create(dir.path().to_str().unwrap())
+ .with_table_name("chunks")
+ .with_record_key_fields(["id"])
+ .with_ordering_fields(["v"])
+ .with_partition_fields(["part"])
+ .create()
+ .await
+ .unwrap();
+ table
+ .append([
+ with_part(list_batch(&[("a", 1.0)], 1), "p1"),
+ with_part(list_batch(&[("b", 2.0)], 1), "p2"),
+ ])
+ .await
+ .unwrap();
+ table
+ .dynamic_partition_overwrite([with_part(fixed_batch(&[("c", 3.0)], 2),
"p2")])
+ .await
+ .unwrap();
+ assert_eq!(
+ read_embeddings(&table).await,
+ vec![("a".into(), 1.0), ("c".into(), 3.0)]
+ );
+}
+
+/// `id`, `payload: struct<vec: list<item: float32>>` (or a fixed-size list of
+/// 2 when `fixed`), `v`.
+fn payload_batch(rows: &[(&str, f32)], version: i64, fixed: bool) ->
RecordBatch {
+ use arrow::array::StructArray;
+ let item = Arc::new(Field::new("item", DataType::Float32, true));
+ let values = Arc::new(embedding_values(rows));
+ let (vec_type, vec): (DataType, ArrayRef) = if fixed {
+ (
+ DataType::FixedSizeList(item.clone(), 2),
+ Arc::new(FixedSizeListArray::new(item, 2, values, None)),
+ )
+ } else {
+ let offsets =
arrow::buffer::OffsetBuffer::from_lengths(rows.iter().map(|_| 2));
+ (
+ DataType::List(item.clone()),
+ Arc::new(ListArray::new(item, offsets, values, None)),
+ )
+ };
+ let vec_field = Arc::new(Field::new("vec", vec_type, true));
+ let payload = StructArray::from(vec![(vec_field.clone(), vec)]);
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Utf8, false),
+ Field::new("payload", DataType::Struct(vec![vec_field].into()), true),
+ Field::new("v", DataType::Int64, false),
+ ]));
+ RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(StringArray::from(
+ rows.iter().map(|(id, _)| *id).collect::<Vec<_>>(),
+ )),
+ Arc::new(payload) as ArrayRef,
+ Arc::new(Int64Array::from(vec![version; rows.len()])),
+ ],
+ )
+ .unwrap()
+}
+
+/// `(id, first payload.vec value)` for every row, sorted by id, read with both
+/// file group readers.
+async fn read_payload_vecs(table: &Table) -> Vec<(String, f32)> {
+ let mut by_version = Vec::new();
+ for reader_version in ["1", "2"] {
+ let options = ReadOptions::new().with_hudi_option(
+ HudiReadConfig::FileGroupReaderVersion.as_ref(),
+ reader_version,
+ );
+ let mut rows = Vec::new();
+ for batch in table.read(&options).await.unwrap() {
+ let ids = batch
+ .column_by_name("id")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap()
+ .clone();
+ let payload = batch
+ .column_by_name("payload")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<arrow::array::StructArray>()
+ .unwrap()
+ .clone();
+ let vecs = payload
+ .column_by_name("vec")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .expect("payload.vec reads back as a variable-size list")
+ .clone();
+ for i in 0..batch.num_rows() {
+ let values = vecs.value(i);
+ let first = values
+ .as_any()
+ .downcast_ref::<Float32Array>()
+ .unwrap()
+ .value(0);
+ rows.push((ids.value(i).to_string(), first));
+ }
+ }
+ rows.sort_by(|a, b| a.0.cmp(&b.0));
+ by_version.push(rows);
+ }
+ assert_eq!(by_version[0], by_version[1]);
+ by_version.remove(0)
+}
+
+#[tokio::test]
+async fn test_upsert_and_update_struct_containing_list() {
+ for table_type in [TableTypeValue::CopyOnWrite,
TableTypeValue::MergeOnRead] {
+ let dir = tempdir().unwrap();
+ let mut table = new_table(&dir, table_type).await;
+ table
+ .append([payload_batch(&[("a", 1.0), ("b", 2.0)], 1, false)])
+ .await
+ .unwrap();
+ let result = table
+ .upsert([payload_batch(&[("b", 20.0), ("c", 3.0)], 2, true)])
+ .await
+ .unwrap();
+ assert_eq!((result.num_updates, result.num_inserts), (1, 1));
+ let set = payload_batch(&[("ignored", 7.0)], 3, false)
+ .project(&[1])
+ .unwrap();
+ let result = table.update("id = 'a'", set).await.unwrap();
+ assert_eq!(result.num_updates, 1);
+ assert_eq!(
+ read_payload_vecs(&table).await,
+ vec![("a".into(), 7.0), ("b".into(), 20.0), ("c".into(), 3.0)]
+ );
+ }
+}
