Add Polars tabular file reader (#238)

Author: damonbayer

cfa/stf/forecasttools/__init__.py

@@ -15,7 +15,7 @@
 from .augment_samples_with_observations import augment_samples_with_observations
 from .create_proportions import create_proportions
 from .location_table import LOCATION_LIST
-from .utils import coalesce_common_columns
+from .utils import coalesce_common_columns, read_tabular
 
 
 def __getattr__(name):
@@ -28,6 +28,7 @@ def __getattr__(name):
 
 __all__ = [
     "coalesce_common_columns",
+    "read_tabular",
     "append_prop_data",
     "augment_samples_with_observations",
     "create_proportions",

cfa/stf/forecasttools/utils.py

@@ -1,7 +1,65 @@
+from pathlib import Path
+from typing import Any
+
 import polars as pl
 import polars.selectors as cs
 
 
+def _read_parquet_with_timezone_correction(
+    path_to_file: str | Path, **kwargs: Any
+) -> pl.DataFrame:
+    """Read Parquet data, interpreting timezone-naive timestamps as UTC."""
+    df = pl.read_parquet(path_to_file, **kwargs)
+    timestamp_columns_without_timezone = [
+        name
+        for name, dtype in df.schema.items()
+        if isinstance(dtype, pl.Datetime) and dtype.time_zone is None
+    ]
+
+    return df.with_columns(
+        pl.col(timestamp_columns_without_timezone).dt.convert_time_zone("UTC")
+    )
+
+
+def read_tabular(path_to_file: str | Path, **kwargs: Any) -> pl.DataFrame:
+    """Read a tabular file, inferring its format from the file extension.
+
+    Parameters
+    ----------
+    path_to_file
+        Path to a ``.csv``, ``.tsv``, or ``.parquet`` file. The extension is
+        matched case-insensitively.
+    **kwargs
+        Additional keyword arguments passed to :func:`polars.read_csv` or
+        :func:`polars.read_parquet`, depending on the inferred format.
+
+    Returns
+    -------
+    pl.DataFrame
+        The contents of the file.
+
+    Raises
+    ------
+    ValueError
+        If the path does not have a supported file extension.
+    """
+    file_format = Path(path_to_file).suffix.removeprefix(".").lower()
+
+    if file_format == "csv":
+        return pl.read_csv(path_to_file, **kwargs)
+    if file_format == "tsv":
+        kwargs.setdefault("separator", "\t")
+        return pl.read_csv(path_to_file, **kwargs)
+    if file_format == "parquet":
+        return _read_parquet_with_timezone_correction(path_to_file, **kwargs)
+
+    supported_formats = ".csv, .tsv, and .parquet"
+    raise ValueError(
+        f"Unsupported file extension {Path(path_to_file).suffix!r}; "
+        f"expected one of {supported_formats}."
+    )
+
+
 def coalesce_common_columns(
     df: pl.DataFrame, suffix: str, new_colname: str | None = None
 ) -> pl.DataFrame:

tests/cfa/stf/forecasttools/test_read_tabular.py

@@ -0,0 +1,87 @@
+from datetime import datetime
+from zoneinfo import ZoneInfo
+
+import polars as pl
+import polars.testing as plt
+import pytest
+
+import cfa.stf.forecasttools as ft
+
+
+@pytest.fixture
+def tabular_data():
+    return pl.DataFrame(
+        {
+            "location": ["US", "01"],
+            "value": [1.5, 2.5],
+        }
+    )
+
+
+@pytest.mark.parametrize("file_format", ["csv", "tsv", "parquet"])
+def test_read_tabular_reads_supported_formats(tmp_path, tabular_data, file_format):
+    path = tmp_path / f"data.{file_format}"
+    if file_format == "csv":
+        tabular_data.write_csv(path)
+    elif file_format == "tsv":
+        tabular_data.write_csv(path, separator="\t")
+    else:
+        tabular_data.write_parquet(path)
+
+    result = ft.read_tabular(path)
+
+    plt.assert_frame_equal(result, tabular_data)
+
+
+def test_read_tabular_matches_extension_case_insensitively(tmp_path, tabular_data):
+    path = tmp_path / "data.CSV"
+    tabular_data.write_csv(path)
+
+    result = ft.read_tabular(path)
+
+    plt.assert_frame_equal(result, tabular_data)
+
+
+def test_read_tabular_forwards_reader_options(tmp_path):
+    path = tmp_path / "data.csv"
+    path.write_text("value\n1\n2\n3\n")
+
+    result = ft.read_tabular(path, n_rows=2)
+
+    assert result.get_column("value").to_list() == [1, 2]
+
+
+def test_read_tabular_corrects_timezone_naive_parquet_timestamps(tmp_path):
+    path = tmp_path / "timestamps.parquet"
+    tokyo = ZoneInfo("Asia/Tokyo")
+    data = pl.DataFrame(
+        {
+            "timestamp_without_timezone": [datetime(2026, 1, 15, 12, 30)],
+            "timestamp_with_timezone": [datetime(2026, 1, 15, 12, 30, tzinfo=tokyo)],
+        },
+        schema={
+            "timestamp_without_timezone": pl.Datetime("us"),
+            "timestamp_with_timezone": pl.Datetime("us", "Asia/Tokyo"),
+        },
+    )
+    data.write_parquet(path)
+
+    result = ft.read_tabular(path)
+
+    assert result.schema == pl.Schema(
+        {
+            "timestamp_without_timezone": pl.Datetime("us", "UTC"),
+            "timestamp_with_timezone": pl.Datetime("us", "Asia/Tokyo"),
+        }
+    )
+    assert result.select(pl.all().to_physical()).row(0) == data.select(
+        pl.all().to_physical()
+    ).row(0)
+
+
+@pytest.mark.parametrize("filename", ["data.json", "data"])
+def test_read_tabular_rejects_unsupported_extensions(tmp_path, filename):
+    path = tmp_path / filename
+
+    with pytest.raises(ValueError, match="Unsupported file extension"):
+        ft.read_tabular(path)