Skip to content

Commit

Permalink
feat: support datetime objects in literal instantiation (#1542)
Browse files Browse the repository at this point in the history
* feat: support datetime objects in literal instantiation

* add integration test

* proper tests

* fix typing
  • Loading branch information
jayceslesar authored Jan 22, 2025
1 parent 9d5638c commit 2cd4e78
Show file tree
Hide file tree
Showing 4 changed files with 38 additions and 1 deletion.
4 changes: 4 additions & 0 deletions pyiceberg/expressions/literals.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

import struct
from abc import ABC, abstractmethod
from datetime import datetime
from decimal import ROUND_HALF_UP, Decimal
from functools import singledispatchmethod
from math import isnan
Expand All @@ -49,6 +50,7 @@
)
from pyiceberg.utils.datetime import (
date_str_to_days,
datetime_to_micros,
micros_to_days,
time_str_to_micros,
timestamp_to_micros,
Expand Down Expand Up @@ -145,6 +147,8 @@ def literal(value: L) -> Literal[L]:
return BinaryLiteral(value)
elif isinstance(value, Decimal):
return DecimalLiteral(value)
elif isinstance(value, datetime):
return TimestampLiteral(datetime_to_micros(value)) # type: ignore
else:
raise TypeError(f"Invalid literal value: {repr(value)}")

Expand Down
3 changes: 2 additions & 1 deletion pyiceberg/typedef.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from __future__ import annotations

from abc import abstractmethod
from datetime import datetime
from decimal import Decimal
from functools import lru_cache
from typing import (
Expand Down Expand Up @@ -78,7 +79,7 @@ def __missing__(self, key: K) -> V:
RecursiveDict = Dict[str, Union[str, "RecursiveDict"]]

# Represents the literal value
L = TypeVar("L", str, bool, int, float, bytes, UUID, Decimal, covariant=True)
L = TypeVar("L", str, bool, int, float, bytes, UUID, Decimal, datetime, covariant=True)


@runtime_checkable
Expand Down
4 changes: 4 additions & 0 deletions tests/expressions/test_literals.py
Original file line number Diff line number Diff line change
Expand Up @@ -906,6 +906,10 @@ def test_uuid_to_binary() -> None:
assert isinstance(binary_literal, BinaryLiteral) # type: ignore


def test_literal_from_datetime() -> None:
assert isinstance(literal(datetime.datetime.now()), TimestampLiteral)


# __ __ ___
# | \/ |_ _| _ \_ _
# | |\/| | || | _/ || |
Expand Down
28 changes: 28 additions & 0 deletions tests/integration/test_reads.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import math
import time
import uuid
from datetime import datetime, timedelta
from pathlib import PosixPath
from urllib.parse import urlparse

Expand Down Expand Up @@ -950,3 +951,30 @@ def test_read_from_s3_and_local_fs(catalog: Catalog, tmp_path: PosixPath) -> Non

result_table = tbl.scan().to_arrow()
assert result_table["colA"].to_pylist() == ["one", "one"]


@pytest.mark.integration
@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
def test_scan_with_datetime(catalog: Catalog) -> None:
table = create_table(catalog)

yesterday = datetime.now() - timedelta(days=1)
table.append(
pa.Table.from_pylist(
[
{
"str": "foo",
"int": 1,
"bool": True,
"datetime": yesterday,
}
],
schema=table.schema().as_arrow(),
),
)

df = table.scan(row_filter=GreaterThanOrEqual("datetime", yesterday)).to_pandas()
assert len(df) == 1

df = table.scan(row_filter=LessThan("datetime", yesterday)).to_pandas()
assert len(df) == 0

0 comments on commit 2cd4e78

Please sign in to comment.