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 aba23877 feat(core): write composite and non-string record keys (#871)
aba23877 is described below
commit aba238779ff1af7ecfb2fbd7053046d46981393c
Author: Y Ethan Guo <[email protected]>
AuthorDate: Fri Oct 30 00:04:33 2026 -0700
feat(core): write composite and non-string record keys (#871)
The native write path accepted only a single `Utf8` record key field, so
tables keyed on several fields or on an integer could not be written. Record
keys are now built the way Java's `KeyGenerator.constructRecordKey` builds them
(`field:value` pairs joined by commas, `__null__` and `__empty__` placeholders,
the same all-null error) for string, boolean and integer key fields. The
writer, `SimpleIndex` and the virtual-key reader share one encoder and one key
generator resolver in `keyg [...]
Date, decimal, timestamp and floating-point key fields stay unsupported:
Spark and Flink store them differently (a date as an ISO `LocalDate` or as the
epoch day, a decimal through `toString` or `toPlainString`), the table does not
record which engine wrote it, and an encoding that does not match the table's
would duplicate rows. A delete or update filter on one field of a composite key
scans instead of using the index, and an `IN` list on `_hoodie_record_key` is
refused for composite [...]
---
.../src/file_group/reader_v2/record_context.rs | 202 +++----------
crates/core/src/index/simple.rs | 50 ++--
crates/core/src/keygen/mod.rs | 218 ++++++++++++--
crates/core/src/keygen/record_key.rs | 296 +++++++++++++++++++
crates/core/src/table/mod.rs | 22 ++
crates/core/src/write/keygen.rs | 261 ++++++----------
crates/core/src/write/rewrite/mod.rs | 21 +-
crates/core/src/write/rewrite/mor.rs | 4 +-
crates/core/src/write/rewrite/rows.rs | 105 ++++---
crates/core/src/write/rewrite/validate.rs | 9 +-
.../core/tests/table_write_composite_keys_tests.rs | 328 +++++++++++++++++++++
11 files changed, 1051 insertions(+), 465 deletions(-)
diff --git a/crates/core/src/file_group/reader_v2/record_context.rs
b/crates/core/src/file_group/reader_v2/record_context.rs
index 283ab160..a33fb5f1 100644
--- a/crates/core/src/file_group/reader_v2/record_context.rs
+++ b/crates/core/src/file_group/reader_v2/record_context.rs
@@ -30,12 +30,14 @@
//! struct.
use crate::Result;
+use crate::config::HudiConfigs;
use crate::config::table::HudiTableConfig;
use crate::error::CoreError;
use crate::file_group::reader_v2::buffered_record::{BufferedRecord,
OrderingValue};
use crate::file_group::reader_v2::delete_context::DeleteContext;
+use crate::keygen::record_key::{construct_record_key, render_key_field};
use crate::metadata::meta_field::MetaField;
-use arrow_array::{Array, ArrayRef, ListArray, RecordBatch, StringArray,
StructArray, UnionArray};
+use arrow_array::{Array, ListArray, RecordBatch, StringArray, StructArray,
UnionArray};
use arrow_schema::{DataType, SchemaRef};
use std::collections::HashMap;
use std::sync::Arc;
@@ -320,16 +322,6 @@ impl Default for RecordContext {
}
impl RecordContext {
- /// Composite-record-key separators + placeholders, matching Java
`KeyGenerator`
- /// (`DEFAULT_COLUMN_VALUE_SEPARATOR` `:`,
`DEFAULT_RECORD_KEY_PARTS_SEPARATOR` `,`,
- /// `NULL_RECORDKEY_PLACEHOLDER`, `EMPTY_RECORDKEY_PLACEHOLDER`). These
MUST stay in sync
- /// with the writer's key generator so a key reconstructed at read time
equals the one
- /// the writer stored.
- const COMPOSITE_KEY_COLUMN_VALUE_SEPARATOR: char = ':';
- const COMPOSITE_KEY_PARTS_SEPARATOR: char = ',';
- const NULL_RECORDKEY_PLACEHOLDER: &'static str = "__null__";
- const EMPTY_RECORDKEY_PLACEHOLDER: &'static str = "__empty__";
-
/// Canonical parse of the table's record-key field list — the single
source of truth
/// shared with [`ReaderContext::record_key_fields_from`] (which delegates
here) so the
/// reader can't disagree with the record context on what the key fields
are.
@@ -422,25 +414,12 @@ impl RecordContext {
}
}
- /// Table version at or above which a `ComplexKeyGenerator` ALWAYS encodes
the
- /// single record-key field name (Java `HoodieTableVersion.NINE`).
- const COMPLEX_KEYGEN_ENCODE_MIN_TABLE_VERSION: i32 = 9;
- /// Write config controlling single-field complex-keygen encoding for table
- /// versions below NINE. Java
`HoodieWriteConfig.COMPLEX_KEYGEN_NEW_ENCODING`
- /// (default `false`). When absent/false, the field name IS encoded.
- const COMPLEX_KEYGEN_NEW_ENCODING_KEY: &'static str =
- "hoodie.write.complex.keygen.new.encoding";
-
/// Whether a single-field VIRTUAL key is stored as `<field>:<value>` by
the
/// writer's key generator, so the reader must reconstruct the same shape.
- ///
- /// Mirrors Java: a `ComplexKeyGenerator`/`ComplexAvroKeyGenerator` over a
SINGLE
- /// record-key field prefixes the field name when
- /// `KeyGenUtils.encodeSingleKeyFieldNameForComplexKeyGen` is true —
- /// `tableVersion >= NINE || !COMPLEX_KEYGEN_NEW_ENCODING` (the config
defaults to
- /// `false`, so v9 always encodes and v6/v8 encode by default). Only
applies to
- /// virtual-key tables; meta-field tables read the already-encoded key from
- /// `_hoodie_record_key`.
+ /// Resolved by the writer's own rule (key generator class, then type, then
+ /// inference, plus
`KeyGenUtils.encodeSingleKeyFieldNameForComplexKeyGen`).
+ /// Only applies to virtual-key tables; meta-field tables read the
+ /// already-encoded key from `_hoodie_record_key`.
fn should_encode_single_complex_key_field(
table_config: &HashMap<String, String>,
populate_meta_fields: bool,
@@ -449,25 +428,10 @@ impl RecordContext {
if populate_meta_fields || num_record_key_fields != 1 {
return false;
}
- let is_complex_keygen = table_config
- .get(HudiTableConfig::KeyGeneratorClass.as_ref())
- .map(|c| c.ends_with("ComplexKeyGenerator") ||
c.ends_with("ComplexAvroKeyGenerator"))
- .unwrap_or(false);
- if !is_complex_keygen {
- return false;
- }
- // `tableVersion >= NINE || !COMPLEX_KEYGEN_NEW_ENCODING`. An
unparseable /
- // absent version falls back to the new-encoding config alone (which
defaults
- // to `false` → encode), matching the common default.
- let table_version = table_config
- .get(HudiTableConfig::TableVersion.as_ref())
- .and_then(|v| v.trim().parse::<i32>().ok());
- let new_encoding = table_config
- .get(Self::COMPLEX_KEYGEN_NEW_ENCODING_KEY)
- .map(|v| v.eq_ignore_ascii_case("true"))
- .unwrap_or(false);
- table_version.is_some_and(|v| v >=
Self::COMPLEX_KEYGEN_ENCODE_MIN_TABLE_VERSION)
- || !new_encoding
+ let configs = HudiConfigs::new(table_config.iter().map(|(k, v)|
(k.as_str(), v.clone())));
+ // The generators the writer does not reproduce (timestamp, custom)
store a
+ // single-field key as the bare value.
+ crate::keygen::encodes_single_key_field_name(&configs).unwrap_or(false)
}
//
=========================================================================
@@ -494,10 +458,9 @@ impl RecordContext {
/// allocation, preserving the vectorized base-merge kernel's hot path).
With
/// VIRTUAL KEYS (`hoodie.populate.meta.fields=false`) the key field is a
real
/// data column whose Arrow type is NOT necessarily `Utf8` — e.g. an
- /// `id INT` primary key. Such a column is stringified via
- /// [`Self::stringify_record_key_column`], mirroring Hudi's
- /// `SimpleKeyGenerator`, which derives the record key from the field's
string
- /// value. Both the base-file side (here) and the log side
+ /// `id INT` primary key. Such a column is stringified the way the Java key
+ /// generator renders the field value (shared with the writer through
+ /// [`render_key_field`]). Both the base-file side (here) and the log side
/// ([`Self::get_record_keys`] via `batch_to_buffered_records`) go through
this
/// method, so their merge keys are stringified identically and match by
/// construction.
@@ -505,8 +468,7 @@ impl RecordContext {
/// # Errors
/// Returns [`CoreError::ReadFileSliceError`] if the key field is absent
from
/// `batch`, and [`CoreError::Unsupported`] if a virtual-key column has a
type
- /// hudi-rs does not derive record keys from (see
- /// [`Self::stringify_record_key_column`]).
+ /// hudi-rs does not derive record keys from (see [`render_key_field`]).
pub fn record_key_array(&self, batch: &RecordBatch) -> Result<StringArray>
{
// Composite virtual key (>1 record-key field): reconstruct
`field:val,field:val`
// per row, mirroring Java `KeyGenerator.constructRecordKey`. Both the
base-file side
@@ -531,41 +493,15 @@ impl RecordContext {
))
})?;
- Self::stringify_key_field_column(batch.column(col_idx),
&self.record_key_field)
- }
-
- /// Stringify one record-key field's column into a [`StringArray`], the
shared building
- /// block for both the single-field key and every field of a composite
key. A `Utf8`
- /// column is cloned (a cheap Arc-buffer bump, not a data copy); any other
type goes
- /// through [`Self::stringify_record_key_column`] (integral types
supported, others
- /// rejected loudly). Using one helper guarantees the single-field,
composite, base, and
- /// log sides all derive keys identically.
- fn stringify_key_field_column(column: &ArrayRef, field: &str) ->
Result<StringArray> {
- match column.data_type() {
- DataType::Utf8 => Ok(column
- .as_any()
- .downcast_ref::<StringArray>()
- .expect("Utf8 column downcasts to StringArray")
- .clone()),
- _ => Self::stringify_record_key_column(column.as_ref(), field),
- }
+ render_key_field(batch.column(col_idx).as_ref(),
&self.record_key_field)
}
- /// Build the merge key for a COMPOSITE virtual key (>1 record-key field),
mirroring Java
- /// `KeyGenerator.constructRecordKey`: for each field append
`<field>:<value>`, join the
- /// parts with `,`, substitute `__null__` / `__empty__` for null / empty
values, and error
- /// if every field is null/empty (matching Java's `HoodieKeyException`).
Each field column
- /// is stringified through the same path as the single-field key (`Utf8`
clone or
- /// [`Self::stringify_record_key_column`]), so a key over e.g. `id INT,
name STRING` is
- /// keyed identically on the base and log sides. All record-key fields are
present in the
- /// batch because the schema handler adds every one to the required schema.
- ///
- /// Note: `,` / `:` occurring INSIDE field values are NOT escaped,
matching Java's
- /// `KeyGenerator.constructRecordKey` bug-for-bug (e.g. `name = "a,b:c"`
yields an
- /// ambiguous key). That ambiguity is inherited by design: both the base
and log sides
- /// construct the key identically through this method, so base/log pairing
is unaffected.
+ /// Build the merge key for a COMPOSITE virtual key (>1 record-key field),
or a single
+ /// field the complex key generator encodes, through Java
+ /// `KeyGenerator.constructRecordKey` (shared with the writer). All
record-key fields are
+ /// present in the batch because the schema handler adds every one to the
required
+ /// schema.
fn build_composite_record_key_array(&self, batch: &RecordBatch) ->
Result<StringArray> {
- // Stringify each key field's column once, preserving nulls.
let mut per_field: Vec<StringArray> =
Vec::with_capacity(self.record_key_fields.len());
for field in &self.record_key_fields {
let col_idx = batch.schema().index_of(field).map_err(|e| {
@@ -573,88 +509,17 @@ impl RecordContext {
"Composite key field '{field}' not found in schema: {e}"
))
})?;
- per_field.push(Self::stringify_key_field_column(
- batch.column(col_idx),
- field,
- )?);
- }
-
- let num_rows = batch.num_rows();
- let mut keys: Vec<String> = Vec::with_capacity(num_rows);
- for row in 0..num_rows {
- let mut key = String::new();
- let mut all_null_or_empty = true;
- for (i, field) in self.record_key_fields.iter().enumerate() {
- if i > 0 {
- key.push(Self::COMPOSITE_KEY_PARTS_SEPARATOR);
- }
- key.push_str(field);
- key.push(Self::COMPOSITE_KEY_COLUMN_VALUE_SEPARATOR);
- let col = &per_field[i];
- if col.is_null(row) {
- key.push_str(Self::NULL_RECORDKEY_PLACEHOLDER);
- } else {
- let value = col.value(row);
- if value.is_empty() {
- key.push_str(Self::EMPTY_RECORDKEY_PLACEHOLDER);
- } else {
- key.push_str(value);
- all_null_or_empty = false;
- }
- }
- }
- if all_null_or_empty {
- return Err(CoreError::ReadFileSliceError(format!(
- "composite record key is entirely null/empty for fields
{:?}; a record \
- key cannot be null or empty (mirrors Java
HoodieKeyException)",
- self.record_key_fields
- )));
- }
- keys.push(key);
+ per_field.push(render_key_field(batch.column(col_idx).as_ref(),
field)?);
}
+ let keys = (0..batch.num_rows())
+ .map(|row| {
+ construct_record_key(&self.record_key_fields, &per_field, row)
+ .map_err(CoreError::ReadFileSliceError)
+ })
+ .collect::<Result<Vec<String>>>()?;
Ok(StringArray::from(keys))
}
- /// Convert a non-`Utf8` virtual record-key column into a [`StringArray`],
- /// mirroring Hudi's `SimpleKeyGenerator` (the key is the field value's
string
- /// form). Supports string and integral key fields — the types Hudi allows
as
- /// a simple record key in practice. Any other type is rejected LOUDLY with
- /// [`CoreError::Unsupported`] rather than guessing a string form that
might
- /// not match the writer's key generator (which would silently mis-key the
- /// merge). For a COMPOSITE virtual key,
[`Self::build_composite_record_key_array`]
- /// calls this per field, so the same per-column type handling is shared.
- fn stringify_record_key_column(column: &dyn Array, field: &str) ->
Result<StringArray> {
- match column.data_type() {
- DataType::LargeUtf8
- | DataType::Int8
- | DataType::Int16
- | DataType::Int32
- | DataType::Int64
- | DataType::UInt8
- | DataType::UInt16
- | DataType::UInt32
- | DataType::UInt64 => {
- let string_column = arrow_cast::cast(column,
&DataType::Utf8).map_err(|e| {
- CoreError::ReadFileSliceError(format!(
- "failed to stringify virtual record-key column
'{field}' of type {:?}: {e}",
- column.data_type()
- ))
- })?;
- Ok(string_column
- .as_any()
- .downcast_ref::<StringArray>()
- .expect("cast to Utf8 yields a StringArray")
- .clone())
- }
- other => Err(CoreError::Unsupported(format!(
- "virtual record-key column '{field}' has type {other:?};
hudi-rs \
- derives record keys only from string and integral key fields
(matching \
- SimpleKeyGenerator). Refusing to guess a string form that may
not match \
- the writer's key generator."
- ))),
- }
- }
-
/// Extract ordering values from the first ordering field column.
///
/// Mirrors Java's `RecordContext.getOrderingValue(T, Schema,
List<String>)`,
@@ -1334,7 +1199,7 @@ impl RecordContext {
mod tests {
use super::*;
use crate::file_group::reader_v2::delete_context::DeleteContext;
- use arrow_array::{BooleanArray, Int32Array, Int64Array, StringArray};
+ use arrow_array::{ArrayRef, BooleanArray, Int32Array, Int64Array,
StringArray};
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
@@ -1562,6 +1427,15 @@ mod tests {
vec!["id:1".to_string(), "id:42".to_string()]
);
+ // Only the type recorded, as hudi-rs and Java 1.x write tables.
+ let complex_type =
make_ctx(&[(HudiTableConfig::KeyGeneratorType.as_ref(), "COMPLEX")]);
+ assert_eq!(
+ keys(&complex_type),
+ vec!["id:1".to_string(), "id:42".to_string()]
+ );
+ let simple_type =
make_ctx(&[(HudiTableConfig::KeyGeneratorType.as_ref(), "SIMPLE")]);
+ assert_eq!(keys(&simple_type), vec!["1".to_string(),
"42".to_string()]);
+
// SimpleKeyGenerator → bare value (no field-name prefix).
let simple = make_ctx(&[(
HudiTableConfig::KeyGeneratorClass.as_ref(),
diff --git a/crates/core/src/index/simple.rs b/crates/core/src/index/simple.rs
index 0f848691..2d5b3609 100644
--- a/crates/core/src/index/simple.rs
+++ b/crates/core/src/index/simple.rs
@@ -22,9 +22,9 @@ use std::collections::{HashMap, HashSet};
use arrow::array::StringArray;
use crate::Result;
-use crate::config::table::HudiTableConfig::RecordKeyFields;
use crate::error::CoreError;
use crate::index::{HoodieIndex, HoodieKey, RecordLocation};
+use crate::keygen::record_key::record_keys_for_batch;
use crate::metadata::meta_field::MetaField;
use crate::table::{ReadOptions, Table};
@@ -50,50 +50,38 @@ impl HoodieIndex for SimpleIndex {
if requested.is_empty() {
return Ok(locations);
}
- let record_key_fields: Vec<String> = table
- .hudi_configs
- .try_get(RecordKeyFields)?
- .map(Into::into)
- .unwrap_or_default();
- if record_key_fields.len() != 1 {
- return Err(CoreError::Unsupported(
- "SimpleIndex currently requires exactly one record key
field".to_string(),
- ));
- }
- let fallback_key = &record_key_fields[0];
let reader = table
.create_file_group_reader_with_options(None,
std::iter::empty::<(&str, &str)>())
.await?;
let mut by_record_key: HashMap<String, RecordLocation> =
HashMap::new();
for slice in table.get_file_slices(&ReadOptions::new()).await? {
let batch = reader.read_file_slice(&slice,
&ReadOptions::new()).await?;
- let key_name = if batch
- .column_by_name(MetaField::RecordKey.as_ref())
- .is_some()
- {
- MetaField::RecordKey.as_ref()
- } else {
- fallback_key
- };
- let key_array = batch
- .column_by_name(key_name)
- .ok_or_else(|| {
- CoreError::Schema(format!("record key field '{key_name}'
is missing"))
- })?
- .as_any()
- .downcast_ref::<StringArray>()
- .ok_or_else(|| {
- CoreError::Unsupported(format!("record key field
'{key_name}' must be Utf8"))
- })?;
let location = RecordLocation {
file_id: slice.file_id().to_string(),
instant_time: table.timeline.get_latest_commit_timestamp()?,
partition_path: slice.partition_path.clone(),
};
- for key in key_array.iter().flatten() {
+ let mut tag = |key: &str| {
if requested.contains(key) {
by_record_key.insert(key.to_string(), location.clone());
}
+ };
+ match batch.column_by_name(MetaField::RecordKey.as_ref()) {
+ Some(column) => column
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .ok_or_else(|| {
+ CoreError::Unsupported(format!(
+ "record key field '{}' must be Utf8",
+ MetaField::RecordKey.as_ref()
+ ))
+ })?
+ .iter()
+ .flatten()
+ .for_each(tag),
+ None => record_keys_for_batch(&table.hudi_configs, &batch)?
+ .iter()
+ .for_each(|key| tag(key)),
}
}
for (hoodie_key, slot) in locations.iter_mut() {
diff --git a/crates/core/src/keygen/mod.rs b/crates/core/src/keygen/mod.rs
index 7cb466f3..3d2c5568 100644
--- a/crates/core/src/keygen/mod.rs
+++ b/crates/core/src/keygen/mod.rs
@@ -17,39 +17,157 @@
* under the License.
*/
-//! Key generator implementations for transforming user filters to partition
filters.
+//! Key generator implementations: record key encoding, and transforms of user
+//! filters to partition filters.
+pub(crate) mod record_key;
pub mod timestamp_based;
use crate::Result;
use crate::config::HudiConfigs;
-use crate::config::table::HudiTableConfig::{KeyGeneratorClass,
KeyGeneratorType};
+use crate::config::table::HudiTableConfig;
+use crate::config::table::HudiTableConfig::{
+ KeyGeneratorClass, KeyGeneratorType, PartitionFields, TableVersion,
+};
+use crate::error::CoreError;
use crate::expr::filter::Filter;
+use crate::keygen::record_key::configured_record_key_fields;
-/// Returns true if the table uses a timestamp-based key generator,
-/// checking both `hoodie.table.keygenerator.class` and
-/// `hoodie.table.keygenerator.type`.
-pub fn is_timestamp_based_keygen(hudi_configs: &HudiConfigs) -> Result<bool> {
- let by_class = hudi_configs
- .try_get(KeyGeneratorClass)?
- .map(|v| {
- let s: String = v.into();
- s.contains("TimestampBasedKeyGenerator")
- })
- .unwrap_or(false);
-
- if by_class {
- return Ok(true);
+/// Java `HoodieWriteConfig.COMPLEX_KEYGEN_NEW_ENCODING`.
+const COMPLEX_KEYGEN_NEW_ENCODING: &str =
"hoodie.write.complex.keygen.new.encoding";
+
+/// The key generators hudi-rs tells apart, as Java `KeyGeneratorType` names
them.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum KeyGeneratorKind {
+ Simple,
+ NonPartitioned,
+ Complex,
+ /// Spark SQL's `SqlKeyGenerator`, which delegates to `ComplexKeyGenerator`
+ /// when the table config names the wrapper itself rather than an origin
class.
+ SparkSql,
+ TimestampBased,
+ /// Custom, global-delete and user-provided generators.
+ Other,
+}
+
+impl KeyGeneratorKind {
+ fn from_class(class: &str) -> Self {
+ match class {
+ "org.apache.hudi.keygen.SimpleKeyGenerator"
+ | "org.apache.hudi.keygen.SimpleAvroKeyGenerator" => Self::Simple,
+ "org.apache.hudi.keygen.NonpartitionedKeyGenerator"
+ | "org.apache.hudi.keygen.NonpartitionedAvroKeyGenerator" =>
Self::NonPartitioned,
+ "org.apache.hudi.keygen.ComplexKeyGenerator"
+ | "org.apache.hudi.keygen.ComplexAvroKeyGenerator" =>
Self::Complex,
+ "org.apache.spark.sql.hudi.command.SqlKeyGenerator" =>
Self::SparkSql,
+ "org.apache.hudi.keygen.TimestampBasedKeyGenerator"
+ | "org.apache.hudi.keygen.TimestampBasedAvroKeyGenerator" =>
Self::TimestampBased,
+ _ => Self::Other,
+ }
}
- Ok(hudi_configs
- .try_get(KeyGeneratorType)?
- .map(|v| {
- let s: String = v.into();
- let upper = s.to_uppercase();
- upper == "TIMESTAMP" || upper == "TIMESTAMP_AVRO"
- })
- .unwrap_or(false))
+ fn from_type(key_generator_type: &str) -> Self {
+ match key_generator_type.to_ascii_uppercase().as_str() {
+ "SIMPLE" | "SIMPLE_AVRO" => Self::Simple,
+ "NON_PARTITION" | "NON_PARTITION_AVRO" => Self::NonPartitioned,
+ "COMPLEX" | "COMPLEX_AVRO" => Self::Complex,
+ "SPARK_SQL" => Self::SparkSql,
+ "TIMESTAMP" | "TIMESTAMP_AVRO" => Self::TimestampBased,
+ _ => Self::Other,
+ }
+ }
+}
+
+fn declared<'a>(hudi_configs: &'a HudiConfigs, config: &HudiTableConfig) ->
Option<&'a str> {
+ hudi_configs
+ .get_raw(config.as_ref())
+ .map(str::trim)
+ .filter(|v| !v.is_empty())
+}
+
+/// The generator the table config names. The class wins over the type, as in
+/// Java `KeyGeneratorType.getKeyGeneratorClassName`.
+fn declared_key_generator(hudi_configs: &HudiConfigs) ->
Option<KeyGeneratorKind> {
+ declared(hudi_configs, &KeyGeneratorClass)
+ .map(KeyGeneratorKind::from_class)
+ .or_else(|| declared(hudi_configs,
&KeyGeneratorType).map(KeyGeneratorKind::from_type))
+}
+
+/// The table's key generator: the declared one, and with none it is inferred
+/// from the field lists (`KeyGenUtils.inferKeyGeneratorType`). Generators
whose
+/// keys or partition paths hudi-rs does not reproduce are rejected, since
+/// writing with a different encoding would split a key across two index
entries.
+pub(crate) fn key_generator_kind(hudi_configs: &HudiConfigs) ->
Result<KeyGeneratorKind> {
+ let kind = match declared_key_generator(hudi_configs) {
+ Some(kind) => kind,
+ None => {
+ let partition_fields: Vec<String> =
hudi_configs.get_or_default(PartitionFields).into();
+ if partition_fields.first().is_some_and(|f| f.contains(':')) {
+ // A `field:type` partition spec is the custom key generator.
+ KeyGeneratorKind::Other
+ } else {
+ match partition_fields.len() {
+ 0 => KeyGeneratorKind::NonPartitioned,
+ 1 if configured_record_key_fields(hudi_configs)?.len() > 1
=> {
+ KeyGeneratorKind::Complex
+ }
+ 1 => KeyGeneratorKind::Simple,
+ _ => KeyGeneratorKind::Complex,
+ }
+ }
+ }
+ };
+ if matches!(
+ kind,
+ KeyGeneratorKind::TimestampBased | KeyGeneratorKind::Other
+ ) {
+ let setting = [KeyGeneratorClass, KeyGeneratorType, PartitionFields]
+ .iter()
+ .find_map(|config| {
+ declared(hudi_configs, config).map(|value|
format!("{}={value}", config.as_ref()))
+ })
+ .unwrap_or_default();
+ return Err(CoreError::Unsupported(format!(
+ "writes do not support {setting}; supported key generators are
SIMPLE, COMPLEX, \
+ NON_PARTITION, and SPARK_SQL"
+ )));
+ }
+ Ok(kind)
+}
+
+/// Whether the table's key generator stores a single-field record key as
+/// `field:value` (Java `ComplexKeyGenerator`) rather than the bare value
+/// (`SimpleKeyGenerator`, `NonpartitionedKeyGenerator`).
+pub(crate) fn encodes_single_key_field_name(hudi_configs: &HudiConfigs) ->
Result<bool> {
+ if !matches!(
+ key_generator_kind(hudi_configs)?,
+ KeyGeneratorKind::Complex | KeyGeneratorKind::SparkSql
+ ) {
+ return Ok(false);
+ }
+ // Java `KeyGenUtils.encodeSingleKeyFieldNameForComplexKeyGen`.
+ let table_version: Option<isize> =
hudi_configs.try_get(TableVersion)?.map(Into::into);
+ let new_encoding = hudi_configs
+ .get_raw(COMPLEX_KEYGEN_NEW_ENCODING)
+ .is_some_and(|v| v.trim().eq_ignore_ascii_case("true"));
+ Ok(table_version.is_some_and(|v| v >= 9) || !new_encoding)
+}
+
+/// Java `KeyGeneratorType.isComplexKeyGenerator`: only a declared type or
+/// class counts, never one inferred from the field lists, and the type is
+/// checked before the class.
+pub(crate) fn declares_complex_key_generator(hudi_configs: &HudiConfigs) ->
bool {
+ let kind = match hudi_configs.get_raw(KeyGeneratorType.as_ref()) {
+ Some(key_generator_type) =>
KeyGeneratorKind::from_type(key_generator_type.trim()),
+ None => declared(hudi_configs, &KeyGeneratorClass)
+ .map_or(KeyGeneratorKind::Other, KeyGeneratorKind::from_class),
+ };
+ kind == KeyGeneratorKind::Complex
+}
+
+/// Returns true if the table uses a timestamp-based key generator.
+pub fn is_timestamp_based_keygen(hudi_configs: &HudiConfigs) -> Result<bool> {
+ Ok(declared_key_generator(hudi_configs) ==
Some(KeyGeneratorKind::TimestampBased))
}
/// Trait for key generators that can transform user filters on data columns
@@ -61,3 +179,55 @@ pub trait KeyGeneratorFilterTransformer {
/// Transforms a filter on the source field to one or more filters on
partition fields.
fn transform_filter(&self, filter: &Filter) -> Result<Vec<Filter>>;
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ fn configs(pairs: &[(&str, &str)]) -> HudiConfigs {
+ HudiConfigs::new(pairs.iter().map(|(k, v)| (*k, v.to_string())))
+ }
+
+ #[test]
+ fn test_declared_key_generator_kinds() {
+ let timestamp = configs(&[(
+ "hoodie.table.keygenerator.class",
+ "org.apache.hudi.keygen.TimestampBasedAvroKeyGenerator",
+ )]);
+ assert!(is_timestamp_based_keygen(×tamp).unwrap());
+ assert!(
+
is_timestamp_based_keygen(&configs(&[("hoodie.table.keygenerator.type",
"timestamp")]))
+ .unwrap()
+ );
+ assert!(!is_timestamp_based_keygen(&configs(&[])).unwrap());
+ // The class wins over the type.
+ assert!(
+ !is_timestamp_based_keygen(&configs(&[
+ ("hoodie.table.keygenerator.type", "TIMESTAMP"),
+ (
+ "hoodie.table.keygenerator.class",
+ "org.apache.hudi.keygen.SimpleKeyGenerator",
+ ),
+ ]))
+ .unwrap()
+ );
+
+ // Java checks the type before the class, and Spark SQL's wrapper is
not complex.
+ let complex_type = configs(&[
+ ("hoodie.table.keygenerator.type", "COMPLEX"),
+ (
+ "hoodie.table.keygenerator.class",
+ "org.apache.hudi.keygen.SimpleKeyGenerator",
+ ),
+ ]);
+ assert!(declares_complex_key_generator(&complex_type));
+ assert!(!declares_complex_key_generator(&configs(&[(
+ "hoodie.table.keygenerator.type",
+ "SPARK_SQL"
+ )])));
+ assert!(!declares_complex_key_generator(&configs(&[(
+ "hoodie.table.partition.fields",
+ "p,q"
+ )])));
+ }
+}
diff --git a/crates/core/src/keygen/record_key.rs
b/crates/core/src/keygen/record_key.rs
new file mode 100644
index 00000000..0a8182cf
--- /dev/null
+++ b/crates/core/src/keygen/record_key.rs
@@ -0,0 +1,296 @@
+/*
+ * 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.
+ */
+//! Record key encoding shared by the writer, the index and the virtual-key
+//! reader, so a key derived at read time equals the one a Java writer stored.
+
+use arrow::array::{
+ Array, BooleanArray, Int8Array, Int16Array, Int32Array, Int64Array,
LargeStringArray,
+ StringArray, StringViewArray, UInt8Array, UInt16Array, UInt32Array,
UInt64Array,
+};
+use arrow::datatypes::DataType;
+use arrow::record_batch::RecordBatch;
+
+use crate::Result;
+use crate::config::HudiConfigs;
+use crate::config::table::HudiTableConfig::RecordKeyFields;
+use crate::error::CoreError;
+use crate::keygen::{KeyGeneratorKind, encodes_single_key_field_name,
key_generator_kind};
+
+/// Java `KeyGenerator.NULL_RECORDKEY_PLACEHOLDER`.
+pub(crate) const NULL_RECORDKEY_PLACEHOLDER: &str = "__null__";
+/// Java `KeyGenerator.EMPTY_RECORDKEY_PLACEHOLDER`.
+pub(crate) const EMPTY_RECORDKEY_PLACEHOLDER: &str = "__empty__";
+/// Java `KeyGenerator.DEFAULT_COLUMN_VALUE_SEPARATOR`.
+const COLUMN_VALUE_SEPARATOR: char = ':';
+/// Java `KeyGenerator.DEFAULT_RECORD_KEY_PARTS_SEPARATOR`.
+const RECORD_KEY_PARTS_SEPARATOR: char = ',';
+
+/// Render a string, boolean or integer record key field column as the Java
+/// key generators do: `toString` of the value. Nulls stay null.
+///
+/// Every other type is [`CoreError::Unsupported`]: Java writers store it in
+/// more than one form that the table does not record, and an encoding that
+/// does not match the table's would duplicate rows. Spark renders a date as
+/// its ISO `LocalDate` and Flink as the epoch day; decimals differ between
+/// `toString` and `toPlainString`; timestamps are stored as a raw epoch count,
+/// a JVM-zone `java.sql.Timestamp` or epoch millis; and float keys depend on
+/// the JDK's `Double.toString` digits.
+pub(crate) fn render_key_field(column: &dyn Array, field: &str) ->
Result<StringArray> {
+ macro_rules! render {
+ ($array:ty, $render:expr) => {
+ downcast::<$array>(column, field)?
+ .iter()
+ .map(|value| value.map($render))
+ .collect::<Vec<Option<String>>>()
+ };
+ }
+ let rendered = match column.data_type() {
+ DataType::Utf8 => return Ok(downcast::<StringArray>(column,
field)?.clone()),
+ DataType::LargeUtf8 => render!(LargeStringArray, str::to_string),
+ DataType::Utf8View => render!(StringViewArray, str::to_string),
+ DataType::Boolean => render!(BooleanArray, |v: bool| v.to_string()),
+ DataType::Int8 => render!(Int8Array, |v: i8| v.to_string()),
+ DataType::Int16 => render!(Int16Array, |v: i16| v.to_string()),
+ DataType::Int32 => render!(Int32Array, |v: i32| v.to_string()),
+ DataType::Int64 => render!(Int64Array, |v: i64| v.to_string()),
+ DataType::UInt8 => render!(UInt8Array, |v: u8| v.to_string()),
+ DataType::UInt16 => render!(UInt16Array, |v: u16| v.to_string()),
+ DataType::UInt32 => render!(UInt32Array, |v: u32| v.to_string()),
+ DataType::UInt64 => render!(UInt64Array, |v: u64| v.to_string()),
+ other => {
+ return Err(CoreError::Unsupported(format!(
+ "record key field '{field}' has type {other}, whose Java key
rendering is not \
+ reproduced; supported record key types are strings, booleans,
and integers"
+ )));
+ }
+ };
+ Ok(StringArray::from(rendered))
+}
+
+/// Java `KeyGenerator.constructRecordKey`: `field:value` parts in field order,
+/// joined by `,`, with null and empty values replaced by their placeholders.
+/// Values containing `,` or `:` are not escaped, as in Java.
+///
+/// Returns Java's `HoodieKeyException` message as the error when every part is
+/// null or empty, so callers can raise it in their own error variant.
+pub(crate) fn construct_record_key(
+ fields: &[String],
+ values: &[StringArray],
+ row: usize,
+) -> std::result::Result<String, String> {
+ let mut key = String::new();
+ let mut entirely_null_or_empty = true;
+ for (i, (field, column)) in fields.iter().zip(values).enumerate() {
+ if i > 0 {
+ key.push(RECORD_KEY_PARTS_SEPARATOR);
+ }
+ key.push_str(field);
+ key.push(COLUMN_VALUE_SEPARATOR);
+ if column.is_null(row) {
+ key.push_str(NULL_RECORDKEY_PLACEHOLDER);
+ } else if column.value(row).is_empty() {
+ key.push_str(EMPTY_RECORDKEY_PLACEHOLDER);
+ } else {
+ key.push_str(column.value(row));
+ entirely_null_or_empty = false;
+ }
+ }
+ if entirely_null_or_empty {
+ return Err(format!(
+ "recordKey values: \"{key}\" for fields: [{}] cannot be entirely
null or empty.",
+ fields.join(", ")
+ ));
+ }
+ Ok(key)
+}
+
+/// Java `KeyGenUtils.getRecordKeyFields`: the configured record key fields,
+/// trimmed, with empty entries dropped. Empty means auto-generated keys.
+pub(crate) fn configured_record_key_fields(hudi_configs: &HudiConfigs) ->
Result<Vec<String>> {
+ let fields: Vec<String> = hudi_configs
+ .try_get(RecordKeyFields)?
+ .map(Into::into)
+ .unwrap_or_default();
+ Ok(fields
+ .iter()
+ .map(|f| f.trim())
+ .filter(|f| !f.is_empty())
+ .map(str::to_string)
+ .collect())
+}
+
+/// Record keys of `batch` derived from the configured record key fields, as
+/// the table's key generator encodes them (Java `KeyGenerator.getRecordKey`).
+/// Fails on a table with auto-generated keys, which no row value determines.
+pub(crate) fn record_keys_for_batch(
+ hudi_configs: &HudiConfigs,
+ batch: &RecordBatch,
+) -> Result<Vec<String>> {
+ let record_key_fields = configured_record_key_fields(hudi_configs)?;
+ if record_key_fields.is_empty() {
+ return Err(CoreError::Unsupported(
+ "auto-generated record keys cannot be derived from row
values".to_string(),
+ ));
+ }
+ encode_record_keys(hudi_configs, &record_key_fields, batch)
+}
+
+/// Several fields, or one field under the complex key generator, encode as
+/// `field:value` parts (Java `KeyGenUtils.getRecordKey(record, fields, ..)`);
+/// one field otherwise is the bare value, which may not be null or empty.
+pub(crate) fn encode_record_keys(
+ hudi_configs: &HudiConfigs,
+ fields: &[String],
+ batch: &RecordBatch,
+) -> Result<Vec<String>> {
+ if fields.len() > 1 && key_generator_kind(hudi_configs)? ==
KeyGeneratorKind::Simple {
+ return Err(CoreError::Unsupported(format!(
+ "the simple key generator expects a single record key field; got
'{}'",
+ fields.join(",")
+ )));
+ }
+ let mut values = Vec::with_capacity(fields.len());
+ for field in fields {
+ let column = batch.column_by_name(field).ok_or_else(|| {
+ CoreError::Schema(format!(
+ "record key field '{field}' is missing from write batch"
+ ))
+ })?;
+ values.push(render_key_field(column.as_ref(), field)?);
+ }
+ if let ([field], [column]) = (fields, values.as_slice())
+ && !encodes_single_key_field_name(hudi_configs)?
+ {
+ return (0..batch.num_rows())
+ .map(|row| {
+ if column.is_null(row) || column.value(row).is_empty() {
+ return Err(CoreError::Write(format!(
+ "record key field '{field}' must be non-null and
non-empty"
+ )));
+ }
+ Ok(column.value(row).to_string())
+ })
+ .collect();
+ }
+ (0..batch.num_rows())
+ .map(|row| construct_record_key(fields, &values,
row).map_err(CoreError::Write))
+ .collect()
+}
+
+fn downcast<'a, A: 'static>(column: &'a dyn Array, field: &str) -> Result<&'a
A> {
+ column.as_any().downcast_ref::<A>().ok_or_else(|| {
+ CoreError::Schema(format!(
+ "record key field '{field}' does not match its declared type {}",
+ column.data_type()
+ ))
+ })
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::{
+ ArrayRef, BinaryArray, Date32Array, Decimal128Array, Float32Array,
Float64Array,
+ TimestampMicrosecondArray, TimestampMillisecondArray,
+ };
+ use std::sync::Arc;
+
+ fn rendered(column: ArrayRef) -> Vec<Option<String>> {
+ render_key_field(column.as_ref(), "k")
+ .unwrap()
+ .iter()
+ .map(|v| v.map(str::to_string))
+ .collect()
+ }
+
+ fn some(values: &[&str]) -> Vec<Option<String>> {
+ values.iter().map(|v| Some(v.to_string())).collect()
+ }
+
+ #[test]
+ fn test_render_key_field_scalar_types() {
+ assert_eq!(
+ rendered(Arc::new(LargeStringArray::from(vec![Some("x"), None]))),
+ vec![Some("x".to_string()), None]
+ );
+ assert_eq!(
+ rendered(Arc::new(BooleanArray::from(vec![true, false]))),
+ some(&["true", "false"])
+ );
+ assert_eq!(
+ rendered(Arc::new(Int64Array::from(vec![i64::MIN]))),
+ some(&["-9223372036854775808"])
+ );
+ assert_eq!(
+ rendered(Arc::new(UInt64Array::from(vec![u64::MAX]))),
+ some(&["18446744073709551615"])
+ );
+ }
+
+ #[test]
+ fn test_render_key_field_rejects_unrendered_types() {
+ let columns: Vec<ArrayRef> = vec![
+ Arc::new(Float32Array::from(vec![1.0f32])),
+ Arc::new(Float64Array::from(vec![1.0])),
+ Arc::new(TimestampMicrosecondArray::from(vec![1])),
+
Arc::new(TimestampMillisecondArray::from(vec![1]).with_timezone("UTC")),
+ Arc::new(BinaryArray::from_vec(vec![b"x"])),
+ Arc::new(Date32Array::from(vec![0])),
+ Arc::new(
+ Decimal128Array::from(vec![1])
+ .with_precision_and_scale(10, 2)
+ .unwrap(),
+ ),
+ ];
+ for column in columns {
+ let err = render_key_field(column.as_ref(), "k").unwrap_err();
+ assert!(matches!(err, CoreError::Unsupported(_)), "{err:?}");
+ let message = err.to_string();
+ assert!(
+ message.contains(&format!(
+ "record key field 'k' has type {}",
+ column.data_type()
+ )),
+ "{message}"
+ );
+ }
+ }
+
+ #[test]
+ fn test_construct_record_key() {
+ let fields = vec!["a".to_string(), "b".to_string()];
+ let values = vec![
+ StringArray::from(vec![Some("1"), None, Some("")]),
+ StringArray::from(vec![Some("x,y:z"), Some("2"), None]),
+ ];
+ assert_eq!(
+ construct_record_key(&fields, &values, 0).unwrap(),
+ "a:1,b:x,y:z"
+ );
+ assert_eq!(
+ construct_record_key(&fields, &values, 1).unwrap(),
+ "a:__null__,b:2"
+ );
+ assert_eq!(
+ construct_record_key(&fields, &values, 2).unwrap_err(),
+ "recordKey values: \"a:__empty__,b:__null__\" for fields: [a, b]
cannot be \
+ entirely null or empty."
+ );
+ }
+}
diff --git a/crates/core/src/table/mod.rs b/crates/core/src/table/mod.rs
index c96ffb59..57d20cae 100644
--- a/crates/core/src/table/mod.rs
+++ b/crates/core/src/table/mod.rs
@@ -114,6 +114,7 @@ use crate::file_group::file_slice::FileSlice;
use crate::file_group::reader::FileGroupReader;
use crate::file_group::reader_v2::reader_context::CompletionGateInputs;
use crate::keygen::is_timestamp_based_keygen;
+use crate::keygen::record_key::configured_record_key_fields;
use crate::metadata::METADATA_TABLE_PARTITION_FIELD;
use crate::metadata::commit::HoodieCommitMetadata;
use crate::metadata::meta_field::MetaField;
@@ -419,6 +420,7 @@ impl Table {
/// A filter that matches no rows succeeds without creating a commit and
reports zero deletes.
pub async fn delete(&mut self, filter: &str) -> Result<WriteResult> {
let filter = parse_write_filter(filter)?;
+ ensure_record_key_list_parsable(&self.hudi_configs, &filter)?;
crate::write::delete_filter(self, filter).await
}
@@ -442,6 +444,7 @@ impl Table {
/// Predicates may target any column. Zero matches succeed without a
commit.
pub async fn update(&mut self, filter: &str, updates: RecordBatch) ->
Result<WriteResult> {
let filter = parse_write_filter(filter)?;
+ ensure_record_key_list_parsable(&self.hudi_configs, &filter)?;
crate::write::update_filter(self, filter, updates).await
}
@@ -1519,6 +1522,25 @@ fn parse_write_filter(filter: &str) -> Result<Filter> {
)))
}
+/// [`parse_write_filter`] splits an `IN` list on every `,`, and every
+/// composite record key contains one, so the list would name key fragments.
+fn ensure_record_key_list_parsable(hudi_configs: &HudiConfigs, filter:
&Filter) -> Result<()> {
+ if filter.field == MetaField::RecordKey.as_ref()
+ && matches!(
+ filter.operator,
+ crate::expr::ExprOperator::In | crate::expr::ExprOperator::NotIn
+ )
+ && configured_record_key_fields(hudi_configs)?.len() > 1
+ {
+ return Err(CoreError::Unsupported(format!(
+ "an IN list on {} cannot hold composite record keys, which contain
','; \
+ use one `=` predicate per key or a predicate on a non-key column",
+ filter.field
+ )));
+ }
+ Ok(())
+}
+
#[cfg(test)]
mod tests {
use super::*;
diff --git a/crates/core/src/write/keygen.rs b/crates/core/src/write/keygen.rs
index 20147228..88fed5c0 100644
--- a/crates/core/src/write/keygen.rs
+++ b/crates/core/src/write/keygen.rs
@@ -24,11 +24,12 @@ use arrow::record_batch::RecordBatch;
use crate::Result;
use crate::config::HudiConfigs;
use crate::config::table::HudiTableConfig::{
- IsHiveStylePartitioning, IsPartitionPathUrlencoded, KeyGeneratorClass,
KeyGeneratorType,
- PartitionFields, RecordKeyFields, TableVersion,
+ IsHiveStylePartitioning, IsPartitionPathUrlencoded, PartitionFields,
TableVersion,
};
use crate::error::CoreError;
use crate::index::HoodieKey;
+use crate::keygen::record_key::{configured_record_key_fields,
encode_record_keys};
+use crate::keygen::{declares_complex_key_generator,
encodes_single_key_field_name};
/// Build per-row [`HoodieKey`]s from a write batch.
///
@@ -54,74 +55,9 @@ pub(crate) const VALIDATION_INSTANT: &str =
"00000000000000000";
/// Java `PartitionPathEncodeUtils.DEFAULT_PARTITION_PATH`.
const DEFAULT_PARTITION_PATH: &str = "__HIVE_DEFAULT_PARTITION__";
-/// Java `HoodieWriteConfig.COMPLEX_KEYGEN_NEW_ENCODING`.
-const COMPLEX_KEYGEN_NEW_ENCODING: &str =
"hoodie.write.complex.keygen.new.encoding";
-
/// Java `HoodieWriteConfig.ENABLE_COMPLEX_KEYGEN_VALIDATION`.
pub(crate) const COMPLEX_KEYGEN_VALIDATION: &str =
"hoodie.write.complex.keygen.validation.enable";
-/// Whether the table's key generator stores a single-field record key as
-/// `field:value` (Java `ComplexKeyGenerator`) rather than the bare value
-/// (`SimpleKeyGenerator`, `NonpartitionedKeyGenerator`).
-///
-/// The generator resolves like Java
`KeyGeneratorType.getKeyGeneratorClassName`:
-/// the class wins over the type, and with neither it is inferred from the
field
-/// lists (`KeyGenUtils.inferKeyGeneratorType`). Generators whose keys or
partition
-/// paths the writer does not reproduce are rejected, since writing with a
-/// different encoding would split a key across two index entries.
-pub(crate) fn encodes_single_key_field_name(hudi_configs: &HudiConfigs) ->
Result<bool> {
- let complex = if let Some(class) = hudi_configs
- .get_raw(KeyGeneratorClass.as_ref())
- .map(str::trim)
- .filter(|c| !c.is_empty())
- {
- match class {
- "org.apache.hudi.keygen.SimpleKeyGenerator"
- | "org.apache.hudi.keygen.SimpleAvroKeyGenerator"
- | "org.apache.hudi.keygen.NonpartitionedKeyGenerator"
- | "org.apache.hudi.keygen.NonpartitionedAvroKeyGenerator" => false,
- // Spark SQL's wrapper delegates to ComplexKeyGenerator when the
table
- // config names the wrapper itself rather than an origin class.
- "org.apache.hudi.keygen.ComplexKeyGenerator"
- | "org.apache.hudi.keygen.ComplexAvroKeyGenerator"
- | "org.apache.spark.sql.hudi.command.SqlKeyGenerator" => true,
- other => {
- return
Err(unsupported_key_generator(KeyGeneratorClass.as_ref(), other));
- }
- }
- } else if let Some(key_generator_type) = hudi_configs
- .get_raw(KeyGeneratorType.as_ref())
- .map(str::trim)
- .filter(|t| !t.is_empty())
- {
- match key_generator_type {
- "SIMPLE" | "SIMPLE_AVRO" | "NON_PARTITION" | "NON_PARTITION_AVRO"
=> false,
- "COMPLEX" | "COMPLEX_AVRO" | "SPARK_SQL" => true,
- other => {
- return
Err(unsupported_key_generator(KeyGeneratorType.as_ref(), other));
- }
- }
- } else {
- let partition_fields: Vec<String> =
hudi_configs.get_or_default(PartitionFields).into();
- if partition_fields.first().is_some_and(|f| f.contains(':')) {
- return Err(unsupported_key_generator(
- PartitionFields.as_ref(),
- &partition_fields.join(","),
- ));
- }
- partition_fields.len() > 1
- };
- if !complex {
- return Ok(false);
- }
- // Java `KeyGenUtils.encodeSingleKeyFieldNameForComplexKeyGen`.
- let table_version: Option<isize> =
hudi_configs.try_get(TableVersion)?.map(Into::into);
- let new_encoding = hudi_configs
- .get_raw(COMPLEX_KEYGEN_NEW_ENCODING)
- .is_some_and(|v| v.trim().eq_ignore_ascii_case("true"));
- Ok(table_version.is_some_and(|v| v >= 9) || !new_encoding)
-}
-
/// Fail a write on a table whose key generator the writer cannot reproduce.
pub(crate) fn ensure_supported_key_generator(hudi_configs: &HudiConfigs) ->
Result<()> {
encodes_single_key_field_name(hudi_configs).map(|_| ())
@@ -134,10 +70,7 @@ pub(crate) fn ensure_supported_key_generator(hudi_configs:
&HudiConfigs) -> Resu
/// refused until the validation is turned off.
pub(crate) fn ensure_complex_keygen_validated(hudi_configs: &HudiConfigs) ->
Result<()> {
let table_version: Option<isize> =
hudi_configs.try_get(TableVersion)?.map(Into::into);
- let record_key_fields: Vec<String> = hudi_configs
- .try_get(RecordKeyFields)?
- .map(Into::into)
- .unwrap_or_default();
+ let record_key_fields = configured_record_key_fields(hudi_configs)?;
if table_version.is_none_or(|v| v >= 9)
|| record_key_fields.len() != 1
|| !complex_keygen_validation_enabled(|key| hudi_configs.get_raw(key))
@@ -162,30 +95,6 @@ pub(crate) fn complex_keygen_validation_enabled<'a>(
!get_raw(COMPLEX_KEYGEN_VALIDATION).is_some_and(|v|
v.trim().eq_ignore_ascii_case("false"))
}
-/// Java `KeyGeneratorType.isComplexKeyGenerator`: only a declared type or
-/// class counts, never one inferred from the field lists.
-fn declares_complex_key_generator(hudi_configs: &HudiConfigs) -> bool {
- if let Some(key_generator_type) =
hudi_configs.get_raw(KeyGeneratorType.as_ref()) {
- return matches!(key_generator_type.trim(), "COMPLEX" | "COMPLEX_AVRO");
- }
- matches!(
- hudi_configs
- .get_raw(KeyGeneratorClass.as_ref())
- .map(str::trim),
- Some(
- "org.apache.hudi.keygen.ComplexKeyGenerator"
- | "org.apache.hudi.keygen.ComplexAvroKeyGenerator"
- )
- )
-}
-
-fn unsupported_key_generator(config: &str, value: &str) -> CoreError {
- CoreError::Unsupported(format!(
- "writes do not support {config}={value}; supported key generators are
SIMPLE, \
- COMPLEX, NON_PARTITION, and SPARK_SQL"
- ))
-}
-
/// Run key/partition derivation for its error side only — call before fencing
/// an instant so invalid input leaves no timeline trace.
pub(crate) fn validate_keygen_inputs(
@@ -214,90 +123,43 @@ pub fn hoodie_keys_for_batch_with_offset(
row_id_offset: usize,
partition_id: u32,
) -> Result<Vec<HoodieKey>> {
- let record_key_fields: Option<Vec<String>> =
- hudi_configs.try_get(RecordKeyFields)?.map(|v| v.into());
- let auto_keys = record_key_fields
- .as_ref()
- .map(|f| f.is_empty() || (f.len() == 1 && f[0].is_empty()))
- .unwrap_or(true);
-
- let encode_field_name = encodes_single_key_field_name(hudi_configs)?;
+ ensure_supported_key_generator(hudi_configs)?;
+ let record_key_fields = configured_record_key_fields(hudi_configs)?;
let partition_fields: Vec<String> =
hudi_configs.get_or_default(PartitionFields).into();
let hive_style: bool =
hudi_configs.get_or_default(IsHiveStylePartitioning).into();
let url_encode: bool = hudi_configs
.get_or_default(IsPartitionPathUrlencoded)
.into();
- if auto_keys {
+ let record_keys = if record_key_fields.is_empty() {
let instant = instant.ok_or_else(|| {
CoreError::Write(
"auto record-key generation requires a commit instant
timestamp".to_string(),
)
})?;
- let mut keys = Vec::with_capacity(batch.num_rows());
- for row in 0..batch.num_rows() {
- // Match Java HoodieRecord.generateSequenceId(instant,
partitionId, rowId).
- let record_key = format!("{instant}_{partition_id}_{}",
row_id_offset + row);
- let partition_path =
- partition_path_for_row(batch, row, &partition_fields,
hive_style, url_encode)?;
- keys.push(HoodieKey {
- record_key,
- partition_path,
- });
- }
- return Ok(keys);
- }
-
- let Some(fields) = record_key_fields else {
- // Unreachable: auto_keys above covers None, but the no-panic rule
- // wants an error, not an unwrap.
- return Err(CoreError::Config(
- crate::config::error::ConfigError::NotFound("record key
fields".to_string()),
- ));
+ // Match Java HoodieRecord.generateSequenceId(instant, partitionId,
rowId).
+ (0..batch.num_rows())
+ .map(|row| format!("{instant}_{partition_id}_{}", row_id_offset +
row))
+ .collect()
+ } else {
+ encode_record_keys(hudi_configs, &record_key_fields, batch)?
};
- if fields.len() != 1 {
- return Err(CoreError::Unsupported(
- "writes currently require exactly one record key field (or none
for auto keys)"
- .to_string(),
- ));
- }
- let key_name = &fields[0];
- let key_array = batch
- .column_by_name(key_name)
- .ok_or_else(|| {
- CoreError::Schema(format!(
- "record key field '{key_name}' is missing from write batch"
- ))
- })?
- .as_any()
- .downcast_ref::<StringArray>()
- .ok_or_else(|| {
- CoreError::Unsupported(format!(
- "record key field '{key_name}' must be Utf8 for writes"
- ))
- })?;
-
- let mut keys = Vec::with_capacity(batch.num_rows());
- for row in 0..batch.num_rows() {
- if key_array.is_null(row) || key_array.value(row).is_empty() {
- return Err(CoreError::Write(format!(
- "record key field '{key_name}' must be non-null and non-empty"
- )));
- }
- let value = key_array.value(row);
- let record_key = if encode_field_name {
- format!("{key_name}:{value}")
- } else {
- value.to_string()
- };
- let partition_path =
- partition_path_for_row(batch, row, &partition_fields, hive_style,
url_encode)?;
- keys.push(HoodieKey {
- record_key,
- partition_path,
- });
- }
- Ok(keys)
+ record_keys
+ .into_iter()
+ .enumerate()
+ .map(|(row, record_key)| {
+ Ok(HoodieKey {
+ record_key,
+ partition_path: partition_path_for_row(
+ batch,
+ row,
+ &partition_fields,
+ hive_style,
+ url_encode,
+ )?,
+ })
+ })
+ .collect()
}
fn partition_path_for_row(
@@ -431,24 +293,71 @@ mod tests {
}
#[test]
- fn test_non_string_record_key_errors() {
+ fn test_non_string_record_key_renders_like_java() {
let cfg = configs(&[("hoodie.table.recordkey.fields", "id")]);
let b = batch(
vec![Field::new("id", DataType::Int64, false)],
- vec![Arc::new(Int64Array::from(vec![7]))],
+ vec![Arc::new(Int64Array::from(vec![7, -3]))],
);
- assert!(hoodie_keys_for_batch(&cfg, &b, None).is_err());
+ let keys = hoodie_keys_for_batch(&cfg, &b, None).unwrap();
+ assert_eq!(keys[0].record_key, "7");
+ assert_eq!(keys[1].record_key, "-3");
+
+ let b = batch(
+ vec![Field::new("id", DataType::Binary, false)],
+ vec![Arc::new(arrow::array::BinaryArray::from_vec(vec![b"x"]))],
+ );
+ let err = hoodie_keys_for_batch(&cfg, &b, None).unwrap_err();
+ assert!(matches!(err, CoreError::Unsupported(_)), "{err:?}");
}
#[test]
- fn test_multi_field_record_key_unsupported() {
- let cfg = configs(&[("hoodie.table.recordkey.fields", "a,b")]);
- let b = batch(
- vec![Field::new("a", DataType::Utf8, false)],
- vec![Arc::new(StringArray::from(vec!["x"]))],
+ fn test_multi_field_record_key_encodes_field_names() {
+ let fields = vec![
+ Field::new("a", DataType::Int32, true),
+ Field::new("b", DataType::Utf8, true),
+ ];
+ let columns: Vec<arrow::array::ArrayRef> = vec![
+ Arc::new(Int32Array::from(vec![Some(1), None, Some(3), None])),
+ Arc::new(StringArray::from(vec![
+ Some("x"),
+ Some("y"),
+ Some(""),
+ Some(""),
+ ])),
+ ];
+ // Java trims the configured field list and drops empty entries.
+ let cfg = configs(&[("hoodie.table.recordkey.fields", " a , b ,")]);
+ let b = batch(fields.clone(), columns.clone());
+ let keys = hoodie_keys_for_batch(&cfg, &b.slice(0, 3), None).unwrap();
+ let keys: Vec<&str> = keys.iter().map(|k|
k.record_key.as_str()).collect();
+ assert_eq!(keys, vec!["a:1,b:x", "a:__null__,b:y", "a:3,b:__empty__"]);
+ let err = hoodie_keys_for_batch(&cfg, &b, None).unwrap_err();
+ assert!(
+ err.to_string().contains(
+ "recordKey values: \"a:__null__,b:__empty__\" for fields: [a,
b] cannot be \
+ entirely null or empty"
+ ),
+ "{err}"
+ );
+
+ // One partition field with several key fields infers the complex
generator.
+ let cfg = configs(&[
+ ("hoodie.table.recordkey.fields", "a,b"),
+ ("hoodie.table.partition.fields", "b"),
+ ]);
+ assert_eq!(
+ hoodie_keys_for_batch(&cfg, &b.slice(0, 1),
None).unwrap()[0].record_key,
+ "a:1,b:x"
);
+
+ // Java's simple key generator rejects several key fields.
+ let cfg = configs(&[
+ ("hoodie.table.recordkey.fields", "a,b"),
+ ("hoodie.table.keygenerator.type", "SIMPLE"),
+ ]);
let err = hoodie_keys_for_batch(&cfg, &b, None).unwrap_err();
- assert!(err.to_string().contains("exactly one record key field"));
+ assert!(err.to_string().contains("single record key field"), "{err}");
}
#[test]
diff --git a/crates/core/src/write/rewrite/mod.rs
b/crates/core/src/write/rewrite/mod.rs
index 28f9ac48..630f0c2d 100644
--- a/crates/core/src/write/rewrite/mod.rs
+++ b/crates/core/src/write/rewrite/mod.rs
@@ -22,17 +22,22 @@ use std::str::FromStr;
use arrow::array::{Array, ArrayRef, BooleanArray, StringArray, UInt32Array};
use arrow::compute::concat_batches;
+use arrow::datatypes::DataType;
use arrow::record_batch::RecordBatch;
use arrow_select::take::take;
use arrow_select::zip::zip;
use crate::Result;
-use crate::config::table::HudiTableConfig::{OrderingFields, RecordKeyFields,
RecordMergeStrategy};
+use crate::config::table::HudiTableConfig::{OrderingFields,
RecordMergeStrategy};
use crate::error::CoreError;
use crate::expr::ExprOperator;
use crate::expr::filter::{Filter, filters_to_row_mask,
validate_fields_against_schemas};
use crate::file_group::record_batches::RecordBatches;
use crate::index::{HoodieIndex, HoodieKey, is_record_index_enabled};
+use crate::keygen::encodes_single_key_field_name;
+use crate::keygen::record_key::{
+ configured_record_key_fields, record_keys_for_batch, render_key_field,
+};
use crate::merge::RecordMergeStrategyValue;
use crate::merge::record_merger::RecordMerger;
use crate::metadata::commit::HoodieWriteStat;
@@ -186,8 +191,8 @@ pub async fn upsert_batches(
let instant = request_rewrite_instant(table, "UPSERT",
RewriteKind::Commit).await?;
let incoming = prepare_batches_for_write(&table.hudi_configs, batches,
&instant, "pending")?;
let incoming = concat(&incoming)?;
- let key_name = record_key_name(table, &incoming)?;
- let incoming = deduplicate_last_by_key(&incoming, &key_name,
&table.hudi_configs)?;
+ let incoming_keys = batch_record_keys(table, &incoming)?;
+ let incoming = deduplicate_last_by_key(&incoming, &incoming_keys,
&table.hudi_configs)?;
let incoming_hoodie_keys = hoodie_keys_for_batch(&table.hudi_configs,
&incoming, None)?;
let delete_markers = delete_marker_rows(&incoming);
let locations = crate::index::for_table_checked(table)
@@ -490,7 +495,7 @@ pub async fn upsert_batches(
pub async fn delete_filter(table: &mut Table, filter: Filter) ->
Result<WriteResult> {
table.reload_timeline_for_write().await?;
// Key equality / IN on the record key (or meta record key) →
RLI/SimpleIndex path.
- if let Some(keys) = keys_from_record_key_filter(table, &filter)?
+ if let Some(keys) = keys_from_record_key_filter(table, &filter).await?
&& !keys.is_empty()
{
return delete_keys(table, &keys).await;
@@ -502,13 +507,7 @@ pub async fn delete_filter(table: &mut Table, filter:
Filter) -> Result<WriteRes
return Ok(WriteResult::default());
};
let mask = filters_to_row_mask(&[filter], &old)?;
- let key_name = record_key_name(table, &old).map_err(|_| {
- CoreError::Unsupported(
- "MOR expression delete requires record keys (or
_hoodie_record_key) in the data"
- .to_string(),
- )
- })?;
- let record_keys = keys(&old, &key_name)?;
+ let record_keys = batch_record_keys(table, &old)?;
let partitions = partition_paths_for_batch(&old);
let delete_key_list = mask
.iter()
diff --git a/crates/core/src/write/rewrite/mor.rs
b/crates/core/src/write/rewrite/mor.rs
index 1febfca2..612d43a5 100644
--- a/crates/core/src/write/rewrite/mor.rs
+++ b/crates/core/src/write/rewrite/mor.rs
@@ -116,8 +116,8 @@ pub(super) async fn mor_upsert_batches(
));
}
}
- let key_name = record_key_name(table, &incoming)?;
- let incoming = deduplicate_last_by_key(&incoming, &key_name,
&table.hudi_configs)?;
+ let incoming_keys = batch_record_keys(table, &incoming)?;
+ let incoming = deduplicate_last_by_key(&incoming, &incoming_keys,
&table.hudi_configs)?;
crate::write::keygen::validate_keygen_inputs(
&table.hudi_configs,
std::slice::from_ref(&incoming),
diff --git a/crates/core/src/write/rewrite/rows.rs
b/crates/core/src/write/rewrite/rows.rs
index e6d628be..b67a862c 100644
--- a/crates/core/src/write/rewrite/rows.rs
+++ b/crates/core/src/write/rewrite/rows.rs
@@ -150,77 +150,90 @@ pub(super) fn concat(batches: &[RecordBatch]) ->
Result<RecordBatch> {
concat_batches(&first.schema(), batches).map_err(Into::into)
}
-pub(super) fn record_key_name(table: &Table, batch: &RecordBatch) ->
Result<String> {
+/// Record keys of `batch`: `_hoodie_record_key` when the batch carries it,
+/// otherwise derived from the configured record key fields.
+pub(super) fn batch_record_keys(table: &Table, batch: &RecordBatch) ->
Result<Vec<String>> {
if batch
.column_by_name(MetaField::RecordKey.as_ref())
.is_some()
{
- return Ok(MetaField::RecordKey.as_ref().to_string());
+ return keys(batch, MetaField::RecordKey.as_ref());
}
- let fields: Vec<String> = table
- .hudi_configs
- .try_get(RecordKeyFields)?
- .map(Into::into)
- .unwrap_or_default();
- if fields.len() != 1 || fields[0].is_empty() {
+ if configured_record_key_fields(&table.hudi_configs)?.is_empty() {
return Err(CoreError::Unsupported(
"upsert/delete require a record key field (or _hoodie_record_key
in the batch); auto-generated keys are insert-only"
.to_string(),
));
}
- Ok(fields[0].clone())
+ record_keys_for_batch(&table.hudi_configs, batch)
}
-pub(super) fn configured_record_key_field(table: &Table) ->
Result<Option<String>> {
- let fields: Vec<String> = table
- .hudi_configs
- .try_get(RecordKeyFields)?
- .map(Into::into)
- .unwrap_or_default();
- if fields.len() == 1 && !fields[0].is_empty() {
- Ok(Some(fields[0].clone()))
- } else {
- Ok(None)
- }
-}
-
-/// If `filter` is `=` / `IN` on the configured record key or
`_hoodie_record_key`,
-/// return keys for the RLI/SimpleIndex delete path.
-pub(super) fn keys_from_record_key_filter(
+/// If `filter` is `=` / `IN` on `_hoodie_record_key`, or on the record key
+/// field of a single-field key, return keys for the RLI/SimpleIndex delete
+/// path. A filter on one field of a composite key does not name whole keys.
+pub(super) async fn keys_from_record_key_filter(
table: &Table,
filter: &Filter,
) -> Result<Option<Vec<HoodieKey>>> {
- let configured = configured_record_key_field(table)?;
- let is_key_field = filter.field == MetaField::RecordKey.as_ref()
- || configured.as_deref() == Some(filter.field.as_str());
- if !is_key_field {
+ if !matches!(filter.operator, ExprOperator::Eq | ExprOperator::In) {
return Ok(None);
}
- if !matches!(filter.operator, ExprOperator::Eq | ExprOperator::In) {
+ let to_keys = |record_keys: Vec<String>| {
+ record_keys
+ .into_iter()
+ .map(|record_key| HoodieKey {
+ record_key,
+ partition_path: String::new(),
+ })
+ .collect()
+ };
+ if filter.field == MetaField::RecordKey.as_ref() {
+ return Ok(Some(to_keys(filter.values.clone())));
+ }
+ if configured_record_key_fields(&table.hudi_configs)? !=
[filter.field.as_str()] {
+ return Ok(None);
+ }
+ let key_type = match table.get_schema().await {
+ Ok(schema) => schema
+ .field_with_name(&filter.field)
+ .map(|field| field.data_type().clone())
+ .unwrap_or(DataType::Utf8),
+ Err(_) => DataType::Utf8,
+ };
+ // The index matches stored key text, so a value names a key only when it
+ // renders back to itself: `id = 07` matches row 7, whose key is `7`.
+ let Ok(values) =
arrow_cast::cast(&StringArray::from(filter.values.clone()), &key_type) else {
+ return Ok(None);
+ };
+ let renders_as_given = render_key_field(values.as_ref(),
&filter.field).is_ok_and(|rendered| {
+ rendered
+ .iter()
+ .zip(&filter.values)
+ .all(|(rendered, given)| rendered == Some(given.as_str()))
+ });
+ // An empty value names no key, as the writer refuses to store one.
+ if !renders_as_given || filter.values.iter().any(String::is_empty) {
return Ok(None);
}
// A filter on the key column holds column values, which the complex key
// generator stores as `field:value`. A generator the writer cannot
// reproduce falls back to the scan path.
- let encode_field_name = filter.field != MetaField::RecordKey.as_ref()
- && match
crate::write::keygen::encodes_single_key_field_name(&table.hudi_configs) {
- Ok(encodes) => encodes,
- Err(_) => return Ok(None),
- };
- Ok(Some(
+ let Ok(encode_field_name) =
encodes_single_key_field_name(&table.hudi_configs) else {
+ return Ok(None);
+ };
+ Ok(Some(to_keys(
filter
.values
.iter()
- .map(|value| HoodieKey {
- record_key: if encode_field_name {
+ .map(|value| {
+ if encode_field_name {
format!("{}:{value}", filter.field)
} else {
value.clone()
- },
- partition_path: String::new(),
+ }
})
.collect(),
- ))
+ )))
}
pub(super) fn partition_paths_for_batch(batch: &RecordBatch) -> Vec<String> {
@@ -350,12 +363,7 @@ pub(super) fn apply_set_updates(
/// Upsert/delete need a configured `hoodie.table.recordkey.fields` (Java auto
keys are insert-only).
pub(super) fn ensure_configured_record_key(table: &Table) -> Result<()> {
- let fields: Vec<String> = table
- .hudi_configs
- .try_get(RecordKeyFields)?
- .map(Into::into)
- .unwrap_or_default();
- if fields.len() != 1 || fields[0].is_empty() {
+ if configured_record_key_fields(&table.hudi_configs)?.is_empty() {
return Err(CoreError::Unsupported(
"upsert/delete require hoodie.table.recordkey.fields;
auto-generated keys are insert-only"
.to_string(),
@@ -377,7 +385,7 @@ pub(super) fn keys(batch: &RecordBatch, field: &str) ->
Result<Vec<String>> {
/// arbitrary by design when no ordering field is configured.
pub(super) fn deduplicate_last_by_key(
batch: &RecordBatch,
- field: &str,
+ keys: &[String],
hudi_configs: &crate::config::HudiConfigs,
) -> Result<RecordBatch> {
// Java routes in-batch dedupe through the merge-mode-derived merger
@@ -405,7 +413,6 @@ pub(super) fn deduplicate_last_by_key(
})
.collect::<std::result::Result<Vec<_>, _>>()?;
let mut winners: HashMap<String, u32> =
HashMap::with_capacity(batch.num_rows());
- let keys = keys(batch, field)?;
for (index, key) in keys.iter().enumerate() {
let index = index as u32;
match winners.get(key) {
diff --git a/crates/core/src/write/rewrite/validate.rs
b/crates/core/src/write/rewrite/validate.rs
index a6bade40..1cedacb3 100644
--- a/crates/core/src/write/rewrite/validate.rs
+++ b/crates/core/src/write/rewrite/validate.rs
@@ -47,14 +47,7 @@ pub(crate) fn ensure_writable_table_version(table: &Table)
-> Result<()> {
/// identity would strand its index entry and diverge the partition value from
/// the path it physically lives in.
pub(super) fn ensure_set_columns_updatable(table: &Table, updates:
&RecordBatch) -> Result<()> {
- let mut identity: Vec<String> = Vec::new();
- if let Some(fields) = table
- .hudi_configs
- .try_get(RecordKeyFields)?
- .map(|v| -> Vec<String> { v.into() })
- {
- identity.extend(fields);
- }
+ let mut identity = configured_record_key_fields(&table.hudi_configs)?;
let partition_fields: Vec<String> = table
.hudi_configs
.get_or_default(crate::config::table::HudiTableConfig::PartitionFields)
diff --git a/crates/core/tests/table_write_composite_keys_tests.rs
b/crates/core/tests/table_write_composite_keys_tests.rs
new file mode 100644
index 00000000..b3bc02d2
--- /dev/null
+++ b/crates/core/tests/table_write_composite_keys_tests.rs
@@ -0,0 +1,328 @@
+/*
+ * 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 to tables keyed on several record key fields, or on a non-string
+//! field, with `_hoodie_record_key` encoded like Java's key generators.
+
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int64Array, StringArray};
+use arrow::datatypes::{DataType, Field, Schema};
+use arrow::record_batch::RecordBatch;
+use hudi_core::config::table::TableTypeValue;
+use hudi_core::index::{HoodieIndex, HoodieKey, SimpleIndex};
+use hudi_core::table::{ReadOptions, Table};
+use tempfile::tempdir;
+
+const KEY_FIELDS: [&str; 2] = ["wr_order_number", "wr_item_sk"];
+
+type OrderRow<'a> = (i64, &'a str, i64);
+
+fn order_batch(rows: &[OrderRow]) -> RecordBatch {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("wr_order_number", DataType::Int64, true),
+ Field::new("wr_item_sk", DataType::Utf8, true),
+ Field::new("wr_quantity", DataType::Int64, false),
+ ]));
+ RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(Int64Array::from(
+ rows.iter().map(|(order, _, _)| *order).collect::<Vec<_>>(),
+ )),
+ Arc::new(StringArray::from(
+ rows.iter().map(|(_, item, _)| *item).collect::<Vec<_>>(),
+ )),
+ Arc::new(Int64Array::from(
+ rows.iter().map(|(_, _, qty)| *qty).collect::<Vec<_>>(),
+ )),
+ ],
+ )
+ .unwrap()
+}
+
+fn strings(batch: &RecordBatch, name: &str) -> Vec<String> {
+ let column = batch
+ .column_by_name(name)
+ .unwrap()
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ (0..batch.num_rows())
+ .map(|i| column.value(i).to_string())
+ .collect()
+}
+
+fn longs(batch: &RecordBatch, name: &str) -> Vec<i64> {
+ let column = batch
+ .column_by_name(name)
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ (0..batch.num_rows()).map(|i| column.value(i)).collect()
+}
+
+/// `(_hoodie_record_key, wr_order_number, wr_item_sk, wr_quantity)`, sorted.
+async fn order_rows(table: &Table) -> Vec<(String, i64, String, i64)> {
+ let mut rows = Vec::new();
+ for batch in table.read(&ReadOptions::new()).await.unwrap() {
+ let keys = strings(&batch, "_hoodie_record_key");
+ let orders = longs(&batch, "wr_order_number");
+ let items = strings(&batch, "wr_item_sk");
+ let quantities = longs(&batch, "wr_quantity");
+ for i in 0..batch.num_rows() {
+ rows.push((keys[i].clone(), orders[i], items[i].clone(),
quantities[i]));
+ }
+ }
+ rows.sort();
+ rows
+}
+
+fn order_row(order: i64, item: &str, quantity: i64) -> (String, i64, String,
i64) {
+ (
+ format!("wr_order_number:{order},wr_item_sk:{item}"),
+ order,
+ item.to_string(),
+ quantity,
+ )
+}
+
+fn key(record_key: &str) -> HoodieKey {
+ HoodieKey {
+ record_key: record_key.to_string(),
+ partition_path: String::new(),
+ }
+}
+
+/// Append, upsert, delete by key, delete by filter, and the SET guard on a
+/// table keyed on `(wr_order_number BIGINT, wr_item_sk STRING)`.
+async fn run_composite_key_lifecycle(table_type: TableTypeValue, record_index:
bool) {
+ let dir = tempdir().unwrap();
+ let mut table = Table::create(dir.path().to_str().unwrap())
+ .with_table_name("web_returns")
+ .with_table_type(table_type)
+ .with_record_key_fields(KEY_FIELDS)
+ .with_record_index(record_index)
+ .create()
+ .await
+ .unwrap();
+ let props =
std::fs::read_to_string(dir.path().join(".hoodie/hoodie.properties")).unwrap();
+
assert!(props.contains("hoodie.table.recordkey.fields=wr_order_number,wr_item_sk"));
+ assert!(props.contains("hoodie.table.keygenerator.type=NON_PARTITION"));
+
+ table
+ .append([order_batch(&[(1, "a", 10), (1, "b", 20), (2, "a", 30)])])
+ .await
+ .unwrap();
+ assert_eq!(
+ order_rows(&table).await,
+ vec![
+ order_row(1, "a", 10),
+ order_row(1, "b", 20),
+ order_row(2, "a", 30)
+ ]
+ );
+
+ let upsert = table
+ .upsert([order_batch(&[(1, "b", 21), (3, "c", 40), (4, "a", 50)])])
+ .await
+ .unwrap();
+ assert_eq!((upsert.num_updates, upsert.num_inserts), (1, 2));
+ assert_eq!(
+ order_rows(&table).await,
+ vec![
+ order_row(1, "a", 10),
+ order_row(1, "b", 21),
+ order_row(2, "a", 30),
+ order_row(3, "c", 40),
+ order_row(4, "a", 50),
+ ]
+ );
+
+ let deleted = table
+ .delete_keys([key("wr_order_number:2,wr_item_sk:a")])
+ .await
+ .unwrap();
+ assert_eq!(deleted.num_deletes, 1);
+
+ let deleted = table
+ .delete("_hoodie_record_key = 'wr_order_number:3,wr_item_sk:c'")
+ .await
+ .unwrap();
+ assert_eq!(deleted.num_deletes, 1);
+ let remaining = vec![
+ order_row(1, "a", 10),
+ order_row(1, "b", 21),
+ order_row(4, "a", 50),
+ ];
+ assert_eq!(order_rows(&table).await, remaining);
+
+ // Write filters split an IN list on every ',', which each composite key
holds.
+ let in_list = "_hoodie_record_key IN ('wr_order_number:1,wr_item_sk:a', \
+ 'wr_order_number:4,wr_item_sk:a')";
+ let set = order_batch(&[(1, "a", 0)]);
+ let errors = [
+ table.delete(in_list).await.unwrap_err(),
+ table.update(in_list, set).await.unwrap_err(),
+ ];
+ for err in errors {
+ assert!(
+ err.to_string()
+ .contains("cannot hold composite record keys"),
+ "{err}"
+ );
+ }
+ assert_eq!(order_rows(&table).await, remaining);
+
+ // One key field is not a whole key: every row with that item goes.
+ let deleted = table.delete("wr_item_sk = 'a'").await.unwrap();
+ assert_eq!(deleted.num_deletes, 2);
+ assert_eq!(order_rows(&table).await, vec![order_row(1, "b", 21)]);
+
+ // The guard runs before the SET type check, so one Int64 column probes
both fields.
+ for field in KEY_FIELDS {
+ let set =
+ RecordBatch::try_from_iter([(field,
Arc::new(Int64Array::from(vec![9])) as ArrayRef)])
+ .unwrap();
+ let err = table.update("wr_quantity = 21", set).await.unwrap_err();
+ assert!(
+ err.to_string().contains(&format!(
+ "cannot SET record-key or partition column '{field}'"
+ )),
+ "{err}"
+ );
+ }
+}
+
+#[tokio::test]
+async fn test_cow_composite_key_lifecycle_with_record_index() {
+ run_composite_key_lifecycle(TableTypeValue::CopyOnWrite, true).await;
+}
+
+#[tokio::test]
+async fn test_cow_composite_key_lifecycle_with_simple_index() {
+ run_composite_key_lifecycle(TableTypeValue::CopyOnWrite, false).await;
+}
+
+#[tokio::test]
+async fn test_mor_composite_key_lifecycle_with_record_index() {
+ run_composite_key_lifecycle(TableTypeValue::MergeOnRead, true).await;
+}
+
+#[tokio::test]
+async fn test_mor_composite_key_lifecycle_with_simple_index() {
+ run_composite_key_lifecycle(TableTypeValue::MergeOnRead, false).await;
+}
+
+/// Without meta fields, SimpleIndex derives the composite keys from the data
columns.
+#[tokio::test]
+async fn test_simple_index_tags_virtual_composite_keys() {
+ let dir = tempdir().unwrap();
+ let mut table = Table::create(dir.path().to_str().unwrap())
+ .with_table_name("web_returns")
+ .with_record_key_fields(KEY_FIELDS)
+ .with_populates_meta_fields(false)
+ .with_record_index(false)
+ .create()
+ .await
+ .unwrap();
+ let append = table
+ .append([order_batch(&[(1, "a", 10), (1, "b", 20)])])
+ .await
+ .unwrap();
+ let stored = key("wr_order_number:1,wr_item_sk:b");
+ let absent = key("wr_order_number:2,wr_item_sk:b");
+ let locations = SimpleIndex
+ .tag_location(&table, &[stored.clone(), absent.clone()])
+ .await
+ .unwrap();
+ let location = locations[&stored].as_ref().expect("stored key is tagged");
+ assert!(append.base_file_path.contains(&location.file_id));
+ assert!(locations[&absent].is_none());
+}
+
+#[tokio::test]
+async fn test_composite_key_partitioned_table_uses_complex_key_generator() {
+ let dir = tempdir().unwrap();
+ let mut table = Table::create(dir.path().to_str().unwrap())
+ .with_table_name("web_returns")
+ .with_record_key_fields(KEY_FIELDS)
+ .with_partition_fields(["wr_item_sk"])
+ .create()
+ .await
+ .unwrap();
+ let props =
std::fs::read_to_string(dir.path().join(".hoodie/hoodie.properties")).unwrap();
+ assert!(props.contains("hoodie.table.keygenerator.type=COMPLEX"));
+
+ table.append([order_batch(&[(1, "a", 10)])]).await.unwrap();
+ let upsert = table
+ .upsert([order_batch(&[(1, "a", 11), (2, "a", 20)])])
+ .await
+ .unwrap();
+ assert_eq!((upsert.num_updates, upsert.num_inserts), (1, 1));
+ assert_eq!(
+ order_rows(&table).await,
+ vec![order_row(1, "a", 11), order_row(2, "a", 20)]
+ );
+}
+
+#[tokio::test]
+async fn test_single_int64_record_key_lifecycle() {
+ let dir = tempdir().unwrap();
+ let mut table = Table::create(dir.path().to_str().unwrap())
+ .with_table_name("web_returns")
+ .with_record_key_fields(["wr_order_number"])
+ .create()
+ .await
+ .unwrap();
+ table
+ .append([order_batch(&[(1, "a", 10), (2, "a", 20)])])
+ .await
+ .unwrap();
+ let upsert = table
+ .upsert([order_batch(&[(2, "a", 21), (3, "a", 30)])])
+ .await
+ .unwrap();
+ assert_eq!((upsert.num_updates, upsert.num_inserts), (1, 1));
+ let key_and_quantity = |rows: Vec<(String, i64, String, i64)>| ->
Vec<(String, i64)> {
+ rows.into_iter()
+ .map(|(key, _, _, qty)| (key, qty))
+ .collect()
+ };
+ assert_eq!(
+ key_and_quantity(order_rows(&table).await),
+ [
+ ("1".to_string(), 10),
+ ("2".to_string(), 21),
+ ("3".to_string(), 30)
+ ]
+ );
+
+ assert_eq!(table.delete_keys([key("1")]).await.unwrap().num_deletes, 1);
+ // `02` is not the stored key text, but it still names row 2.
+ for filter in ["wr_order_number = 3", "wr_order_number = 02"] {
+ assert_eq!(
+ table.delete(filter).await.unwrap().num_deletes,
+ 1,
+ "{filter}"
+ );
+ }
+ assert!(order_rows(&table).await.is_empty());
+}
