Fix Dataset.to_json() truncating/corrupting non-millisecond temporal columns

[aakif-ms]

Fixes #8390

Problem

Dataset.to_json() writes every timestamp/duration column through pandas.DataFrame.to_json(), which only supports a single global date_unit for the entire call (defaulting to milliseconds).

This ignores each column's actual Arrow unit (s/ms/us/ns), causing a to_json() → from_json() round trip to either:

• raise OverflowError for timestamp[s], or
• silently return the wrong value for timestamp[us]/timestamp[ns] (off by tens of years).

Root Cause

JsonDatasetWriter._batch_json converts the Arrow batch to pandas via to_pandas() before calling to_json().

That conversion collapses every temporal column to a generic datetime64[ns] dtype. By the time to_json() runs, the original per-column Arrow unit is already gone, so there is no way to recover it. Pandas then falls back to its default millisecond date_unit for every temporal column.

Fix

Cast timestamp/duration columns to their own native-unit int64 representation before calling to_pandas(), while the original Arrow unit is still available.

This bypasses pandas' date handling entirely for these columns, so they are written as plain integers.

The fix is symmetric with the JSON loader, which already reads a raw integer column and casts it directly to the declared timestamp/duration unit via table_cast, treating the integer as already being in that unit.

As a result, writing the native-unit integer makes the to_json() → from_json() round trip exact, with no changes needed on the read side.

Testing

Added parametrized round-trip tests covering:

• timestamp[s]
• timestamp[ms]
• timestamp[us]
• timestamp[ns]
• duration[us]
• null handling

Confirmed that the new tests fail on unpatched main (3 of 4 timestamp units) and pass with the fix.

The full tests/io/test_json.py suite passes: 54 passed. There are also 3 pre-existing, unrelated errors caused by a missing pytest-datadir fixture in my local environment.

[ebarkhordar]

The top-level fix works. Running the same script against 8bbcd703 and against the merge-base 48b7ee7b, in a clean python:3.11-slim container with pyarrow 25.0.1 and pandas 3.0.5:

timestamp[s], top-level column
  48b7ee7b  {"t":1704112245000}   -> OverflowError: date value out of range
  8bbcd703  {"t":1704112245}      -> datetime.datetime(2024, 1, 1, 12, 30, 45)

A timestamp nested inside a struct is still written at pandas' millisecond default, and still fails the same way at this head:

import datetime
from datasets import Dataset, Features, Value

f = Features({"t": {"inner": Value("timestamp[s]")}})
ds = Dataset.from_dict({"t": [{"inner": datetime.datetime(2024, 1, 1, 12, 30, 45)}]}, features=f)
ds.to_json("/tmp/t.jsonl")
print(open("/tmp/t.jsonl").read().strip())        # {"t":{"inner":1704112245000}}
Dataset.from_json("/tmp/t.jsonl", features=f)[0]  # OverflowError: date value out of range

pa.types.is_timestamp is only asked about the entries of batch.schema.types, so for this column it is asked about struct<inner: timestamp[s]>, gets False, and the column reaches to_pandas() unconverted. Lists are not affected: List(Value("timestamp[s]")) and List(Value("timestamp[us]")) both round trip on the merge-base as well as on this head, so the remaining gap is specifically struct fields.

Whether that belongs in this PR is your call, and #8390's reproducer is top-level only. If you keep the scope as it is, it would help to say in the new test that nested temporal fields still take the old path, so the next reader does not read the regression test as covering them.

One minor thing, and only because the checks have not run: both workflow runs on this head sha are sitting at action_required, so ruff has not seen the diff yet. ruff check reports I001 (the new import datetime sorts above import io) and W292 (no newline at the end of tests/io/test_json.py), and ruff format wants the trailing whitespace removed from the blank line at src/datasets/io/json.py:135. All three are --fix-able.

[aakif-ms]

Thanks for the thorough review, and for pressure-testing this against a struct case — good catch.

Confirmed the struct issue locally: the loop only inspects batch.schema.types at the top level, so a struct<inner: timestamp[s]> column type fails pa.types.is_timestamp() and passes through to to_pandas() unconverted, hitting the same overflow this PR fixes for top-level columns.

