feat: hive partitioning for id/var (TODO: make it optional?)

Author: martibosch

tstore/archive/io.py

@@ -3,6 +3,8 @@
 import glob
 import os
 import shutil
+from pathlib import Path
+from typing import Union
 
 
 def check_tstore_structure(tstore_structure):
@@ -31,12 +33,43 @@ def define_attributes_filepath(base_dir):
     return fpath
 
 
-def define_tsarray_filepath(base_dir, tstore_id, ts_variable, tstore_structure):
-    """Define filepath of a TStore TS."""
+def define_tsarray_filepath(
+    base_dir: Union[Path, str],
+    tstore_id: str,
+    ts_variable: str,
+    tstore_structure: str,
+    id_prefix: str,
+    var_prefix: str,
+) -> str:
+    """
+    Define filepath of a TStore TS.
+
+    Parameters
+    ----------
+    base_dir : path-like
+        Base directory of the TStore.
+    tstore_id : str
+        Value of the time series ID.
+    ts_variable : str
+        Name of the time series variable.
+    ts_structure : ["id-var", "var-id"]
+        TStore structure, either "id-var" or "var-id".
+    id_prefix : str
+        Prefix for the ID directory.
+    var_prefix : str
+        Prefix for the variable directory.
+
+    Returns
+    -------
+    fpath : str
+        Filepath for the time series.
+    """
+    id_dir_basename = f"{id_prefix}={tstore_id}"
+    var_dir_basename = f"{var_prefix}={ts_variable}"
     if tstore_structure == "id-var":
-        fpath = os.path.join(base_dir, tstore_id, ts_variable)
+        fpath = os.path.join(base_dir, id_dir_basename, var_dir_basename)
     elif tstore_structure == "var-id":
-        fpath = os.path.join(base_dir, ts_variable, tstore_id)
+        fpath = os.path.join(base_dir, var_dir_basename, id_dir_basename)
     else:
         raise ValueError("Valid tstore_structure are 'id-var' and 'var-id'.")
     return fpath
@@ -82,16 +115,48 @@ def get_partitions(base_dir, ts_variable):
     return partitions
 
 
-def get_ts_info(base_dir, ts_variable):
-    """Retrieve filepaths and tstore_ids for a specific ts_variable."""
+def get_ts_info(
+    base_dir: Union[Path, str],
+    ts_variable: str,
+    var_prefix: str,
+):
+    """
+    Retrieve filepaths and tstore_ids for a specific ts_variable.
+
+    Parameters
+    ----------
+    base_dir : path-like
+        Base directory of the TStore.
+    ts_variable : str
+        Name of the time series variable.
+    var_prefix : str
+        Prefix for the variable directory.
+
+    Returns
+    -------
+    fpaths : list of str
+        List of filepaths for the time series.
+    tstore_ids : list of str
+        List of time series IDs.
+    partitions : list of str
+        List of partitions.
+    """
     tstore_structure = get_tstore_structure(base_dir)
+
+    # TODO: DRY with `define_tsarray_filepath`?
+    var_dir_basename = f"{var_prefix}={ts_variable}"
+
     if tstore_structure == "id-var":
-        fpaths = glob.glob(os.path.join(base_dir, "*", ts_variable))
+        fpaths = glob.glob(os.path.join(base_dir, "*", var_dir_basename))
         tstore_ids = [os.path.basename(os.path.dirname(fpath)) for fpath in fpaths]
     elif tstore_structure == "var-id":
-        fpaths = glob.glob(os.path.join(base_dir, ts_variable, "*"))
+        fpaths = glob.glob(os.path.join(base_dir, var_dir_basename, "*"))
         tstore_ids = [os.path.basename(fpath) for fpath in fpaths]
     else:
         raise ValueError("Valid tstore_structure are 'id-var' and 'var-id'.")
+    # get only id values (remove prefix from hive prefix=value notation)
+    tstore_ids = [tstore_id.split("=")[1] for tstore_id in tstore_ids]
+
     partitions = get_partitions(base_dir, ts_variable)
