Add parquet skim reading support (dense row/col-major + sparse)

Author: Copilot

activitysim/core/configuration/network.py

@@ -125,6 +125,18 @@ class TAZ_Settings(PydanticBase):
 
     This is treated as a fallback for the raw input data, if ZARR format data
     is not available.
+
+    As an alternative to OMX, skim files can instead be provided in Parquet
+    format (using a ``.parquet`` file extension). The input format is
+    auto-detected from the file extension, so no other settings need to
+    change to use Parquet input. Parquet skim files should have an origin
+    column and a destination column (the first two columns in the file),
+    followed by one column for each named skim matrix (matching the naming
+    conventions used for OMX skims, including double-underscore delimited
+    time periods). Parquet skim data may be dense (one row for every
+    origin-destination combination, sorted in row-major or column-major
+    order) or sparse (only some origin-destination combinations present, in
+    any order).
     """
 
     zarr: str = None

activitysim/core/skim_dict_factory.py

@@ -15,6 +15,7 @@
 
 from activitysim.core import skim_dictionary, util
 from activitysim.core.exceptions import TableTypeError
+from activitysim.core.skim_parquet import ParquetSkimFile, is_parquet_file
 
 logger = logging.getLogger(__name__)
 
@@ -60,18 +61,24 @@ def __init__(self, state, skim_tag, network_los):
 
         skim_tag:           str             (e.g. 'TAZ')
         dtype_name:         str             (e.g. 'float32')
-        omx_manifest:       dict            dict mapping { omx_key: omx_file_name }
-        omx_shape:          2D tuple        shape of omx matrix: (<number_of_zones>, <number_of_zones>)
-        num_skims:          int             total number of individual skim matrices in omx files
+        omx_manifest:       dict            dict mapping { omx_key: skim_file_name }, whether the skim
+                                            file is an omx file or a parquet file
+        omx_shape:          2D tuple        shape of skim matrix: (<number_of_zones>, <number_of_zones>)
+        num_skims:          int             total number of individual skim matrices in omx/parquet files
         skim_data_shape:    3D tuple        (num_skims, omx_shape[0], omx_shape[1]) if ROW_MAJOR_LAYOUT
-        offset_map:         dict or None    1D ndarray as returned by omx_file.mapentries, if omx file has mappings
+        offset_map:         dict or None    1D ndarray as returned by omx_file.mapentries, if omx file has
+                                            mappings, or the (sorted) zone ids found in a parquet skim file
         offset_map_name:    str             name of offset_map in omx_filecorresponding to offset_map, if there was one
         omx_keys:       dict            dict mapping skim key (str or tuple) to skim key in omx file
                                             {DISTWALK: DISTWALK,
                                             ('DRV_COM_WLK_BOARDS', 'AM'): DRV_COM_WLK_BOARDS__AM, ...}
         base_keys:          list of str     e.g. 'BIKEDIST' or 'SOVTOLL_VTOLL' (base key of 3d skim)
         block_offsets:      dict            dict mapping skim key tuple to offset
 
+        Skim files can be in either OMX or Parquet format; the format is auto-detected
+        from each file's extension (``.omx`` vs ``.parquet``/``.pq``), and OMX and Parquet
+        files can be freely mixed within the list of files for a single skim_tag.
+
         Parameters
         ----------
         skim_tag
@@ -92,6 +99,10 @@ def __init__(self, state, skim_tag, network_los):
         self.base_keys = None
         self.block_offsets = None
 
+        # cache of ParquetSkimFile instances, keyed by file path, so files
+        # opened during load_skim_info are not re-parsed when reading data
+        self.parquet_files = {}
+
         if skim_tag:
             self.load_skim_info(state, skim_tag)
 
@@ -119,6 +130,44 @@ def load_skim_info(self, state, skim_tag):
         for omx_file_path in self.omx_file_paths:
             logger.debug(f"load_skim_info {skim_tag} reading {omx_file_path}")
 
+            if is_parquet_file(omx_file_path):
+                # Skim data provided as a parquet file (auto-detected by extension)
+                # instead of an omx file. The file is inspected here (and cached)
+                # to determine its zone list, shape, and dense/sparse layout.
+                parquet_skim_file = ParquetSkimFile(omx_file_path)
+                self.parquet_files[omx_file_path] = parquet_skim_file
+
+                # Check the shape of the skims, same as is done for omx files below.
+                if self.omx_shape is None:
+                    self.omx_shape = parquet_skim_file.shape
+                else:
+                    assert (
+                        self.omx_shape == parquet_skim_file.shape
+                    ), f"Mismatch shape {self.omx_shape} != {parquet_skim_file.shape}"
+
+                for skim_name in parquet_skim_file.data_cols:
+                    if skim_name in self.omx_manifest:
+                        warnings.warn(
+                            f"duplicate skim '{skim_name}' found in {self.omx_manifest[skim_name]} and {omx_file_path}"
+                        )
+                    self.omx_manifest[skim_name] = omx_file_path
+
+                # The origin/destination (zone id) values found in the parquet file
+                # serve the same purpose as an omx file's offset mapping.  Each parquet
+                # file is checked independently (it need not have zones in the same
+                # order as other files) but the set of zone ids found must match.
+                if self.offset_map is None:
+                    self.offset_map_name = f"{omx_file_path} zone ids"
+                    self.offset_map = parquet_skim_file.zone_ids
+                    assert len(self.offset_map) == self.omx_shape[0]
+                else:
+                    if not np.array_equal(self.offset_map, parquet_skim_file.zone_ids):
+                        raise RuntimeError(
+                            f"Mismatched zone ids in parquet skim file {omx_file_path}: "
+                            f"expected zone ids consistent with {self.offset_map_name}"
+                        )
+                continue
+
             with omx.open_file(omx_file_path, mode="r") as omx_file:
 
                 # Check the shape of the skims. All skim files loaded within this
@@ -322,7 +371,7 @@ def load_skim_info(self, state, skim_tag):
 
     def _read_skims_from_omx(self, skim_info, skim_data):
         """
