diff --git a/.github/workflows/unittests.yml b/.github/workflows/unittests.yml new file mode 100644 index 0000000..9680e66 --- /dev/null +++ b/.github/workflows/unittests.yml @@ -0,0 +1,105 @@ +name: Unit tests + +on: + push: + paths-ignore: + - "docs/**" + - "*.md" + pull_request: + paths-ignore: + - "docs/**" + - "*.md" + +jobs: + format: + if: | + (github.event_name == 'push' || github.event.pull_request.head.repo.fork) && + (!contains(github.event.head_commit.message, 'skip ci')) + + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Set up Python 3.13 + uses: actions/setup-python@v5 + with: + python-version: 3.13 + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install black toml + - name: Format source w/ black + # fail when file(s) would be formatted + run: black --check arrow_expts tests + + lint: + if: | + (github.event_name == 'push' || github.event.pull_request.head.repo.fork) && + (!contains(github.event.head_commit.message, 'skip ci')) + + strategy: + matrix: + pyver: ["py310", "py313"] + + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Ruff with Python target version ${{ matrix.pyver }} + uses: chartboost/ruff-action@v1 + with: + src: arrow_expts + args: "check --target-version ${{ matrix.pyver }}" + + test: + if: | + (github.event_name == 'push' || github.event.pull_request.head.repo.fork) && + (!contains(github.event.head_commit.message, 'skip ci')) + needs: [format, lint] + + strategy: + matrix: + python-version: ["3.10", "3.13"] + os: [ubuntu-latest, windows-latest, macos-latest] + # include: + # - python-version: "3.14" + # os: ubuntu-latest + + runs-on: ${{ matrix.os }} + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + allow-prereleases: true + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install . + - name: Run pytest + run: | + pytest -vv tests + + type-hints: + if: | + (github.event_name == 'push' || github.event.pull_request.head.repo.fork) && + (!contains(github.event.head_commit.message, 'skip ci')) + + strategy: + matrix: + python-version: ["3.10", "3.13", "3.14"] + + runs-on: ubuntu-latest + continue-on-error: true + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + allow-prereleases: true + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install mypy{,_extensions} + - name: Type check w/ mypy + run: echo y | mypy --install-types --check-untyped-defs --pretty arrow_expts diff --git a/arrow_expts/schema/json-blob-schema.json b/arrow_expts/schema/json-blob-schema.json index 672b1a6..222d14a 100644 --- a/arrow_expts/schema/json-blob-schema.json +++ b/arrow_expts/schema/json-blob-schema.json @@ -61,6 +61,20 @@ }, "type": "array" }, + { + "items": { + "anyOf": [ + { + "format": "binary", + "type": "string" + }, + { + "type": "null" + } + ] + }, + "type": "array" + }, { "items": { "anyOf": [ @@ -112,6 +126,7 @@ "integer", "number", "boolean", + "bytes", "date-time", "duration", "time-pattern" @@ -280,6 +295,20 @@ }, "type": "array" }, + { + "items": { + "anyOf": [ + { + "format": "binary", + "type": "string" + }, + { + "type": "null" + } + ] + }, + "type": "array" + }, { "items": { "anyOf": [ @@ -331,6 +360,7 @@ "integer", "number", "boolean", + "bytes", "date-time", "duration", "time-pattern" @@ -501,6 +531,20 @@ }, "type": "array" }, + { + "items": { + "anyOf": [ + { + "format": "binary", + "type": "string" + }, + { + "type": "null" + } + ] + }, + "type": "array" + }, { "items": { "anyOf": [ @@ -552,6 +596,7 @@ "integer", "number", "boolean", + "bytes", "date-time", "duration", "time-pattern" diff --git a/arrow_expts/schema/models.py b/arrow_expts/schema/models.py index ba28c92..e68630b 100755 --- a/arrow_expts/schema/models.py +++ b/arrow_expts/schema/models.py @@ -2,65 +2,100 @@ # /// script # requires-python = ">=3.10" # dependencies = [ +# "python-dateutil", # "pandas>=2", # "pydantic>=2", # ] # /// -"""Write JSON schema for JSON blob in SpineDB +"""Write JSON schema for JSON blob in SpineDB""" -""" - -# from dataclasses import dataclass -# from dataclasses import field from datetime import datetime, timedelta -from typing import Annotated, Literal, Type, TypeAlias +from typing import Annotated, Literal, TypeAlias +from dateutil.relativedelta import relativedelta +import numpy as np import pandas as pd -from pydantic import RootModel -from pydantic.dataclasses import dataclass -from pydantic.dataclasses import Field as field -from pydantic.types import StringConstraints + +if __name__ == "__main__": + from pydantic.dataclasses import dataclass + from pydantic.dataclasses import Field as field +else: + from dataclasses import dataclass + from dataclasses import field Floats: TypeAlias = list[float] Integers: TypeAlias = list[int] Strings: TypeAlias = list[str] Booleans: TypeAlias = list[bool] +BytesList: TypeAlias = list[bytes] # array of bytes to support mixed types Datetimes: TypeAlias = list[datetime] Timedeltas: TypeAlias = list[timedelta] # FIXME: how to do w/o Pydantic? time_pat_re = r"(Y|M|D|WD|h|m|s)[0-9]+-[0-9]+" -TimePattern: TypeAlias = Annotated[str, StringConstraints(pattern=time_pat_re)] + +# generate schema w/ Pydantic by running as a script +if __name__ == "__main__": + from pydantic.types import StringConstraints + + TimePattern: TypeAlias = Annotated[str, StringConstraints(pattern=time_pat_re)] +else: + from re import Pattern + + @dataclass(frozen=True) + class TimePattern: + pattern: str + re: str | Pattern[str] = time_pat_re + + TimePatterns: TypeAlias = list[TimePattern] +# nullable variant of arrays NullableIntegers: TypeAlias = list[int | None] NullableFloats: TypeAlias = list[float | None] NullableStrings: TypeAlias = list[str | None] NullableBooleans: TypeAlias = list[bool | None] +NullableBytesList: TypeAlias = list[bytes | None] NullableDatetimes: TypeAlias = list[datetime | None] NullableTimedeltas: TypeAlias = list[timedelta | None] NullableTimePatterns: TypeAlias = list[TimePattern | None] +# sets of types used to define array schemas below IndexTypes: TypeAlias = Integers | Strings | Datetimes | Timedeltas | TimePatterns ValueTypes: TypeAlias = ( - Integers | Strings | Floats | Booleans | Datetimes | Timedeltas | TimePatterns + Integers + | Strings + | Floats + | Booleans + | Datetimes + | Timedeltas + | TimePatterns + | BytesList ) NullableValueTypes: TypeAlias = ( NullableIntegers | NullableStrings | NullableFloats | NullableBooleans + | NullableBytesList | NullableDatetimes | NullableTimedeltas | NullableTimePatterns ) - +# names of types used in the schema ValueTypeNames: TypeAlias = Literal[ - "string", "integer", "number", "boolean", "date-time", "duration", "time-pattern" + "string", + "integer", + "number", + "boolean", + "bytes", + "date-time", + "duration", + "time-pattern", ] IndexValueTypeNames: TypeAlias = Literal[ "string", "integer", "date-time", "duration", "time-pattern" @@ -69,13 +104,25 @@ type_map: dict[type, ValueTypeNames] = { str: "string", int: "integer", + np.int8: "integer", + np.int16: "integer", + np.int32: "integer", + np.int64: "integer", float: "number", + np.float16: "number", + np.float32: "number", + np.float64: "number", + # np.float128: "number", # not available on macos bool: "boolean", + np.bool: "boolean", datetime: "date-time", pd.Timestamp: "date-time", timedelta: "duration", pd.Timedelta: "duration", + relativedelta: "duration", + pd.DateOffset: "duration", TimePattern: "time-pattern", + bytes: "string", } @@ -184,7 +231,7 @@ class Array(_TypeInferMixin): # NOTE: To add run-length encoding to the schema, add it to the # following type union following which, we need to implement a -# converter to an Arrow array type +# converter to a compatible pyarrow array type Table: TypeAlias = list[ RunEndIndex | DictEncodedIndex | ArrayIndex | RunEndArray | DictEncodedArray | Array ] @@ -195,6 +242,8 @@ class Array(_TypeInferMixin): import json from pathlib import Path + from pydantic import RootModel + parser = ArgumentParser(__doc__) parser.add_argument("json_file", help="Path of JSON schema file to write") opts = parser.parse_args() diff --git a/arrow_expts/schema/reencode.py b/arrow_expts/schema/reencode.py index 5251e0e..43901ff 100755 --- a/arrow_expts/schema/reencode.py +++ b/arrow_expts/schema/reencode.py @@ -8,37 +8,39 @@ # ] # /// -"""Reencode old map type JSON to new table/tables type JSON - -""" +"""Reencode old map type JSON to new table/tables type JSON""" from datetime import datetime, timedelta import json from pathlib import Path +from typing import overload +import numpy as np import pandas as pd -import pyarrow as pa from pydantic import RootModel -from rich.pretty import pprint +from rich.pretty import pprint # noqa: F401, keep for debugging from ..spine.dbmap import make_records from .models import ( Array, ArrayIndex, - DEArray, - DEIndex, - REArray, - REIndex, - RLIndex, + DictEncodedArray, + DictEncodedIndex, + RunEndArray, + RunEndIndex, + RunLengthArray, + RunLengthIndex, Table, ) def to_df(json_doc: dict): - data = make_records(json_doc, {}, [], idx_name="metric") - tbl = pa.Table.from_pylist(data) - df = tbl.to_pandas(types_mapper=pd.ArrowDtype) + data = make_records(json_doc, {}, []) + # NOTE: don't use pyarrow, difficult to support mixed types + # tbl = pa.Table.from_pylist(data) + # df = tbl.to_pandas(types_mapper=pd.ArrowDtype) + df = pd.DataFrame.from_records(data) return df @@ -46,70 +48,107 @@ class _sentinel: pass -def rl_encode(arr: ArrayIndex) -> RLIndex: - last = _sentinel() - values = [] - run_len = [] +SENTINEL = _sentinel() + + +@overload +def rl_encode(arr: Array) -> RunLengthArray: ... + + +@overload +def rl_encode(arr: ArrayIndex) -> RunLengthIndex: ... + + +def rl_encode(arr): + last = SENTINEL + values, run_len = [], [] for val in arr.values: if val != last: values.append(val) run_len.append(1) + last = val else: run_len[-1] += 1 - return RLIndex(name=arr.name, values=values, run_len=run_len) + if isinstance(arr, ArrayIndex): + return RunLengthIndex(name=arr.name, values=values, run_len=run_len) + else: + return RunLengthArray(name=arr.name, values=values, run_len=run_len) -def re_encode(arr: ArrayIndex) -> REIndex: - last = arr.values[0] # _sentinel() - values = [last] - run_end = [] - for idx, val in enumerate(arr.values[1:], start=1): - if val != last: - last = val +@overload +def re_encode(arr: Array) -> RunEndArray: ... + + +@overload +def re_encode(arr: ArrayIndex) -> RunEndIndex: ... + + +def re_encode(arr): + last = SENTINEL + values, run_end = [], [] + for idx, val in enumerate(arr.values, start=1): + if last != val: values.append(val) run_end.append(idx) - run_end.append(len(arr.values)) - return REIndex(name=arr.name, values=values, run_end=run_end) + else: + run_end[-1] = idx + last = val + if isinstance(arr, ArrayIndex): + return RunEndIndex(name=arr.name, values=values, run_end=run_end) + else: + return RunEndArray(name=arr.name, values=values, run_end=run_end) + +@overload +def de_encode(arr: Array) -> DictEncodedArray: ... -def de_encode(arr: ArrayIndex) -> DEIndex: + +@overload +def de_encode(arr: ArrayIndex) -> DictEncodedIndex: ... + + +def de_encode(arr): # not using list(set(...)) to preserve order values = list(dict.fromkeys(arr.values)) indices = list(map(values.index, arr.values)) - return DEIndex(name=arr.name, values=values, indices=indices) + if isinstance(arr, ArrayIndex): + return DictEncodedIndex(name=arr.name, values=values, indices=indices) + else: + return DictEncodedArray(name=arr.name, values=values, indices=indices) -def series_to_col(col: pd.Series) -> ArrayIndex | DEIndex | Array | DEArray: +def series_to_col( + col: pd.Series, +) -> ArrayIndex | DictEncodedIndex | Array | DictEncodedArray: match col.name, col.dtype.type: - case "value", t if issubclass(t, bool | str) or t is object: - print(f"idx_type: {t}, value: {col.iloc[:3]}") + case "value", t if issubclass(t, str) or t in (object, np.object_): col = col.astype("category") - arr = DEArray( + return DictEncodedArray( name=col.name, values=col.cat.categories, indices=col.cat.codes, ) - case _, t if issubclass(t, int | float): - arr = Array(name=col.name, values=col.values) - case _, t if issubclass(t, str) or t is object: - print(f"type: {t}, value: {col.iloc[:3]}") + case "value", t if issubclass(t, int): + return Array(name=col.name, values=col.to_list()) + case _, t if issubclass(t, (bool, float, bytes)): + return Array(name=col.name, values=col.to_list()) + case _, t if issubclass(t, str) or t in (object, np.object_): col = col.astype("category") - arr = DEIndex( + return DictEncodedIndex( name=col.name, - values=col.cat.categories, - indices=col.cat.codes, + values=col.cat.categories.to_list(), + indices=col.cat.codes.to_list(), ) - case _, t if issubclass(t, int | str | datetime | timedelta) or t is object: - arr = ArrayIndex(name=col.name, values=col.values) - case _, _: - raise NotImplementedError(f"unknown type {t}") - return arr + case _, t if issubclass(t, (int, datetime, timedelta)) or t is object: + return ArrayIndex(name=col.name, values=col.to_list()) + case n, t: + raise NotImplementedError(f"{n}: unknown type {t}") def to_tables(df: pd.DataFrame) -> Table: if df.empty: return [] - return [series_to_col(df[colname]) for colname in df.columns] + return [series_to_col(col) for _, col in df.items()] if __name__ == "__main__": @@ -117,20 +156,12 @@ def to_tables(df: pd.DataFrame) -> Table: parser = ArgumentParser(__doc__) parser.add_argument("old_json") - #parser.add_argument("new_json") + parser.add_argument("new_json") opts = parser.parse_args() - new_json = opts.old_json.replace("/0_orig/", "/1_reencoded/") # FIXME brittle df = to_df(json.loads(Path(opts.old_json).read_text())) - print("== DF ==") - #print(df) tbls = to_tables(df) - print("== TBLS ==") - print(tbls) - model = RootModel[Table](tbls).model_dump_json(indent=2) - print("== MODEL ==") - print(model) # NOTE: using pydantic is optional here, we can also write our own # JSON serialisation if we want, probably all we need is `asdict(...)`. - #Path(opts.new_json).write_text(RootModel[Table](tbls).model_dump_json()) - Path(new_json).write_text(model) + json_blob = RootModel[Table](tbls).model_dump_json(indent=2) + Path(opts.new_json).write_text(json_blob) diff --git a/arrow_expts/spine/dbmap.py b/arrow_expts/spine/dbmap.py index 7b9e163..8053dc8 100644 --- a/arrow_expts/spine/dbmap.py +++ b/arrow_expts/spine/dbmap.py @@ -1,13 +1,14 @@ from argparse import ArgumentParser -from datetime import datetime -from enum import Enum, auto import json import re -from typing import cast +from typing import cast, Any, Callable, Iterable, TypeAlias +from warnings import warn import weakref import pandas as pd +import numpy as np +from arrow_expts.schema.models import TimePattern from spinedb_api import DatabaseMapping from spinedb_api.temp_id import TempId @@ -16,39 +17,175 @@ def json_loads_ts(json_str: str | bytes): return pd.Series(json.loads(json_str)["data"]) -SEQ_PAT = re.compile(r"(t|p)([0-9]+)") +# Regex pattern to indentify numerical sequences encoded as string +SEQ_PAT = re.compile(r"^(t|p)([0-9]+)$") +# Regex pattern to identify a number encoded as a string +FREQ_PAT = re.compile("^[0-9]+$") +# Regex pattern to duration strings +DUR_PAT = re.compile(r"([0-9]+) *(Y|M|W|D|h|min|s)") -def filter_frequencies(freq: str) -> str: - # not very robust yet - filtered_freq = freq \ - .replace("years", "Y") \ - .replace("year", "Y") \ - .replace("months", "M") \ - .replace("month", "M") \ - .replace("quarters", "Q") \ - .replace("quarter", "Q") \ - .replace("weeks", "W") \ - .replace("week", "W") \ - .replace("hours", "h") \ - .replace("hour", "h") \ - .replace("minutes", "min") \ - .replace("minute", "min") \ - .replace("seconds", "s") \ - .replace("second", "s") \ - .replace("microseconds", "us") \ - .replace("microsecond", "us") \ - .replace("nanoseconds", "ns") \ - .replace("nanosecond", "ns") - if re.compile("^[0-9]+$").match(filtered_freq): + +def _normalise_freq(freq: int | str): + """Normalise integer/string to frequency. + + The frequency value is as understood by `pandas.Timedelta`. Note + that ambiguous values such as month or year are still retained + with the intention to handle later in the pipeline. + + """ + if isinstance(freq, int): + return str(freq) + "min" + if FREQ_PAT.match(freq): # If frequency is an integer, the implied unit is "minutes" - return freq + "m" + return freq + "min" + # not very robust yet + return ( + freq.replace("years", "Y") + .replace("year", "Y") + .replace("months", "M") + .replace("month", "M") + .replace("weeks", "W") + .replace("week", "W") + .replace("days", "D") + .replace("day", "D") + .replace("hours", "h") + .replace("hour", "h") + .replace("minutes", "min") + .replace("minute", "min") + .replace("seconds", "s") + .replace("second", "s") + ) + + +_to_numpy_time_units = { + "Y": "Y", + "M": "M", + "W": "W", + "D": "D", + "h": "h", + "min": "m", + "s": "s", +} + + +def _low_res_datetime(start: str, freq: str, periods: int) -> pd.DatetimeIndex: + """Create pd.DatetimeIndex with lower time resolution. + + The default resolution of pd.date_time is [ns], which puts + boundaries on allowed start- and end-dates due to limited storage + capacity. Choosing a resolution of [s] instead opens up that range + considerably. + + "For nanosecond resolution, the time span that can be represented + using a 64-bit integer is limited to approximately 584 years." - + https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#timestamp-limitations + + You can check the available ranges with `pd.Timestamp.min` and + `pd.Timestamp.max`. + + """ + if re_match := DUR_PAT.match(_normalise_freq(freq)): + number_str, unit = re_match.groups() + else: + raise ValueError(f"invalid frequency: {freq!r}") + + start_date_np = np.datetime64(start, "s") + freq_np = np.timedelta64(int(number_str), _to_numpy_time_units[unit]) + freq_pd = pd.Timedelta(freq_np) + + date_array = np.arange(start_date_np, start_date_np + periods * freq_np, freq_np) + date_array_with_frequency = pd.DatetimeIndex( + date_array, freq=freq_pd, dtype="datetime64[s]" + ) + + return date_array_with_frequency + + +def _to_dateoffset(val: str) -> pd.DateOffset: + if (m := DUR_PAT.match(val)) is None: + raise ValueError(f"{val}: bad duration value") + num_str, freq = m.groups() + num = int(num_str) + match freq: + case "Y": + return pd.DateOffset(years=num) + case "M": + return pd.DateOffset(months=num) + case "W": + return pd.DateOffset(weeks=num) + case "D": + return pd.DateOffset(days=num) + case "h": + return pd.DateOffset(hours=num) + case "min": + return pd.DateOffset(minutes=num) + case "s": + return pd.DateOffset(seconds=num) + case _: + # should not get here + raise ValueError(f"{val}: unknown duration") + + +def _atoi(val: str) -> int | str: + """Convert string to number if it matches `t0001` or `p2001`. + + If a match is found, also override the name to "time" or "period" + respectively. + + """ + if m := SEQ_PAT.match(val): + return int(m.group(2)) else: - return freq + return val + + +_FmtIdx: TypeAlias = Callable[[str, str | Any], dict[str, Any]] + + +def _formatter(index_type: str) -> _FmtIdx: + """Get a function that formats the values of a name value pair. + + The name is the column name. The function returned depends on the + `index_type`. An unknown `index_type` returns a noop formatter, + but it also issues a warning. A noop formatter can be requested + explicitly by passing the type "noop"; no warning is issued in + this case. -class IndexType(Enum): - Timestamp = auto() - Sequence = auto() - Generic = auto() + Index types: + ============ + + - "date_time" :: converts value to `datetime` + + - "duration" :: converts string to `pandas.DateOffset`; this + allows for ambiguous units like month or year. + + - "str" :: convert the value to integer if it matches `t0001` or + `p2002`, and the name to "time" and "period" respectively; + without a match it is a noop. + + - "float" | "time_pattern" | "noop" :: noop + + - fallback :: noop with a warning + + """ + match index_type: + case "date_time" | "datetime": + return lambda name, key: {name: pd.Timestamp(key)} + case "duration": + return lambda name, key: {name: _to_dateoffset(_normalise_freq(key))} + case "str": + # don't use lambda, can't add type hints + def _atoi_dict(name: str, val: str) -> dict[str, int | str]: + return {name: _atoi(val)} + + return _atoi_dict + case "time_pattern" | "timepattern": + return lambda name, key: {name: TimePattern(key)} + case "float" | "noop": + return lambda name, key: {name: key} + case _: # fallback to noop w/ a warning + warn(f"{index_type}: unknown type, fallback to noop formatter") + return lambda name, key: {name: key} def make_records( @@ -56,149 +193,131 @@ def make_records( idx_lvls: dict, res: list[dict], *, - idx_name: str = "default", + lvlname_base: str = "default", ) -> list[dict]: - """Parse time-series w/ a multi-index stored as a nested map + """Parse parameter value into a list of records + + Spine db stores parameter_value as JSON. After the JSON blob has + been decoded to a Python dict, this function can transform it into + a list of records (dict) like a table. These records can then be + consumed by Pandas to create a dataframe. - Ask Suvayu for the example DB + The parsing logic works recursively by traversing depth first. + Each call incrementally accumulates a cell/level of a record in + the `idx_lvls` dictionary, once the traversal reaches a leaf node, + the final record is appended to the list `res`. The final result + is also returned by the function, allowing for composition. + + If at any level, the index level name is missing, a default base + name can be provided by setting a default `lvlname_base`. The + level name is derived by concatenating the base name with depth + level. """ + lvlname = lvlname_base + f"{len(idx_lvls)}" + + # NOTE: The private functions below are closures, defined early in + # the function such that they have the original arguments to + # `make_records` available to them, but nothing more. They either + # help with some computation, raise a warning, or are helpers to + # append to the result. + _msg_assert = ( + "for the type checker: rest of the function expects `json_doc` to be a dict" + ) + + def _uniquify_index_name(default: str) -> str: + assert isinstance(json_doc, dict), _msg_assert + index_name = json_doc.get("index_name", default) + return index_name + f"{len(idx_lvls)}" if index_name in idx_lvls else index_name + + def _from_pairs(data: Iterable[Iterable], fmt: _FmtIdx): + index_name = _uniquify_index_name(lvlname) + for key, val in data: + _lvls = {**idx_lvls, **fmt(index_name, key)} + make_records(val, _lvls, res, lvlname_base=lvlname_base) + + def _deprecated(var: str, val: Any): + assert isinstance(json_doc, dict), _msg_assert + index_name = json_doc.get("index_name", lvlname) + msg = f"{index_name}: {var}={val} is deprecated, handle in model, defaulting to time index from 0001-01-01." + warn(msg, DeprecationWarning) + + def _time_index(idx: dict, length: int): + start = idx.get("start", "0001-01-01T00:00:00") + resolution = idx.get("resolution", "1h") + return _low_res_datetime(start=start, freq=resolution, periods=length) + + def _append_arr(arr: Iterable, fmt: _FmtIdx): + index_name = _uniquify_index_name("i") + for value in arr: + res.append({**idx_lvls, **fmt(index_name, value)}) match json_doc: # maps - case {"data": list() as data, "type": "map", **_r}: - index_name = json_doc.get("index_name", "time") # use "time" if "index_name" does not exist - index_type = json_doc.get("index_type") - for key, val in data: - if index_type == "date_time": - key = datetime.fromisoformat(key) - if index_type == "duration": - key = pd.Timedelta(key) - make_records(val, {**idx_lvls, index_name: key}, res) - case {"data": dict() as data, "type": "map", **_r}: - index_name = json_doc.get("index_name", "time") # use "time" if "index_name" does not exist - index_type = json_doc.get("index_type") - for key, val in data.items(): - if index_type == "date_time": - key = datetime.fromisoformat(key) - if index_type == "duration": - key = pd.Timedelta(key) - make_records(val, {**idx_lvls, index_name: key}, res) + case {"data": dict() as data, "type": "map"}: + # NOTE: is "index_type" mandatory? In case it's not, we + # check for it separately, and fallback in a way that + # raises a warning but doesn't crash; same for the + # 2-column array variant below. + index_type = json_doc.get("index_type", "undefined-index_type-in-map") + _from_pairs(data.items(), _formatter(index_type)) + case {"data": dict() as data, "index_type": index_type}: + # NOTE: relies on other types not having "index_type"; + # same for the 2-column array variant below. + _from_pairs(data.items(), _formatter(index_type)) + case {"data": [[_, _], *_] as data, "type": "map"}: + index_type = json_doc.get("index_type", "undefined-index_type-in-map") + _from_pairs(data, _formatter(index_type)) + case {"data": [[_, _], *_] as data, "index_type": index_type}: + _from_pairs(data, _formatter(index_type)) # time series - case {"data": dict() as data, "type": "time_series", **_r}: - index_name = json_doc.get("index_name", "time") # use "time" if "index_name" does not exist - for key, val in data.items(): - key = datetime.fromisoformat(key) - make_records(val, {**idx_lvls, index_name: key}, res) + case {"data": dict() as data, "type": "time_series"}: + _from_pairs(data.items(), _formatter("date_time")) + case {"data": [[str(), float() | int()], *_] as data, "type": "time_series"}: + _from_pairs(data, _formatter("date_time")) case { - "data": [[str(), float() | int()], *_] as data, + "data": [float() | int(), *_] as data, "type": "time_series", - **_r, + "index": dict() as idx, }: - index_name = json_doc.get("index_name", "time") # use "time" if "index_name" does not exist - for key, val in data: - key = datetime.fromisoformat(key) - make_records(val, {**idx_lvls, index_name: key}, res) - case {"data": [float() | int(), *_] as data, "type": "time_series", **_r}: - index_name = json_doc.get("index_name", "time") # use "time" if "index_name" does not exist - ignore_year = json_doc.get(index, {}).get("ignore_year", None) - repeat = json_doc.get(index, {}).get("repeat", None) - if ignore_year == False: - raise ValueError('Can\'t handle `ignore_year == False`. Please re-format your dataset.') - if repeat == False: - raise ValueError('Can\'t handle `repeat == False`. Please re-format your dataset.') - match json_doc: - case { - "index": { - "start": start, - "resolution": freq, - "ignore_year": bool(), - "repeat": bool(), - }, - **_r, - }: - freq = filter_frequencies(freq) - index = pd.date_range(start=start, freq=freq, periods=len(data)) - case _: - if json_doc.get(index, {}) != {}: - raise NotImplementedError('Can\'t handle a partially set `index` value. Please re-format your dataset.') - index = pd.date_range( - start="0001-01-01", periods=len(data), freq="1h" - ) - for time, val in zip(index, data): - make_records(val, {**idx_lvls, index_name: time, "value": val}, res) + match idx: + case {"ignore_year": ignore_year}: + _deprecated("ignore_year", ignore_year) + case {"repeat": repeat}: + _deprecated("repeat", repeat) + + index = _time_index(idx, len(data)) + _from_pairs(zip(index, data), _formatter("noop")) + case {"type": "time_series", "data": [float() | int(), *_] as data}: + msg = "array-like 'time_series' without time-stamps, relies on 'ignore_year' and 'repeat' implicitly" + warn(msg, DeprecationWarning) + updated = {**json_doc, "index": {"ignore_year": True, "repeat": True}} + make_records(updated, idx_lvls, res, lvlname_base=lvlname_base) + # time_pattern + case {"type": "time_pattern", "data": dict() as data}: + _from_pairs(data.items(), _formatter("time_pattern")) + # arrays case { - "data": [[str(), dict() | float() | int()], *_] as data, - "type": "time_series", - **_r, + "type": "array", + "value_type": value_type, + "data": [str() | float() | int(), *_] as data, }: - if m := SEQ_PAT.match(data[0][0]): - idx_type = IndexType.Sequence - index_name = json_doc.get( - "index_name", "period" if "p" == m.group(1) else "seq" - ) - else: - idx_type = IndexType.Generic - index_name = json_doc.get("index_name", idx_name) # use idx_name if "index_name" does not exist - for key, val in data: - if idx_type == IndexType.Sequence: - m = SEQ_PAT.match(key) - assert m is not None - key = int(m.group(2)) - make_records(val, {**idx_lvls, index_name: key}, res) - - # arrays - case {"type": "array", "data": [str() | float() | int(), *_] as data, **_r}: - value_type = json_doc.get("value_type", "float") # use "float" if "value_type" does not exist - index_name = json_doc.get("index_name", "i") # use "i" if "index_name" does not exist - - if value_type == "duration": - try: - data = [pd.Timedelta(filter_frequencies(value)) for value in data] - except ValueError as err: - if "invalid unit abbreviation" in repr(err): - raise ValueError('"year" and "month" are ambiguous time units. Please convert them to days.') - else: - raise err - elif value_type == "date_time": - data = [datetime.fromisoformat(key) for key in data] - elif value_type == "float": - data = [pd.to_numeric(value) for value in data] - elif value_type == "str": - pass - else: - raise NotImplementedError(f"Can't match {value_type} arrays.") - for value in data: - res.append({index_name: value}) - - # date-time - case {"type": "date_time", "data": str() as data, **_r}: - idx_lvls["value"] = datetime.fromisoformat(data) - res.append(idx_lvls) - - # duration - case {"type": "duration", "data": int() | str() as data, **_r}: - if type(data) == int: - data = str(data) + "m" # integer time unit is "minutes" - try: - data = pd.Timedelta(filter_frequencies(data)) - except ValueError as err: - if "invalid unit abbreviation" in repr(err): - raise ValueError('"year" and "month" are ambiguous time units. Please convert them to days.') - else: - raise err - res.append({"value": data}) - - # time_pattern - case {"type": "time_pattern", **_r}: - raise NotImplementedError("Can't convert `time_pattern`. Please convert it to a `time_series`.") - + _append_arr(data, _formatter(value_type)) + case {"type": "array", "data": [float() | int(), *_] as data}: + _append_arr(data, _formatter("float")) + # date_time | duration + case { + "type": "date_time" | "duration" as data_t, + "data": str() | int() as data, + }: + _fmt = _formatter(data_t) + res.append({**idx_lvls, **_fmt("value", data)}) # values - case int() | float() | str() | bool(): - idx_lvls["value"] = json_doc - res.append(idx_lvls) + case int() | float() | str() | bool() as data: + _fmt = _formatter("noop") + res.append({**idx_lvls, **_fmt("value", data)}) case _: - raise NotImplementedError("Can't match this JSON structure yet.") + raise ValueError(f"match not found: {json_doc}") return res diff --git a/arrow_expts/utils.py b/arrow_expts/utils.py index db22730..8f1b69e 100644 --- a/arrow_expts/utils.py +++ b/arrow_expts/utils.py @@ -7,9 +7,9 @@ from typing import Any, Callable -def execute_on_sqlite(database_uri: str, - my_function: Callable, - *args: Any, **kwargs: Any) -> Any: +def execute_on_sqlite( + database_uri: str, my_function: Callable, *args: Any, **kwargs: Any +) -> Any: """Execute a function on an SQlite database. Returns: @@ -21,10 +21,13 @@ def execute_on_sqlite(database_uri: str, connection.commit() return result -def _insert_data_into_table(cursor: adbc_driver_manager.dbapi.Cursor, - table_name: str, - data: pa.lib.Table, - mode: str) -> int: + +def _insert_data_into_table( + cursor: adbc_driver_manager.dbapi.Cursor, + table_name: str, + data: pa.lib.Table, + mode: str, +) -> int: """Create a table and fill it with data or append data to an existing one. Args: @@ -44,11 +47,9 @@ def _insert_data_into_table(cursor: adbc_driver_manager.dbapi.Cursor, return result - -def write_data_to_db(database_uri: str, - table_name: str, - data: pa.lib.Table, - mode: str) -> int: +def write_data_to_db( + database_uri: str, table_name: str, data: pa.lib.Table, mode: str +) -> int: """Create a table and fill it with data or append data to an existing one. Args: @@ -60,16 +61,13 @@ def write_data_to_db(database_uri: str, Returns: Number of rows inserted """ - rows_inserted: int = execute_on_sqlite(database_uri, - _insert_data_into_table, - table_name, - data, - mode) + rows_inserted: int = execute_on_sqlite( + database_uri, _insert_data_into_table, table_name, data, mode + ) return rows_inserted -def database_equality(uri_database_1: str, - uri_database_2: str) -> bool: +def database_equality(uri_database_1: str, uri_database_2: str) -> bool: """Compare equality of two SQlite databases.""" connection: adbc_driver_manager.dbapi.Connection diff --git a/pyproject.toml b/pyproject.toml index e3dc54a..b360043 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,10 +14,16 @@ dependencies = [ "pyarrow", "pydantic", "pytest", + "python-dateutil", "rich", "spinedb-api", "sqlalchemy", ] +[project.optional-dependencies] +dev = [ + "mypy", + "pytest", +] [tool.hatch.metadata.source] include = ["arrow_expts"] diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..1781ccf --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,3 @@ +from pathlib import Path + +JSONDIR = Path("tests/json") diff --git a/arrow_expts/schema/json/array.durations.json b/tests/json/array.durations.json similarity index 100% rename from arrow_expts/schema/json/array.durations.json rename to tests/json/array.durations.json diff --git a/arrow_expts/schema/json/array.numbers.json b/tests/json/array.numbers.json similarity index 100% rename from arrow_expts/schema/json/array.numbers.json rename to tests/json/array.numbers.json diff --git a/arrow_expts/schema/json/array.strings.json b/tests/json/array.strings.json similarity index 75% rename from arrow_expts/schema/json/array.strings.json rename to tests/json/array.strings.json index fa93e6f..7990ce6 100644 --- a/arrow_expts/schema/json/array.strings.json +++ b/tests/json/array.strings.json @@ -1,5 +1,6 @@ { "type": "array", "data": ["one", "two"], + "value_type": "str", "index_name": "step" } diff --git a/arrow_expts/schema/json/date-time.json b/tests/json/date-time.json similarity index 100% rename from arrow_expts/schema/json/date-time.json rename to tests/json/date-time.json diff --git a/arrow_expts/schema/json/duration.compact.json b/tests/json/duration.compact.json similarity index 100% rename from arrow_expts/schema/json/duration.compact.json rename to tests/json/duration.compact.json diff --git a/arrow_expts/schema/json/duration.integer.json b/tests/json/duration.integer.json similarity index 100% rename from arrow_expts/schema/json/duration.integer.json rename to tests/json/duration.integer.json diff --git a/arrow_expts/schema/json/duration.verbose.json b/tests/json/duration.verbose.json similarity index 100% rename from arrow_expts/schema/json/duration.verbose.json rename to tests/json/duration.verbose.json diff --git a/arrow_expts/schema/json/map.dictionary.json b/tests/json/map.dictionary.json similarity index 100% rename from arrow_expts/schema/json/map.dictionary.json rename to tests/json/map.dictionary.json diff --git a/arrow_expts/schema/json/map.stochastic-time-series-antti.json b/tests/json/map.stochastic-time-series-antti.json similarity index 100% rename from arrow_expts/schema/json/map.stochastic-time-series-antti.json rename to tests/json/map.stochastic-time-series-antti.json diff --git a/arrow_expts/schema/json/map.stochastic-time-series.json b/tests/json/map.stochastic-time-series.json similarity index 100% rename from arrow_expts/schema/json/map.stochastic-time-series.json rename to tests/json/map.stochastic-time-series.json diff --git a/arrow_expts/schema/json/map.two-column-array.json b/tests/json/map.two-column-array.json similarity index 100% rename from arrow_expts/schema/json/map.two-column-array.json rename to tests/json/map.two-column-array.json diff --git a/arrow_expts/schema/json/time-pattern.json b/tests/json/time-pattern.json similarity index 100% rename from arrow_expts/schema/json/time-pattern.json rename to tests/json/time-pattern.json diff --git a/arrow_expts/schema/json/time-series.dict.json b/tests/json/time-series.dict.json similarity index 100% rename from arrow_expts/schema/json/time-series.dict.json rename to tests/json/time-series.dict.json diff --git a/arrow_expts/schema/json/time-series.one-column-array-custom-indices.json b/tests/json/time-series.one-column-array-custom-indices.json similarity index 100% rename from arrow_expts/schema/json/time-series.one-column-array-custom-indices.json rename to tests/json/time-series.one-column-array-custom-indices.json diff --git a/arrow_expts/schema/json/time-series.one-column-array.json b/tests/json/time-series.one-column-array.json similarity index 100% rename from arrow_expts/schema/json/time-series.one-column-array.json rename to tests/json/time-series.one-column-array.json diff --git a/arrow_expts/schema/json/time-series.two-column-array-named-indices.json b/tests/json/time-series.two-column-array-named-indices.json similarity index 100% rename from arrow_expts/schema/json/time-series.two-column-array-named-indices.json rename to tests/json/time-series.two-column-array-named-indices.json diff --git a/arrow_expts/schema/json/time-series.two-column-array.json b/tests/json/time-series.two-column-array.json similarity index 100% rename from arrow_expts/schema/json/time-series.two-column-array.json rename to tests/json/time-series.two-column-array.json diff --git a/tests/test_dbmap.py b/tests/test_dbmap.py new file mode 100644 index 0000000..b79202d --- /dev/null +++ b/tests/test_dbmap.py @@ -0,0 +1,136 @@ +import json +from types import NoneType +from typing import Callable + +import numpy as np +import pandas as pd +import pytest + +from arrow_expts.schema.models import TimePattern +from arrow_expts.spine.dbmap import ( + _atoi, + _formatter, + _low_res_datetime, + _normalise_freq, + _to_dateoffset, + make_records, +) + +from .conftest import JSONDIR + + +def test_filter_frequencies(): + assert _normalise_freq("5 years") == "5 Y" + assert _normalise_freq("5 year") == "5 Y" + assert _normalise_freq("3s") == "3s" + assert _normalise_freq("60") == "60min" + + +def test_low_res_datetime(): + target = pd.date_range(start="2001-01-01", freq="7D", periods=5) + low_res = _low_res_datetime(start="2001-01-01", freq="1W", periods=5) + + assert low_res.dtype == "datetime64[s]" + assert target.freq == low_res.freq + for date_target, date_low_res in zip(target, low_res): + assert date_target == date_low_res + + res = _low_res_datetime(start="0001-01-01", freq="1s", periods=1) + assert res[0] == pd.Timestamp("0001-01-01") + + +@pytest.mark.parametrize( + "val, expect", + [("foo", "foo"), ("t0012", 12), ("p2023", 2023), ("t0012.1", "t0012.1")], +) +def test_atoi(val: str, expect: int | str): + assert expect == _atoi(val) + + +@pytest.mark.parametrize("part", ["durations", "numbers", "strings"]) +def test_arrays(part: str): + data = json.loads((JSONDIR / f"array.{part}.json").read_text()) + fmt = _formatter(data.get("value_type", "float")) + index_name = data.get("index_name", "i") + expect = [fmt(index_name, i) for i in data["data"]] + assert expect == make_records(data, {}, []) + + +def fmt_durations(val: str | int): + return _to_dateoffset(_normalise_freq(val)) + + +@pytest.mark.parametrize( + "type_,value,fmt", + [ + ("date_time", "2019-06-01T22:15:00+01:00", pd.Timestamp), + ("duration", "1h", fmt_durations), + ("duration", 60, fmt_durations), + ("duration", "1 hour", fmt_durations), + ], +) +def test_values(type_: str, value, fmt: Callable): + res = make_records({"type": type_, "data": value}, {}, []) + assert res == [{"value": fmt(value)}] + + +def test_time_pattern(): + data = json.loads((JSONDIR / "time-pattern.json").read_text()) + index_name = data.get("index_name", "default0") + expect = [{index_name: TimePattern(k), "value": v} for k, v in data["data"].items()] + res = make_records(data, {}, []) + assert res == expect + + +@pytest.mark.parametrize( + "part, interval", + [ + ("dict", "30 min"), + ("one-column-array-custom-indices", "30 min"), + ("one-column-array", "1 h"), + ("two-column-array", "30 min"), + ("two-column-array-named-indices", "30 min"), + ], +) +def test_time_series(part: str, interval: str): + data = json.loads((JSONDIR / f"time-series.{part}.json").read_text()) + res = make_records(data, {}, []) + + # index column name matches + index_name = data.get("index_name", "default0") + assert index_name in res[0] + assert isinstance(res[0][index_name], pd.Timestamp) + assert res[1][index_name] - res[0][index_name] == pd.Timedelta(interval) + + +@pytest.mark.parametrize( + "part, lvls, types", + [ + ( + "dictionary", + ["default0", "default1", "value"], + [pd.Timestamp, pd.DateOffset, (float, int)], + ), + ( + "stochastic-time-series-antti", + ["stochastic_scenario", "analysis_time", "default2", "value"], + [str, pd.Timestamp, pd.Timestamp, (float, int)], + ), + ( + "stochastic-time-series", + ["Forecast time", "Target time", "Stochastic scenario", "value"], + [pd.Timestamp, pd.Timestamp, int, (float, int)], + ), + ("two-column-array", ["default0", "value"], [str, (float, int)]), + ], +) +def test_map(part: str, lvls: list[str], types: list[type]): + data = json.loads((JSONDIR / f"map.{part}.json").read_text()) + res = make_records(data, {}, []) + + ncols = len(lvls) + assert all(len(i) <= ncols for i in res) + # fallback to `NoneType` to account for missing values + assert all( + isinstance(i.get(k), (t, NoneType)) for i in res for k, t in zip(lvls, types) + ) diff --git a/tests/test_integration.py b/tests/test_integration.py new file mode 100644 index 0000000..323e3fe --- /dev/null +++ b/tests/test_integration.py @@ -0,0 +1,52 @@ +from dataclasses import astuple +import json + +import numpy as np +import pandas as pd +import pytest + +from arrow_expts.schema.models import Array +from arrow_expts.schema.reencode import series_to_col, to_df, to_tables + +from .conftest import JSONDIR + + +def np_eq(i, j): + if isinstance(cmp := (i == j), (np.ndarray, pd.Series, pd.DataFrame)): + return cmp.all() + else: + return cmp + + +@pytest.mark.parametrize("part", ["numbers"]) +def test_series_to_col(part): + input_json = JSONDIR / f"array.{part}.json" + data = json.loads(input_json.read_text()) + + exp = Array(name=data.get("index_name", "i"), values=data["data"]) + df = to_df(data) + + res = series_to_col(df[exp.name]) + for i, j in zip(astuple(res), astuple(exp)): + assert np_eq(i, j) + + tbl = to_tables(df) + assert len(tbl) == 1 + assert isinstance(tbl[0], Array) + + +@pytest.mark.parametrize("part", ["numbers"]) +def test_to_tables(part): + input_json = JSONDIR / f"array.{part}.json" + data = json.loads(input_json.read_text()) + + exp = Array(name=data.get("index_name", "i"), values=data["data"]) + df = to_df(data) + + res = series_to_col(df[exp.name]) + for i, j in zip(astuple(res), astuple(exp)): + assert np_eq(i, j) + + tbl = to_tables(df) + assert len(tbl) == 1 + assert isinstance(tbl[0], Array)