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/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/types.py b/python/python/lance/types.py index 3b5cc5b0683..8641214fd4d 100644 --- a/python/python/lance/types.py +++ b/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 diff --git a/python/python/lance/util.py b/python/python/lance/util.py index 180e2441b43..6169e99514b 100644 --- a/python/python/lance/util.py +++ b/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: diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index f48138f2bae..d5274f012e9 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 @@ -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}] 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..806f9eb0a96 --- /dev/null +++ b/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]