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 e5a2552c feat(python): add write APIs to the Python binding (#878)
e5a2552c is described below
commit e5a2552cabf154b76d989ff4050d58392d6c587b
Author: Y Ethan Guo <[email protected]>
AuthorDate: Fri Oct 30 00:16:39 2026 -0700
feat(python): add write APIs to the Python binding (#878)
The Python package exposed only the read APIs, so `pip install hudi` could
not create or change a table. `HudiTable` gains `create` (wrapping
`TableCreateBuilder`, forwarding only the arguments a caller passes so the Rust
defaults stay the defaults), `append`, `upsert` (with optional `update_columns`
on copy-on-write), `update`, `delete`, `delete_keys`, `overwrite` and
`dynamic_partition_overwrite`, each returning a `HudiWriteResult`. Every write
releases the GIL and maps errors to Py [...]
Write input is any `pyarrow.Table`, `RecordBatch` or `RecordBatchReader`, a
list of batches, or any object implementing the Arrow PyCapsule stream
interface, which covers Polars and pandas 2.2+ without adding either as a
dependency. View, large and dictionary types are normalized to the types the
writer accepts (moving that into core is #864), nanosecond timestamps from
pandas and Polars become microseconds (a value with sub-microsecond precision
is rejected rather than truncated), an [...]
---
README.md | 82 ++-
python/hudi/__init__.py | 2 +
python/hudi/_internal.pyi | 204 +++++++-
python/pyproject.toml | 2 +
python/src/internal.rs | 1046 +++++++++++++++++++++++++++++++++++++-
python/src/lib.rs | 3 +-
python/tests/test_table_write.py | 774 ++++++++++++++++++++++++++++
7 files changed, 2101 insertions(+), 12 deletions(-)
diff --git a/README.md b/README.md
index ce078cd1..35ad9df0 100644
--- a/README.md
+++ b/README.md
@@ -322,7 +322,7 @@ std::vector<std::string> log_file_paths{};
ArrowArrayStream* stream_ptr =
reader->read_file_slice_from_paths("relative/path.parquet", log_file_paths);
```
-## Writing Tables (Rust)
+## Writing Tables
hudi-rs includes a native, single-node write path that is storage-compatible
with the Apache Hudi Spark writer: tables written here are readable — and
@@ -333,12 +333,30 @@ metadata table (file listings, column/partition stats,
record-level index) on
every commit. See [docs/writer-design.md](docs/writer-design.md) for the full
design.
-> [!NOTE]
-> Write APIs are currently Rust-only; Python and C++ bindings expose the read
-> APIs shown above.
+Python exposes the same write APIs on `HudiTable`, accepting PyArrow tables,
+record batches and readers, or any object implementing the Arrow PyCapsule
+stream interface (`__arrow_c_stream__`, e.g. polars and pandas >= 2.2
+DataFrames). Each write returns a `HudiWriteResult`. The C++ binding exposes
+the read APIs only.
### Create a Table
+#### Python
+
+```python
+from hudi import HudiTable
+
+table = HudiTable.create(
+ "/tmp/trips_table",
+ record_key="uuid",
+ partition_by="city",
+ ordering_fields="ts", # event-time ordering
+ table_type="cow", # or "mor"
+)
+```
+
+#### Rust
+
```rust
use hudi::table::Table;
use hudi::config::table::TableTypeValue;
@@ -361,6 +379,20 @@ an index lookup, splitting input into files bounded by
path: unpartitioned append-only tables with `hoodie.populate.meta.fields=false`
and no metadata table, one base file per call.
+#### Python
+
+```python
+import pyarrow as pa
+
+# `table` is the HudiTable created above; the input is plain Arrow data.
+batch = pa.record_batch(
+ {"uuid": ["a", "b"], "city": ["san_francisco"] * 2, "ts": [1, 1], "fare":
[30.0, 45.0]}
+)
+table.append(batch)
+```
+
+#### Rust
+
```rust
use arrow_array::RecordBatch;
@@ -376,6 +408,18 @@ update their file group in place (a rewritten base file on
COW, a log file on
MOR); new keys pack into existing small files first and then form new file
groups.
+#### Python
+
+```python
+result = table.upsert(batch)
+print(f"instant={result.instant_time} updates={result.num_updates}
inserts={result.num_inserts}")
+
+# Partial upsert (copy-on-write only): matched rows take only these columns
from the input.
+table.upsert(batch, update_columns=["fare"])
+```
+
+#### Rust
+
```rust
let result = table.upsert([batch]).await?;
println!(
@@ -386,6 +430,24 @@ println!(
### Update and Delete
+A write whose filter or keys match nothing succeeds without creating a commit
+and reports zero rows.
+
+#### Python
+
+```python
+# SQL-style update: each value must fit its column's type exactly.
+table.update("city = 'san_francisco'", {"fare": 99.0})
+
+# Delete by predicate (single comparison, or IN for multiple keys).
+table.delete("uuid IN ('a')")
+
+# Delete by key, optionally with the partition path.
+table.delete_keys([("b", "city=san_francisco")])
+```
+
+#### Rust
+
```rust
use hudi::table::HoodieKey;
@@ -405,6 +467,18 @@ table.delete_keys([HoodieKey {
### Overwrite
+#### Python
+
+```python
+# INSERT_OVERWRITE_TABLE: replace all file groups.
+table.overwrite(batch)
+
+# INSERT_OVERWRITE: replace only the partitions present in the input.
+table.dynamic_partition_overwrite(batch)
+```
+
+#### Rust
+
```rust
// INSERT_OVERWRITE_TABLE: replace all file groups.
table.overwrite([batch]).await?;
diff --git a/python/hudi/__init__.py b/python/hudi/__init__.py
index 09a9e2fe..81c10e94 100644
--- a/python/hudi/__init__.py
+++ b/python/hudi/__init__.py
@@ -27,6 +27,7 @@ from hudi._internal import (
HudiRecordBatchStream,
HudiTable,
HudiTimeline,
+ HudiWriteResult,
)
from hudi._internal import __version__ as __version__
from hudi.table.builder import HudiTableBuilder
@@ -45,4 +46,5 @@ __all__ = [
"HudiTableBuilder",
"HudiTableConfig",
"HudiTimeline",
+ "HudiWriteResult",
]
diff --git a/python/hudi/_internal.pyi b/python/hudi/_internal.pyi
index 2ec130f8..c5d3d9df 100644
--- a/python/hudi/_internal.pyi
+++ b/python/hudi/_internal.pyi
@@ -15,12 +15,24 @@
# specific language governing permissions and limitations
# under the License.
from dataclasses import dataclass
-from typing import ClassVar, Dict, List, Optional, Tuple
+from typing import Any, ClassVar, Dict, List, Optional, Protocol, Tuple, Union
import pyarrow # type: ignore[import-untyped]
__version__: str
+class _ArrowStreamExportable(Protocol):
+ def __arrow_c_stream__(self, requested_schema: Any = None) -> Any: ...
+
+WriteData = Union[
+ "pyarrow.Table",
+ "pyarrow.RecordBatch",
+ "pyarrow.RecordBatchReader",
+ List["pyarrow.RecordBatch"],
+ _ArrowStreamExportable,
+]
+"""Tabular input accepted by the :class:`HudiTable` write methods."""
+
def _config_keys() -> Dict[str, List[Tuple[str, str]]]: ...
class HudiQueryType:
@@ -255,6 +267,10 @@ class HudiTable:
"""
Represents a Hudi table and provides methods to interact with it.
+ A ``HudiTable`` must not be shared across threads while a write is running:
+ each write method takes exclusive access to the instance, so a concurrent
+ call on it raises ``RuntimeError``. Use one instance per thread.
+
Attributes:
base_uri (str): The base URI of the Hudi table.
options (Optional[Dict[str, str]]): Additional options for table
operations.
@@ -436,6 +452,192 @@ class HudiTable:
is not yet supported and raises ``HudiCoreError``.
"""
...
+ @staticmethod
+ def create(
+ base_uri: str,
+ *,
+ record_key: Union[str, List[str]],
+ partition_by: Optional[Union[str, List[str]]] = None,
+ ordering_fields: Optional[Union[str, List[str]]] = None,
+ table_type: str = "cow",
+ table_name: Optional[str] = None,
+ table_version: Optional[int] = None,
+ record_index: Optional[bool] = None,
+ metadata_table: Optional[bool] = None,
+ column_stats: Optional[bool] = None,
+ partition_stats: Optional[bool] = None,
+ hive_style_partitioning: Optional[bool] = None,
+ hudi_options: Optional[Dict[str, str]] = None,
+ storage_options: Optional[Dict[str, str]] = None,
+ ) -> "HudiTable":
+ """
+ Create a new Hudi table at ``base_uri`` and open it for reads and
writes.
+
+ Arguments left as ``None`` keep the native defaults: table version 9,
+ metadata table on with the record index, column stats and partition
+ stats, and hive-style partition paths. Turning ``metadata_table`` off
also
+ turns off the indexes that live in it.
+
+ Parameters:
+ base_uri (str): Location of the new table.
+ record_key (Union[str, List[str]]): Record key field(s).
+ partition_by (Optional[Union[str, List[str]]]): Partition
field(s); unpartitioned if omitted.
+ ordering_fields (Optional[Union[str, List[str]]]): Field(s) that
pick the winning
+ record when a key is written more than once; commit order if
omitted.
+ table_type (str): ``"cow"``/``"COPY_ON_WRITE"`` or
``"mor"``/``"MERGE_ON_READ"``,
+ case-insensitive.
+ table_name (Optional[str]): Defaults to the last path segment of
``base_uri``.
+ table_version (Optional[int]): Table version to write, 8 or 9.
+ record_index (Optional[bool]): Record-level index in the metadata
table.
+ metadata_table (Optional[bool]): Metadata table.
+ column_stats (Optional[bool]): Column stats index in the metadata
table.
+ partition_stats (Optional[bool]): Partition stats index in the
metadata table.
+ hive_style_partitioning (Optional[bool]): Use ``field=value``
partition paths.
+ hudi_options (Optional[Dict[str, str]]): Extra ``hoodie.*`` table
configs, written into
+ ``hoodie.properties``. Read options belong on the reader, not
here.
+ storage_options (Optional[Dict[str, str]]): Storage options such
as cloud credentials.
+
+ Returns:
+ HudiTable: The new table.
+
+ Raises:
+ ValueError: If ``table_type`` is not recognized.
+ HudiCoreError: If a table already exists at ``base_uri`` or the
+ options are inconsistent.
+ """
+ ...
+ def append(self, data: WriteData) -> "HudiWriteResult":
+ """
+ Insert ``data`` as one commit without looking up existing keys.
+
+ Parameters:
+ data (WriteData): A pyarrow ``Table``, ``RecordBatch``,
``RecordBatchReader``,
+ a list of ``RecordBatch``, or any object implementing
``__arrow_c_stream__``
+ (for example a polars or pandas >= 2.2 DataFrame).
+
+ Returns:
+ HudiWriteResult: ``num_writes`` and ``num_inserts`` are the rows
written.
+ """
+ ...
+ def upsert(
+ self,
+ data: WriteData,
+ update_columns: Optional[List[str]] = None,
+ ) -> "HudiWriteResult":
+ """
+ Update rows whose record key exists and insert the rest.
+
+ Parameters:
+ data (WriteData): Rows to upsert; see :meth:`append` for accepted
types.
+ update_columns (Optional[List[str]]): When set, matched rows only
take these
+ columns from ``data`` and keep their other values.
Copy-on-write tables only.
+
+ Returns:
+ HudiWriteResult: The commit and its update/insert counts.
+ """
+ ...
+ def update(
+ self,
+ predicate: str,
+ updates: Union[
+ Dict[str, Any],
+ "pyarrow.RecordBatch",
+ "pyarrow.Table",
+ _ArrowStreamExportable,
+ ],
+ ) -> "HudiWriteResult":
+ """
+ Set columns on the rows matching ``predicate``.
+
+ Parameters:
+ predicate (str): A single binary predicate such as ``"city =
'nyc'"``.
+ updates (Union[Dict[str, Any], pyarrow.RecordBatch,
pyarrow.Table]): Column values to
+ set, either as a dict or as a one-row batch carrying only
those columns.
+ Dict values must fit their column's type exactly, except that
a float32
+ column rounds to the nearest float32; a bool fits only a bool
column. Date
+ columns take ``datetime.date`` and timestamp columns take
+ ``datetime.datetime``; a tz-aware datetime needs a timestamp
column with a
+ time zone. Struct fields missing from a dict value are set to
null.
+
+ Returns:
+ HudiWriteResult: No commit is created when nothing matches.
+
+ Raises:
+ ValueError: For dict input, if a key is not a str column of the
table schema
+ or a value does not fit its column; for batch input, if it
yields no record
+ batches (such as an empty ``pyarrow.Table``).
+ HudiCoreError: For batch input that does not match the table
schema or does not
+ hold exactly one row.
+ """
+ ...
+ def delete(self, predicate: str) -> "HudiWriteResult":
+ """
+ Delete the rows matching ``predicate``.
+
+ Parameters:
+ predicate (str): A single binary predicate such as ``"id IN ('t1',
't2')"``.
+ Predicates on the record key use the index; others scan the
table.
+
+ Returns:
+ HudiWriteResult: No commit is created when nothing matches.
+ """
+ ...
+ def delete_keys(self, keys: List[Union[str, Tuple[str, str]]]) ->
"HudiWriteResult":
+ """
+ Delete records by record key.
+
+ Parameters:
+ keys (List[Union[str, Tuple[str, str]]]): Record keys, or
+ ``(record_key, partition_path)`` tuples. ``partition_path`` is
the
+ table-relative path as stored, e.g. ``city=nyc`` with
hive-style
+ partitioning. A bare key is looked up in every partition.
+
+ Returns:
+ HudiWriteResult: No commit is created when no key exists.
+ """
+ ...
+ def overwrite(self, data: WriteData) -> "HudiWriteResult":
+ """
+ Replace all rows in the table with ``data``.
+
+ Parameters:
+ data (WriteData): Replacement rows; see :meth:`append` for
accepted types.
+
+ Returns:
+ HudiWriteResult: The replace commit.
+ """
+ ...
+ def dynamic_partition_overwrite(self, data: WriteData) ->
"HudiWriteResult":
+ """
+ Replace only the partitions present in ``data`` in a partitioned table.
+
+ Parameters:
+ data (WriteData): Replacement rows; see :meth:`append` for
accepted types.
+
+ Returns:
+ HudiWriteResult: The replace commit.
+ """
+ ...
+
+@dataclass(init=False)
+class HudiWriteResult:
+ """
+ Outcome of a write on :class:`HudiTable`.
+
+ Attributes:
+ instant_time (Optional[str]): Instant time of the completed commit, or
``None``
+ when no commit was made.
+ num_writes (int): Rows written to data files.
+ num_updates (int): Input rows that replaced existing records.
+ num_inserts (int): Input rows that were new records.
+ num_deletes (int): Records deleted.
+ """
+
+ instant_time: Optional[str]
+ num_writes: int
+ num_updates: int
+ num_inserts: int
+ num_deletes: int
@dataclass(init=False)
class HudiTimeline:
diff --git a/python/pyproject.toml b/python/pyproject.toml
index 16c5d1b3..4c25fec2 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -48,6 +48,8 @@ devel = [
"ruff==0.17.1",
"mypy==2.4.0",
"pre-commit",
+ "pandas>=2.2",
+ "polars>=1.0",
]
lint = [
"ruff==0.17.1",
diff --git a/python/src/internal.rs b/python/src/internal.rs
index 98639e12..8c3adbf7 100644
--- a/python/src/internal.rs
+++ b/python/src/internal.rs
@@ -17,13 +17,21 @@
* under the License.
*/
-use arrow::pyarrow::ToPyArrow;
-use arrow::record_batch::RecordBatch;
+use arrow::array::{
+ Array, ArrayData, ArrayRef, AsArray, GenericListArray, OffsetSizeTrait,
make_array,
+};
+use arrow::compute::{cast, concat_batches};
+use arrow::datatypes::TimestampNanosecondType;
+use arrow::datatypes::{DataType, FieldRef, Schema, TimeUnit};
+use arrow::ffi_stream::ArrowArrayStreamReader;
+use arrow::pyarrow::{FromPyArrow, ToPyArrow};
+use arrow::record_batch::{RecordBatch, RecordBatchOptions};
use futures::StreamExt;
use futures::stream::BoxStream;
use std::collections::HashMap;
use std::convert::From;
use std::path::PathBuf;
+use std::str::FromStr;
use std::sync::{Arc, OnceLock};
use tokio::runtime::Runtime;
use tokio::sync::Mutex;
@@ -33,7 +41,7 @@ use datafusion::error::DataFusionError;
use hudi::config::plan::HudiPlanConfig;
use hudi::config::read::HudiReadConfig;
-use hudi::config::table::HudiTableConfig;
+use hudi::config::table::{HudiTableConfig, TableTypeValue};
use hudi::error::CoreError;
use hudi::error::Result as HudiResult;
use hudi::file_group::FileGroup;
@@ -41,12 +49,15 @@ use hudi::file_group::file_slice::FileSlice;
use hudi::file_group::reader::FileGroupReader;
use hudi::storage::error::StorageError;
use hudi::table::builder::TableBuilder;
-use hudi::table::{QueryType, ReadOptions, Table};
+use hudi::table::{AppendResult, QueryType, ReadOptions, Table};
use hudi::timeline::Timeline;
use hudi::timeline::instant::Instant;
+use hudi::{HoodieKey, UpsertOptions, WriteResult};
use pyo3::exceptions::PyException;
+use pyo3::exceptions::{PyTypeError, PyValueError};
use pyo3::prelude::*;
-use pyo3::{PyErr, PyResult, Python, create_exception, pyclass, pyfunction,
pymethods};
+use pyo3::types::{PyBool, PyDict, PyList, PyString, PyTuple};
+use pyo3::{PyErr, PyResult, Python, create_exception, intern, pyclass,
pyfunction, pymethods};
use std::error::Error;
type RecordBatchBoxStream = BoxStream<'static, HudiResult<RecordBatch>>;
@@ -849,6 +860,707 @@ impl HudiTable {
})?;
Ok(HudiRecordBatchStream::from_stream(stream))
}
+
+ #[staticmethod]
+ #[pyo3(signature = (
+ base_uri,
+ *,
+ record_key,
+ partition_by=None,
+ ordering_fields=None,
+ table_type="cow",
+ table_name=None,
+ table_version=None,
+ record_index=None,
+ metadata_table=None,
+ column_stats=None,
+ partition_stats=None,
+ hive_style_partitioning=None,
+ hudi_options=None,
+ storage_options=None,
+ ))]
+ #[allow(clippy::too_many_arguments)]
+ fn create(
+ py: Python,
+ base_uri: &str,
+ record_key: OneOrMany,
+ partition_by: Option<OneOrMany>,
+ ordering_fields: Option<OneOrMany>,
+ table_type: &str,
+ table_name: Option<String>,
+ table_version: Option<isize>,
+ record_index: Option<bool>,
+ metadata_table: Option<bool>,
+ column_stats: Option<bool>,
+ partition_stats: Option<bool>,
+ hive_style_partitioning: Option<bool>,
+ hudi_options: Option<HashMap<String, String>>,
+ storage_options: Option<HashMap<String, String>>,
+ ) -> PyResult<Self> {
+ let table_type = TableTypeValue::from_str(table_type).map_err(|_| {
+ PyValueError::new_err(format!(
+ "invalid table_type '{table_type}': expected 'cow'
(COPY_ON_WRITE) or 'mor' (MERGE_ON_READ)"
+ ))
+ })?;
+ let table_name = table_name.unwrap_or_else(||
last_path_segment(base_uri));
+ let mut builder = Table::create(base_uri)
+ .with_table_name(table_name)
+ .with_table_type(table_type)
+ .with_record_key_fields(record_key.into_vec());
+ if let Some(fields) = partition_by {
+ builder = builder.with_partition_fields(fields.into_vec());
+ }
+ if let Some(fields) = ordering_fields {
+ builder = builder.with_ordering_fields(fields.into_vec());
+ }
+ // Order matters: the version and metadata setters adjust the index
+ // flags that follow them, so explicit index choices must come last.
+ if let Some(version) = table_version {
+ builder = builder.with_table_version(version);
+ }
+ if let Some(enabled) = metadata_table {
+ builder = builder.with_metadata(enabled);
+ }
+ if let Some(enabled) = record_index {
+ builder = builder.with_record_index(enabled);
+ }
+ if let Some(enabled) = column_stats {
+ builder = builder.with_column_stats(enabled);
+ }
+ if let Some(enabled) = partition_stats {
+ builder = builder.with_partition_stats(enabled);
+ }
+ if let Some(enabled) = hive_style_partitioning {
+ builder = builder.with_hive_style_partitioning(enabled);
+ }
+ for (k, v) in hudi_options.unwrap_or_default() {
+ builder = builder.with_option(k, v);
+ }
+ for (k, v) in storage_options.unwrap_or_default() {
+ builder = builder.with_storage_option(k, v);
+ }
+ let inner = py.detach(||
rt().block_on(builder.create()).map_err(PythonError::from))?;
+ Ok(HudiTable { inner })
+ }
+
+ fn append(&mut self, data: &Bound<'_, PyAny>, py: Python) ->
PyResult<HudiWriteResult> {
+ let batches = to_batches(data)?;
+ let result = py.detach(|| {
+ rt().block_on(self.inner.append(batches))
+ .map_err(PythonError::from)
+ })?;
+ Ok(HudiWriteResult::from(result))
+ }
+
+ #[pyo3(signature = (data, update_columns=None))]
+ fn upsert(
+ &mut self,
+ data: &Bound<'_, PyAny>,
+ update_columns: Option<Vec<String>>,
+ py: Python,
+ ) -> PyResult<HudiWriteResult> {
+ let batches = to_batches(data)?;
+ let options = UpsertOptions { update_columns };
+ let result = py.detach(|| {
+ rt().block_on(self.inner.upsert_with(batches, options))
+ .map_err(PythonError::from)
+ })?;
+ Ok(HudiWriteResult::from(result))
+ }
+
+ fn update(
+ &mut self,
+ predicate: &str,
+ updates: &Bound<'_, PyAny>,
+ py: Python,
+ ) -> PyResult<HudiWriteResult> {
+ let updates = if let Ok(dict) = updates.cast::<PyDict>() {
+ let schema = py.detach(|| {
+ rt().block_on(self.inner.get_schema())
+ .map_err(PythonError::from)
+ })?;
+ one_row_batch_from_dict(py, &schema, dict)?
+ } else {
+ let batches = to_batches(updates)?;
+ let Some(first) = batches.first() else {
+ return Err(PyValueError::new_err(
+ "`updates` must contain exactly one row, got no batches",
+ ));
+ };
+ concat_batches(&first.schema(), &batches).map_err(|e| {
+ PyValueError::new_err(format!("failed to combine `updates`
batches: {e}"))
+ })?
+ };
+ let result = py.detach(|| {
+ rt().block_on(self.inner.update(predicate, updates))
+ .map_err(PythonError::from)
+ })?;
+ Ok(HudiWriteResult::from(result))
+ }
+
+ fn delete(&mut self, predicate: &str, py: Python) ->
PyResult<HudiWriteResult> {
+ let result = py.detach(|| {
+ rt().block_on(self.inner.delete(predicate))
+ .map_err(PythonError::from)
+ })?;
+ Ok(HudiWriteResult::from(result))
+ }
+
+ fn delete_keys(&mut self, keys: &Bound<'_, PyAny>, py: Python) ->
PyResult<HudiWriteResult> {
+ let keys = to_hoodie_keys(keys)?;
+ let result = py.detach(|| {
+ rt().block_on(self.inner.delete_keys(keys))
+ .map_err(PythonError::from)
+ })?;
+ Ok(HudiWriteResult::from(result))
+ }
+
+ fn overwrite(&mut self, data: &Bound<'_, PyAny>, py: Python) ->
PyResult<HudiWriteResult> {
+ let batches = to_batches(data)?;
+ let result = py.detach(|| {
+ rt().block_on(self.inner.overwrite(batches))
+ .map_err(PythonError::from)
+ })?;
+ Ok(HudiWriteResult::from(result))
+ }
+
+ fn dynamic_partition_overwrite(
+ &mut self,
+ data: &Bound<'_, PyAny>,
+ py: Python,
+ ) -> PyResult<HudiWriteResult> {
+ let batches = to_batches(data)?;
+ let result = py.detach(|| {
+ rt().block_on(self.inner.dynamic_partition_overwrite(batches))
+ .map_err(PythonError::from)
+ })?;
+ Ok(HudiWriteResult::from(result))
+ }
+}
+
+/// Result of a table write, shared by every write verb.
+#[cfg(not(tarpaulin_include))]
+#[pyclass(frozen, get_all)]
+pub struct HudiWriteResult {
+ instant_time: Option<String>,
+ num_writes: usize,
+ num_updates: usize,
+ num_inserts: usize,
+ num_deletes: usize,
+}
+
+#[cfg(not(tarpaulin_include))]
+#[pymethods]
+impl HudiWriteResult {
+ fn __repr__(&self) -> String {
+ let instant_time = self
+ .instant_time
+ .as_ref()
+ .map_or_else(|| "None".to_string(), |t| format!("'{t}'"));
+ format!(
+ "HudiWriteResult(instant_time={instant_time}, num_writes={},
num_updates={}, num_inserts={}, num_deletes={})",
+ self.num_writes, self.num_updates, self.num_inserts,
self.num_deletes,
+ )
+ }
+}
+
+impl From<WriteResult> for HudiWriteResult {
+ fn from(r: WriteResult) -> Self {
+ Self {
+ instant_time: committed(r.instant),
+ num_writes: r.num_writes,
+ num_updates: r.num_updates,
+ num_inserts: r.num_inserts,
+ num_deletes: r.num_deletes,
+ }
+ }
+}
+
+impl From<AppendResult> for HudiWriteResult {
+ fn from(r: AppendResult) -> Self {
+ Self {
+ instant_time: committed(r.instant),
+ num_writes: r.num_rows,
+ num_updates: 0,
+ num_inserts: r.num_rows,
+ num_deletes: 0,
+ }
+ }
+}
+
+/// The core reports a write that made no commit with an empty instant.
+fn committed(instant: String) -> Option<String> {
+ (!instant.is_empty()).then_some(instant)
+}
+
+/// A Python argument that is either one string or a list of strings.
+#[cfg(not(tarpaulin_include))]
+#[derive(FromPyObject)]
+enum OneOrMany {
+ One(String),
+ Many(Vec<String>),
+}
+
+#[cfg(not(tarpaulin_include))]
+impl OneOrMany {
+ fn into_vec(self) -> Vec<String> {
+ match self {
+ OneOrMany::One(s) => vec![s],
+ OneOrMany::Many(v) => v,
+ }
+ }
+}
+
+fn last_path_segment(base_uri: &str) -> String {
+ base_uri
+ .trim_end_matches('/')
+ .rsplit('/')
+ .next()
+ .unwrap_or_default()
+ .to_string()
+}
+
+/// Convert any supported Python tabular object into record batches.
+///
+/// Objects implementing the Arrow PyCapsule interface (`__arrow_c_stream__`,
+/// `__arrow_c_array__`) are imported zero-copy, which covers pyarrow, polars
+/// and pandas >= 2.2. The pyarrow class checks serve pyarrow releases that
+/// predate the PyCapsule interface.
+#[cfg(not(tarpaulin_include))]
+fn to_batches(data: &Bound<'_, PyAny>) -> PyResult<Vec<RecordBatch>> {
+ collect_batches(data)?
+ .into_iter()
+ .map(|batch| normalize_batch(batch).map_err(PyValueError::new_err))
+ .collect()
+}
+
+/// Rewrite a batch into types the writer accepts: drop the columns pandas
+/// adds for an unnamed index, and cast view and large string, binary and
+/// list types (Polars exports `Utf8View`, pandas exports `LargeUtf8`) and
+/// dictionaries to their plain forms at any nesting depth.
+///
+/// Errors are plain strings so the function runs without an interpreter.
+fn normalize_batch(batch: RecordBatch) -> Result<RecordBatch, String> {
+ let schema = batch.schema();
+ let from_pandas = schema.metadata().contains_key("pandas");
+ 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 schema.fields().iter().zip(batch.columns()) {
+ if from_pandas && is_unnamed_pandas_index(field.name()) {
+ changed = true;
+ continue;
+ }
+ match writable_field(field) {
+ Some(target) => {
+ // The cast to microseconds truncates silently.
+ if has_sub_microsecond(column.as_ref()) {
+ return Err(format!(
+ "column '{}' has timestamps with sub-microsecond
precision, \
+ which Hudi tables cannot store",
+ field.name()
+ ));
+ }
+ let casted = cast(column.as_ref(),
target.data_type()).map_err(|e| {
+ format!(
+ "failed to cast column '{}' from {} to {}: {e}",
+ field.name(),
+ field.data_type(),
+ target.data_type()
+ )
+ })?;
+ fields.push(target);
+ columns.push(casted);
+ changed = true;
+ }
+ None => {
+ fields.push(field.clone());
+ columns.push(column.clone());
+ }
+ }
+ }
+ if !changed {
+ return Ok(batch);
+ }
+ let options =
RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
+ RecordBatch::try_new_with_options(
+ Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())),
+ columns,
+ &options,
+ )
+ .map_err(|e| format!("failed to rebuild record batch: {e}"))
+}
+
+/// Whether a nanosecond timestamp anywhere in `array` is not a whole
microsecond.
+///
+/// Only values the array references are checked: a sliced list keeps its
+/// whole child buffer and a dictionary may hold unused values.
+fn has_sub_microsecond(array: &dyn Array) -> bool {
+ match array.data_type() {
+ DataType::Timestamp(TimeUnit::Nanosecond, _) => array
+ .as_primitive::<TimestampNanosecondType>()
+ .iter()
+ .flatten()
+ .any(|v| v % 1000 != 0),
+ DataType::List(_) => {
+
has_sub_microsecond(referenced_values(array.as_list::<i32>()).as_ref())
+ }
+ DataType::LargeList(_) => {
+
has_sub_microsecond(referenced_values(array.as_list::<i64>()).as_ref())
+ }
+ DataType::Struct(_) => array
+ .as_struct()
+ .columns()
+ .iter()
+ .any(|c| has_sub_microsecond(c.as_ref())),
+ DataType::Dictionary(_, value) => {
+ cast(array, value).is_ok_and(|plain|
has_sub_microsecond(plain.as_ref()))
+ }
+ _ => false,
+ }
+}
+
+fn referenced_values<O: OffsetSizeTrait>(list: &GenericListArray<O>) ->
ArrayRef {
+ let offsets = list.value_offsets();
+ let start = offsets[0].as_usize();
+ let end = offsets[offsets.len() - 1].as_usize();
+ list.values().slice(start, end - start)
+}
+
+/// pandas exports an unnamed index as `__index_level_N__`; a named index
+/// keeps its name and may be the record key.
+fn is_unnamed_pandas_index(name: &str) -> bool {
+ name.strip_prefix("__index_level_")
+ .and_then(|rest| rest.strip_suffix("__"))
+ .is_some_and(|n| !n.is_empty() && n.bytes().all(|b|
b.is_ascii_digit()))
+}
+
+/// The writable form of `field`, or `None` when it is already writable.
+/// Names, nullability and metadata are kept, including list item names.
+fn writable_field(field: &FieldRef) -> Option<FieldRef> {
+ writable_type(field.data_type()).map(|dt|
Arc::new(field.as_ref().clone().with_data_type(dt)))
+}
+
+fn writable_type(dt: &DataType) -> Option<DataType> {
+ match dt {
+ DataType::Utf8View | DataType::LargeUtf8 => Some(DataType::Utf8),
+ DataType::BinaryView | DataType::LargeBinary => Some(DataType::Binary),
+ DataType::Timestamp(TimeUnit::Nanosecond, tz) => {
+ Some(DataType::Timestamp(TimeUnit::Microsecond, tz.clone()))
+ }
+ DataType::Dictionary(_, value) => {
+ Some(writable_type(value).unwrap_or_else(||
value.as_ref().clone()))
+ }
+ DataType::List(item) => writable_field(item).map(DataType::List),
+ DataType::LargeList(item) => Some(DataType::List(
+ writable_field(item).unwrap_or_else(|| item.clone()),
+ )),
+ DataType::Struct(children) => {
+ let mapped: Vec<Option<FieldRef>> =
children.iter().map(writable_field).collect();
+ if mapped.iter().all(Option::is_none) {
+ return None;
+ }
+ Some(DataType::Struct(
+ children
+ .iter()
+ .zip(mapped)
+ .map(|(child, m)| m.unwrap_or_else(|| child.clone()))
+ .collect(),
+ ))
+ }
+ _ => None,
+ }
+}
+
+#[cfg(not(tarpaulin_include))]
+fn collect_batches(data: &Bound<'_, PyAny>) -> PyResult<Vec<RecordBatch>> {
+ let py = data.py();
+ if let Ok(list) = data.cast::<PyList>() {
+ return list
+ .iter()
+ .enumerate()
+ .map(|(i, item)| {
+ RecordBatch::from_pyarrow_bound(&item).map_err(|e| {
+ PyTypeError::new_err(format!(
+ "data[{i}] is not a pyarrow.RecordBatch or an Arrow
array-capsule object: {e}"
+ ))
+ })
+ })
+ .collect();
+ }
+ if data.hasattr(intern!(py, "__arrow_c_stream__"))? {
+ return collect_stream(data);
+ }
+ if data.hasattr(intern!(py, "__arrow_c_array__"))? {
+ return Ok(vec![RecordBatch::from_pyarrow_bound(data)?]);
+ }
+ let pa = py.import(intern!(py, "pyarrow"))?;
+ if data.is_instance(&pa.getattr(intern!(py, "RecordBatchReader"))?)? {
+ return collect_stream(data);
+ }
+ if data.is_instance(&pa.getattr(intern!(py, "RecordBatch"))?)? {
+ return Ok(vec![RecordBatch::from_pyarrow_bound(data)?]);
+ }
+ if data.is_instance(&pa.getattr(intern!(py, "Table"))?)? {
+ return collect_batches(&data.call_method0(intern!(py, "to_batches"))?);
+ }
+ Err(PyTypeError::new_err(format!(
+ "unsupported data type '{}': expected a pyarrow Table, RecordBatch,
RecordBatchReader, \
+ a list of RecordBatch, or an object implementing __arrow_c_stream__",
+ data.get_type().name()?
+ )))
+}
+
+#[cfg(not(tarpaulin_include))]
+fn collect_stream(data: &Bound<'_, PyAny>) -> PyResult<Vec<RecordBatch>> {
+ ArrowArrayStreamReader::from_pyarrow_bound(data)?
+ .collect::<Result<Vec<_>, _>>()
+ .map_err(|e| PyValueError::new_err(format!("failed to read Arrow
stream: {e}")))
+}
+
+/// Build the one-row `SET` batch for `update`, typing each value by the
+/// table schema so that e.g. a Python int lands in a float64 column.
+#[cfg(not(tarpaulin_include))]
+fn one_row_batch_from_dict(
+ py: Python,
+ schema: &Schema,
+ updates: &Bound<'_, PyDict>,
+) -> PyResult<RecordBatch> {
+ if updates.is_empty() {
+ return Err(PyValueError::new_err(
+ "`updates` must name at least one column",
+ ));
+ }
+ let pa = py.import(intern!(py, "pyarrow"))?;
+ let mut fields = Vec::with_capacity(updates.len());
+ let mut columns = Vec::with_capacity(updates.len());
+ for (key, value) in updates.iter() {
+ let name: String = key.extract().map_err(|_| {
+ PyValueError::new_err(format!(
+ "update column names must be str, got {}",
+ key.repr().map(|r| r.to_string()).unwrap_or_default()
+ ))
+ })?;
+ let field = schema.field_with_name(&name).map_err(|_| {
+ PyValueError::new_err(format!("update column '{name}' is not in
the table schema"))
+ })?;
+ let value = unwrap_scalars(&value, &pa.getattr(intern!(py,
"Scalar"))?)?;
+ let kwargs = PyDict::new(py);
+ kwargs.set_item(intern!(py, "type"),
field.data_type().to_pyarrow(py)?)?;
+ let array = pa
+ .call_method(
+ intern!(py, "array"),
+ (PyList::new(py, [&value])?,),
+ Some(&kwargs),
+ )
+ .map_err(|e| {
+ PyValueError::new_err(format!(
+ "cannot convert update value for column '{name}' to {}:
{e}",
+ field.data_type()
+ ))
+ })?;
+ let stored = array
+ .get_item(0)
+ .and_then(|v| v.call_method0(intern!(py, "as_py")))
+ .map_err(|e| {
+ PyValueError::new_err(format!(
+ "cannot read back update value for column '{name}' as {}:
{e}",
+ field.data_type()
+ ))
+ })?;
+ if let Some(reason) = misfit(field.data_type(), &value, &stored)? {
+ return Err(PyValueError::new_err(format!(
+ "update value {} for column '{name}' does not fit {} exactly:
{reason} \
+ (stored as {})",
+ value.repr()?,
+ field.data_type(),
+ stored.repr()?
+ )));
+ }
+ columns.push(make_array(ArrayData::from_pyarrow_bound(&array)?));
+ fields.push(field.clone());
+ }
+ RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
+ .map_err(|e| PyValueError::new_err(format!("invalid `updates`: {e}")))
+}
+
+/// `value` with every `pyarrow.Scalar`, at any depth of dicts, lists and
+/// tuples, replaced by its `as_py()`.
+#[cfg(not(tarpaulin_include))]
+fn unwrap_scalars<'py>(
+ value: &Bound<'py, PyAny>,
+ scalar: &Bound<'py, PyAny>,
+) -> PyResult<Bound<'py, PyAny>> {
+ let py = value.py();
+ if value.is_instance(scalar)? {
+ return value.call_method0(intern!(py, "as_py"));
+ }
+ if let Ok(dict) = value.cast::<PyDict>() {
+ let out = PyDict::new(py);
+ for (k, v) in dict.iter() {
+ out.set_item(k, unwrap_scalars(&v, scalar)?)?;
+ }
+ return Ok(out.into_any());
+ }
+ let is_tuple = value.is_instance_of::<PyTuple>();
+ if is_tuple || value.is_instance_of::<PyList>() {
+ let items = value
+ .try_iter()?
+ .map(|v| unwrap_scalars(&v?, scalar))
+ .collect::<PyResult<Vec<_>>>()?;
+ return Ok(if is_tuple {
+ PyTuple::new(py, items)?.into_any()
+ } else {
+ PyList::new(py, items)?.into_any()
+ });
+ }
+ Ok(value.clone())
+}
+
+/// Why `value` does not survive the typed `pa.array` constructor for `dt`,
+/// given `stored`, the `as_py()` of what it produced; `None` when it fits.
+///
+/// The constructor coerces silently: 1.5 into int64 stores 1, a numpy bool
+/// into int64 stores 1, a datetime into date32 drops the time, a number into
+/// a timestamp is read as epoch units, a coarser timestamp unit drops
+/// sub-unit digits, and unknown struct keys are dropped. Python equality
+/// accepts exact conversions (3 == 3.0);
+/// timestamps are not compared whole because a naive input never equals its
+/// UTC-aware stored form.
+#[cfg(not(tarpaulin_include))]
+fn misfit(
+ dt: &DataType,
+ value: &Bound<'_, PyAny>,
+ stored: &Bound<'_, PyAny>,
+) -> PyResult<Option<String>> {
+ if value.is_none() || stored.is_none() {
+ return Ok(None);
+ }
+ let py = value.py();
+ match dt {
+ DataType::Struct(children) => {
+ let by_name = value.cast::<PyDict>().ok();
+ if let Some(d) = &by_name {
+ for key in d.keys() {
+ let known = key
+ .extract::<String>()
+ .is_ok_and(|k| children.iter().any(|c| c.name() ==
&k));
+ if !known {
+ return Ok(Some(format!("unknown struct field {}",
key.repr()?)));
+ }
+ }
+ }
+ let mut positional = match by_name {
+ Some(_) => None,
+ None => Some(value.try_iter()?),
+ };
+ for child in children {
+ let v = match (&by_name, positional.as_mut()) {
+ (Some(d), _) => d.get_item(child.name())?,
+ (None, Some(it)) => it.next().transpose()?,
+ (None, None) => None,
+ };
+ if let Some(v) = v {
+ let s = stored.get_item(child.name())?;
+ if let Some(reason) = misfit(child.data_type(), &v, &s)? {
+ return Ok(Some(reason));
+ }
+ }
+ }
+ Ok(None)
+ }
+ DataType::List(item) | DataType::LargeList(item) => {
+ for (v, s) in value.try_iter()?.zip(stored.try_iter()?) {
+ if let Some(reason) = misfit(item.data_type(), &v?, &s?)? {
+ return Ok(Some(reason));
+ }
+ }
+ Ok(None)
+ }
+ dt if dt.is_numeric() && is_bool(value)? => Ok(Some("a bool is not a
number".into())),
+ DataType::Float16 | DataType::Float32 => {
+ let (v, s): (f64, f64) = (value.extract()?, stored.extract()?);
+ Ok((v.is_finite() && !s.is_finite())
+ .then(|| "the value overflows the float type".into()))
+ }
+ dt if dt.is_numeric() => {
+ let both_nan = stored.ne(stored)? && value.ne(value)?;
+ Ok((!(stored.eq(value)? || both_nan)).then(|| "the value would
change".into()))
+ }
+ DataType::Date32 | DataType::Date64 => {
+ let datetime = py.import(intern!(py, "datetime"))?;
+ let is_date = value.is_instance(&datetime.getattr(intern!(py,
"date"))?)?;
+ let is_datetime = value.is_instance(&datetime.getattr(intern!(py,
"datetime"))?)?;
+ Ok((!is_date || is_datetime).then(|| "expected a
datetime.date".into()))
+ }
+ DataType::Timestamp(unit, tz) => {
+ let datetime = py
+ .import(intern!(py, "datetime"))?
+ .getattr(intern!(py, "datetime"))?;
+ if !value.is_instance(&datetime)? {
+ return Ok(Some("expected a datetime.datetime".into()));
+ }
+ if tz.is_none() && !value.getattr(intern!(py,
"tzinfo"))?.is_none() {
+ return Ok(Some(
+ "a tz-aware datetime is ambiguous for a timestamp without
a time zone".into(),
+ ));
+ }
+ let nanos = intern!(py, "nanosecond");
+ let micros = intern!(py, "microsecond");
+ let truncated = (*unit != TimeUnit::Nanosecond
+ && value.hasattr(nanos)?
+ && !value.getattr(nanos)?.eq(0)?)
+ || !stored.getattr(micros)?.eq(value.getattr(micros)?)?;
+ Ok(truncated.then(|| "the timestamp unit would truncate the
value".into()))
+ }
+ _ => Ok(None),
+ }
+}
+
+#[cfg(not(tarpaulin_include))]
+fn is_bool(value: &Bound<'_, PyAny>) -> PyResult<bool> {
+ if value.is_instance_of::<PyBool>() {
+ return Ok(true);
+ }
+ let ty = value.get_type();
+ let module: String = ty.getattr(intern!(value.py(),
"__module__"))?.extract()?;
+ let name = ty.name()?;
+ Ok(module == "numpy" && (name == "bool" || name == "bool_"))
+}
+
+/// Accept record keys as `str` (empty partition path) or `(record_key,
partition_path)`.
+#[cfg(not(tarpaulin_include))]
+fn to_hoodie_keys(keys: &Bound<'_, PyAny>) -> PyResult<Vec<HoodieKey>> {
+ if keys.is_instance_of::<PyString>() {
+ return Err(PyTypeError::new_err(
+ "keys must be a list of record keys, not a single str",
+ ));
+ }
+ keys.try_iter()?
+ .map(|item| {
+ let item = item?;
+ if let Ok(record_key) = item.extract::<String>() {
+ return Ok(HoodieKey {
+ record_key,
+ partition_path: String::new(),
+ });
+ }
+ let (record_key, partition_path) =
+ item.extract::<(String, String)>().map_err(|_| {
+ PyTypeError::new_err(format!(
+ "each key must be a str or a (record_key,
partition_path) tuple, got {}",
+ item.get_type()
+ .name()
+ .map(|n| n.to_string())
+ .unwrap_or_default()
+ ))
+ })?;
+ Ok(HoodieKey {
+ record_key,
+ partition_path,
+ })
+ })
+ .collect()
}
#[cfg(not(tarpaulin_include))]
@@ -1030,7 +1742,329 @@ pub fn rt() -> &'static Runtime {
#[cfg(test)]
mod tests {
- use super::pascal_to_screaming_snake;
+ use super::{
+ HudiWriteResult, has_sub_microsecond, is_unnamed_pandas_index,
last_path_segment,
+ normalize_batch, pascal_to_screaming_snake, writable_field,
writable_type,
+ };
+ use arrow::array::{
+ Array, ArrayRef, DictionaryArray, Int8Array, Int64Array,
LargeListArray, LargeStringArray,
+ ListArray, StringArray, StructArray, TimestampMicrosecondArray,
TimestampNanosecondArray,
+ };
+ use arrow::buffer::OffsetBuffer;
+ use arrow::datatypes::{DataType, Field, Fields, Schema, TimeUnit};
+ use arrow::record_batch::RecordBatch;
+ use hudi::WriteResult;
+ use hudi::table::AppendResult;
+ use std::collections::HashMap;
+ use std::sync::Arc;
+
+ #[test]
+ fn writable_field_rewrites_nested_types_and_keeps_names_and_metadata() {
+ let metadata = HashMap::from([("k".to_string(), "v".to_string())]);
+ let item =
+ Arc::new(Field::new("item", DataType::Utf8View,
true).with_metadata(metadata.clone()));
+ let field = Arc::new(
+ Field::new(
+ "f",
+ DataType::Struct(Fields::from(vec![
+ Field::new("tags", DataType::LargeList(item), true),
+ Field::new(
+ "kind",
+ DataType::Dictionary(
+ Box::new(DataType::Int8),
+ Box::new(DataType::LargeUtf8),
+ ),
+ true,
+ ),
+ Field::new("n", DataType::Int64, true),
+ ])),
+ false,
+ )
+ .with_metadata(metadata.clone()),
+ );
+
+ let written = writable_field(&field).unwrap();
+ assert_eq!(written.name(), "f");
+ assert!(!written.is_nullable());
+ assert_eq!(written.metadata(), &metadata);
+ let DataType::Struct(children) = written.data_type() else {
+ panic!("expected struct, got {}", written.data_type());
+ };
+ let DataType::List(item) = children[0].data_type() else {
+ panic!("expected list, got {}", children[0].data_type());
+ };
+ assert_eq!(item.name(), "item");
+ assert_eq!(item.data_type(), &DataType::Utf8);
+ assert_eq!(item.metadata(), &metadata);
+ assert_eq!(children[1].data_type(), &DataType::Utf8);
+ assert_eq!(children[2].data_type(), &DataType::Int64);
+
+ let plain = Arc::new(Field::new(
+ "n",
+ DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
+ true,
+ ));
+ assert!(writable_field(&plain).is_none());
+ }
+
+ #[test]
+ fn unnamed_pandas_index_detection() {
+ assert!(is_unnamed_pandas_index("__index_level_0__"));
+ assert!(is_unnamed_pandas_index("__index_level_12__"));
+ assert!(!is_unnamed_pandas_index("__index_level___"));
+ assert!(!is_unnamed_pandas_index("__index_level_x__"));
+ assert!(!is_unnamed_pandas_index("id"));
+ }
+
+ #[test]
+ fn writable_type_maps_leaf_types() {
+ let utc: Option<Arc<str>> = Some("UTC".into());
+ for (from, to) in [
+ (DataType::LargeBinary, Some(DataType::Binary)),
+ (DataType::BinaryView, Some(DataType::Binary)),
+ (
+ DataType::Timestamp(TimeUnit::Nanosecond, utc.clone()),
+ Some(DataType::Timestamp(TimeUnit::Microsecond, utc.clone())),
+ ),
+ (
+ DataType::Dictionary(Box::new(DataType::Int8),
Box::new(DataType::Int64)),
+ Some(DataType::Int64),
+ ),
+ (DataType::Timestamp(TimeUnit::Microsecond, utc), None),
+ (DataType::Binary, None),
+ (DataType::Int64, None),
+ (
+ DataType::Struct(Fields::from(vec![Field::new("n",
DataType::Int64, true)])),
+ None,
+ ),
+ ] {
+ assert_eq!(writable_type(&from), to, "for {from}");
+ }
+
+ // A large list narrows to a list even when its item is already
writable.
+ let item = Arc::new(Field::new("element", DataType::Int64, false));
+ assert_eq!(
+ writable_type(&DataType::LargeList(item.clone())),
+ Some(DataType::List(item))
+ );
+ }
+
+ fn nanos(values: Vec<Option<i64>>) -> TimestampNanosecondArray {
+ TimestampNanosecondArray::from(values)
+ }
+
+ fn list_of(values: TimestampNanosecondArray) -> ArrayRef {
+ let field = Arc::new(Field::new("item", values.data_type().clone(),
true));
+ let offsets = OffsetBuffer::from_lengths([values.len()]);
+ Arc::new(ListArray::new(field, offsets, Arc::new(values), None))
+ }
+
+ fn large_list_of(values: TimestampNanosecondArray) -> ArrayRef {
+ let field = Arc::new(Field::new("item", values.data_type().clone(),
true));
+ let offsets = OffsetBuffer::from_lengths([values.len()]);
+ Arc::new(LargeListArray::new(field, offsets, Arc::new(values), None))
+ }
+
+ fn struct_of(values: TimestampNanosecondArray) -> ArrayRef {
+ let field = Arc::new(Field::new("ts", values.data_type().clone(),
true));
+ Arc::new(StructArray::from(vec![(
+ field,
+ Arc::new(values) as ArrayRef,
+ )]))
+ }
+
+ fn dictionary_of(values: TimestampNanosecondArray) -> ArrayRef {
+ let keys = Int8Array::from_iter_values(0..values.len() as i8);
+ Arc::new(DictionaryArray::new(keys, Arc::new(values)))
+ }
+
+ #[test]
+ fn has_sub_microsecond_checks_values_at_any_depth() {
+ let whole = || nanos(vec![Some(2_000), None]);
+ let partial = || nanos(vec![Some(2_000), Some(2_001)]);
+ assert!(!has_sub_microsecond(&whole()));
+ assert!(has_sub_microsecond(&partial()));
+ for wrap in [list_of, large_list_of, struct_of, dictionary_of] {
+ assert!(!has_sub_microsecond(wrap(whole()).as_ref()));
+ assert!(has_sub_microsecond(wrap(partial()).as_ref()));
+ }
+ // Values outside a slice and unused dictionary values are ignored.
+ let lists = ListArray::new(
+ Arc::new(Field::new(
+ "item",
+ DataType::Timestamp(TimeUnit::Nanosecond, None),
+ true,
+ )),
+ OffsetBuffer::from_lengths([1, 1]),
+ Arc::new(nanos(vec![Some(1), Some(2_000)])),
+ None,
+ );
+ assert!(!has_sub_microsecond(&lists.slice(1, 1)));
+ assert!(has_sub_microsecond(&lists.slice(0, 1)));
+ let unused = DictionaryArray::new(
+ Int8Array::from(vec![1]),
+ Arc::new(nanos(vec![Some(1), Some(2_000)])),
+ );
+ assert!(!has_sub_microsecond(&unused));
+ // Only nanosecond timestamps can hold sub-microsecond digits.
+ assert!(!has_sub_microsecond(&TimestampMicrosecondArray::from(
+ vec![1]
+ )));
+ assert!(!has_sub_microsecond(&Int64Array::from(vec![1])));
+ }
+
+ fn batch(columns: Vec<(&str, ArrayRef)>, pandas: bool) -> RecordBatch {
+ let fields: Vec<Field> = columns
+ .iter()
+ .map(|(name, c)| Field::new(*name, c.data_type().clone(), true))
+ .collect();
+ let metadata = if pandas {
+ HashMap::from([("pandas".to_string(), "{}".to_string())])
+ } else {
+ HashMap::new()
+ };
+ let schema = Arc::new(Schema::new_with_metadata(fields, metadata));
+ RecordBatch::try_new(schema, columns.into_iter().map(|(_, c)|
c).collect()).unwrap()
+ }
+
+ #[test]
+ fn normalize_batch_returns_a_writable_batch_unchanged() {
+ let input = batch(
+ vec![
+ ("id", Arc::new(StringArray::from(vec!["a"]))),
+ ("__index_level_0__", Arc::new(Int64Array::from(vec![0]))),
+ ],
+ false,
+ );
+ let output = normalize_batch(input.clone()).unwrap();
+ assert!(Arc::ptr_eq(&input.schema(), &output.schema()));
+ }
+
+ #[test]
+ fn normalize_batch_drops_pandas_index_and_casts_columns() {
+ let input = batch(
+ vec![
+ ("__index_level_0__", Arc::new(Int64Array::from(vec![5]))),
+ ("id", Arc::new(LargeStringArray::from(vec!["a"]))),
+ ("ts", list_of(nanos(vec![Some(3_000), None]))),
+ ],
+ true,
+ );
+ let output = normalize_batch(input.clone()).unwrap();
+ assert_eq!(output.schema().metadata(), input.schema().metadata());
+ let names: Vec<_> = output
+ .schema()
+ .fields()
+ .iter()
+ .map(|f| f.name().clone())
+ .collect();
+ assert_eq!(names, ["id", "ts"]);
+ assert_eq!(output.column(0).data_type(), &DataType::Utf8);
+ let DataType::List(item) = output.column(1).data_type() else {
+ panic!("expected list, got {}", output.column(1).data_type());
+ };
+ assert_eq!(
+ item.data_type(),
+ &DataType::Timestamp(TimeUnit::Microsecond, None)
+ );
+ let values = output
+ .column(1)
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .unwrap();
+ let values = values
+ .values()
+ .as_any()
+ .downcast_ref::<TimestampMicrosecondArray>()
+ .unwrap();
+ assert_eq!(values.iter().collect::<Vec<_>>(), [Some(3), None]);
+ }
+
+ #[test]
+ fn normalize_batch_keeps_row_count_when_only_the_index_was_present() {
+ let input = batch(
+ vec![(
+ "__index_level_0__",
+ Arc::new(Int64Array::from(vec![0, 1, 2])),
+ )],
+ true,
+ );
+ let output = normalize_batch(input).unwrap();
+ assert_eq!((output.num_columns(), output.num_rows()), (0, 3));
+ }
+
+ #[test]
+ fn normalize_batch_rejects_sub_microsecond_timestamps() {
+ for column in [
+ Arc::new(nanos(vec![Some(1)])) as ArrayRef,
+ struct_of(nanos(vec![Some(1_000_001)])),
+ ] {
+ let err = normalize_batch(batch(vec![("ts", column)],
false)).unwrap_err();
+ assert!(
+ err.contains("column 'ts' has timestamps with
sub-microsecond"),
+ "{err}"
+ );
+ }
+ }
+
+ #[test]
+ fn write_results_convert_to_python_counts() {
+ let rewrite = HudiWriteResult::from(WriteResult {
+ instant: "20240101000000000".to_string(),
+ completion_time: None,
+ num_writes: 4,
+ num_updates: 1,
+ num_inserts: 2,
+ num_deletes: 3,
+ });
+ assert_eq!(rewrite.instant_time.as_deref(), Some("20240101000000000"));
+ assert_eq!(
+ (
+ rewrite.num_writes,
+ rewrite.num_updates,
+ rewrite.num_inserts,
+ rewrite.num_deletes
+ ),
+ (4, 1, 2, 3)
+ );
+
+ let append = HudiWriteResult::from(AppendResult {
+ instant: "20240101000000001".to_string(),
+ completion_time: None,
+ commit_relative_path: String::new(),
+ base_file_path: String::new(),
+ num_rows: 5,
+ });
+ assert_eq!(append.instant_time.as_deref(), Some("20240101000000001"));
+ assert_eq!(
+ (
+ append.num_writes,
+ append.num_updates,
+ append.num_inserts,
+ append.num_deletes
+ ),
+ (5, 0, 5, 0)
+ );
+
+ // A write that matched nothing reports no instant.
+ let noop = HudiWriteResult::from(WriteResult {
+ instant: String::new(),
+ completion_time: None,
+ num_writes: 0,
+ num_updates: 0,
+ num_inserts: 0,
+ num_deletes: 0,
+ });
+ assert_eq!(noop.instant_time, None);
+ }
+
+ #[test]
+ fn last_path_segment_names_the_table_directory() {
+ assert_eq!(last_path_segment("/data/trips"), "trips");
+ assert_eq!(last_path_segment("file:///data/trips//"), "trips");
+ assert_eq!(last_path_segment("s3://bucket/warehouse/trips"), "trips");
+ assert_eq!(last_path_segment("trips"), "trips");
+ }
#[test]
fn pascal_to_screaming_snake_basic() {
diff --git a/python/src/lib.rs b/python/src/lib.rs
index 6308ff37..aa796387 100644
--- a/python/src/lib.rs
+++ b/python/src/lib.rs
@@ -33,7 +33,7 @@ fn _internal(_py: Python, m: &Bound<'_, PyModule>) ->
PyResult<()> {
use internal::{
HudiFileGroupReader, HudiFileSlice, HudiInstant, HudiQueryType,
HudiReadOptions,
- HudiRecordBatchStream, HudiTable, HudiTimeline,
+ HudiRecordBatchStream, HudiTable, HudiTimeline, HudiWriteResult,
};
m.add_class::<HudiFileGroupReader>()?;
m.add_class::<HudiFileSlice>()?;
@@ -43,6 +43,7 @@ fn _internal(_py: Python, m: &Bound<'_, PyModule>) ->
PyResult<()> {
m.add_class::<HudiRecordBatchStream>()?;
m.add_class::<HudiTable>()?;
m.add_class::<HudiTimeline>()?;
+ m.add_class::<HudiWriteResult>()?;
m.add_function(wrap_pyfunction!(internal::_config_keys, m)?)?;
diff --git a/python/tests/test_table_write.py b/python/tests/test_table_write.py
new file mode 100644
index 00000000..64da9426
--- /dev/null
+++ b/python/tests/test_table_write.py
@@ -0,0 +1,774 @@
+# 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.
+
+"""Tests for the HudiTable write API."""
+
+import datetime
+import math
+from decimal import Decimal
+
+import pyarrow as pa
+import pytest
+
+from hudi import HudiTable, HudiWriteResult
+
+SCHEMA = pa.schema(
+ [
+ ("id", pa.string()),
+ ("city", pa.string()),
+ ("ts", pa.int64()),
+ ("fare", pa.float64()),
+ ]
+)
+
+
+def trips(*rows):
+ ids, cities, ts, fares = zip(*rows)
+ return pa.table({"id": ids, "city": cities, "ts": ts, "fare": fares},
schema=SCHEMA)
+
+
+def snapshot(table):
+ """Map record key to (city, ts, fare) for the current snapshot."""
+ read = pa.Table.from_batches(table.read())
+ return {
+ r["id"]: (r["city"], r["ts"], r["fare"])
+ for r in read.select(["id", "city", "ts", "fare"]).to_pylist()
+ }
+
+
+def new_table(tmp_path, table_type="cow"):
+ table = HudiTable.create(
+ str(tmp_path / "trips"),
+ record_key="id",
+ partition_by="city",
+ ordering_fields="ts",
+ table_type=table_type,
+ )
+ table.append(
+ trips(("t1", "sf", 1, 30.0), ("t2", "sf", 1, 45.0), ("t3", "nyc", 1,
20.0))
+ )
+ return table
+
+
+class StreamOnly:
+ """Exposes only the Arrow PyCapsule stream interface, like polars or
pandas."""
+
+ def __init__(self, table):
+ self._table = table
+
+ def __arrow_c_stream__(self, requested_schema=None):
+ return self._table.__arrow_c_stream__(requested_schema)
+
+
+class ArrayOnly:
+ """Exposes only the Arrow PyCapsule array interface for one record
batch."""
+
+ def __init__(self, batch):
+ self._batch = batch
+
+ def __arrow_c_array__(self, requested_schema=None):
+ return self._batch.__arrow_c_array__(requested_schema)
+
+
[email protected](
+ "table_type, expected_type",
+ [
+ ("cow", "COPY_ON_WRITE"),
+ ("Merge_On_Read", "MERGE_ON_READ"),
+ ],
+)
+def test_create_table_types(tmp_path, table_type, expected_type):
+ base = tmp_path / "trips"
+ table = HudiTable.create(str(base), record_key=["id"],
table_type=table_type)
+ assert table.table_type == expected_type
+ assert table.is_mor == (expected_type == "MERGE_ON_READ")
+ assert table.table_name == "trips"
+ assert (base / ".hoodie" / "hoodie.properties").is_file()
+ assert (base / ".hoodie" / "metadata" / "record_index").is_dir()
+
+
+def test_create_with_explicit_options(tmp_path):
+ base = tmp_path / "plain"
+ table = HudiTable.create(
+ str(base) + "/",
+ record_key="id",
+ partition_by=["city"],
+ table_name="renamed",
+ table_version=8,
+ metadata_table=False,
+ hive_style_partitioning=False,
+ hudi_options={"hoodie.table.timeline.timezone": "UTC"},
+ )
+ assert table.table_name == "renamed"
+ assert table.hudi_options()["hoodie.table.version"] == "8"
+ assert not (base / ".hoodie" / "metadata").exists()
+ table.append(trips(("t1", "sf", 1, 30.0)))
+ assert (base / "sf").is_dir()
+
+
+def test_create_rejects_bad_table_type(tmp_path):
+ with pytest.raises(ValueError, match="invalid table_type 'parquet'"):
+ HudiTable.create(str(tmp_path / "t"), record_key="id",
table_type="parquet")
+ assert not (tmp_path / "t").exists()
+
+
+def test_create_on_existing_table_raises(tmp_path):
+ HudiTable.create(str(tmp_path / "t"), record_key="id")
+ with pytest.raises(Exception, match="already exists"):
+ HudiTable.create(str(tmp_path / "t"), record_key="id")
+
+
[email protected](
+ "flags, error",
+ [
+ (
+ dict(metadata_table=False, record_index=True),
+ "record index requires the metadata",
+ ),
+ (
+ dict(metadata_table=False, column_stats=True),
+ "column_stats requires the metadata",
+ ),
+ # The explicit partition_stats choice is applied after column_stats,
+ # so it is reported rather than silently turned off.
+ (dict(column_stats=False, partition_stats=True), "partition_stats
requires"),
+ ],
+)
+def test_create_rejects_inconsistent_index_options(tmp_path, flags, error):
+ with pytest.raises(Exception, match=error):
+ HudiTable.create(str(tmp_path / "t"), record_key="id", **flags)
+ assert not (tmp_path / "t").exists()
+
+
+def test_create_index_options(tmp_path):
+ base = tmp_path / "t"
+ table = HudiTable.create(
+ str(base),
+ record_key="id",
+ partition_by="city",
+ record_index=False,
+ column_stats=True,
+ partition_stats=False,
+ )
+ partitions =
table.hudi_options()["hoodie.table.metadata.partitions"].split(",")
+ assert "column_stats" in partitions
+ assert "partition_stats" not in partitions
+ assert "record_index" not in partitions
+ # Without the record index, upserts tag records through the simple index.
+ table.append(trips(("t1", "sf", 1, 30.0)))
+ assert table.upsert(trips(("t1", "sf", 2, 31.0))).num_updates == 1
+ assert snapshot(table) == {"t1": ("sf", 2, 31.0)}
+
+
+def test_create_with_field_lists_and_storage_options(tmp_path):
+ table = HudiTable.create(
+ str(tmp_path / "t"),
+ record_key=["id"],
+ partition_by=["city", "ts"],
+ ordering_fields=["ts"],
+ storage_options={"allow_http": "true"},
+ )
+ options = table.hudi_options()
+ assert options["hoodie.table.recordkey.fields"] == "id"
+ assert options["hoodie.table.partition.fields"] == "city,ts"
+ assert options["hoodie.table.ordering.fields"] == "ts"
+ assert table.storage_options()["allow_http"] == "true"
+ table.append(trips(("t1", "sf", 1, 30.0)))
+ assert (tmp_path / "t" / "city=sf" / "ts=1").is_dir()
+
+
[email protected]("table_type", ["cow", "mor"])
+def test_append_then_read(tmp_path, table_type):
+ table = HudiTable.create(
+ str(tmp_path / "trips"),
+ record_key="id",
+ partition_by="city",
+ table_type=table_type,
+ )
+ res = table.append(trips(("t1", "sf", 1, 30.0), ("t2", "nyc", 1, 45.0)))
+ assert isinstance(res, HudiWriteResult)
+ assert res.instant_time
+ assert (res.num_writes, res.num_inserts, res.num_updates, res.num_deletes)
== (
+ 2,
+ 2,
+ 0,
+ 0,
+ )
+ assert "num_inserts=2" in repr(res)
+ assert f"instant_time='{res.instant_time}'" in repr(res)
+ with pytest.raises(AttributeError):
+ res.num_inserts = 0
+ assert snapshot(table) == {"t1": ("sf", 1, 30.0), "t2": ("nyc", 1, 45.0)}
+ # A fresh handle sees the same committed data.
+ assert snapshot(HudiTable(str(tmp_path / "trips"))) == snapshot(table)
+
+
[email protected]("table_type", ["cow", "mor"])
+def test_upsert_updates_and_inserts(tmp_path, table_type):
+ table = new_table(tmp_path, table_type)
+ res = table.upsert(trips(("t2", "sf", 2, 50.0), ("t4", "nyc", 2, 15.0)))
+ assert (res.num_updates, res.num_inserts, res.num_deletes) == (1, 1, 0)
+ assert snapshot(table) == {
+ "t1": ("sf", 1, 30.0),
+ "t2": ("sf", 2, 50.0),
+ "t3": ("nyc", 1, 20.0),
+ "t4": ("nyc", 2, 15.0),
+ }
+
+
+def test_upsert_keeps_higher_ordering_value(tmp_path):
+ table = new_table(tmp_path)
+ table.upsert(trips(("t1", "sf", 0, 99.0)))
+ assert snapshot(table)["t1"] == ("sf", 1, 30.0)
+
+
+def test_upsert_update_columns(tmp_path):
+ table = new_table(tmp_path)
+ res = table.upsert(trips(("t1", "sf", 5, 31.0)), update_columns=["fare"])
+ assert res.num_updates == 1
+ assert snapshot(table)["t1"] == ("sf", 1, 31.0)
+
+
[email protected]("table_type", ["cow", "mor"])
+def test_update_with_dict(tmp_path, table_type):
+ table = new_table(tmp_path, table_type)
+ # An int value lands in the float64 column.
+ res = table.update(predicate="city = 'nyc'", updates={"fare": 99})
+ assert res.num_updates == 1
+ assert snapshot(table)["t3"] == ("nyc", 1, 99.0)
+ assert snapshot(table)["t1"] == ("sf", 1, 30.0)
+
+
[email protected]("table_type", ["cow", "mor"])
+def test_update_with_batch(tmp_path, table_type):
+ table = new_table(tmp_path, table_type)
+ set_batch = pa.record_batch(
+ {"fare": pa.array([7.5]), "ts": pa.array([3], pa.int64())}
+ )
+ table.update("id = 't1'", set_batch)
+ assert snapshot(table)["t1"] == ("sf", 3, 7.5)
+ table.update("id = 't2'", pa.table({"fare": [8.5]}))
+ assert snapshot(table)["t2"] == ("sf", 1, 8.5)
+ table.update("id = 't3'", [pa.record_batch({"fare": [9.5]})])
+ assert snapshot(table)["t3"] == ("nyc", 1, 9.5)
+
+
+def test_update_set_batches_must_hold_one_row(tmp_path):
+ table = new_table(tmp_path)
+ with pytest.raises(ValueError, match="got no batches"):
+ table.update("id = 't1'", [])
+ two = [pa.record_batch({"fare": [1.0]}), pa.record_batch({"fare": [2.0]})]
+ with pytest.raises(Exception, match="single-row"):
+ table.update("id = 't1'", two)
+ assert snapshot(table)["t1"] == ("sf", 1, 30.0)
+
+
+def test_update_dict_null_and_nan(tmp_path):
+ table = new_table(tmp_path)
+ table.update("id = 't1'", {"fare": None})
+ assert snapshot(table)["t1"] == ("sf", 1, None)
+ table.update("id = 't2'", {"fare": float("nan")})
+ assert math.isnan(snapshot(table)["t2"][2])
+
+
+def test_update_no_match_creates_no_commit(tmp_path):
+ table = new_table(tmp_path)
+ commits = len(table.get_timeline().get_completed_commits())
+ res = table.update("city = 'la'", {"fare": 1.0})
+ assert (res.instant_time, res.num_updates) == (None, 0)
+ assert "instant_time=None" in repr(res)
+ assert len(table.get_timeline().get_completed_commits()) == commits
+
+
+def test_update_rejects_bad_set(tmp_path):
+ table = new_table(tmp_path)
+ with pytest.raises(
+ ValueError, match="update column 'nope' is not in the table schema"
+ ):
+ table.update("city = 'nyc'", {"nope": 1})
+ with pytest.raises(
+ ValueError, match="cannot convert update value for column 'fare'"
+ ):
+ table.update("city = 'nyc'", {"fare": "not a number"})
+ with pytest.raises(ValueError, match="at least one column"):
+ table.update("city = 'nyc'", {})
+
+
+def test_update_dict_values_must_fit_column_type(tmp_path):
+ table = HudiTable.create(str(tmp_path / "t"), record_key="id")
+ table.append(
+ pa.table(
+ {
+ "id": ["t1"],
+ "n": pa.array([1], pa.int64()),
+ "x": pa.array([1.0], pa.float64()),
+ "f": pa.array([1.0], pa.float32()),
+ "amount": pa.array([Decimal("1.00")], pa.decimal128(10, 2)),
+ }
+ )
+ )
+
+ def row():
+ return (
+ pa.Table.from_batches(table.read())
+ .select(["n", "x", "amount"])
+ .to_pylist()[0]
+ )
+
+ table.update("id = 't1'", {"x": 3, "amount": 2})
+ table.update("id = 't1'", {"n": 4.0})
+ # float32 rounding is accepted; only finite to inf is not.
+ table.update("id = 't1'", {"f": 0.1})
+ assert row() == {"n": 4, "x": 3.0, "amount": Decimal("2.00")}
+
+ for set_, column in [
+ ({"n": 1.5}, "n"),
+ ({"amount": 1.234}, "amount"),
+ ({"amount": Decimal("1.234")}, "amount"),
+ ({"n": "7"}, "n"),
+ ({"f": 1e300}, "f"),
+ ({"x": True}, "x"),
+ ]:
+ with pytest.raises(ValueError, match=f"column '{column}'"):
+ table.update("id = 't1'", set_)
+ assert row() == {"n": 4, "x": 3.0, "amount": Decimal("2.00")}
+
+
+def test_update_dict_nested_and_numpy_values_must_fit(tmp_path):
+ np = pytest.importorskip("numpy")
+ table = HudiTable.create(str(tmp_path / "t"), record_key="id")
+ table.append(
+ pa.table(
+ {
+ "id": ["t1"],
+ "n": pa.array([1], pa.int64()),
+ "st": pa.array([{"k": 1}], pa.struct([("k", pa.int64())])),
+ "xs": pa.array([[1]], pa.list_(pa.int64())),
+ }
+ )
+ )
+ table.update("id = 't1'", {"st": {"k": 3}, "xs": [np.int64(4), 5.0]})
+
+ for set_, column in [
+ ({"st": {"k": 2.5}}, "st"),
+ ({"xs": [np.float64(7.9)]}, "xs"),
+ ({"n": np.True_}, "n"),
+ ({"xs": [np.True_]}, "xs"),
+ ]:
+ with pytest.raises(ValueError, match=f"column '{column}'"):
+ table.update("id = 't1'", set_)
+ row = pa.Table.from_batches(table.read()).select(["n", "st",
"xs"]).to_pylist()
+ assert row == [{"n": 1, "st": {"k": 3}, "xs": [4, 5]}]
+
+
+def _np():
+ return pytest.importorskip("numpy")
+
+
+def _pd():
+ return pytest.importorskip("pandas")
+
+
+_AWARE = datetime.datetime(
+ 2024, 1, 3, 2, tzinfo=datetime.timezone(datetime.timedelta(hours=2))
+)
+
+
[email protected](
+ "make_set, error",
+ [
+ # Accepted.
+ (lambda: {"d": datetime.date(2024, 2, 2)}, None),
+ (lambda: {"ms": datetime.datetime(2024, 1, 2, 0, 0, 0, 5000)}, None),
+ (lambda: {"utc": datetime.datetime(2024, 1, 2)}, None),
+ (lambda: {"utc": _AWARE}, None),
+ (lambda: {"n": pa.scalar(5)}, None),
+ (lambda: {"st": {"k": pa.scalar(3), "inner": {"j": pa.scalar(2)}}},
None),
+ (
+ lambda: {
+ "tl": [pa.scalar(datetime.datetime(2024, 1, 1),
pa.timestamp("ms"))]
+ },
+ None,
+ ),
+ (lambda: {"st": {"k": 3}}, None),
+ # Temporal columns take only dates and datetimes, at any depth.
+ (lambda: {"d": datetime.datetime(2024, 3, 3, 12)}, "column 'd'"),
+ (lambda: {"d": 5}, "column 'd'"),
+ (lambda: {"utc": 1.7e9}, "column 'utc'"),
+ (lambda: {"ms": Decimal(5)}, "column 'ms'"),
+ (lambda: {"ms": _np().datetime64("2024-01-01")}, "column 'ms'"),
+ (
+ lambda: {"tl": _np().array(["2024-01-01"], dtype="datetime64[s]")},
+ "column 'tl'",
+ ),
+ # Sub-unit truncation, including pandas nanoseconds.
+ (lambda: {"ms": datetime.datetime(2024, 1, 2, 0, 0, 0, 123456)},
"column 'ms'"),
+ (
+ lambda: {"utc": _pd().Timestamp("2024-01-01 00:00:00.000000001")},
+ "column 'utc'",
+ ),
+ # A tz-aware datetime into a timestamp without a time zone.
+ (lambda: {"ms": _AWARE}, "column 'ms'"),
+ # Unknown struct keys, at any depth.
+ (lambda: {"st": {"k": 1, "x": 2}}, "column 'st'.*'x'"),
+ (lambda: {"st": {"inner": {"j": 1, "y": 2}}}, "column 'st'.*'y'"),
+ # pyarrow scalars are unwrapped before the check.
+ (lambda: {"n": pa.scalar(1.5)}, "column 'n'"),
+ (lambda: {"st": {"k": pa.scalar(2.5)}}, "column 'st'"),
+ (lambda: {"tl": [pa.scalar(5)]}, "column 'tl'"),
+ # A value pyarrow stores but cannot read back.
+ (lambda: {"ms": 10**15}, "read back .*column 'ms'"),
+ (lambda: {1: 2}, "must be str"),
+ ],
+)
+def test_update_dict_value_checks(tmp_path, make_set, error):
+ set_ = make_set()
+ table = HudiTable.create(str(tmp_path / "t"), record_key="id")
+ inner = pa.struct([("j", pa.int64())])
+ table.append(
+ pa.table(
+ {
+ "id": ["t1"],
+ "n": pa.array([1], pa.int64()),
+ "d": pa.array([datetime.date(2024, 1, 1)], pa.date32()),
+ "ms": pa.array([datetime.datetime(2024, 1, 1)],
pa.timestamp("ms")),
+ "utc": pa.array(
+ [datetime.datetime(2024, 1, 1)], pa.timestamp("us", "UTC")
+ ),
+ "tl": pa.array(
+ [[datetime.datetime(2024, 1, 1)]],
pa.list_(pa.timestamp("ms"))
+ ),
+ "st": pa.array(
+ [{"k": 1, "inner": {"j": 1}}],
+ pa.struct([("k", pa.int64()), ("inner", inner)]),
+ ),
+ }
+ )
+ )
+ if error is None:
+ assert table.update("id = 't1'", set_).num_updates == 1
+ return
+ before = pa.Table.from_batches(table.read()).drop_columns(
+ ["_hoodie_commit_time", "_hoodie_commit_seqno"]
+ )
+ with pytest.raises(ValueError, match=error):
+ table.update("id = 't1'", set_)
+ after = pa.Table.from_batches(table.read()).drop_columns(
+ ["_hoodie_commit_time", "_hoodie_commit_seqno"]
+ )
+ assert after.equals(before)
+
+
[email protected]("table_type", ["cow", "mor"])
+def test_delete_by_key_predicate(tmp_path, table_type):
+ table = new_table(tmp_path, table_type)
+ res = table.delete("id IN ('t1', 't3')")
+ assert res.num_deletes == 2
+ assert set(snapshot(table)) == {"t2"}
+
+
+def test_delete_by_non_key_predicate(tmp_path):
+ table = new_table(tmp_path)
+ res = table.delete(predicate="fare > 40")
+ assert res.num_deletes == 1
+ assert set(snapshot(table)) == {"t1", "t3"}
+ res = table.delete("fare > 1000")
+ assert (res.instant_time, res.num_deletes) == (None, 0)
+
+
[email protected]("table_type", ["cow", "mor"])
+def test_delete_keys(tmp_path, table_type):
+ table = new_table(tmp_path, table_type)
+ res = table.delete_keys(["t1", ("t3", "city=nyc")])
+ assert res.num_deletes == 2
+ assert set(snapshot(table)) == {"t2"}
+ res = table.delete_keys(["missing"])
+ assert (res.instant_time, res.num_deletes) == (None, 0)
+ res = table.delete_keys(k for k in ("t2",))
+ assert res.num_deletes == 1
+ assert snapshot(table) == {}
+
+
+def test_delete_keys_rejects_bad_input(tmp_path):
+ table = new_table(tmp_path)
+ with pytest.raises(TypeError, match="not a single str"):
+ table.delete_keys("t1")
+ for bad in ([1], [("t1",)], [("t1", "city=sf", "x")]):
+ with pytest.raises(TypeError, match="each key must be a str"):
+ table.delete_keys(bad)
+ assert len(snapshot(table)) == 3
+
+
[email protected]("table_type", ["cow", "mor"])
+def test_overwrite(tmp_path, table_type):
+ table = new_table(tmp_path, table_type)
+ res = table.overwrite(trips(("t9", "sf", 9, 10.0)))
+ assert res.instant_time
+ assert snapshot(table) == {"t9": ("sf", 9, 10.0)}
+
+
[email protected]("table_type", ["cow", "mor"])
+def test_dynamic_partition_overwrite(tmp_path, table_type):
+ table = new_table(tmp_path, table_type)
+ res = table.dynamic_partition_overwrite(trips(("t8", "sf", 9, 11.0)))
+ assert res.instant_time
+ assert snapshot(table) == {"t8": ("sf", 9, 11.0), "t3": ("nyc", 1, 20.0)}
+
+
[email protected](
+ "make_data",
+ [
+ pytest.param(lambda t: t, id="table"),
+ pytest.param(lambda t: t.to_batches()[0], id="record_batch"),
+ pytest.param(lambda t: t.to_reader(), id="record_batch_reader"),
+ pytest.param(lambda t: t.to_batches(), id="list"),
+ pytest.param(StreamOnly, id="arrow_c_stream"),
+ pytest.param(
+ lambda t: ArrayOnly(t.combine_chunks().to_batches()[0]),
id="arrow_c_array"
+ ),
+ pytest.param(
+ lambda t: [ArrayOnly(b) for b in t.to_batches()],
id="list_of_arrow_c_array"
+ ),
+ pytest.param(lambda t: polars().from_arrow(t), id="polars"),
+ pytest.param(
+ lambda t: t.to_pandas(types_mapper=pandas().ArrowDtype),
id="pandas_arrow"
+ ),
+ pytest.param(lambda t: t.to_pandas(), id="pandas_numpy"),
+ ],
+)
+def test_data_input_types(tmp_path, make_data):
+ table = HudiTable.create(str(tmp_path / "t"), record_key="id")
+ data = trips(("t1", "sf", 1, 30.0), ("t2", "nyc", 1, 45.0))
+ table.append(make_data(data))
+ res = table.upsert(make_data(trips(("t2", "nyc", 2, 46.0))))
+ assert res.num_updates == 1
+ assert snapshot(table) == {"t1": ("sf", 1, 30.0), "t2": ("nyc", 2, 46.0)}
+
+
+def polars():
+ return pytest.importorskip("polars")
+
+
+def pandas():
+ return pytest.importorskip("pandas")
+
+
+def test_nested_view_large_and_dictionary_columns_are_cast(tmp_path):
+ table = HudiTable.create(str(tmp_path / "t"), record_key="id")
+ addr = pa.field(
+ "addr",
+ pa.struct([("city", pa.large_string()), ("zip", pa.int64())]),
+ metadata={b"k": b"v"},
+ )
+ schema = pa.schema(
+ [
+ ("id", pa.string_view()),
+ addr,
+ ("kind", pa.dictionary(pa.int8(), pa.large_string())),
+ ]
+ )
+
+ def rows(*values):
+ ids, cities, kinds = zip(*values)
+ return pa.table(
+ {
+ "id": ids,
+ "addr": [{"city": c, "zip": 94000} for c in cities],
+ "kind": list(kinds),
+ },
+ schema=schema,
+ )
+
+ table.append(rows(("t1", "sf", "x"), ("t2", "nyc", "y")))
+ res = table.upsert(rows(("t2", "la", "x")))
+ assert res.num_updates == 1
+ read = pa.Table.from_batches(table.read()).sort_by("id")
+ assert [a["city"] for a in read.column("addr").to_pylist()] == ["sf", "la"]
+ assert read.column("kind").to_pylist() == ["x", "x"]
+
+
+def polars_frame(ids, tags, cities):
+ return polars().DataFrame(
+ {"id": ids, "tags": tags, "addr": [{"city": c} for c in cities]}
+ )
+
+
+def test_polars_nested_columns_append(tmp_path):
+ table = HudiTable.create(str(tmp_path / "t"), record_key="id")
+ table.append(polars_frame(["t1", "t2"], [["a", "b"], ["c"]], ["sf",
"nyc"]))
+ written = table.get_schema()
+ assert written.field("tags").type.value_type == pa.string()
+ assert written.field("addr").type.field("city").type == pa.string()
+ read = pa.Table.from_batches(table.read()).sort_by("id")
+ assert read.column("tags").to_pylist() == [["a", "b"], ["c"]]
+
+
+def test_polars_nested_columns_upsert(tmp_path):
+ table = HudiTable.create(str(tmp_path / "t"), record_key="id")
+ table.append(polars_frame(["t1", "t2"], [["a", "b"], ["c"]], ["sf",
"nyc"]))
+ res = table.upsert(polars_frame(["t2"], [["d"]], ["la"]))
+ assert res.num_updates == 1
+ read = pa.Table.from_batches(table.read()).sort_by("id")
+ assert read.column("tags").to_pylist() == [["a", "b"], ["d"]]
+ assert [a["city"] for a in read.column("addr").to_pylist()] == ["sf", "la"]
+
+
[email protected]("source", ["pandas", "polars"])
+def test_nanosecond_timestamps_are_checked(tmp_path, source):
+ # 2024-01-02 03:04:05 plus `nanos`, as a naive nanosecond timestamp.
+ def frame(nanos):
+ epoch_ns = 1_704_164_645 * 10**9 + nanos
+ if source == "pandas":
+ pd = pandas()
+ ts = pd.Series(pd.to_datetime([epoch_ns], unit="ns")).astype(
+ "datetime64[ns]"
+ )
+ return pd.DataFrame({"id": ["t1"], "ts": ts})
+ pl = polars()
+ ts = pl.Series([epoch_ns]).cast(pl.Datetime("ns"))
+ return pl.DataFrame({"id": ["t1"], "ts": ts})
+
+ assert pa.table(frame(1)).schema.field("ts").type == pa.timestamp("ns")
+ table = HudiTable.create(str(tmp_path / "t"), record_key="id")
+ table.append(frame(123_456_000))
+ read = pa.Table.from_batches(table.read())
+ assert read.column("ts").to_pylist() == [
+ datetime.datetime(2024, 1, 2, 3, 4, 5, 123456)
+ ]
+ with pytest.raises(
+ ValueError, match="column 'ts' has timestamps with sub-microsecond"
+ ):
+ table.upsert(frame(1))
+
+
+def test_unreferenced_nanosecond_values_are_ignored(tmp_path):
+ ns = pa.timestamp("ns")
+ sliced = pa.table(
+ {"id": ["a", "b"], "ts": pa.array([[1], [1_000]], pa.list_(ns))}
+ ).slice(1)
+ unused = pa.DictionaryArray.from_arrays(
+ pa.array([1], pa.int8()), pa.array([1, 1_000], ns)
+ )
+ table = HudiTable.create(str(tmp_path / "t"), record_key="id")
+ table.append(sliced.append_column("d", unused))
+ read = pa.Table.from_batches(table.read()).select(["id", "ts",
"d"]).to_pylist()
+ one_us = datetime.datetime(1970, 1, 1, 0, 0, 0, 1)
+ assert read == [{"id": "b", "ts": [one_us], "d": one_us}]
+
+
+def test_pandas_category_and_index_columns(tmp_path):
+ pd = pandas()
+ table = HudiTable.create(str(tmp_path / "t"), record_key="id")
+ frame = pd.DataFrame(
+ {
+ "id": ["t2", "t1", "t3"],
+ "kind": pd.Categorical(["x", "y", "x"]),
+ "fare": [2.0, 1.0, 3.0],
+ }
+ )
+ # A sorted frame keeps a non-default index, which pandas exports as an
+ # unnamed `__index_level_0__` column.
+ table.append(frame.sort_values("id"))
+ assert "__index_level_0__" not in table.get_schema().names
+ res = table.upsert(
+ pd.DataFrame({"id": ["t1"], "kind": pd.Categorical(["z"]), "fare":
[9.0]})
+ )
+ assert res.num_updates == 1
+ read = pa.Table.from_batches(table.read()).sort_by("id")
+ assert read.column("kind").to_pylist() == ["z", "x", "x"]
+ assert read.column("fare").to_pylist() == [9.0, 2.0, 3.0]
+
+
+def test_index_like_column_without_pandas_metadata_is_kept(tmp_path):
+ table = HudiTable.create(str(tmp_path / "t"), record_key="id")
+ table.append(pa.table({"id": ["t1"], "__index_level_0__": [7]}))
+ assert "__index_level_0__" in table.get_schema().names
+
+
+def test_pandas_named_index_is_kept(tmp_path):
+ pd = pandas()
+ table = HudiTable.create(str(tmp_path / "t"), record_key="id")
+ frame = pd.DataFrame({"id": ["t1", "t2"], "fare": [1.0,
2.0]}).set_index("id")
+ table.append(frame)
+ res = table.upsert(pd.DataFrame({"fare": [5.0]}, index=pd.Index(["t2"],
name="id")))
+ assert res.num_updates == 1
+ read = pa.Table.from_batches(table.read()).sort_by("id")
+ assert read.column("fare").to_pylist() == [1.0, 5.0]
+
+
+def test_write_errors_leave_the_table_unchanged(tmp_path):
+ table = new_table(tmp_path)
+ timeline = tmp_path / "trips" / ".hoodie" / "timeline"
+ instants = sorted(p.name for p in timeline.iterdir())
+ no_key = pa.table({"city": ["sf"], "ts": [2], "fare": [1.0]})
+ with pytest.raises(Exception, match="schema does not match"):
+ table.upsert(no_key)
+ with pytest.raises(Exception, match="nope"):
+ table.upsert(trips(("t1", "sf", 2, 1.0)), update_columns=["nope"])
+ with pytest.raises(Exception, match="'no_such_column' not found"):
+ table.delete("no_such_column = 1")
+ assert sorted(p.name for p in timeline.iterdir()) == instants
+ assert snapshot(table)["t1"] == ("sf", 1, 30.0)
+
+
+def test_binary_large_list_and_dictionary_columns_are_cast(tmp_path):
+ table = HudiTable.create(str(tmp_path / "t"), record_key="id")
+ schema = pa.schema(
+ [
+ ("id", pa.string()),
+ ("blob", pa.large_binary()),
+ ("view", pa.binary_view()),
+ ("xs", pa.large_list(pa.int64())),
+ ("code", pa.dictionary(pa.int8(), pa.int64())),
+ ]
+ )
+
+ def rows(id_, blob, xs, code):
+ return pa.table(
+ {"id": [id_], "blob": [blob], "view": [blob], "xs": [xs], "code":
[code]},
+ schema=schema,
+ )
+
+ table.append(rows("t1", b"a", [1, 2], 7))
+ written = table.get_schema()
+ assert written.field("blob").type == pa.binary()
+ assert written.field("view").type == pa.binary()
+ assert pa.types.is_list(written.field("xs").type)
+ assert written.field("xs").type.value_type == pa.int64()
+ assert written.field("code").type == pa.int64()
+ assert table.upsert(rows("t1", b"b", [3], 8)).num_updates == 1
+ read = pa.Table.from_batches(table.read()).select(["blob", "view", "xs",
"code"])
+ assert read.to_pylist() == [{"blob": b"b", "view": b"b", "xs": [3],
"code": 8}]
+
+
+def test_update_dict_struct_value_by_position(tmp_path):
+ table = HudiTable.create(str(tmp_path / "t"), record_key="id")
+ st = pa.struct([("k", pa.int64()), ("v", pa.string())])
+ table.append(pa.table({"id": ["t1"], "st": pa.array([{"k": 1, "v": "a"}],
st)}))
+ table.update("id = 't1'", {"st": (2, "b")})
+ with pytest.raises(ValueError, match="column 'st'"):
+ table.update("id = 't1'", {"st": (2.5, "c")})
+ read = pa.Table.from_batches(table.read()).select(["st"]).to_pylist()
+ assert read == [{"st": {"k": 2, "v": "b"}}]
+
+
+def test_unsupported_data_type_raises(tmp_path):
+ table = HudiTable.create(str(tmp_path / "t"), record_key="id")
+ with pytest.raises(TypeError, match="unsupported data type 'dict'"):
+ table.append({"id": ["t1"]})
+ with pytest.raises(TypeError, match=r"data\[0\] is not a
pyarrow.RecordBatch"):
+ table.append([{"id": "t1"}])