-        read skims from omx file into skim_data
+        read skims from omx and/or parquet files into skim_data
         """
 
         skim_tag = skim_info.skim_tag
@@ -334,6 +383,34 @@ def _read_skims_from_omx(self, skim_info, skim_data):
 
             logger.info(f"_read_skims_from_omx {omx_file_path}")
 
+            if is_parquet_file(omx_file_path):
+                parquet_skim_file = skim_info.parquet_files.get(omx_file_path)
+                if parquet_skim_file is None:
+                    parquet_skim_file = ParquetSkimFile(omx_file_path)
+                for skim_key, omx_key in omx_keys.items():
+                    if omx_manifest[omx_key] == omx_file_path:
+                        offset = skim_info.block_offsets[skim_key]
+                        logger.debug(
+                            f"_read_skims_from_omx (parquet) file {omx_file_path} "
+                            f"omx_key {omx_key} skim_key {skim_key} to offset {offset}"
+                        )
+
+                        if skim_dictionary.ROW_MAJOR_LAYOUT:
+                            a = skim_data[offset, :, :]
+                        else:
+                            a = skim_data[:, :, offset]
+
+                        a[:] = parquet_skim_file.read_matrix(
+                            omx_key, dtype=skim_info.dtype_name
+                        )
+
+                        num_skims_loaded += 1
+
+                logger.info(
+                    f"_read_skims_from_omx loaded {num_skims_loaded} skims from {omx_file_path}"
+                )
+                continue
+
             # read skims into skim_data
             with omx.open_file(omx_file_path, mode="r") as omx_file:
                 for skim_key, omx_key in omx_keys.items():

activitysim/core/skim_parquet.py

@@ -0,0 +1,151 @@
+# ActivitySim
+# See full license in LICENSE.txt.
+from __future__ import annotations
+
+import logging
+import os
+
+import numpy as np
+import pyarrow.parquet as pq
+
+logger = logging.getLogger(__name__)
+
+PARQUET_SUFFIXES = (".parquet", ".pq")
+
+# layout tags
+ROW_MAJOR = "row_major"
+COL_MAJOR = "col_major"
+SPARSE = "sparse"
+
+
+def is_parquet_file(file_path):
+    """
+    Return True if file_path appears to be a parquet skim file, based on its extension.
+
+    Parameters
+    ----------
+    file_path : str or Path
+
+    Returns
+    -------
+    bool
+    """
+    return os.fspath(file_path).lower().endswith(PARQUET_SUFFIXES)
+
+
+class ParquetSkimFile:
+    """
+    Inspect and read a single skim matrix table stored in parquet format.
+
+    The parquet file is expected to have an origin column and a destination column
+    (the first two columns in the file, whatever their names), followed by one
+    column per named skim matrix (analogous to the matrices in an omx file). The
+    origin/destination values are used to determine the (square) shape of the
+    matrices, and whether the data is arranged 'densely' (i.e. every combination
+    of origin and destination is present exactly once) in row-major or column-major
+    order, or is instead 'sparse' (i.e. not every combination is present, or the
+    dense data is not sorted in row-major or column-major order).
+    """
+
+    def __init__(self, file_path):
+        self.file_path = file_path
+
+        parquet_file = pq.ParquetFile(file_path)
+        column_names = list(parquet_file.schema_arrow.names)
+        if len(column_names) < 3:
+            raise ValueError(
+                f"parquet skim file {file_path} must have at least 3 columns "
+                f"(origin, destination, and at least one data column), "
+                f"found {len(column_names)}: {column_names}"
+            )
+
+        self.orig_col = column_names[0]
+        self.dest_col = column_names[1]
+        self.data_cols = column_names[2:]
+
+        od_table = parquet_file.read(columns=[self.orig_col, self.dest_col])
+        origins = od_table.column(self.orig_col).to_numpy(zero_copy_only=False)
+        destinations = od_table.column(self.dest_col).to_numpy(zero_copy_only=False)
+
+        zone_ids = np.unique(np.concatenate([origins, destinations]))
+        self.zone_ids = zone_ids
+        self.n_zones = len(zone_ids)
+
+        self.shape = (self.n_zones, self.n_zones)
+
+        n_rows = len(origins)
+        self.is_dense = n_rows == self.n_zones * self.n_zones
+
+        zone_index = {z: i for i, z in enumerate(zone_ids)}
+        orig_idx = np.fromiter(
+            (zone_index[o] for o in origins), dtype=np.int64, count=n_rows
+        )
+        dest_idx = np.fromiter(
+            (zone_index[d] for d in destinations), dtype=np.int64, count=n_rows
+        )
+        self._orig_idx = orig_idx
+        self._dest_idx = dest_idx
+
+        if self.is_dense:
+            self.layout = self._detect_dense_layout(orig_idx, dest_idx)
+        else:
+            self.layout = SPARSE
+
+    def _detect_dense_layout(self, orig_idx, dest_idx):
+        """
+        Determine whether dense data is arranged in row-major or column-major
+        order.  Raises a ValueError if the data is dense but not sorted in
+        either of these orders.
+        """
+        n = self.n_zones
+
+        row_major_orig = np.repeat(np.arange(n), n)
+        row_major_dest = np.tile(np.arange(n), n)
+        if np.array_equal(orig_idx, row_major_orig) and np.array_equal(
+            dest_idx, row_major_dest
+        ):
+            return ROW_MAJOR
+
+        col_major_orig = np.tile(np.arange(n), n)
+        col_major_dest = np.repeat(np.arange(n), n)
+        if np.array_equal(orig_idx, col_major_orig) and np.array_equal(
+            dest_idx, col_major_dest
+        ):
+            return COL_MAJOR
+
+        raise ValueError(
+            f"parquet skim file {self.file_path} appears to contain dense data "
+            f"(one row for every origin-destination pair) but the rows are not "
+            f"sorted in row-major or column-major order. Dense parquet skim data "
+            f"must be sorted so it can be read efficiently; alternatively, omit "
+            f"rows to store the data in (unsorted or sorted) sparse format."
+        )
+
+    def read_matrix(self, column_name, dtype=None):
+        """
+        Read a single named skim matrix from the parquet file as a dense 2D array.
+
+        Parameters
+        ----------
+        column_name : str
+        dtype : dtype convertible, optional
+
+        Returns
+        -------
+        np.ndarray, shape (n_zones, n_zones)
+        """
+        table = pq.read_table(self.file_path, columns=[column_name])
+        values = table.column(column_name).to_numpy(zero_copy_only=False)
+        if dtype is not None:
+            values = values.astype(dtype, copy=False)
+
+        n = self.n_zones
+        if self.layout == ROW_MAJOR:
+            return np.ascontiguousarray(values.reshape(n, n))
+        elif self.layout == COL_MAJOR:
+            return np.ascontiguousarray(values.reshape(n, n, order="F"))
+        else:
+            # sparse layout (may or may not be sorted); scatter into dense matrix
+            matrix = np.zeros((n, n), dtype=values.dtype)
+            matrix[self._orig_idx, self._dest_idx] = values
+            return matrix

activitysim/core/test/los/configs_1z_parquet/network_los.yaml

@@ -0,0 +1,9 @@
+zone_system: 1
+
+taz_skims: z1_taz_skims.parquet
+
+skim_time_periods:
+    time_window: 1440
+    period_minutes: 60
+    periods: [0, 6, 11, 16, 20, 24]
+    labels: ['EA', 'AM', 'MD', 'PM', 'EV']

activitysim/core/test/los/configs_1z_parquet/settings.yaml

@@ -0,0 +1,2 @@
+
+multiprocess: False

activitysim/core/test/los/configs_1z_parquet_multi/network_los.yaml

@@ -0,0 +1,11 @@
+zone_system: 1
+
+taz_skims:
+  - z1_taz_skims_part1.parquet
+  - z1_taz_skims_part2.parquet
+
+skim_time_periods:
+    time_window: 1440
+    period_minutes: 60
+    periods: [0, 6, 11, 16, 20, 24]
+    labels: ['EA', 'AM', 'MD', 'PM', 'EV']

activitysim/core/test/los/configs_1z_parquet_multi/settings.yaml

@@ -0,0 +1,2 @@
+
+multiprocess: False