+
     return fpaths, tstore_ids, partitions

tstore/tests/test_tslong.py

@@ -101,8 +101,10 @@ def test_store(
     assert dirpath.is_dir()
 
     # Check directory content
-    assert sorted(os.listdir(dirpath)) == ["1", "2", "3", "4", "_attributes.parquet", "tstore_metadata.yaml"]
-    assert os.listdir(dirpath / "1" / "ts_variable" / "year=2000" / "month=1") == ["part-0.parquet"]
+    assert sorted(os.listdir(dirpath)) == ["_attributes.parquet"] + [f"store_id={i}" for i in ["1", "2", "3", "4"]] + [
+        "tstore_metadata.yaml",
+    ]
+    assert os.listdir(dirpath / "store_id=1" / "variable=ts_variable" / "year=2000" / "month=1") == ["part-0.parquet"]
 
 
 class TestLoad:

tstore/tslong/pandas.py

@@ -31,6 +31,7 @@ def to_tstore(
         # TSTORE options
         partitioning=None,
         tstore_structure="id-var",
+        var_prefix="variable",
         overwrite=True,
     ):
         """Write the wrapped dataframe as a TStore structure."""
@@ -117,6 +118,8 @@ def to_tstore(
                     tstore_id=tstore_id,
                     ts_variable=ts_variable,
                     tstore_structure=tstore_structure,
+                    id_prefix=id_var,
+                    var_prefix=var_prefix,
                 )
 
                 # -----------------------------------------------

tstore/tslong/polars.py

@@ -31,6 +31,7 @@ def to_tstore(
         # TSTORE options
         partitioning=None,
         tstore_structure="id-var",
+        var_prefix="variable",
         overwrite=True,
     ) -> None:
         """Write the wrapped dataframe as a TStore structure."""
@@ -110,6 +111,8 @@ def to_tstore(
                     tstore_id=tstore_id,
                     ts_variable=ts_variable,
                     tstore_structure=tstore_structure,
+                    id_prefix=id_var,
+                    var_prefix=var_prefix,
                 )
 
                 # TODO; Maybe create TS object and use TS.to_parquet() once implemented

tstore/tslong/pyarrow.py

@@ -28,6 +28,7 @@ def to_tstore(self) -> None:
     def from_tstore(
         base_dir,
         ts_variables=None,
+        var_prefix="variable",
         start_time=None,
         end_time=None,
         tstore_ids=None,
@@ -50,6 +51,7 @@ def from_tstore(
             id_var=id_var,
             time_var=time_var,
             ts_variables=ts_variables,
+            var_prefix=var_prefix,
             start_time=start_time,
             end_time=end_time,
             columns=columns,
@@ -104,6 +106,7 @@ def _read_ts_variable(
     base_dir,
     id_var,
     ts_variable,
+    var_prefix,
     start_time=None,
     end_time=None,
     columns=None,
@@ -112,7 +115,8 @@ def _read_ts_variable(
 ):
     """Read a TStore ts_variable into pyarrow long-format."""
     # Find TS and associated TStore IDs
-    fpaths, tstore_ids, partitions = get_ts_info(base_dir=base_dir, ts_variable=ts_variable)
+    fpaths, tstore_ids, partitions = get_ts_info(base_dir=base_dir, ts_variable=ts_variable, var_prefix=var_prefix)
+
     # Read each TS
     list_tables = [
         _read_ts(
@@ -147,6 +151,7 @@ def _read_ts_variables(
     id_var,
     time_var,
     ts_variables,
+    var_prefix,
     start_time=None,
     end_time=None,
     columns=None,
@@ -160,6 +165,7 @@ def _read_ts_variables(
             base_dir=base_dir,
             id_var=id_var,
             ts_variable=ts_variable,
+            var_prefix=var_prefix,
             start_time=start_time,
             end_time=end_time,
             columns=columns,  # columns[ts_variable] in future