fix(python): correct the asof and reader input annotations

python/pyproject.toml

@@ -119,6 +119,7 @@ include = [
     "python/lance/arrow.py",
     "python/tests/test_arrow.py",
     "python/tests/test_fragment_typing.py",
+    "python/tests/test_optional_types_typing.py",
     "python/tests/test_udf.py",
 ]
 # Dependencies like pyarrow make this difficult to enforce strictly.

python/python/lance/__init__.py

@@ -72,13 +72,10 @@
 from .util import sanitize_ts
 
 if TYPE_CHECKING:
-    from datetime import datetime
     from pathlib import Path
 
     from lance.commit import CommitLock
-    from lance.dependencies import pandas as pd
-
-    ts_types = Union[datetime, pd.Timestamp, str]
+    from lance.util import ts_types
 
 
 __all__ = [

python/python/lance/dataset.py

@@ -7766,7 +7766,11 @@ def write_dataset(
     ----------
     data_obj: Reader-like
         The data to be written. Acceptable types are:
-        - Pandas DataFrame, Pyarrow Table, Dataset, Scanner, or RecordBatchReader
+        - Pandas DataFrame, Polars DataFrame
+        - Pyarrow Table, RecordBatch, Dataset, Scanner, or RecordBatchReader
+        - An iterable of Pyarrow RecordBatch (requires ``schema``)
+        - A dict of columns, or a list of row dicts
+        - A list of Pydantic model instances
         - Huggingface dataset
     uri: str, Path, LanceDataset, or None
         Where to write the dataset to (directory). If a LanceDataset is passed,

python/python/lance/types.py

@@ -3,30 +3,45 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Iterable, Optional, Union
+from typing import TYPE_CHECKING, Any, Iterable, Mapping, Optional, Sequence, Union
 
 import pyarrow as pa
 from pyarrow import RecordBatch
+from pyarrow.dataset import Dataset as ArrowDataset
+from pyarrow.dataset import Scanner as ArrowScanner
 
-from . import dataset
 from .dependencies import (
     _check_for_hugging_face,
     _check_for_pandas,
+    _check_for_polars,
     _is_pydantic_base_model,
     _validate_pydantic_list,
     model_to_dict,
 )
 from .dependencies import pandas as pd
+from .dependencies import polars as pl
 
 if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+    from .dependencies import datasets
+
+    # Keep in step with the branches of ``_coerce_reader``.
     ReaderLike = Union[
-        pd.Timestamp,
+        pd.DataFrame,
+        pl.DataFrame,
         pa.Table,
-        pa.dataset.Dataset,
-        pa.dataset.Scanner,
         pa.RecordBatch,
-        Iterable[RecordBatch],
         pa.RecordBatchReader,
+        # ``LanceDataset`` is an ``ArrowDataset`` subclass.
+        ArrowDataset,
+        ArrowScanner,
+        datasets.Dataset,
+        datasets.IterableDataset,
+        Mapping[str, Any],
+        Sequence[Mapping[str, Any]],
+        Sequence[BaseModel],
+        Iterable[RecordBatch],
     ]
 
 
@@ -70,10 +85,7 @@ def _is_materialized(data_obj: ReaderLike) -> bool:
         return True
     if isinstance(data_obj, (pa.Table, pa.RecordBatch)):
         return True
-    if (
-        type(data_obj).__module__.startswith("polars")
-        and data_obj.__class__.__name__ == "DataFrame"
-    ):
+    if _check_for_polars(data_obj) and isinstance(data_obj, pl.DataFrame):
         return True
     if isinstance(data_obj, dict):
         return True
@@ -89,24 +101,24 @@ def _is_materialized(data_obj: ReaderLike) -> bool:
 def _coerce_reader(
     data_obj: ReaderLike, schema: Optional[pa.Schema] = None
 ) -> pa.RecordBatchReader:
+    # Local import: ``lance.dataset`` imports this module.
+    from .dataset import LanceDataset
+
     if _check_for_pandas(data_obj) and isinstance(data_obj, pd.DataFrame):
         return pa.Table.from_pandas(data_obj, schema=schema).to_reader()
     elif isinstance(data_obj, pa.Table):
         return data_obj.to_reader()
     elif isinstance(data_obj, pa.RecordBatch):
         return pa.Table.from_batches([data_obj]).to_reader()
-    elif isinstance(data_obj, dataset.LanceDataset):
+    elif isinstance(data_obj, LanceDataset):
         return data_obj.scanner().to_reader()
-    elif isinstance(data_obj, pa.dataset.Dataset):
-        return pa.dataset.Scanner.from_dataset(data_obj).to_reader()
-    elif isinstance(data_obj, pa.dataset.Scanner):
+    elif isinstance(data_obj, ArrowDataset):
+        return ArrowScanner.from_dataset(data_obj).to_reader()
+    elif isinstance(data_obj, ArrowScanner):
         return data_obj.to_reader()
     elif isinstance(data_obj, pa.RecordBatchReader):
         return data_obj
-    elif (
-        type(data_obj).__module__.startswith("polars")
-        and data_obj.__class__.__name__ == "DataFrame"
-    ):
+    elif _check_for_polars(data_obj) and isinstance(data_obj, pl.DataFrame):
         return data_obj.to_arrow().to_reader()
     elif _check_for_hugging_face(data_obj):
         from .dependencies import datasets as hf_datasets

python/python/lance/util.py

@@ -18,13 +18,17 @@
 
 import pyarrow as pa
 
-from .dependencies import _check_for_numpy, _check_for_pandas
+from .dependencies import _PANDAS_AVAILABLE, _check_for_numpy, _check_for_pandas
 from .dependencies import numpy as np
 from .dependencies import pandas as pd
 from .lance import _Hnsw, _KMeans
 
 if TYPE_CHECKING:
-    ts_types = Union[datetime, pd.Timestamp, str]
+    from pandas.api.typing import NaTType
+
+    # ``pd.Timestamp`` is a ``datetime``, but pandas types its constructor as
+    # ``Timestamp | NaTType``; ``sanitize_ts`` rejects ``NaT`` at runtime.
+    ts_types = Union[datetime, NaTType, str]
 
 MetricType = Literal["l2", "euclidean", "dot", "cosine"]
 
@@ -40,20 +44,24 @@ def _normalize_metric_type(metric_type: str) -> MetricType:
 
 def sanitize_ts(ts: ts_types) -> datetime:
     """Returns a python datetime object from various timestamp input types."""
-    if _check_for_pandas(ts) and isinstance(ts, str):
-        ts = pd.to_datetime(ts).to_pydatetime()
-    elif isinstance(ts, str):
-        try:
-            ts = datetime.strptime(ts, "%Y-%m-%d %H:%M:%S")
-        except ValueError:
-            raise ValueError(
-                f"Failed to parse timestamp string {ts}. Try installing Pandas."
-            )
-    elif _check_for_pandas(ts) and isinstance(ts, pd.Timestamp):
-        ts = ts.to_pydatetime()
-    elif not isinstance(ts, datetime):
-        raise TypeError(f"Unrecognized version timestamp {ts} of type {type(ts)}")
-    return ts
+    if isinstance(ts, str):
+        if not _PANDAS_AVAILABLE:
+            try:
+                return datetime.strptime(ts, "%Y-%m-%d %H:%M:%S")
+            except ValueError:
+                raise ValueError(
+                    f"Failed to parse timestamp string {ts}. Try installing Pandas."
+                ) from None
+        ts = pd.to_datetime(ts)
+    if _check_for_pandas(ts):
+        # ``NaT`` subclasses ``datetime`` but compares false against every version.
+        if ts is pd.NaT:
+            raise ValueError("NaT is not a valid version timestamp")
+        if isinstance(ts, pd.Timestamp):
+            return ts.to_pydatetime()
+    if isinstance(ts, datetime):
+        return ts
+    raise TypeError(f"Unrecognized version timestamp {ts} of type {type(ts)}")
 
 
 def td_to_micros(td: timedelta) -> int:

python/python/tests/test_dataset.py

@@ -33,7 +33,7 @@
 from lance.debug import format_fragment
 from lance.file import LanceFileWriter, stable_version
 from lance.schema import LanceSchema
-from lance.util import validate_vector_index
+from lance.util import sanitize_ts, validate_vector_index
 
 BaseModel = pytest.importorskip("pydantic").BaseModel
 
@@ -70,6 +70,8 @@ class _InputModel(BaseModel):
     ),
     # Pydantic model instances are auto-converted
     (None, [_InputModel(a=1.0, b=20), _InputModel(a=2.0, b=30)]),
+    (None, {"a": [1.0, 2.0], "b": [20, 30]}),
+    (None, [{"a": 1.0, "b": 20}, {"a": 2.0, "b": 30}]),
 ]
 
 
