diff --git a/pyiceberg/expressions/literals.py b/pyiceberg/expressions/literals.py index 39922fde33..153e94eb17 100644 --- a/pyiceberg/expressions/literals.py +++ b/pyiceberg/expressions/literals.py @@ -45,7 +45,9 @@ IntegerType, LongType, StringType, + TimestampNanoType, TimestampType, + TimestamptzNanoType, TimestamptzType, TimeType, UUIDType, @@ -56,11 +58,17 @@ datetime_to_micros, days_to_date, micros_to_days, + micros_to_nanos, micros_to_timestamp, + nanos_to_days, + nanos_to_micros, + nanos_to_timestamp, time_str_to_micros, time_to_micros, timestamp_to_micros, + timestamp_to_nanos, timestamptz_to_micros, + timestamptz_to_nanos, ) from pyiceberg.utils.decimal import decimal_to_unscaled, unscaled_to_decimal from pyiceberg.utils.singleton import Singleton @@ -347,6 +355,16 @@ def _(self, _: TimestampType) -> Literal[int]: def _(self, _: TimestamptzType) -> Literal[int]: return TimestampLiteral(self.value) + @to.register(TimestampNanoType) + def _(self, type_var: TimestampNanoType) -> Literal[int]: + # The value is assumed to be in microseconds, to match the TimestampType case above + return TimestampLiteral(self.value).to(type_var) + + @to.register(TimestamptzNanoType) + def _(self, type_var: TimestamptzNanoType) -> Literal[int]: + # The value is assumed to be in microseconds, to match the TimestamptzType case above + return TimestampLiteral(self.value).to(type_var) + @to.register(DecimalType) def _(self, type_var: DecimalType) -> Literal[Decimal]: unscaled = Decimal(self.value) @@ -495,6 +513,54 @@ def _(self, _: TimestamptzType) -> Literal[int]: def _(self, _: DateType) -> Literal[int]: return DateLiteral(micros_to_days(self.value)) + @to.register(TimestampNanoType) + def _(self, _: TimestampNanoType) -> Literal[int]: + return TimestampNanoLiteral(micros_to_nanos(self.value)) + + @to.register(TimestamptzNanoType) + def _(self, _: TimestamptzNanoType) -> Literal[int]: + return TimestampNanoLiteral(micros_to_nanos(self.value)) + + +class TimestampNanoLiteral(Literal[int]): + def __init__(self, value: int) -> None: + super().__init__(value, int) + + @model_serializer + def ser_model(self) -> str: + # Python datetime only goes down to microseconds, so the last three digits are appended + return f"{nanos_to_timestamp(self.root).isoformat(timespec='microseconds')}{self.root % 1000:03d}" + + def increment(self) -> Literal[int]: + return TimestampNanoLiteral(self.value + 1) + + def decrement(self) -> Literal[int]: + return TimestampNanoLiteral(self.value - 1) + + @singledispatchmethod + def to(self, type_var: IcebergType) -> Literal: # type: ignore + raise TypeError(f"Cannot convert TimestampNanoLiteral into {type_var}") + + @to.register(TimestampNanoType) + def _(self, _: TimestampNanoType) -> Literal[int]: + return self + + @to.register(TimestamptzNanoType) + def _(self, _: TimestamptzNanoType) -> Literal[int]: + return self + + @to.register(TimestampType) + def _(self, _: TimestampType) -> Literal[int]: + return TimestampLiteral(nanos_to_micros(self.value)) + + @to.register(TimestamptzType) + def _(self, _: TimestamptzType) -> Literal[int]: + return TimestampLiteral(nanos_to_micros(self.value)) + + @to.register(DateType) + def _(self, _: DateType) -> Literal[int]: + return DateLiteral(nanos_to_days(self.value)) + class DecimalLiteral(Literal[Decimal]): def __init__(self, value: Decimal) -> None: @@ -615,6 +681,14 @@ def _(self, _: TimestampType) -> Literal[int]: def _(self, _: TimestamptzType) -> Literal[int]: return TimestampLiteral(timestamptz_to_micros(self.value)) + @to.register(TimestampNanoType) + def _(self, _: TimestampNanoType) -> Literal[int]: + return TimestampNanoLiteral(timestamp_to_nanos(self.value)) + + @to.register(TimestamptzNanoType) + def _(self, _: TimestamptzNanoType) -> Literal[int]: + return TimestampNanoLiteral(timestamptz_to_nanos(self.value)) + @to.register(UUIDType) def _(self, _: UUIDType) -> Literal[bytes]: return UUIDLiteral(UUID(self.value).bytes) diff --git a/pyiceberg/transforms.py b/pyiceberg/transforms.py index 5e0027a829..46ef38a69b 100644 --- a/pyiceberg/transforms.py +++ b/pyiceberg/transforms.py @@ -65,6 +65,7 @@ Literal, LongLiteral, TimestampLiteral, + TimestampNanoLiteral, literal, ) from pyiceberg.typedef import IcebergRootModel, L @@ -1049,7 +1050,7 @@ def _truncate_number( ) -> UnboundPredicate | None: boundary = pred.literal - if not isinstance(boundary, (LongLiteral, DecimalLiteral, DateLiteral, TimestampLiteral)): + if not isinstance(boundary, (LongLiteral, DecimalLiteral, DateLiteral, TimestampLiteral, TimestampNanoLiteral)): raise ValueError(f"Expected a numeric literal, got: {type(boundary)}") if isinstance(pred, BoundLessThan): @@ -1071,7 +1072,7 @@ def _truncate_number_strict( ) -> UnboundPredicate | None: boundary = pred.literal - if not isinstance(boundary, (LongLiteral, DecimalLiteral, DateLiteral, TimestampLiteral)): + if not isinstance(boundary, (LongLiteral, DecimalLiteral, DateLiteral, TimestampLiteral, TimestampNanoLiteral)): raise ValueError(f"Expected a numeric literal, got: {type(boundary)}") if isinstance(pred, BoundLessThan): diff --git a/pyiceberg/utils/datetime.py b/pyiceberg/utils/datetime.py index ea7329ea20..276e1593e0 100644 --- a/pyiceberg/utils/datetime.py +++ b/pyiceberg/utils/datetime.py @@ -26,6 +26,8 @@ timedelta, ) +from pyiceberg.types import LongType + EPOCH_DATE = date.fromisoformat("1970-01-01") EPOCH_TIMESTAMP = datetime.fromisoformat("1970-01-01T00:00:00.000000") ISO_TIMESTAMP = re.compile(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(.\d{1,6})?") @@ -35,6 +37,14 @@ ISO_TIMESTAMPTZ_NANO = re.compile(r"(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})(.\d{1,6})?(\d{1,3})?([-+]\d{2}:\d{2})") +def _check_nanos_range(nanos: int, source: int | str) -> int: + """Reject a nanosecond timestamp that does not fit in a signed 64-bit integer.""" + # Python integers are unbounded, so unlike Java this does not overflow on its own + if not LongType.min <= nanos <= LongType.max: + raise OverflowError(f"Timestamp cannot be converted to nanoseconds, out of range: {source}") + return nanos + + def micros_to_days(timestamp: int) -> int: """Convert a timestamp in microseconds to a date in days.""" return timedelta(microseconds=timestamp).days @@ -126,7 +136,8 @@ def timestamp_to_nanos(timestamp_str: str) -> int: ns_str = (match.group(3) or "0").ljust(3, "0") ms_str = match.group(2) if match.group(2) else "" timestamp_str_without_ns_str = match.group(1) + ms_str - return datetime_to_nanos(datetime.fromisoformat(timestamp_str_without_ns_str)) + int(ns_str) + nanos = datetime_to_nanos(datetime.fromisoformat(timestamp_str_without_ns_str)) + int(ns_str) + return _check_nanos_range(nanos, timestamp_str) if ISO_TIMESTAMPTZ_NANO.fullmatch(timestamp_str): # When we can match a timestamp without a zone, we can give a more specific error raise ValueError(f"Zone offset provided, but not expected: {timestamp_str}") @@ -143,7 +154,8 @@ def timestamptz_to_nanos(timestamptz_str: str) -> int: ns_str = (match.group(3) or "0").ljust(3, "0") ms_str = match.group(2) if match.group(2) else "" timestamptz_str_without_ns_str = match.group(1) + ms_str + match.group(4) - return datetime_to_nanos(datetime.fromisoformat(timestamptz_str_without_ns_str)) + int(ns_str) + nanos = datetime_to_nanos(datetime.fromisoformat(timestamptz_str_without_ns_str)) + int(ns_str) + return _check_nanos_range(nanos, timestamptz_str) if ISO_TIMESTAMP_NANO.fullmatch(timestamptz_str): # When we can match a timestamp without a zone, we can give a more specific error raise ValueError(f"Missing zone offset: {timestamptz_str} (must be ISO-8601)") @@ -283,3 +295,8 @@ def nanos_to_hours(nanos: int) -> int: def nanos_to_micros(nanos: int) -> int: """Convert a nanoseconds timestamp to microsecond timestamp by dropping precision.""" return nanos // 1000 + + +def micros_to_nanos(micros: int) -> int: + """Convert a microseconds timestamp to a nanosecond timestamp.""" + return _check_nanos_range(micros * 1000, micros) diff --git a/tests/expressions/test_expressions.py b/tests/expressions/test_expressions.py index 8ce48a6897..78315e62d4 100644 --- a/tests/expressions/test_expressions.py +++ b/tests/expressions/test_expressions.py @@ -77,6 +77,8 @@ NestedField, StringType, StructType, + TimestampNanoType, + TimestamptzNanoType, ) @@ -110,6 +112,20 @@ def test_invert_not_nan_bind() -> None: assert ~NotNaN(Reference("a")).bind(schema) == IsNaN(Reference("a")).bind(schema) +def test_bind_timestamp_nano() -> None: + schema = Schema(NestedField(2, "a", TimestampNanoType(), required=False), schema_id=1) + + assert GreaterThan("a", "2017-08-18T14:21:01.919234567").bind(schema).literal.value == 1503066061919234567 + # A long is read as microseconds, matching the plain timestamp case + assert EqualTo("a", 1503066061919234).bind(schema).literal.value == 1503066061919234000 + + +def test_bind_timestamptz_nano() -> None: + schema = Schema(NestedField(2, "a", TimestamptzNanoType(), required=False), schema_id=1) + + assert GreaterThan("a", "2017-08-18T14:21:01.919234567+00:00").bind(schema).literal.value == 1503066061919234567 + + def test_bind_expr_does_not_exists() -> None: schema = Schema(NestedField(2, "a", IntegerType()), schema_id=1) with pytest.raises(ValueError) as exc_info: diff --git a/tests/expressions/test_literals.py b/tests/expressions/test_literals.py index 9251e79a7d..d19051d4dc 100644 --- a/tests/expressions/test_literals.py +++ b/tests/expressions/test_literals.py @@ -45,6 +45,7 @@ StringLiteral, TimeLiteral, TimestampLiteral, + TimestampNanoLiteral, literal, ) from pyiceberg.types import ( @@ -59,7 +60,9 @@ LongType, PrimitiveType, StringType, + TimestampNanoType, TimestampType, + TimestamptzNanoType, TimestamptzType, TimeType, UUIDType, @@ -88,6 +91,7 @@ def test_literal_from_nan_error() -> None: DateLiteral, TimeLiteral, TimestampLiteral, + TimestampNanoLiteral, DecimalLiteral, StringLiteral, FixedLiteral, @@ -226,6 +230,16 @@ def test_long_to_timestamp() -> None: assert timestamp_lit.value == long_lit.value +@pytest.mark.parametrize("timestamp_nano_type", [TimestampNanoType(), TimestamptzNanoType()]) +def test_long_to_timestamp_nano(timestamp_nano_type: PrimitiveType) -> None: + # A long is read as microseconds, matching the plain timestamp case + long_lit = literal(1647305201).to(LongType()) + timestamp_nano_lit = long_lit.to(timestamp_nano_type) + + assert isinstance(timestamp_nano_lit, TimestampNanoLiteral) + assert timestamp_nano_lit.value == 1647305201 * 1_000 + + @pytest.mark.parametrize( "decimal_type, decimal_value", [(DecimalType(9, 0), "34"), (DecimalType(9, 2), "34.00"), (DecimalType(9, 4), "34.0000")] ) @@ -300,6 +314,69 @@ def test_timestamp_to_date() -> None: assert date_lit.value == 0 +@pytest.mark.parametrize("timestamp_nano_type", [TimestampNanoType(), TimestamptzNanoType()]) +def test_timestamp_to_timestamp_nano(timestamp_nano_type: PrimitiveType) -> None: + timestamp_lit = TimestampLiteral(1503066061919234) + timestamp_nano_lit = timestamp_lit.to(timestamp_nano_type) + + assert isinstance(timestamp_nano_lit, TimestampNanoLiteral) + assert timestamp_nano_lit.value == 1503066061919234000 + + +@pytest.mark.parametrize("timestamp_nano_type", [TimestampNanoType(), TimestamptzNanoType()]) +def test_timestamp_to_timestamp_nano_out_of_range(timestamp_nano_type: PrimitiveType) -> None: + # A value that is already in nanoseconds no longer fits once it is scaled up by a thousand + nanos = 1503066061919234567 + + with pytest.raises(OverflowError, match=f"Timestamp cannot be converted to nanoseconds, out of range: {nanos}"): + _ = TimestampLiteral(nanos).to(timestamp_nano_type) + + with pytest.raises(OverflowError, match=f"Timestamp cannot be converted to nanoseconds, out of range: {nanos}"): + _ = literal(nanos).to(LongType()).to(timestamp_nano_type) + + +@pytest.mark.parametrize("timestamp_type", [TimestampType(), TimestamptzType()]) +def test_timestamp_nano_to_timestamp(timestamp_type: PrimitiveType) -> None: + # Sub-microsecond precision is truncated towards negative infinity + assert TimestampNanoLiteral(1503066061919234567).to(timestamp_type).value == 1503066061919234 + assert TimestampNanoLiteral(-1).to(timestamp_type).value == -1 + + +@pytest.mark.parametrize("timestamp_nano_type", [TimestampNanoType(), TimestamptzNanoType()]) +def test_timestamp_nano_to_timestamp_nano(timestamp_nano_type: PrimitiveType) -> None: + timestamp_nano_lit = TimestampNanoLiteral(1503066061919234567) + + assert timestamp_nano_lit.to(timestamp_nano_type) is timestamp_nano_lit + + +def test_timestamp_nano_to_date() -> None: + assert TimestampNanoLiteral(0).to(DateType()).value == 0 + assert TimestampNanoLiteral(1503066061919234567).to(DateType()).value == 17396 + # One nanosecond before the epoch still falls on the previous day + assert TimestampNanoLiteral(-1).to(DateType()).value == -1 + + +def test_timestamp_nano_increment_decrement() -> None: + timestamp_nano_lit = TimestampNanoLiteral(1503066061919234567) + + assert timestamp_nano_lit.increment() == TimestampNanoLiteral(1503066061919234568) + assert timestamp_nano_lit.decrement() == TimestampNanoLiteral(1503066061919234566) + + +@pytest.mark.parametrize( + "nanos, expected", + [ + (1503066061919234567, "2017-08-18T14:21:01.919234567"), + # The microsecond part is zero, so the fraction still has to be padded out + (1_000_000_500, "1970-01-01T00:00:01.000000500"), + (0, "1970-01-01T00:00:00.000000000"), + (-1, "1969-12-31T23:59:59.999999999"), + ], +) +def test_timestamp_nano_serialization(nanos: int, expected: str) -> None: + assert TimestampNanoLiteral(nanos).model_dump() == expected + + def test_string_literal() -> None: sqrt2 = literal("1.414").to(StringType()) pi = literal("3.141").to(StringType()) @@ -359,6 +436,40 @@ def test_string_to_timestamp_literal() -> None: assert avro_val == timestamp.value +def test_string_to_timestamp_nano_literal() -> None: + timestamp_str = literal("2017-08-18T14:21:01.919234567+00:00") + timestamp_nano = timestamp_str.to(TimestamptzNanoType()) + + assert isinstance(timestamp_nano, TimestampNanoLiteral) + assert timestamp_nano.value == 1503066061919234567 + + timestamp_str = literal("2017-08-18T14:21:01.919234567") + timestamp_nano = timestamp_str.to(TimestampNanoType()) + assert timestamp_nano.value == 1503066061919234567 + + timestamp_str = literal("2017-08-18T14:21:01.919234567-07:00") + timestamp_nano = timestamp_str.to(TimestamptzNanoType()) + assert timestamp_nano.value == 1503091261919234567 + + +def test_string_to_timestamp_nano_out_of_range() -> None: + with pytest.raises(OverflowError, match="Timestamp cannot be converted to nanoseconds, out of range"): + _ = literal("2300-01-01T00:00:00").to(TimestampNanoType()) + + with pytest.raises(OverflowError, match="Timestamp cannot be converted to nanoseconds, out of range"): + _ = literal("2300-01-01T00:00:00+00:00").to(TimestamptzNanoType()) + + +def test_string_to_timestamp_nano_zone_mismatch() -> None: + with pytest.raises(ValueError) as e: + _ = literal("2017-08-18T14:21:01.919234567").to(TimestamptzNanoType()) + assert "Missing zone offset: 2017-08-18T14:21:01.919234567 (must be ISO-8601)" in str(e.value) + + with pytest.raises(ValueError) as e: + _ = literal("2017-08-18T14:21:01.919234567+07:00").to(TimestampNanoType()) + assert "Zone offset provided, but not expected: 2017-08-18T14:21:01.919234567+07:00" in str(e.value) + + def test_timestamp_with_zone_without_zone_in_literal() -> None: timestamp_str = literal("2017-08-18T14:21:01.919234") with pytest.raises(ValueError) as e: @@ -764,6 +875,25 @@ def test_invalid_timestamp_conversions() -> None: ) +def test_invalid_timestamp_nano_conversions() -> None: + assert_invalid_conversions( + TimestampNanoLiteral(1503066061919234567), + [ + BooleanType(), + IntegerType(), + LongType(), + FloatType(), + DoubleType(), + TimeType(), + DecimalType(9, 2), + StringType(), + UUIDType(), + FixedType(1), + BinaryType(), + ], + ) + + def test_invalid_decimal_conversion_scale() -> None: lit = literal(Decimal("34.11")) with pytest.raises(ValueError) as e: diff --git a/tests/test_transforms.py b/tests/test_transforms.py index c977fcea14..7f97e11b0c 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -72,6 +72,7 @@ DateLiteral, DecimalLiteral, TimestampLiteral, + TimestampNanoLiteral, literal, ) from pyiceberg.partitioning import _to_partition_representation @@ -670,6 +671,13 @@ def bound_reference_timestamp() -> BoundReference: ) +@pytest.fixture +def bound_reference_timestamp_ns() -> BoundReference: + return BoundReference( + field=NestedField(1, "field", TimestampNanoType(), required=False), accessor=Accessor(position=0, inner=None) + ) + + @pytest.fixture def bound_reference_decimal() -> BoundReference: return BoundReference( @@ -774,6 +782,37 @@ def test_projection_day_month_not_in(bound_reference_date: BoundReference) -> No ) +@pytest.mark.parametrize( + "transform, expected", + [ + (YearTransform(), 52), + (MonthTransform(), 634), + (DayTransform(), 19302), + (HourTransform(), 463249), + ], +) +def test_projection_time_transform_literal_ns( + transform: TimeTransform[Any], + expected: int, + bound_reference_timestamp_ns: BoundReference, + bound_reference_timestamp: BoundReference, +) -> None: + # A nanosecond literal must project onto the same partition value as the equivalent microsecond one + micros, nanos = 1667696874_000_000, 1667696874_000_000_000 + + assert transform.project( + "name", BoundGreaterThan(term=bound_reference_timestamp_ns, literal=TimestampNanoLiteral(nanos)) + ) == transform.project("name", BoundGreaterThan(term=bound_reference_timestamp, literal=TimestampLiteral(micros))) + + assert transform.project( + "name", BoundEqualTo(term=bound_reference_timestamp_ns, literal=TimestampNanoLiteral(nanos)) + ) == EqualTo(term="name", literal=expected) + + assert transform.strict_project( + "name", BoundNotEqualTo(term=bound_reference_timestamp_ns, literal=TimestampNanoLiteral(nanos)) + ) == NotEqualTo(term="name", literal=expected) + + def test_projection_day_unary(bound_reference_timestamp: BoundReference) -> None: assert DayTransform().project("name", BoundNotNull(term=bound_reference_timestamp)) == NotNull(term="name") diff --git a/tests/utils/test_datetime.py b/tests/utils/test_datetime.py index d7a6f431ee..eceeee2627 100644 --- a/tests/utils/test_datetime.py +++ b/tests/utils/test_datetime.py @@ -22,6 +22,7 @@ from pyiceberg.utils.datetime import ( datetime_to_millis, datetime_to_nanos, + micros_to_nanos, millis_to_datetime, nanos_to_hours, nanos_to_micros, @@ -158,3 +159,46 @@ def test_nanos_to_micros(nanos: int, micros: int) -> None: ) def test_nanos_to_hours(nanos: int, hours: int) -> None: assert hours == nanos_to_hours(nanos) + + +@pytest.mark.parametrize("micros, nanos", [(1510871468000001, 1510871468000001000), (-1510871468000001, -1510871468000001000)]) +def test_micros_to_nanos(micros: int, nanos: int) -> None: + assert nanos == micros_to_nanos(micros) + + +@pytest.mark.parametrize("micros", [9223372036854776, -9223372036854776]) +def test_micros_to_nanos_out_of_range(micros: int) -> None: + with pytest.raises(OverflowError, match=f"Timestamp cannot be converted to nanoseconds, out of range: {micros}"): + micros_to_nanos(micros) + + +@pytest.mark.parametrize( + "timestamp, nanos", + [ + # 2262-04-11T23:47:16.854775807 is the last timestamp that fits in a signed 64-bit integer + ("2262-04-11T23:47:16.854775807", 9223372036854775807), + ("1677-09-21T00:12:43.145224192", -9223372036854775808), + ], +) +def test_timestamp_to_nanos_at_boundary(timestamp: str, nanos: int) -> None: + assert nanos == timestamp_to_nanos(timestamp) + + +@pytest.mark.parametrize( + "timestamp", + [ + "2300-01-01T00:00:00", + "1600-01-01T00:00:00", + # One nanosecond past either boundary, which only the sub-microsecond digits push over + "2262-04-11T23:47:16.854775808", + "1677-09-21T00:12:43.145224191", + ], +) +def test_timestamp_to_nanos_out_of_range(timestamp: str) -> None: + with pytest.raises(OverflowError, match="Timestamp cannot be converted to nanoseconds, out of range"): + timestamp_to_nanos(timestamp) + + +def test_timestamptz_to_nanos_out_of_range() -> None: + with pytest.raises(OverflowError, match="Timestamp cannot be converted to nanoseconds, out of range"): + timestamptz_to_nanos("2300-01-01T00:00:00+00:00")