I'm keeping this PR scoped to top-level columns, since that's what #8390 reports and reproduces — recursing through nested struct/list fields felt like a large enough change to deserve its own follow-up rather than folding into this one. Added a note in the new test making that scope explicit so it isn't misread as covering the nested case.

Also pushed the ruff fixes (import order, trailing whitespace, missing newline) — thanks for flagging those, the workflow approval gate meant CI hadn't caught them yet.

[shashvat-singham]

Reproduced the bug and confirmed the fix. Round-tripping datetime(2024, 1, 1, 12, 30, 45) through to_json(lines=True) → from_json, per declared unit:

unit	main	this branch
timestamp[s]	OverflowError: date value out of range	2024-01-01 12:30:45
timestamp[ms]	2024-01-01 12:30:45	2024-01-01 12:30:45
timestamp[us]	1970-01-20 17:21:52	2024-01-01 12:30:45
timestamp[ns]	1970-01-01 00:28:24	2024-01-01 12:30:45

The silent-wrong-answer cases are the nasty ones, and they are fixed. The diagnosis in the description is right: on main all four units write the same 1704112245000, because pandas applies one global date_unit="ms" after to_pandas() has already flattened everything to datetime64[ns].

CI will fail as-is. There is a blank line with trailing whitespace at src/datasets/io/json.py:135:

W293 Blank line contains whitespace
1 file would be reformatted

ruff format src/datasets/io/json.py fixes it.

The part I would want stated in the PR: this changes the on-disk JSON for non-ms columns.

timestamp[s], main : 1704112245000
timestamp[s], here : 1704112245

Both are integers, so nothing breaks structurally — but they mean different things, and the old value was the one a non-datasets reader would get right. Pandas' documented default is epoch-milliseconds, so a consumer decoding these files with datetime.utcfromtimestamp(v / 1000) reads the correct instant from main and gets 1970-01-20 from this branch. The file is now only interpretable if you also know the column's declared unit, which is not in the JSON.

That is a real trade: internal round-trip correctness bought with external self-describability. I think it is the right trade — the current behaviour is broken or lossy for three of the four units, and from_json is the primary consumer — but it is a behaviour change for anyone whose pipeline is datasets.to_json → something else, and it deserves a line in the PR description and probably the release notes rather than being implied by "fix".

Worth at least considering the alternative: date_format="iso" writes ISO-8601 strings, which are unambiguous to both datasets and outside readers and sidestep the unit question entirely. Bigger change, different file size, and it would need the loader to parse strings — so not a request, just noting it was not mentioned as an option.

Two small things:

• The test covers to_json → from_json. Worth also asserting the written value for one unit, so a future change to the loader that quietly re-breaks the writer cannot keep the test green by cancelling out.
• The comment block is excellent, but ~20 lines inside _batch_json is a lot in a hot loop; the "why" half would read well just above the method.

Verified on Windows 11 / Python 3.11.9 / pyarrow 25.0.1, main vs pr-8459.

[aakif-ms]

Thanks for the detailed reviews, both — really useful.

@ebarkhordar, confirmed the struct case locally, and you're right about why — my loop only checks the top-level column types, so a struct wrapping a timestamp just sails through untouched. I'm going to leave that out of this PR though, since #8390's original report is only about top-level columns and properly handling nested structs (structs inside structs, structs inside lists, etc.) feels like its own separate piece of work. Added a note in the test so it's clear that's not covered here.

Fixed the ruff stuff too — thanks for catching that, guessing CI hadn't run yet because of the workflow approval gate.

@shashvat-singham, really appreciate you spelling out the on-disk format point, that's a fair thing to flag and I hadn't stated it clearly. You're right that it changes what a plain-integer-reading consumer gets back for the non-ms units. I still think it's the right call given from_json is the main way people read these back and 3 of 4 units are just broken today, but agreed it shouldn't be an implicit side effect — added a line about it in the PR description.

And thanks for the ISO note too, wasn't trying to dodge that option, just went with int64 since it's a smaller change and doesn't touch the loader.

Added the raw-value assertion you suggested as well, good catch that the round-trip-only test could pass for the wrong reasons.