@@ -558,6 +560,28 @@ def test_asof_checkout(tmp_path: Path):
     assert len(ds.to_table()) == 9
 
 
+@pytest.mark.parametrize(
+    "ts", ["2026-01-01", pd.Timestamp("2026-01-01"), datetime(2026, 1, 1)]
+)
+def test_sanitize_ts(ts):
+    assert sanitize_ts(ts) == datetime(2026, 1, 1)
+
+
+@pytest.mark.parametrize(
+    "ts,error", [(pd.NaT, ValueError), (object(), TypeError)], ids=["nat", "object"]
+)
+def test_sanitize_ts_rejects(ts, error):
+    with pytest.raises(error):
+        sanitize_ts(ts)
+
+
+def test_sanitize_ts_without_pandas(monkeypatch):
+    monkeypatch.setattr("lance.util._PANDAS_AVAILABLE", False)
+    assert sanitize_ts("2026-01-01 00:00:00") == datetime(2026, 1, 1)
+    with pytest.raises(ValueError, match="Try installing Pandas"):
+        sanitize_ts("2026-01-01")
+
+
 def test_enable_stable_row_ids(tmp_path: Path):
     table = pa.Table.from_pylist(
         [{"name": "Alice", "age": 20}, {"name": "Bob", "age": 30}]

python/python/tests/test_optional_types_typing.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright The Lance Authors
+
+"""Type-checking regressions for ``asof``; part of the pyright target.
+
+Rejected cases use ``pyright: ignore``, so an overly permissive annotation fails
+too. Like ``test_fragment_typing.py``, this module does not import ``pytest``.
+"""
+
+# pyright: reportUnnecessaryTypeIgnoreComment=true
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from datetime import datetime
+
+    import lance
+    import pandas as pd
+    from lance.util import sanitize_ts
+
+    def _check_asof_types() -> None:
+        _timestamp: datetime = sanitize_ts(pd.Timestamp("2026-01-01"))
+        lance.dataset("memory://", asof=pd.Timestamp("2026-01-01"))
+        lance.dataset("memory://", asof=pd.NaT)
+        lance.dataset("memory://", asof=datetime(2026, 1, 1))
+        lance.dataset("memory://", asof="2026-01-01")
+        lance.dataset("memory://", asof=object())  # pyright: ignore[reportArgumentType]