Skip to content

fix: allow reading pyarrow timestamp as iceberg timestamptz #2333

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Aug 19, 2025
Merged

fix: allow reading pyarrow timestamp as iceberg timestamptz #2333

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f38d79b
add unit test for projecting timestamp to timestamptz
kevinjqliu Aug 14, 2025
deb56d9
reorder
kevinjqliu Aug 19, 2025
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
20 changes: 11 additions & 9 deletions pyiceberg/io/pyarrow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1776,14 +1776,7 @@ def _cast_if_needed(self, field: NestedField, values: pa.Array) -> pa.Array:
file_field = self._file_schema.find_field(field.field_id)

if field.field_type.is_primitive:
if field.field_type != file_field.field_type:
target_schema = schema_to_pyarrow(
promote(file_field.field_type, field.field_type), include_field_ids=self._include_field_ids
)
if self._use_large_types is False:
target_schema = _pyarrow_schema_ensure_small_types(target_schema)
return values.cast(target_schema)
elif (target_type := schema_to_pyarrow(field.field_type, include_field_ids=self._include_field_ids)) != values.type:
if (target_type := schema_to_pyarrow(field.field_type, include_field_ids=self._include_field_ids)) != values.type:
if field.field_type == TimestampType():
# Downcasting of nanoseconds to microseconds
if (
Expand All @@ -1802,13 +1795,22 @@ def _cast_if_needed(self, field: NestedField, values: pa.Array) -> pa.Array:
pa.types.is_timestamp(target_type)
and target_type.tz == "UTC"
and pa.types.is_timestamp(values.type)
and values.type.tz in UTC_ALIASES
and (values.type.tz in UTC_ALIASES or values.type.tz is None)
Copy link
Contributor Author

@kevinjqliu kevinjqliu Aug 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

also allow when type.tz is None, mirrors the logic here

iceberg-python/pyiceberg/io/pyarrow.py

Lines 1350 to 1353 in 8b43eb8

if primitive.tz in UTC_ALIASES:
return TimestamptzType()
elif primitive.tz is None:
return TimestampType()

):
if target_type.unit == "us" and values.type.unit == "ns" and self._downcast_ns_timestamp_to_us:
return values.cast(target_type, safe=False)
elif target_type.unit == "us" and values.type.unit in {"s", "ms", "us"}:
return values.cast(target_type)
raise ValueError(f"Unsupported schema projection from {values.type} to {target_type}")

if field.field_type != file_field.field_type:
target_schema = schema_to_pyarrow(
promote(file_field.field_type, field.field_type), include_field_ids=self._include_field_ids
)
if self._use_large_types is False:
target_schema = _pyarrow_schema_ensure_small_types(target_schema)
return values.cast(target_schema)

return values

def _construct_field(self, field: NestedField, arrow_type: pa.DataType) -> pa.Field:
Expand Down
39 changes: 39 additions & 0 deletions tests/io/test_pyarrow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2521,6 +2521,45 @@ def test_initial_value() -> None:
assert val.as_py() == 22


def test__to_requested_schema_timestamp_to_timestamptz_projection() -> None:
from datetime import datetime, timezone

# file is written with timestamp without timezone
file_schema = Schema(NestedField(1, "ts_field", TimestampType(), required=False))
batch = pa.record_batch(
[
pa.array(
[
datetime(2025, 8, 14, 12, 0, 0),
datetime(2025, 8, 14, 13, 0, 0),
],
type=pa.timestamp("us"),
)
],
names=["ts_field"],
)

# table is written with timestamp with timezone
table_schema = Schema(NestedField(1, "ts_field", TimestamptzType(), required=False))

actual_result = _to_requested_schema(table_schema, file_schema, batch, downcast_ns_timestamp_to_us=True)
expected = pa.record_batch(
[
pa.array(
[
datetime(2025, 8, 14, 12, 0, 0),
datetime(2025, 8, 14, 13, 0, 0),
],
type=pa.timestamp("us", tz=timezone.utc),
)
],
names=["ts_field"],
)

# expect actual_result to have timezone
assert expected.equals(actual_result)


def test__to_requested_schema_timestamps(
arrow_table_schema_with_all_timestamp_precisions: pa.Schema,
arrow_table_with_all_timestamp_precisions: pa.Table,
Expand Down