Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion .github/workflows/test-python.yml
Original file line number Diff line number Diff line change
Expand Up @@ -105,4 +105,7 @@ jobs:
PY

- name: Run Iceberg compatibility tests
run: uv run --no-sync pytest tests/tower/test_tables.py tests/tower/test_storage.py
run: >-
uv run --no-sync pytest
tests/tower/test_table*.py
tests/tower/test_storage.py
149 changes: 71 additions & 78 deletions src/tower/_tables.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,16 @@
from __future__ import annotations

import math
import os
import random
import time
from dataclasses import dataclass
from typing import List, Optional, TypeVar, Union
from typing import Callable, List, Optional, TypeVar, Union

from pyiceberg.exceptions import CommitFailedException, NoSuchTableError

TTable = TypeVar("TTable", bound="Table")

import random
import time
TRetryResult = TypeVar("TRetryResult")

import polars as pl
import pyarrow as pa
Expand Down Expand Up @@ -37,6 +38,8 @@
namespace_or_default,
)

_MAX_COMMIT_RETRY_DELAY_SECONDS = 30.0


@dataclass
class RowsAffectedInformation:
Expand Down Expand Up @@ -249,8 +252,30 @@ def rows_affected(self) -> RowsAffectedInformation:
def _validate_retry_args(max_retries: int, retry_delay_seconds: float) -> None:
if max_retries < 0:
raise ValueError("max_retries must be >= 0")
if retry_delay_seconds < 0:
raise ValueError("retry_delay_seconds must be >= 0")
if not math.isfinite(retry_delay_seconds) or retry_delay_seconds < 0:
raise ValueError("retry_delay_seconds must be finite and >= 0")

def _commit_with_retry(
self,
operation: Callable[[], TRetryResult],
max_retries: int,
retry_delay_seconds: float,
) -> TRetryResult:
retry_ceiling = min(retry_delay_seconds, _MAX_COMMIT_RETRY_DELAY_SECONDS)

for attempt in range(max_retries + 1):
try:
return operation()
except CommitFailedException:
if attempt == max_retries:
raise

delay = random.uniform(0.0, retry_ceiling)
time.sleep(delay)
self._table.refresh()
retry_ceiling = min(retry_ceiling * 2, _MAX_COMMIT_RETRY_DELAY_SECONDS)

raise AssertionError("unreachable")

def insert(
self,
Expand All @@ -270,8 +295,8 @@ def insert(
must match the schema of the target table.
max_retries (int): Maximum number of retry attempts on commit conflicts.
Defaults to 5.
retry_delay_seconds (float): Wait time in seconds between retries.
Defaults to 0.5 seconds.
retry_delay_seconds (float): Base delay in seconds for bounded exponential
backoff with full jitter. Defaults to 0.5 seconds.

Returns:
TTable: The table instance with the newly inserted rows, allowing for method chaining.
Expand All @@ -296,23 +321,11 @@ def insert(
self._validate_retry_args(max_retries, retry_delay_seconds)
self._ensure_read_write_table()

last_exception = None

for attempt in range(max_retries + 1):
try:
if attempt > 0:
self._table.refresh()

self._table.append(data)
self._stats.inserts += data.num_rows
return self

except CommitFailedException as e:
last_exception = e
if attempt < max_retries:
time.sleep(retry_delay_seconds)

raise last_exception
self._commit_with_retry(
lambda: self._table.append(data), max_retries, retry_delay_seconds
)
self._stats.inserts += data.num_rows
return self

def upsert(
self,
Expand All @@ -337,8 +350,8 @@ def upsert(
If not provided, all columns will be used for matching.
max_retries (int): Maximum number of retry attempts on commit conflicts.
Defaults to 5.
retry_delay_seconds (float): Wait time in seconds between retries.
Defaults to 0.5 seconds.
retry_delay_seconds (float): Base delay in seconds for bounded exponential
backoff with full jitter. Defaults to 0.5 seconds.

Returns:
TTable: The table instance with the upserted rows, allowing for method chaining.
Expand Down Expand Up @@ -370,34 +383,24 @@ def upsert(
self._validate_retry_args(max_retries, retry_delay_seconds)
self._ensure_read_write_table()

last_exception = None
res = self._commit_with_retry(
lambda: self._table.upsert(
data,
join_cols=join_cols,
# All upserts will always be case sensitive. Perhaps we'll add this
# as a parameter in the future?
case_sensitive=True,
# These are the defaults, but we're including them to be complete.
when_matched_update_all=True,
when_not_matched_insert_all=True,
),
max_retries,
retry_delay_seconds,
)

for attempt in range(max_retries + 1):
try:
if attempt > 0:
self._table.refresh()

res = self._table.upsert(
data,
join_cols=join_cols,
# All upserts will always be case sensitive. Perhaps we'll add this
# as a parameter in the future?
case_sensitive=True,
# These are the defaults, but we're including them to be complete.
when_matched_update_all=True,
when_not_matched_insert_all=True,
)

self._stats.updates += res.rows_updated
self._stats.inserts += res.rows_inserted
return self

except CommitFailedException as e:
last_exception = e
if attempt < max_retries:
time.sleep(retry_delay_seconds)

raise last_exception
self._stats.updates += res.rows_updated
self._stats.inserts += res.rows_inserted
return self

def delete(
self,
Expand All @@ -421,8 +424,8 @@ def delete(
- A string expression
max_retries (int): Maximum number of retry attempts on commit conflicts.
Defaults to 5.
retry_delay_seconds (float): Wait time in seconds between retries.
Defaults to 0.5 seconds.
retry_delay_seconds (float): Base delay in seconds for bounded exponential
backoff with full jitter. Defaults to 0.5 seconds.

Returns:
TTable: The table instance with the deleted rows, allowing for method chaining.
Expand Down Expand Up @@ -455,30 +458,20 @@ def delete(
next_filters = convert_pyarrow_expressions(filters)
filters = next_filters

last_exception = None

for attempt in range(max_retries + 1):
try:
if attempt > 0:
self._table.refresh()

self._table.delete(
delete_filter=filters,
# We want this to always be the case. Not sure why you wouldn't?
case_sensitive=True,
)

# NOTE: There is, unfortunately, no way to get the number of rows
# deleted besides comparing the two snapshots that were created.

return self
self._commit_with_retry(
lambda: self._table.delete(
delete_filter=filters,
# We want this to always be the case. Not sure why you wouldn't?
case_sensitive=True,
),
max_retries,
retry_delay_seconds,
)

except CommitFailedException as e:
last_exception = e
if attempt < max_retries:
time.sleep(retry_delay_seconds)
# NOTE: There is, unfortunately, no way to get the number of rows
# deleted besides comparing the two snapshots that were created.

raise last_exception
return self

def schema(self) -> pa.Schema:
"""
Expand Down
Loading
Loading