chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,968 @@
|
||||
"""
|
||||
Module to write a Ray Dataset into an iceberg table, by using the Ray Datasink API.
|
||||
"""
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Union
|
||||
|
||||
import ray
|
||||
from ray._common.retry import call_with_retry
|
||||
from ray.data._internal.datasource.parquet_datasource import (
|
||||
PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT,
|
||||
)
|
||||
from ray.data._internal.execution.interfaces import TaskContext
|
||||
from ray.data._internal.savemode import SaveMode
|
||||
from ray.data._internal.util import MiB
|
||||
from ray.data.block import Block, BlockAccessor
|
||||
from ray.data.context import DataContext
|
||||
from ray.data.datasource.datasink import Datasink, WriteResult
|
||||
from ray.data.expressions import Expr
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pyarrow as pa
|
||||
from pyiceberg.catalog import Catalog
|
||||
from pyiceberg.expressions import BooleanExpression
|
||||
from pyiceberg.io import FileIO
|
||||
from pyiceberg.manifest import DataFile
|
||||
from pyiceberg.schema import Schema
|
||||
from pyiceberg.table import DataScan, FileScanTask, Table
|
||||
from pyiceberg.table.metadata import TableMetadata
|
||||
from pyiceberg.table.update.schema import UpdateSchema
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_REWRITE_STALL_TIMEOUT_S = 600
|
||||
|
||||
|
||||
@ray.remote
|
||||
def _rewrite_iceberg_file(
|
||||
file_scan_task: "FileScanTask",
|
||||
keys_ref: "pa.Table",
|
||||
upsert_cols: List[str],
|
||||
table_metadata: "TableMetadata",
|
||||
io: "FileIO",
|
||||
) -> "tuple[Optional[DataFile], List[DataFile]]":
|
||||
"""Read one Iceberg file, anti-join against upsert keys, write preserved rows.
|
||||
|
||||
Preserved rows are rows in the file that are not in the upsert batch. The
|
||||
coarse range filter would delete them (see ``IcebergDatasink._build_coarse_range_filter``),
|
||||
so we preserve them by writing them as new data files before the delete.
|
||||
|
||||
The file is read in streaming fashion via ``ArrowScan.to_record_batches()``
|
||||
so the full file is never materialised at once. The anti-join is applied
|
||||
per RecordBatch and preserved rows are accumulated, then concatenated and
|
||||
written as a single output once the stream is exhausted.
|
||||
|
||||
Returns (original DataFile to delete, list of new preserved DataFiles).
|
||||
If the entire file is matched (no preserved rows), returns (file, []).
|
||||
If the file has no matched rows at all, returns (None, []), leave it untouched.
|
||||
"""
|
||||
import hashlib
|
||||
import time as _time
|
||||
import uuid as _uuid
|
||||
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
from pyiceberg.expressions import AlwaysTrue
|
||||
from pyiceberg.io.pyarrow import ArrowScan, _dataframe_to_data_files
|
||||
|
||||
file_path = file_scan_task.file.file_path
|
||||
file_name = file_path.split("/")[-1]
|
||||
file_size_mb = file_scan_task.file.file_size_in_bytes / MiB
|
||||
t_start = _time.perf_counter()
|
||||
|
||||
# Cast target pulled from keys_ref once. Applied per batch so PyArrow's join
|
||||
# doesn't raise ArrowInvalid on utf8/large_utf8 or similar width mismatches.
|
||||
target_key_schema = pa.schema([keys_ref.schema.field(c) for c in upsert_cols])
|
||||
|
||||
record_batches = ArrowScan(
|
||||
table_metadata=table_metadata,
|
||||
io=io,
|
||||
projected_schema=table_metadata.schema(),
|
||||
row_filter=AlwaysTrue(),
|
||||
).to_record_batches(tasks=[file_scan_task])
|
||||
|
||||
preserved_rows: Optional["pa.Table"] = None
|
||||
total_in_rows = 0
|
||||
total_preserved_rows = 0
|
||||
n_batches = 0
|
||||
|
||||
for rb in record_batches:
|
||||
n_batches += 1
|
||||
batch_table = pa.Table.from_batches([rb])
|
||||
if len(batch_table) == 0:
|
||||
continue
|
||||
total_in_rows += len(batch_table)
|
||||
|
||||
batch_keys = batch_table.select(upsert_cols).cast(target_key_schema)
|
||||
|
||||
idx_col = pa.array(np.arange(len(batch_table), dtype=np.int64))
|
||||
preserved_keys = batch_keys.append_column("__row_idx__", idx_col).join(
|
||||
keys_ref, keys=upsert_cols, join_type="left anti"
|
||||
)
|
||||
|
||||
if len(preserved_keys) > 0:
|
||||
new_rows = batch_table.take(preserved_keys["__row_idx__"])
|
||||
if preserved_rows is None:
|
||||
preserved_rows = new_rows
|
||||
else:
|
||||
preserved_rows = pa.concat_tables(
|
||||
[preserved_rows, new_rows], promote_options="permissive"
|
||||
)
|
||||
total_preserved_rows += len(preserved_keys)
|
||||
|
||||
t_read = _time.perf_counter()
|
||||
logger.debug(
|
||||
"[rewrite] stream-read+join %d rows / %.1f MB (compressed) from %s "
|
||||
"across %d batch(es) in %.2fs",
|
||||
total_in_rows,
|
||||
file_size_mb,
|
||||
file_name,
|
||||
n_batches,
|
||||
t_read - t_start,
|
||||
)
|
||||
|
||||
if total_in_rows == 0:
|
||||
return (None, [])
|
||||
|
||||
if total_preserved_rows == 0:
|
||||
# Every row in this file is being upserted — delete the whole file, no preserved file needed.
|
||||
logger.debug(
|
||||
"[rewrite] %s: all %d rows matched -> whole-file delete",
|
||||
file_name,
|
||||
total_in_rows,
|
||||
)
|
||||
return (file_scan_task.file, [])
|
||||
|
||||
if total_preserved_rows == total_in_rows:
|
||||
# No rows in this file match any upsert key — leave it alone entirely.
|
||||
logger.debug("[rewrite] %s: 0 rows matched -> untouched", file_name)
|
||||
return (None, [])
|
||||
|
||||
# Derive a deterministic write_uuid from the source file path so that
|
||||
# task retries overwrite the same object rather than leaking orphan files.
|
||||
preserved_write_uuid = _uuid.UUID(hashlib.md5(file_path.encode()).hexdigest())
|
||||
preserved_files = list(
|
||||
_dataframe_to_data_files(
|
||||
table_metadata=table_metadata,
|
||||
df=preserved_rows,
|
||||
io=io,
|
||||
write_uuid=preserved_write_uuid,
|
||||
)
|
||||
)
|
||||
logger.debug(
|
||||
"[rewrite] %s: %d/%d rows preserved -> wrote %d preserved file(s) in %.2fs",
|
||||
file_name,
|
||||
total_preserved_rows,
|
||||
total_in_rows,
|
||||
len(preserved_files),
|
||||
_time.perf_counter() - t_read,
|
||||
)
|
||||
return (file_scan_task.file, preserved_files)
|
||||
|
||||
|
||||
@dataclass
|
||||
class IcebergWriteResult:
|
||||
"""Result from writing blocks to Iceberg storage.
|
||||
|
||||
Attributes:
|
||||
data_files: List of DataFile objects containing metadata about written Parquet files.
|
||||
upsert_keys: PyArrow table containing key columns for upsert operations.
|
||||
schemas: List of PyArrow schemas from all non-empty blocks.
|
||||
"""
|
||||
|
||||
data_files: List["DataFile"] = field(default_factory=list)
|
||||
upsert_keys: Optional["pa.Table"] = None
|
||||
schemas: List["pa.Schema"] = field(default_factory=list)
|
||||
|
||||
|
||||
_UPSERT_COLS_ID = "join_cols"
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class IcebergDatasink(Datasink[IcebergWriteResult]):
|
||||
"""
|
||||
Iceberg datasink to write a Ray Dataset into an existing Iceberg table.
|
||||
This datasink handles concurrent writes by:
|
||||
- Each worker writes Parquet files to storage and returns DataFile metadata
|
||||
- The driver collects all DataFile objects and performs a single commit
|
||||
|
||||
Schema evolution is supported:
|
||||
- New columns in incoming data are automatically added to the table schema
|
||||
- Type promotion across blocks is handled via schema reconciliation on the driver
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
table_identifier: str,
|
||||
catalog_kwargs: Optional[Dict[str, Any]] = None,
|
||||
snapshot_properties: Optional[Dict[str, str]] = None,
|
||||
mode: SaveMode = SaveMode.APPEND,
|
||||
overwrite_filter: Optional["Expr"] = None,
|
||||
upsert_kwargs: Optional[Dict[str, Any]] = None,
|
||||
overwrite_kwargs: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
"""
|
||||
Initialize the IcebergDatasink
|
||||
|
||||
Args:
|
||||
table_identifier: The identifier of the table such as `default.taxi_dataset`
|
||||
catalog_kwargs: Optional arguments to use when setting up the Iceberg catalog
|
||||
snapshot_properties: Custom properties to write to snapshot summary
|
||||
mode: Write mode - APPEND, UPSERT, or OVERWRITE. Defaults to APPEND.
|
||||
- APPEND: Add new data without checking for duplicates
|
||||
- UPSERT: Update existing rows or insert new ones based on a join condition
|
||||
- OVERWRITE: Replace table data (all data or filtered subset)
|
||||
overwrite_filter: Optional filter for OVERWRITE mode to perform partial overwrites.
|
||||
Must be a Ray Data expression from `ray.data.expressions`. Only rows matching
|
||||
this filter are replaced. If None with OVERWRITE mode, replaces all table data.
|
||||
upsert_kwargs: Optional arguments for upsert operations.
|
||||
Supported parameters: join_cols (List[str]), case_sensitive (bool),
|
||||
branch (str). Note: This implementation uses a copy-on-write strategy
|
||||
that always updates all columns for matched keys and inserts all new keys.
|
||||
overwrite_kwargs: Optional arguments to pass through to PyIceberg's table.overwrite()
|
||||
method. Supported parameters include case_sensitive (bool) and branch (str).
|
||||
See PyIceberg documentation for details.
|
||||
|
||||
Note:
|
||||
Schema evolution is automatically enabled. New columns in the incoming data
|
||||
are automatically added to the table schema. The schema is extracted from
|
||||
the first input bundle when on_write_start is called.
|
||||
"""
|
||||
self.table_identifier = table_identifier
|
||||
self._catalog_kwargs = (catalog_kwargs or {}).copy()
|
||||
self._snapshot_properties = (snapshot_properties or {}).copy()
|
||||
self._mode = mode
|
||||
self._overwrite_filter = overwrite_filter
|
||||
self._upsert_kwargs = (upsert_kwargs or {}).copy()
|
||||
self._overwrite_kwargs = (overwrite_kwargs or {}).copy()
|
||||
|
||||
# Validate kwargs are only set for relevant modes
|
||||
if self._upsert_kwargs and self._mode != SaveMode.UPSERT:
|
||||
raise ValueError(
|
||||
f"upsert_kwargs can only be specified when mode is SaveMode.UPSERT, but mode is {self._mode}"
|
||||
)
|
||||
if self._overwrite_kwargs and self._mode != SaveMode.OVERWRITE:
|
||||
raise ValueError(
|
||||
f"overwrite_kwargs can only be specified when mode is SaveMode.OVERWRITE, but mode is {self._mode}"
|
||||
)
|
||||
if self._overwrite_filter and self._mode != SaveMode.OVERWRITE:
|
||||
raise ValueError(
|
||||
f"overwrite_filter can only be specified when mode is SaveMode.OVERWRITE, but mode is {self._mode}"
|
||||
)
|
||||
|
||||
# Remove invalid parameters from overwrite_kwargs if present
|
||||
for invalid_param, reason in [
|
||||
(
|
||||
"overwrite_filter",
|
||||
"should be passed as a separate parameter to write_iceberg()",
|
||||
),
|
||||
(
|
||||
"delete_filter",
|
||||
"is an internal PyIceberg parameter; use 'overwrite_filter' instead",
|
||||
),
|
||||
]:
|
||||
if self._overwrite_kwargs.pop(invalid_param, None) is not None:
|
||||
logger.warning(
|
||||
f"Removed '{invalid_param}' from overwrite_kwargs: {reason}"
|
||||
)
|
||||
|
||||
if "name" in self._catalog_kwargs:
|
||||
self._catalog_name = self._catalog_kwargs.pop("name")
|
||||
else:
|
||||
self._catalog_name = "default"
|
||||
|
||||
self._table: "Table" = None
|
||||
self._io: "FileIO" = None
|
||||
self._table_metadata: "TableMetadata" = None
|
||||
self._data_context = DataContext.get_current()
|
||||
|
||||
def __getstate__(self) -> dict:
|
||||
"""Exclude `_table` during pickling."""
|
||||
state = self.__dict__.copy()
|
||||
state.pop("_table", None)
|
||||
return state
|
||||
|
||||
def __setstate__(self, state: dict) -> None:
|
||||
self.__dict__.update(state)
|
||||
self._table = None
|
||||
|
||||
def _with_retry(self, func: Callable, description: str) -> Any:
|
||||
"""Execute a function with retry logic.
|
||||
|
||||
This helper encapsulates the common retry pattern for Iceberg catalog
|
||||
operations, using the configured retry parameters from DataContext.
|
||||
|
||||
Args:
|
||||
func: The callable to execute with retry logic.
|
||||
description: Human-readable description for logging/error messages.
|
||||
|
||||
Returns:
|
||||
The result of calling func.
|
||||
"""
|
||||
iceberg_config = self._data_context.iceberg_config
|
||||
return call_with_retry(
|
||||
func,
|
||||
description=description,
|
||||
match=iceberg_config.catalog_retried_errors,
|
||||
max_attempts=iceberg_config.catalog_max_attempts,
|
||||
max_backoff_s=iceberg_config.catalog_retry_max_backoff_s,
|
||||
)
|
||||
|
||||
def _get_catalog(self) -> "Catalog":
|
||||
from pyiceberg import catalog
|
||||
|
||||
return self._with_retry(
|
||||
lambda: catalog.load_catalog(self._catalog_name, **self._catalog_kwargs),
|
||||
description=f"load Iceberg catalog '{self._catalog_name}'",
|
||||
)
|
||||
|
||||
def _reload_table(self) -> None:
|
||||
"""Reload the Iceberg table from the catalog."""
|
||||
cat = self._get_catalog()
|
||||
self._table = self._with_retry(
|
||||
lambda: cat.load_table(self.table_identifier),
|
||||
description=f"load Iceberg table '{self.table_identifier}'",
|
||||
)
|
||||
self._io = self._table.io
|
||||
self._table_metadata = self._table.metadata
|
||||
|
||||
def _get_upsert_cols(self) -> List[str]:
|
||||
"""Get join columns for upsert, using table identifier fields as fallback."""
|
||||
upsert_cols = self._upsert_kwargs.get(_UPSERT_COLS_ID, [])
|
||||
if not upsert_cols:
|
||||
# Use table's identifier fields as fallback
|
||||
identifier_cols = []
|
||||
schema = self._table_metadata.schema()
|
||||
for field_id in schema.identifier_field_ids:
|
||||
col_name = schema.find_column_name(field_id)
|
||||
if col_name:
|
||||
identifier_cols.append(col_name)
|
||||
return identifier_cols
|
||||
|
||||
case_sensitive = self._upsert_kwargs.get("case_sensitive", True)
|
||||
|
||||
# To support case insensitivity, we need to define a mapping of
|
||||
# provided (possibly case-modified) names to their original names in the schema
|
||||
if not case_sensitive:
|
||||
schema = self._table_metadata.schema()
|
||||
lower_to_original_mapping = {
|
||||
col.name.lower(): col.name for col in schema.fields
|
||||
}
|
||||
resolved_upsert_cols = []
|
||||
for upsert_col in upsert_cols:
|
||||
resolved_col = lower_to_original_mapping.get(upsert_col.lower())
|
||||
if resolved_col is None:
|
||||
raise ValueError(
|
||||
f"Upsert join column {upsert_col!r} does not match any column in "
|
||||
f"table schema (case-insensitive)."
|
||||
)
|
||||
resolved_upsert_cols.append(resolved_col)
|
||||
upsert_cols = resolved_upsert_cols
|
||||
|
||||
return upsert_cols
|
||||
|
||||
def _build_coarse_range_filter(
|
||||
self,
|
||||
keys_table: "pa.Table",
|
||||
upsert_cols: List[str],
|
||||
) -> "BooleanExpression":
|
||||
"""Build an O(1) coarse range filter covering all upsert key values.
|
||||
|
||||
For each upsert column computes AND(GTE(col, min), LTE(col, max)).
|
||||
The filter may match rows outside the upsert batch (filter overshoot);
|
||||
callers must anti-join to identify and preserve those rows.
|
||||
"""
|
||||
import pyarrow.compute as pc
|
||||
from pyiceberg.expressions import (
|
||||
AlwaysTrue,
|
||||
And,
|
||||
GreaterThanOrEqual,
|
||||
LessThanOrEqual,
|
||||
)
|
||||
|
||||
expr = None
|
||||
for col_name in upsert_cols:
|
||||
mm = pc.min_max(keys_table[col_name])
|
||||
min_val = mm["min"].as_py()
|
||||
max_val = mm["max"].as_py()
|
||||
if min_val is None:
|
||||
continue
|
||||
col_expr = And(
|
||||
GreaterThanOrEqual(col_name, min_val),
|
||||
LessThanOrEqual(col_name, max_val),
|
||||
)
|
||||
expr = col_expr if expr is None else And(expr, col_expr)
|
||||
|
||||
return expr if expr is not None else AlwaysTrue()
|
||||
|
||||
def _commit_upsert_scan_merge(
|
||||
self,
|
||||
txn: "Table.transaction",
|
||||
data_files: List["DataFile"],
|
||||
keys_table: "pa.Table",
|
||||
upsert_cols: List[str],
|
||||
) -> None:
|
||||
"""Upsert commit using coarse range filter + per-file distributed anti-join.
|
||||
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ Stage 1: Build coarse filter (driver) │
|
||||
│ keys_table ──► min/max per col ──► coarse_filter │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ Stage 2: Plan candidate files (driver) │
|
||||
│ table.scan(coarse_filter).plan_files() │
|
||||
│ ──► file_scan_tasks │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ Stage 3: Rewrite (one _rewrite_iceberg_file task per file) │
|
||||
│ read file ─► anti-join keys ─► write preserved rows │
|
||||
│ returns (old_file, preserved_files) │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ Stage 4: Atomic overwrite (driver) │
|
||||
│ delete old_file (each rewritten candidate) │
|
||||
│ append preserved_files (preserved rows kept) │
|
||||
│ append data_files (new upsert payload) │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
commit_transaction
|
||||
|
||||
1. Build an O(1) coarse range filter using min-max covering upsert key values (for each column).
|
||||
2. plan_files() on the driver to find candidate files that could be updated
|
||||
3. Dispatch one Ray task per candidate file. Each task reads its file,
|
||||
anti-joins against the upsert keys to find preserved rows (rows that
|
||||
the coarse delete would remove but that are NOT being upserted), and
|
||||
writes them as new data files directly to storage.
|
||||
4. Commit atomically via txn.update_snapshot().overwrite(): delete each
|
||||
original candidate file and append preserved files + new upsert data files.
|
||||
"""
|
||||
import time
|
||||
|
||||
case_sensitive = self._upsert_kwargs.get("case_sensitive", True)
|
||||
branch = self._upsert_kwargs.get("branch", "main")
|
||||
unknown = set(self._upsert_kwargs) - {
|
||||
_UPSERT_COLS_ID,
|
||||
"case_sensitive",
|
||||
"branch",
|
||||
}
|
||||
if unknown:
|
||||
logger.warning(
|
||||
"[scan-merge] ignoring unsupported upsert_kwargs: %s", sorted(unknown)
|
||||
)
|
||||
|
||||
# Dedup keys to minimise per-task anti-join hash table size.
|
||||
keys_table = keys_table.group_by(upsert_cols).aggregate([])
|
||||
|
||||
coarse_filter = self._build_coarse_range_filter(keys_table, upsert_cols)
|
||||
logger.debug("[scan-merge] coarse_filter=%s", coarse_filter)
|
||||
|
||||
# plan_files() reads only manifest metadata, no Parquet data on the driver.
|
||||
t0 = time.perf_counter()
|
||||
scan: "DataScan" = self._table.scan(
|
||||
row_filter=coarse_filter, case_sensitive=case_sensitive
|
||||
)
|
||||
# Use the specific branch for the scan
|
||||
scan = scan.use_ref(branch)
|
||||
file_scan_tasks: List["FileScanTask"] = list(scan.plan_files())
|
||||
|
||||
logger.info(
|
||||
"[scan-merge] planned %d candidate file(s) in %.2fs",
|
||||
len(file_scan_tasks),
|
||||
time.perf_counter() - t0,
|
||||
)
|
||||
|
||||
if not file_scan_tasks:
|
||||
# No existing files match the coarse filter, so it's a pure insert.
|
||||
self._append_and_commit(txn, data_files, branch=branch)
|
||||
return
|
||||
|
||||
# Put the deduped keys in the object store once; all tasks share one copy.
|
||||
keys_ref = ray.put(keys_table)
|
||||
|
||||
t0 = time.perf_counter()
|
||||
refs = [
|
||||
_rewrite_iceberg_file.options(
|
||||
memory=int(
|
||||
task.file.file_size_in_bytes
|
||||
* PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT
|
||||
* 3 # Bump memory estimate to account for the anti-join and the preserved rows (also since to_record_batches materializes the entire table in memory, see https://github.com/apache/iceberg-python/issues/3036)
|
||||
),
|
||||
num_cpus=1,
|
||||
).remote(task, keys_ref, upsert_cols, self._table_metadata, self._io)
|
||||
for task in file_scan_tasks
|
||||
]
|
||||
logger.info("[scan-merge] dispatched %d rewrite task(s)", len(refs))
|
||||
|
||||
# Collect results with periodic progress logs so long rewrites aren't silent.
|
||||
results = []
|
||||
pending = list(refs)
|
||||
_LOG_INTERVAL = max(1, len(refs) // 10) # log ~10 times total
|
||||
while pending:
|
||||
done, pending = ray.wait(
|
||||
pending,
|
||||
num_returns=min(_LOG_INTERVAL, len(pending)),
|
||||
timeout=_REWRITE_STALL_TIMEOUT_S,
|
||||
fetch_local=True,
|
||||
)
|
||||
results.extend(ray.get(done))
|
||||
logger.debug(
|
||||
"[scan-merge] rewrite progress: %d/%d file(s) done (%.1fs elapsed)",
|
||||
len(results),
|
||||
len(refs),
|
||||
time.perf_counter() - t0,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"[scan-merge] all %d file(s) rewritten in %.2fs",
|
||||
len(refs),
|
||||
time.perf_counter() - t0,
|
||||
)
|
||||
|
||||
# Count how many files were wholly deleted vs partially rewritten.
|
||||
n_whole_delete = n_partial = n_untouched = 0
|
||||
for old, preserved_files in results:
|
||||
if old is None:
|
||||
n_untouched += 1
|
||||
elif preserved_files:
|
||||
n_partial += 1
|
||||
else:
|
||||
n_whole_delete += 1
|
||||
logger.info(
|
||||
"[scan-merge] files: %d whole-delete, %d partial-rewrite, %d untouched",
|
||||
n_whole_delete,
|
||||
n_partial,
|
||||
n_untouched,
|
||||
)
|
||||
|
||||
# Single atomic commit: schema update (already staged in txn), and overwrite.
|
||||
# _OverwriteFiles handles both file-level deletes and appends in one snapshot.
|
||||
t0 = time.perf_counter()
|
||||
with txn.update_snapshot(
|
||||
snapshot_properties=self._snapshot_properties, branch=branch
|
||||
).overwrite() as snap:
|
||||
for old_file, preserved_files in results:
|
||||
if old_file is not None:
|
||||
snap.delete_data_file(old_file)
|
||||
for preserved_file in preserved_files:
|
||||
snap.append_data_file(preserved_file)
|
||||
for df in data_files:
|
||||
snap.append_data_file(df)
|
||||
|
||||
self._with_retry(
|
||||
txn.commit_transaction,
|
||||
description=f"commit upsert transaction to Iceberg table '{self.table_identifier}'",
|
||||
)
|
||||
logger.info("[scan-merge] committed in %.2fs", time.perf_counter() - t0)
|
||||
|
||||
def _append_and_commit(
|
||||
self,
|
||||
txn: "Table.transaction",
|
||||
data_files: List["DataFile"],
|
||||
branch: str = "main",
|
||||
) -> None:
|
||||
"""Append data files to a transaction and commit.
|
||||
|
||||
Args:
|
||||
txn: PyIceberg transaction object
|
||||
data_files: List of DataFile objects to append
|
||||
branch: Iceberg branch to commit the snapshot to. Defaults to "main"
|
||||
to match pyiceberg's default
|
||||
"""
|
||||
with txn._append_snapshot_producer(
|
||||
self._snapshot_properties, branch=branch
|
||||
) as append_files:
|
||||
for data_file in data_files:
|
||||
append_files.append_data_file(data_file)
|
||||
|
||||
self._with_retry(
|
||||
txn.commit_transaction,
|
||||
description=f"commit transaction to Iceberg table '{self.table_identifier}'",
|
||||
)
|
||||
|
||||
def _commit_upsert(
|
||||
self,
|
||||
txn: "Table.transaction",
|
||||
data_files: List["DataFile"],
|
||||
upsert_keys: Optional["pa.Table"],
|
||||
) -> None:
|
||||
"""
|
||||
Commit upsert transaction with copy-on-write strategy.
|
||||
|
||||
Args:
|
||||
txn: PyIceberg transaction object
|
||||
data_files: List of DataFile objects to commit
|
||||
upsert_keys: PyArrow table containing upsert key columns
|
||||
"""
|
||||
import functools
|
||||
import time
|
||||
|
||||
import pyarrow as pa
|
||||
|
||||
# Create delete filter if we have join keys
|
||||
if upsert_keys is not None and len(upsert_keys) > 0:
|
||||
# Filter out rows with any NULL values in join columns
|
||||
# (NULL != NULL in SQL semantics)
|
||||
upsert_cols = self._get_upsert_cols()
|
||||
logger.info(
|
||||
"[upsert commit] Filtering NULL keys from %d rows on cols %s",
|
||||
len(upsert_keys),
|
||||
upsert_cols,
|
||||
)
|
||||
t0 = time.perf_counter()
|
||||
masks = (pa.compute.is_valid(upsert_keys[col]) for col in upsert_cols)
|
||||
mask = functools.reduce(pa.compute.and_, masks)
|
||||
keys_table = upsert_keys.filter(mask)
|
||||
logger.info(
|
||||
"[upsert commit] NULL filter done in %.2fs: %d -> %d rows (dropped %d NULLs)",
|
||||
time.perf_counter() - t0,
|
||||
len(upsert_keys),
|
||||
len(keys_table),
|
||||
len(upsert_keys) - len(keys_table),
|
||||
)
|
||||
|
||||
# Only delete if we have non-NULL keys
|
||||
if len(keys_table) > 0:
|
||||
self._commit_upsert_scan_merge(txn, data_files, keys_table, upsert_cols)
|
||||
return
|
||||
else:
|
||||
logger.info("[upsert commit] No upsert keys — skipping delete phase")
|
||||
|
||||
# No non-NULL keys — just append new data files and commit
|
||||
logger.info(
|
||||
"[upsert commit] Appending %d data files and committing ...",
|
||||
len(data_files),
|
||||
)
|
||||
t0 = time.perf_counter()
|
||||
branch = self._upsert_kwargs.get("branch", "main")
|
||||
self._append_and_commit(txn, data_files, branch=branch)
|
||||
logger.info(
|
||||
"[upsert commit] Append+commit done in %.2fs",
|
||||
time.perf_counter() - t0,
|
||||
)
|
||||
|
||||
def _preserve_identifier_field_requirements(
|
||||
self, update: "UpdateSchema", table_schema: "Schema"
|
||||
) -> None:
|
||||
"""Ensure identifier fields remain required after schema union.
|
||||
|
||||
When union_by_name is called with a schema that has nullable fields,
|
||||
PyIceberg may make identifier fields optional. Since identifier fields
|
||||
must be required, this helper ensures they remain required after union.
|
||||
|
||||
Example:
|
||||
Table schema: id: int (required, identifier), val: string
|
||||
Input schema: id: int (optional), val: string
|
||||
|
||||
`union_by_name` merges them to:
|
||||
id: int (optional), val: string
|
||||
|
||||
This violates the identifier constraint. This function forces `id`
|
||||
back to required in the pending update.
|
||||
|
||||
Args:
|
||||
update: The UpdateSchema object from update_schema() context manager
|
||||
table_schema: The current table schema to get identifier field IDs from
|
||||
"""
|
||||
from pyiceberg.types import NestedField
|
||||
|
||||
identifier_field_ids = table_schema.identifier_field_ids
|
||||
for field_id in identifier_field_ids:
|
||||
# Check if this field has a pending update
|
||||
if field_id in update._updates:
|
||||
updated_field = update._updates[field_id]
|
||||
# If it was made optional (likely by union_by_name), force it back to required
|
||||
if not updated_field.required:
|
||||
# Directly update the pending change to enforce required=True.
|
||||
# We create a new NestedField because it might be immutable.
|
||||
# We bypass _set_column_requirement because it has a check that
|
||||
# incorrectly returns early if the original field is already required,
|
||||
# ignoring the fact that we are overwriting a pending update.
|
||||
update._updates[field_id] = NestedField(
|
||||
field_id=updated_field.field_id,
|
||||
name=updated_field.name,
|
||||
field_type=updated_field.field_type,
|
||||
doc=updated_field.doc,
|
||||
required=True,
|
||||
initial_default=updated_field.initial_default,
|
||||
write_default=updated_field.write_default,
|
||||
)
|
||||
|
||||
def _update_schema_with_union(
|
||||
self,
|
||||
update: "UpdateSchema",
|
||||
new_schema: Union["pa.Schema", "Schema"],
|
||||
table_schema: "Schema",
|
||||
) -> None:
|
||||
"""Update schema using union_by_name while preserving identifier field requirements.
|
||||
|
||||
Args:
|
||||
update: The UpdateSchema object.
|
||||
new_schema: The new schema to union with the table schema.
|
||||
table_schema: The current table schema.
|
||||
"""
|
||||
update.union_by_name(new_schema)
|
||||
self._preserve_identifier_field_requirements(update, table_schema)
|
||||
|
||||
def on_write_start(self, schema: Optional["pa.Schema"] = None) -> None:
|
||||
"""Initialize table for writing and create a shared write UUID.
|
||||
|
||||
Args:
|
||||
schema: The PyArrow schema of the data being written. This is
|
||||
automatically extracted from the first input bundle by the
|
||||
Write operator. Used to evolve the table schema before writing
|
||||
to avoid PyIceberg name mapping errors.
|
||||
"""
|
||||
self._reload_table()
|
||||
|
||||
# Evolve schema BEFORE any files are written
|
||||
# This prevents PyIceberg name mapping errors when incoming data has new columns
|
||||
if schema is not None:
|
||||
table_schema = self._table.metadata.schema()
|
||||
|
||||
def _update_schema():
|
||||
with self._table.update_schema() as update:
|
||||
self._update_schema_with_union(update, schema, table_schema)
|
||||
|
||||
self._with_retry(
|
||||
_update_schema,
|
||||
description=f"update schema for Iceberg table '{self.table_identifier}'",
|
||||
)
|
||||
# Succeeded, reload to get latest table version and exit.
|
||||
self._reload_table()
|
||||
|
||||
# Validate join_cols for UPSERT mode before writing any files
|
||||
if self._mode == SaveMode.UPSERT:
|
||||
upsert_cols = self._upsert_kwargs.get(_UPSERT_COLS_ID, [])
|
||||
if not upsert_cols:
|
||||
# Check if table has identifier fields as fallback
|
||||
identifier_field_ids = (
|
||||
self._table.metadata.schema().identifier_field_ids
|
||||
)
|
||||
if not identifier_field_ids:
|
||||
raise ValueError(
|
||||
"join_cols must be specified in upsert_kwargs for UPSERT mode "
|
||||
"when table has no identifier fields"
|
||||
)
|
||||
|
||||
def write(self, blocks: Iterable[Block], ctx: TaskContext) -> IcebergWriteResult:
|
||||
"""
|
||||
Write blocks to Parquet files in storage and return DataFile metadata with schemas.
|
||||
|
||||
This runs on each worker in parallel. Files are written directly to storage
|
||||
(S3, HDFS, etc.) and only metadata is returned to the driver.
|
||||
Schema updates are NOT performed here - they happen on the driver.
|
||||
|
||||
Args:
|
||||
blocks: Iterable of Ray Data blocks to write
|
||||
ctx: TaskContext object containing task-specific information
|
||||
|
||||
Returns:
|
||||
IcebergWriteResult containing DataFile objects, upsert keys, and schemas.
|
||||
"""
|
||||
from pyiceberg.io.pyarrow import _dataframe_to_data_files
|
||||
|
||||
all_data_files = []
|
||||
upsert_keys_tables = []
|
||||
block_schemas = []
|
||||
use_copy_on_write_upsert = self._mode == SaveMode.UPSERT
|
||||
|
||||
for block in blocks:
|
||||
pa_table = BlockAccessor.for_block(block).to_arrow()
|
||||
if pa_table.num_rows > 0:
|
||||
block_schemas.append(pa_table.schema)
|
||||
|
||||
# Extract join key values for copy-on-write upsert
|
||||
if use_copy_on_write_upsert:
|
||||
upsert_cols = self._get_upsert_cols()
|
||||
if len(upsert_cols) > 0:
|
||||
upsert_keys_tables.append(pa_table.select(upsert_cols))
|
||||
|
||||
# Write data files to storage with retry for transient errors
|
||||
def _write_data_files():
|
||||
return list(
|
||||
_dataframe_to_data_files(
|
||||
table_metadata=self._table_metadata,
|
||||
df=pa_table,
|
||||
io=self._io,
|
||||
)
|
||||
)
|
||||
|
||||
iceberg_config = self._data_context.iceberg_config
|
||||
data_files = call_with_retry(
|
||||
_write_data_files,
|
||||
description=f"write data files to Iceberg table '{self.table_identifier}'",
|
||||
match=self._data_context.retried_io_errors,
|
||||
max_attempts=iceberg_config.write_file_max_attempts,
|
||||
max_backoff_s=iceberg_config.write_file_retry_max_backoff_s,
|
||||
)
|
||||
all_data_files.extend(data_files)
|
||||
|
||||
# Combine all upsert key tables into one
|
||||
from ray.data._internal.arrow_ops.transform_pyarrow import concat
|
||||
|
||||
upsert_keys = concat(upsert_keys_tables) if upsert_keys_tables else None
|
||||
|
||||
return IcebergWriteResult(
|
||||
data_files=all_data_files,
|
||||
upsert_keys=upsert_keys,
|
||||
schemas=block_schemas,
|
||||
)
|
||||
|
||||
def _commit_overwrite(
|
||||
self, txn: "Table.transaction", data_files: List["DataFile"]
|
||||
) -> None:
|
||||
"""Commit data files using OVERWRITE mode."""
|
||||
from pyiceberg.expressions import AlwaysTrue
|
||||
|
||||
# Default - Full overwrite - delete all
|
||||
pyi_filter = AlwaysTrue()
|
||||
|
||||
# Delete matching data if filter provided
|
||||
if self._overwrite_filter is not None:
|
||||
from ray.data._internal.datasource.iceberg_datasource import (
|
||||
_IcebergExpressionVisitor,
|
||||
)
|
||||
|
||||
visitor = _IcebergExpressionVisitor()
|
||||
pyi_filter = visitor.visit(self._overwrite_filter)
|
||||
|
||||
txn.delete(
|
||||
delete_filter=pyi_filter,
|
||||
snapshot_properties=self._snapshot_properties,
|
||||
**self._overwrite_kwargs,
|
||||
)
|
||||
|
||||
# Append on the same branch the delete targeted (defaults to "main").
|
||||
branch = self._overwrite_kwargs.get("branch", "main")
|
||||
self._append_and_commit(txn, data_files, branch=branch)
|
||||
|
||||
def on_write_complete(self, write_result: WriteResult) -> None:
|
||||
"""
|
||||
Complete the write by reconciling schemas and committing all data files.
|
||||
|
||||
This runs on the driver after all workers finish writing files.
|
||||
Collects all DataFile objects and schemas from all workers, reconciles schemas
|
||||
(allowing type promotion), updates table schema if needed, then performs a single
|
||||
atomic commit.
|
||||
"""
|
||||
import time
|
||||
|
||||
t_start = time.perf_counter()
|
||||
logger.info("[on_write_complete] Starting commit phase (mode=%s)", self._mode)
|
||||
|
||||
# Collect all data files and schemas from all workers
|
||||
all_data_files: List["DataFile"] = []
|
||||
all_schemas: List["pa.Schema"] = []
|
||||
upsert_keys_tables: List["pa.Table"] = []
|
||||
|
||||
for write_return in write_result.write_returns:
|
||||
if not write_return:
|
||||
continue
|
||||
|
||||
if write_return.data_files: # Only add schema if we have data files
|
||||
all_data_files.extend(write_return.data_files)
|
||||
all_schemas.extend(write_return.schemas)
|
||||
if write_return.upsert_keys is not None:
|
||||
upsert_keys_tables.append(write_return.upsert_keys)
|
||||
|
||||
logger.info(
|
||||
"[on_write_complete] Collected results: %d data files, %d schema blocks, "
|
||||
"%d upsert key batches from workers (%.2fs)",
|
||||
len(all_data_files),
|
||||
len(all_schemas),
|
||||
len(upsert_keys_tables),
|
||||
time.perf_counter() - t_start,
|
||||
)
|
||||
|
||||
if not all_data_files:
|
||||
logger.info("[on_write_complete] No data files written, nothing to commit")
|
||||
return
|
||||
|
||||
# Concatenate all upsert keys from all workers into a single table
|
||||
from ray.data._internal.arrow_ops.transform_pyarrow import concat
|
||||
|
||||
if upsert_keys_tables:
|
||||
total_key_rows = sum(len(t) for t in upsert_keys_tables)
|
||||
logger.info(
|
||||
"[on_write_complete] Concatenating %d upsert key batches (%d total rows) ...",
|
||||
len(upsert_keys_tables),
|
||||
total_key_rows,
|
||||
)
|
||||
t0 = time.perf_counter()
|
||||
upsert_keys = concat(upsert_keys_tables)
|
||||
logger.info(
|
||||
"[on_write_complete] upsert key concat done in %.2fs: %d rows, cols=%s",
|
||||
time.perf_counter() - t0,
|
||||
len(upsert_keys),
|
||||
upsert_keys.column_names,
|
||||
)
|
||||
else:
|
||||
upsert_keys = None
|
||||
|
||||
# Reconcile all schemas from all blocks across all workers
|
||||
# Get table schema and union with reconciled schema using unify_schemas with promotion
|
||||
from pyiceberg.io import pyarrow as pyi_pa_io
|
||||
|
||||
from ray.data._internal.arrow_ops.transform_pyarrow import unify_schemas
|
||||
|
||||
logger.info("[on_write_complete] Reconciling %d schemas ...", len(all_schemas))
|
||||
t0 = time.perf_counter()
|
||||
table_schema = pyi_pa_io.schema_to_pyarrow(self._table.schema())
|
||||
final_reconciled_schema = unify_schemas(
|
||||
[table_schema] + all_schemas, promote_types=True
|
||||
)
|
||||
logger.info(
|
||||
"[on_write_complete] Schema reconciliation done in %.2fs",
|
||||
time.perf_counter() - t0,
|
||||
)
|
||||
|
||||
# Create transaction and commit schema update + data files atomically
|
||||
txn = self._table.transaction()
|
||||
|
||||
# Update table schema within the transaction if it differs
|
||||
if not final_reconciled_schema.equals(table_schema):
|
||||
logger.info(
|
||||
"[on_write_complete] Schema changed — updating table schema ..."
|
||||
)
|
||||
t0 = time.perf_counter()
|
||||
current_table_schema = self._table.metadata.schema()
|
||||
with txn.update_schema() as update:
|
||||
self._update_schema_with_union(
|
||||
update, final_reconciled_schema, current_table_schema
|
||||
)
|
||||
logger.info(
|
||||
"[on_write_complete] Schema update done in %.2fs",
|
||||
time.perf_counter() - t0,
|
||||
)
|
||||
else:
|
||||
logger.info("[on_write_complete] Schema unchanged, skipping update")
|
||||
|
||||
# Create transaction and commit based on mode
|
||||
logger.info(
|
||||
"[on_write_complete] Starting %s commit for %d data files ...",
|
||||
self._mode,
|
||||
len(all_data_files),
|
||||
)
|
||||
t0 = time.perf_counter()
|
||||
if self._mode == SaveMode.APPEND:
|
||||
self._append_and_commit(txn, all_data_files)
|
||||
elif self._mode == SaveMode.OVERWRITE:
|
||||
self._commit_overwrite(txn, all_data_files)
|
||||
elif self._mode == SaveMode.UPSERT:
|
||||
self._commit_upsert(txn, all_data_files, upsert_keys)
|
||||
else:
|
||||
raise ValueError(f"Unsupported mode: {self._mode}")
|
||||
logger.info(
|
||||
"[on_write_complete] Commit complete in %.2fs (total on_write_complete=%.2fs)",
|
||||
time.perf_counter() - t0,
|
||||
time.perf_counter() - t_start,
|
||||
)
|
||||
Reference in New Issue
Block a user