From 8595af14c06f731a9b9e685094b3eafb6fc39b62 Mon Sep 17 00:00:00 2001 From: Jonas Dedden Date: Wed, 2 Sep 2026 22:31:59 +0200 Subject: [PATCH 1/3] fix(python): correct the asof timestamp annotation The stubs bundled with pandas declare `Timestamp.__new__` as returning `Self | NaTType`, so `lance.dataset(asof=pd.Timestamp(...))` was rejected by strict type checkers. Name `NaTType` in `ts_types`; `pd.Timestamp` needs no member of its own because it subclasses `datetime`. `NaT` subclasses `datetime` too, so it previously passed through `sanitize_ts` untouched and then compared false against every version timestamp, surfacing as a misleading "earlier than the first version" error. Reject it at the boundary instead. The pandas string-parsing branch was guarded by `_check_for_pandas(ts)`, which inspects the argument's MRO and is therefore always false for a `str`. That branch had never run: `asof="2026-01-01"` raised "Try installing Pandas" on machines that had pandas installed. Guard on `_PANDAS_AVAILABLE` instead. The new typing regression file joins the pyright target and pins the accepted and rejected `asof` types in both directions, so the annotation also fails the type check if it becomes too permissive. --- python/pyproject.toml | 1 + python/python/lance/__init__.py | 5 +- python/python/lance/util.py | 45 +++++++++----- python/python/tests/test_dataset.py | 28 ++++++++- .../tests/test_optional_types_typing.py | 60 +++++++++++++++++++ 5 files changed, 118 insertions(+), 21 deletions(-) create mode 100644 python/python/tests/test_optional_types_typing.py diff --git a/python/pyproject.toml b/python/pyproject.toml index f6972e03af5..5a793d2e23b 100644 --- a/python/pyproject.toml +++ b/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. diff --git a/python/python/lance/__init__.py b/python/python/lance/__init__.py index 981140858c7..df98ccf69de 100644 --- a/python/python/lance/__init__.py +++ b/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__ = [ diff --git a/python/python/lance/util.py b/python/python/lance/util.py index 180e2441b43..78577258d8d 100644 --- a/python/python/lance/util.py +++ b/python/python/lance/util.py @@ -18,13 +18,20 @@ 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 + + # ``pandas.Timestamp`` is a ``datetime`` subclass, so it needs no member of + # its own. ``NaTType`` does: the stubs bundled with pandas declare + # ``Timestamp.__new__`` as returning ``Self | NaTType``, so every caller + # writing ``asof=pd.Timestamp(...)`` hands us that union. We accept it here + # and reject the ``NaT`` half in ``sanitize_ts``. + ts_types = Union[datetime, NaTType, str] MetricType = Literal["l2", "euclidean", "dot", "cosine"] @@ -40,20 +47,26 @@ 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): + # pandas spells a missing timestamp ``NaT``, which subclasses ``datetime`` + # and compares false against every version timestamp, so it would + # otherwise be reported as being older than the first 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: diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index f48138f2bae..2f0bbd1fb1d 100644 --- a/python/python/tests/test_dataset.py +++ b/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 @@ -558,6 +558,32 @@ def test_asof_checkout(tmp_path: Path): assert len(ds.to_table()) == 9 +def test_sanitize_ts_parses_strings_with_pandas(): + # pandas accepts timestamp strings that the pandas-free fallback format + # rejects, so a date without a time of day has to work here. + assert sanitize_ts("2026-01-01") == datetime(2026, 1, 1) + + +def test_sanitize_ts_rejects_nat(): + # `NaT` is a `datetime` subclass, so it would otherwise pass through and + # compare false against every version timestamp. + with pytest.raises(ValueError, match="NaT is not a valid version timestamp"): + sanitize_ts(pd.NaT) + + +def test_sanitize_ts_rejects_unknown_type(): + with pytest.raises(TypeError, match="Unrecognized version timestamp"): + sanitize_ts(object()) + + +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}] diff --git a/python/python/tests/test_optional_types_typing.py b/python/python/tests/test_optional_types_typing.py new file mode 100644 index 00000000000..89431e3b448 --- /dev/null +++ b/python/python/tests/test_optional_types_typing.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""Type-checking regression tests for optional dependency input types. + +This module is part of the pyright target configured in ``pyproject.toml``, so +a regression in the annotations below fails the repository type check and not +only the runtime suite. + +The rejected cases are pinned with ``pyright: ignore`` comments and +``reportUnnecessaryTypeIgnoreComment``, so the file also fails if an annotation +becomes *too* permissive and one of them stops being an error. They sit in a +``TYPE_CHECKING`` block because they are invalid at runtime; the matching +runtime assertions live in ``test_dataset.py``. + +Like ``test_fragment_typing.py``, this module does not import ``pytest``: the +lint workflow installs pyright without the test dependencies. +""" + +# pyright: reportUnnecessaryTypeIgnoreComment=true + +from datetime import datetime +from typing import TYPE_CHECKING + +import pandas as pd +from lance.util import sanitize_ts + +if TYPE_CHECKING: + import lance + + def _check_accepted_asof_types() -> None: + # ``sanitize_ts`` is exercised at runtime below; this pins the public + # entry point that forwards to it. + lance.dataset("memory://unused", asof=pd.Timestamp("2026-01-01")) + lance.dataset("memory://unused", asof=pd.NaT) + lance.dataset("memory://unused", asof=datetime(2026, 1, 1)) + lance.dataset("memory://unused", asof="2026-01-01") + + def _check_rejected_asof_types() -> None: + # `DatetimeIndex` has a `to_pydatetime` method, so it satisfies a + # structural timestamp annotation without being a valid instant. + index = pd.DatetimeIndex(["2026-01-01"]) + lance.dataset( + "memory://unused", + asof=index, # pyright: ignore[reportArgumentType] + ) + sanitize_ts(object()) # pyright: ignore[reportArgumentType] + + +def test_sanitize_ts_accepts_pandas_timestamp() -> None: + # The stubs bundled with pandas type this constructor as + # ``Timestamp | NaTType``, so both halves have to satisfy ``ts_types``. + result: datetime = sanitize_ts(pd.Timestamp("2026-01-01")) + + assert result == datetime(2026, 1, 1) + + +def test_sanitize_ts_accepts_datetime_and_str() -> None: + assert sanitize_ts(datetime(2026, 1, 1)) == datetime(2026, 1, 1) + assert sanitize_ts("2026-01-01 00:00:00") == datetime(2026, 1, 1) From 9775b5149c161c3b4ec82e606225aba253468ed0 Mon Sep 17 00:00:00 2001 From: Jonas Dedden Date: Wed, 2 Sep 2026 22:38:26 +0200 Subject: [PATCH 2/3] fix(python): complete the reader input union `ReaderLike` listed `pd.Timestamp` rather than a dataframe type, and reached the Arrow classes through `pa.dataset.Dataset`, which type checkers cannot resolve because `pyarrow` does not expose `dataset` as an attribute. The unresolved member made the whole union accept anything, so neither defect surfaced and the union drifted out of step with `_coerce_reader`: pandas and Polars dataframes, HuggingFace datasets, column dicts, row dicts and Pydantic model instances are all coerced at runtime but were absent from the annotation. List every input `_coerce_reader` handles, import the Arrow classes directly, and dispatch on the optional dependencies the way the rest of the codebase does: `_check_for_polars` plus a real `isinstance` replaces the duplicated `__module__.startswith("polars")` string matching, which also lets a type checker narrow the branch. `LanceDataset` moves to a function-local import because `lance.dataset` names both this module and a function on the package. `types.py` cannot join the pyright target yet: without `pyarrow-stubs`, `isinstance(x, pa.Table)` does not narrow and every branch reports an error, so that is left for a follow-up along with the stub dependency. The union members are pinned in the typing regression file instead. --- python/python/lance/dataset.py | 6 ++- python/python/lance/dependencies.py | 4 +- python/python/lance/types.py | 51 ++++++++++++------- python/python/tests/test_dataset.py | 2 + .../tests/test_optional_types_typing.py | 38 ++++++++++++++ 5 files changed, 81 insertions(+), 20 deletions(-) diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 4739b8ff579..901cbbe706d 100644 --- a/python/python/lance/dataset.py +++ b/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, diff --git a/python/python/lance/dependencies.py b/python/python/lance/dependencies.py index e1435e6fca1..d9dc0cb1dad 100644 --- a/python/python/lance/dependencies.py +++ b/python/python/lance/dependencies.py @@ -162,6 +162,7 @@ def _lazy_import(module_name: str) -> tuple[ModuleType, bool]: import numpy import pandas import polars + import pydantic import torch # type: ignore[reportMissingImports] else: # heavy/optional third party libs @@ -170,7 +171,7 @@ def _lazy_import(module_name: str) -> tuple[ModuleType, bool]: polars, _POLARS_AVAILABLE = _lazy_import("polars") torch, _TORCH_AVAILABLE = _lazy_import("torch") datasets, _HUGGING_FACE_AVAILABLE = _lazy_import("datasets") - _, _PYDANTIC_AVAILABLE = _lazy_import("pydantic") + pydantic, _PYDANTIC_AVAILABLE = _lazy_import("pydantic") @lru_cache(maxsize=None) @@ -266,6 +267,7 @@ def _validate_pydantic_list(data: Any, model_class: type) -> None: "numpy", "pandas", "polars", + "pydantic", "torch", # lazy utilities "_check_for_hugging_face", diff --git a/python/python/lance/types.py b/python/python/lance/types.py index 3b5cc5b0683..019288ba7e8 100644 --- a/python/python/lance/types.py +++ b/python/python/lance/types.py @@ -3,30 +3,47 @@ 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 .dependencies import datasets, pydantic + + # Keep in step with the branches of ``_coerce_reader``: every input coerced + # there needs a member here, and every member here needs a branch there. + # The container members are the covariant spellings so that, say, a + # ``list[MyModel]`` is accepted; ``_coerce_reader`` narrows them to ``dict`` + # and ``list`` and reports anything else it cannot read. 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[pydantic.BaseModel], + Iterable[RecordBatch], ] @@ -70,10 +87,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 +103,25 @@ def _is_materialized(data_obj: ReaderLike) -> bool: def _coerce_reader( data_obj: ReaderLike, schema: Optional[pa.Schema] = None ) -> pa.RecordBatchReader: + # Imported here because ``lance.dataset`` imports this module, and because + # the ``lance.dataset`` name is also bound to a function in ``lance``. + 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 diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index 2f0bbd1fb1d..9bd734f2b31 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -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}]), ] diff --git a/python/python/tests/test_optional_types_typing.py b/python/python/tests/test_optional_types_typing.py index 89431e3b448..170596db815 100644 --- a/python/python/tests/test_optional_types_typing.py +++ b/python/python/tests/test_optional_types_typing.py @@ -27,6 +27,44 @@ if TYPE_CHECKING: import lance + import polars as pl + import pyarrow as pa + from lance.types import ReaderLike + from pyarrow.dataset import Dataset as ArrowDataset + from pyarrow.dataset import Scanner as ArrowScanner + from pydantic import BaseModel + + def _accept_reader(reader: ReaderLike) -> None: + pass + + def _check_reader_types( + pandas_dataframe: pd.DataFrame, + polars_dataframe: pl.DataFrame, + arrow_dataset: ArrowDataset, + arrow_scanner: ArrowScanner, + lance_dataset: lance.LanceDataset, + table: pa.Table, + batch: pa.RecordBatch, + reader: pa.RecordBatchReader, + batches: list[pa.RecordBatch], + models: list[BaseModel], + ) -> None: + # One case per branch of `lance.types._coerce_reader`. + _accept_reader(pandas_dataframe) + _accept_reader(polars_dataframe) + _accept_reader(arrow_dataset) + _accept_reader(arrow_scanner) + _accept_reader(lance_dataset) + _accept_reader(table) + _accept_reader(batch) + _accept_reader(reader) + _accept_reader(batches) + _accept_reader(models) + _accept_reader({"a": [1.0, 2.0]}) + _accept_reader([{"a": 1.0}, {"a": 2.0}]) + # No rejected case here: `ReaderLike` names pyarrow types, which are + # unresolved without `pyarrow-stubs`, and a union with an unresolved + # member accepts anything. The `asof` pins below have no such member. def _check_accepted_asof_types() -> None: # ``sanitize_ts`` is exercised at runtime below; this pins the public From 0d763568c9f8b8604c82e24d195591d84886c61b Mon Sep 17 00:00:00 2001 From: Jonas Dedden Date: Thu, 17 Sep 2026 17:06:20 +0200 Subject: [PATCH 3/3] test(python): drop no-op ReaderLike typing checks and trim comments ReaderLike accepts anything under the repo's pyright setup (no pyarrow stubs), so its acceptance checks could not fail. Fold sanitize_ts runtime tests into parametrized cases and import pydantic only for type checking. --- python/python/lance/dependencies.py | 4 +- python/python/lance/types.py | 15 ++- python/python/lance/util.py | 11 +- python/python/tests/test_dataset.py | 26 ++--- .../tests/test_optional_types_typing.py | 101 +++--------------- 5 files changed, 36 insertions(+), 121 deletions(-) diff --git a/python/python/lance/dependencies.py b/python/python/lance/dependencies.py index d9dc0cb1dad..e1435e6fca1 100644 --- a/python/python/lance/dependencies.py +++ b/python/python/lance/dependencies.py @@ -162,7 +162,6 @@ def _lazy_import(module_name: str) -> tuple[ModuleType, bool]: import numpy import pandas import polars - import pydantic import torch # type: ignore[reportMissingImports] else: # heavy/optional third party libs @@ -171,7 +170,7 @@ def _lazy_import(module_name: str) -> tuple[ModuleType, bool]: polars, _POLARS_AVAILABLE = _lazy_import("polars") torch, _TORCH_AVAILABLE = _lazy_import("torch") datasets, _HUGGING_FACE_AVAILABLE = _lazy_import("datasets") - pydantic, _PYDANTIC_AVAILABLE = _lazy_import("pydantic") + _, _PYDANTIC_AVAILABLE = _lazy_import("pydantic") @lru_cache(maxsize=None) @@ -267,7 +266,6 @@ def _validate_pydantic_list(data: Any, model_class: type) -> None: "numpy", "pandas", "polars", - "pydantic", "torch", # lazy utilities "_check_for_hugging_face", diff --git a/python/python/lance/types.py b/python/python/lance/types.py index 019288ba7e8..8641214fd4d 100644 --- a/python/python/lance/types.py +++ b/python/python/lance/types.py @@ -22,13 +22,11 @@ from .dependencies import polars as pl if TYPE_CHECKING: - from .dependencies import datasets, pydantic + from pydantic import BaseModel - # Keep in step with the branches of ``_coerce_reader``: every input coerced - # there needs a member here, and every member here needs a branch there. - # The container members are the covariant spellings so that, say, a - # ``list[MyModel]`` is accepted; ``_coerce_reader`` narrows them to ``dict`` - # and ``list`` and reports anything else it cannot read. + from .dependencies import datasets + + # Keep in step with the branches of ``_coerce_reader``. ReaderLike = Union[ pd.DataFrame, pl.DataFrame, @@ -42,7 +40,7 @@ datasets.IterableDataset, Mapping[str, Any], Sequence[Mapping[str, Any]], - Sequence[pydantic.BaseModel], + Sequence[BaseModel], Iterable[RecordBatch], ] @@ -103,8 +101,7 @@ def _is_materialized(data_obj: ReaderLike) -> bool: def _coerce_reader( data_obj: ReaderLike, schema: Optional[pa.Schema] = None ) -> pa.RecordBatchReader: - # Imported here because ``lance.dataset`` imports this module, and because - # the ``lance.dataset`` name is also bound to a function in ``lance``. + # Local import: ``lance.dataset`` imports this module. from .dataset import LanceDataset if _check_for_pandas(data_obj) and isinstance(data_obj, pd.DataFrame): diff --git a/python/python/lance/util.py b/python/python/lance/util.py index 78577258d8d..6169e99514b 100644 --- a/python/python/lance/util.py +++ b/python/python/lance/util.py @@ -26,11 +26,8 @@ if TYPE_CHECKING: from pandas.api.typing import NaTType - # ``pandas.Timestamp`` is a ``datetime`` subclass, so it needs no member of - # its own. ``NaTType`` does: the stubs bundled with pandas declare - # ``Timestamp.__new__`` as returning ``Self | NaTType``, so every caller - # writing ``asof=pd.Timestamp(...)`` hands us that union. We accept it here - # and reject the ``NaT`` half in ``sanitize_ts``. + # ``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"] @@ -57,9 +54,7 @@ def sanitize_ts(ts: ts_types) -> datetime: ) from None ts = pd.to_datetime(ts) if _check_for_pandas(ts): - # pandas spells a missing timestamp ``NaT``, which subclasses ``datetime`` - # and compares false against every version timestamp, so it would - # otherwise be reported as being older than the first version. + # ``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): diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index 9bd734f2b31..d5274f012e9 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -560,27 +560,23 @@ def test_asof_checkout(tmp_path: Path): assert len(ds.to_table()) == 9 -def test_sanitize_ts_parses_strings_with_pandas(): - # pandas accepts timestamp strings that the pandas-free fallback format - # rejects, so a date without a time of day has to work here. - assert sanitize_ts("2026-01-01") == datetime(2026, 1, 1) - - -def test_sanitize_ts_rejects_nat(): - # `NaT` is a `datetime` subclass, so it would otherwise pass through and - # compare false against every version timestamp. - with pytest.raises(ValueError, match="NaT is not a valid version timestamp"): - sanitize_ts(pd.NaT) +@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) -def test_sanitize_ts_rejects_unknown_type(): - with pytest.raises(TypeError, match="Unrecognized version timestamp"): - sanitize_ts(object()) +@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") diff --git a/python/python/tests/test_optional_types_typing.py b/python/python/tests/test_optional_types_typing.py index 170596db815..806f9eb0a96 100644 --- a/python/python/tests/test_optional_types_typing.py +++ b/python/python/tests/test_optional_types_typing.py @@ -1,98 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The Lance Authors -"""Type-checking regression tests for optional dependency input types. +"""Type-checking regressions for ``asof``; part of the pyright target. -This module is part of the pyright target configured in ``pyproject.toml``, so -a regression in the annotations below fails the repository type check and not -only the runtime suite. - -The rejected cases are pinned with ``pyright: ignore`` comments and -``reportUnnecessaryTypeIgnoreComment``, so the file also fails if an annotation -becomes *too* permissive and one of them stops being an error. They sit in a -``TYPE_CHECKING`` block because they are invalid at runtime; the matching -runtime assertions live in ``test_dataset.py``. - -Like ``test_fragment_typing.py``, this module does not import ``pytest``: the -lint workflow installs pyright without the test dependencies. +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 datetime import datetime from typing import TYPE_CHECKING -import pandas as pd -from lance.util import sanitize_ts - if TYPE_CHECKING: - import lance - import polars as pl - import pyarrow as pa - from lance.types import ReaderLike - from pyarrow.dataset import Dataset as ArrowDataset - from pyarrow.dataset import Scanner as ArrowScanner - from pydantic import BaseModel - - def _accept_reader(reader: ReaderLike) -> None: - pass - - def _check_reader_types( - pandas_dataframe: pd.DataFrame, - polars_dataframe: pl.DataFrame, - arrow_dataset: ArrowDataset, - arrow_scanner: ArrowScanner, - lance_dataset: lance.LanceDataset, - table: pa.Table, - batch: pa.RecordBatch, - reader: pa.RecordBatchReader, - batches: list[pa.RecordBatch], - models: list[BaseModel], - ) -> None: - # One case per branch of `lance.types._coerce_reader`. - _accept_reader(pandas_dataframe) - _accept_reader(polars_dataframe) - _accept_reader(arrow_dataset) - _accept_reader(arrow_scanner) - _accept_reader(lance_dataset) - _accept_reader(table) - _accept_reader(batch) - _accept_reader(reader) - _accept_reader(batches) - _accept_reader(models) - _accept_reader({"a": [1.0, 2.0]}) - _accept_reader([{"a": 1.0}, {"a": 2.0}]) - # No rejected case here: `ReaderLike` names pyarrow types, which are - # unresolved without `pyarrow-stubs`, and a union with an unresolved - # member accepts anything. The `asof` pins below have no such member. - - def _check_accepted_asof_types() -> None: - # ``sanitize_ts`` is exercised at runtime below; this pins the public - # entry point that forwards to it. - lance.dataset("memory://unused", asof=pd.Timestamp("2026-01-01")) - lance.dataset("memory://unused", asof=pd.NaT) - lance.dataset("memory://unused", asof=datetime(2026, 1, 1)) - lance.dataset("memory://unused", asof="2026-01-01") + from datetime import datetime - def _check_rejected_asof_types() -> None: - # `DatetimeIndex` has a `to_pydatetime` method, so it satisfies a - # structural timestamp annotation without being a valid instant. - index = pd.DatetimeIndex(["2026-01-01"]) - lance.dataset( - "memory://unused", - asof=index, # pyright: ignore[reportArgumentType] - ) - sanitize_ts(object()) # pyright: ignore[reportArgumentType] - - -def test_sanitize_ts_accepts_pandas_timestamp() -> None: - # The stubs bundled with pandas type this constructor as - # ``Timestamp | NaTType``, so both halves have to satisfy ``ts_types``. - result: datetime = sanitize_ts(pd.Timestamp("2026-01-01")) - - assert result == datetime(2026, 1, 1) - - -def test_sanitize_ts_accepts_datetime_and_str() -> None: - assert sanitize_ts(datetime(2026, 1, 1)) == datetime(2026, 1, 1) - assert sanitize_ts("2026-01-01 00:00:00") == datetime(2026, 1, 1) + 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]