Skip to content

Commit d4c9865

Browse files
lidavidmlidavidm
authored
fix: compare nanosecond timestamps properly (#29)
## What's Changed Use `whenever` to manually compare timestamps instead of relying on PyArrow which randomly breaks based on what libraries are installed. (We could directly compare the raw timestamp, but this gives us a nicer display.)
1 parent 1cf7e8b commit d4c9865

4 files changed

Lines changed: 171 additions & 45 deletions

File tree

‎adbc_drivers_validation/compare.py‎

Lines changed: 85 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import typing
2020

2121
import pyarrow
22+
import whenever
2223

2324
T = typing.TypeVar("T", pyarrow.Schema, pyarrow.Table)
2425

@@ -67,6 +68,89 @@ def make_nullable(value: T) -> T:
6768
raise TypeError(f"Unsupported type: {type(value)}")
6869

6970

71+
def scalar_to_py_smart(value: pyarrow.Scalar) -> typing.Any:
72+
"""
73+
Convert a PyArrow Scalar to a Python object.
74+
75+
Unlike `pyarrow.Scalar.as_py()`, this function aims to properly handle
76+
times, timestamps, and intervals.
77+
"""
78+
79+
if value is None or not value.is_valid:
80+
return None
81+
elif isinstance(value, (pyarrow.Time32Scalar, pyarrow.Time64Scalar)):
82+
# Needed because to_pylist uses the wrong object for times, losing precision...
83+
match value.type.unit:
84+
case "s":
85+
return NaiveTime(
86+
hour=value.value // 3600,
87+
minute=(value.value % 3600) // 60,
88+
second=value.value % 60,
89+
nanosecond=0,
90+
)
91+
case "ms":
92+
return NaiveTime(
93+
hour=value.value // 3600000,
94+
minute=(value.value % 3600000) // 60000,
95+
second=(value.value % 60000) // 1000,
96+
nanosecond=(value.value % 1000) * 1000000,
97+
)
98+
case "us":
99+
return NaiveTime(
100+
hour=value.value // 3600000000,
101+
minute=(value.value % 3600000000) // 60000000,
102+
second=(value.value % 60000000) // 1000000,
103+
nanosecond=(value.value % 1000000) * 1000,
104+
)
105+
case "ns":
106+
return NaiveTime(
107+
hour=value.value // 3600000000000,
108+
minute=(value.value % 3600000000000) // 60000000000,
109+
second=(value.value % 60000000000) // 1000000000,
110+
nanosecond=value.value % 1000000000,
111+
)
112+
case _:
113+
raise NotImplementedError(str(value))
114+
elif isinstance(value, pyarrow.TimestampScalar):
115+
match value.type.unit:
116+
case "s":
117+
nanos = value.value * 1_000_000_000
118+
case "ms":
119+
nanos = value.value * 1_000_000
120+
case "us":
121+
nanos = value.value * 1000
122+
case "ns":
123+
nanos = value.value
124+
125+
instant = whenever.Instant.from_timestamp_nanos(nanos)
126+
if value.type.tz is None or value.type.tz == "":
127+
# A bit sketch
128+
naive = whenever.PlainDateTime.parse_common_iso(
129+
instant.format_common_iso()[:-1]
130+
)
131+
return str(naive)
132+
elif value.type.tz == "UTC":
133+
return str(instant)
134+
elif value.type.tz[0] in ("+", "-"):
135+
negative = value.type.tz[0] == "-"
136+
hours, _, minutes = value.type.tz[1:].partition(":")
137+
tz_offset = whenever.TimeDelta(hours=int(hours), minutes=int(minutes))
138+
if negative:
139+
tz_offset = -tz_offset
140+
offset = instant.to_fixed_offset(tz_offset)
141+
return str(offset)
142+
zoned = instant.to_tz(value.type.tz)
143+
return str(zoned)
144+
elif isinstance(value, pyarrow.MonthDayNanoIntervalScalar):
145+
mdn = value.as_py()
146+
if mdn is None:
147+
return None
148+
else:
149+
return f"{mdn[0]}M{mdn[1]}d{mdn[2]}ns"
150+
151+
return value.as_py()
152+
153+
70154
def to_pylist(table: pyarrow.Table) -> list[dict[str, typing.Any]]:
71155
"""Convert a PyArrow Table to a list of dictionaries."""
72156
rows = []
@@ -75,51 +159,8 @@ def to_pylist(table: pyarrow.Table) -> list[dict[str, typing.Any]]:
75159
for col_idx in range(table.num_columns):
76160
value = table.column(col_idx)[row_idx]
77161
col_name = table.schema[col_idx].name
78-
if value is None:
79-
row[col_name] = None
80-
elif isinstance(value, (pyarrow.Time32Scalar, pyarrow.Time64Scalar)):
81-
# Needed because to_pylist uses the wrong object for times, losing precision...
82-
match value.type.unit:
83-
case "s":
84-
row[col_name] = NaiveTime(
85-
hour=value.value // 3600,
86-
minute=(value.value % 3600) // 60,
87-
second=value.value % 60,
88-
nanosecond=0,
89-
)
90-
case "ms":
91-
row[col_name] = NaiveTime(
92-
hour=value.value // 3600000,
93-
minute=(value.value % 3600000) // 60000,
94-
second=(value.value % 60000) // 1000,
95-
nanosecond=(value.value % 1000) * 1000000,
96-
)
97-
case "us":
98-
row[col_name] = NaiveTime(
99-
hour=value.value // 3600000000,
100-
minute=(value.value % 3600000000) // 60000000,
101-
second=(value.value % 60000000) // 1000000,
102-
nanosecond=(value.value % 1000000) * 1000,
103-
)
104-
case "ns":
105-
row[col_name] = NaiveTime(
106-
hour=value.value // 3600000000000,
107-
minute=(value.value % 3600000000000) // 60000000000,
108-
second=(value.value % 60000000000) // 1000000000,
109-
nanosecond=value.value % 1000000000,
110-
)
111-
case _:
112-
raise NotImplementedError(str(value))
113-
elif isinstance(value, pyarrow.MonthDayNanoIntervalScalar):
114-
mdn = value.as_py()
115-
if mdn is None:
116-
row[col_name] = None
117-
else:
118-
row[col_name] = f"{mdn[0]}M{mdn[1]}d{mdn[2]}ns"
119-
else:
120-
row[col_name] = value.as_py()
162+
row[col_name] = scalar_to_py_smart(value)
121163
rows.append(row)
122-
123164
return rows
124165

125166

‎pixi.lock‎

Lines changed: 18 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ dependencies = [
2323
"jinja2>=3.1.6,<4",
2424
"pyarrow>=20.0.0,<21",
2525
"pytest>=8.4.1,<9",
26+
"whenever>=0.8.8,<0.9",
2627
]
2728

2829
[build-system]

‎tests/test_compare.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
import typing
16+
1517
import pyarrow
1618
import pytest
1719

@@ -54,3 +56,68 @@ def test_compare_fields():
5456
compare.compare_fields(f1, f4)
5557
with pytest.raises(AssertionError):
5658
compare.compare_fields(f1, f6)
59+
60+
61+
@pytest.mark.parametrize(
62+
"value, expected",
63+
[
64+
pytest.param(pyarrow.scalar(None), None, id="null"),
65+
pytest.param(
66+
pyarrow.scalar(3661, pyarrow.time32("s")),
67+
compare.NaiveTime(hour=1, minute=1, second=1, nanosecond=0),
68+
id="time32[s]",
69+
),
70+
pytest.param(
71+
pyarrow.scalar(3661001, pyarrow.time32("ms")),
72+
compare.NaiveTime(hour=1, minute=1, second=1, nanosecond=1000000),
73+
id="time32[ms]",
74+
),
75+
pytest.param(
76+
pyarrow.scalar(3661001001, pyarrow.time64("us")),
77+
compare.NaiveTime(hour=1, minute=1, second=1, nanosecond=1001000),
78+
id="time64[us]",
79+
),
80+
pytest.param(
81+
pyarrow.scalar(3661001001001, pyarrow.time64("ns")),
82+
compare.NaiveTime(hour=1, minute=1, second=1, nanosecond=1001001),
83+
id="time64[ns]",
84+
),
85+
pytest.param(
86+
pyarrow.scalar(1, pyarrow.timestamp("s")),
87+
"1970-01-01T00:00:01",
88+
id="timestamp[s]",
89+
),
90+
pytest.param(
91+
pyarrow.scalar(1, pyarrow.timestamp("s", "UTC")),
92+
"1970-01-01T00:00:01Z",
93+
id="timestamp[s, UTC]",
94+
),
95+
pytest.param(
96+
pyarrow.scalar(1, pyarrow.timestamp("s", "Asia/Tokyo")),
97+
"1970-01-01T09:00:01+09:00[Asia/Tokyo]",
98+
id="timestamp[s, Asia/Tokyo]",
99+
),
100+
pytest.param(
101+
pyarrow.scalar(1, pyarrow.timestamp("s", "+09:00")),
102+
"1970-01-01T09:00:01+09:00",
103+
id="timestamp[s, +09:00]",
104+
),
105+
pytest.param(
106+
pyarrow.scalar(1, pyarrow.timestamp("s", "-00:30")),
107+
"1969-12-31T23:30:01-00:30",
108+
id="timestamp[s, -00:30]",
109+
),
110+
pytest.param(
111+
pyarrow.scalar(1, pyarrow.timestamp("ns")),
112+
"1970-01-01T00:00:00.000000001",
113+
id="timestamp[ns]",
114+
),
115+
pytest.param(
116+
pyarrow.scalar(1, pyarrow.timestamp("ns", "UTC")),
117+
"1970-01-01T00:00:00.000000001Z",
118+
id="timestamp[ns, UTC]",
119+
),
120+
],
121+
)
122+
def test_scalar_to_py_smart(value: pyarrow.Scalar, expected: typing.Any) -> None:
123+
assert compare.scalar_to_py_smart(value) == expected

0 commit comments

Comments
 (0)