Support reading skims from Parquet files, in addition to OMX

activitysim/core/configuration/network.py

@@ -125,6 +125,20 @@ 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`` or ``.pq`` 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
+    using any stable zone-ID order) or sparse (only some origin-destination
+    combinations present, in any order). Parquet inputs are supported by both
+    the legacy skim-dictionary loaders and Sharrow when Sharrow 2.16 or newer
+    is installed.
     """
 
     zarr: str = None
@@ -219,10 +233,11 @@ class NetworkSettings(PydanticReadable, extra="forbid"):
     """Instructions for how to load and pre-process skim matrices.
 
     If given as a string or a list of strings, it is interpreted as the location
-    for OMX file(s), either as a single file or as a glob-matching pattern for
-    multiple files. The time period for the matrix must be represented at the end
-    of the matrix name and be seperated by a double_underscore (e.g. `BUS_IVT__AM`
-    indicates base skim BUS_IVT with a time period of AM.
+    for OMX or Parquet skim file(s), either as a single file or as a glob-matching
+    pattern for multiple files. Formats are detected from each file's extension
+    and may be mixed. The time period for the matrix must be represented at the
+    end of the matrix name and be separated by a double underscore (e.g.
+    `BUS_IVT__AM` indicates base skim BUS_IVT with a time period of AM).
 
     Alternatively, this can be given as a nested dictionary defined via the
     TAZ_Settings class, which allows for ZARR transformation and pre-processing.

activitysim/core/skim_dataset.py

@@ -18,6 +18,7 @@
 from activitysim.core import flow as __flow  # noqa: 401
 from activitysim.core import workflow
 from activitysim.core.input import read_input_file
+from activitysim.core.skim_parquet import SPARSE, ParquetSkimFile, is_parquet_file
 
 logger = logging.getLogger(__name__)
 
@@ -698,6 +699,235 @@ def load_sparse_maz_skims(
     return dataset
 
 
+def _matrix_time_periods(matrix_names, ignore):
+    """Collect the time-period pages physically present in a source group."""
+    if isinstance(ignore, str):
+        ignore = [ignore]
+
+    available_periods = {}
+    for matrix_name in matrix_names:
+        if ignore and any(re.match(pattern, matrix_name) for pattern in ignore):
+            continue
+        base_name, separator, period_name = matrix_name.partition("__")
+        if separator:
+            available_periods.setdefault(base_name, set()).add(period_name)
+    return available_periods
+
+
+def _mask_synthetic_time_periods(dataset, available_periods):
+    """Replace loader-created zero pages with missing values before merging."""
+
+    for base_name, periods in available_periods.items():
+        if base_name not in dataset or "time_period" not in dataset[base_name].dims:
+            continue
+        if not set(dataset.time_period.values).issubset(periods):
+            dataset[base_name] = dataset[base_name].where(
+                dataset.time_period.isin(list(periods))
+            )
+    return dataset
+
+
+def _restore_synthetic_time_periods(dataset, available_periods):
+    """Restore zero pages for periods absent from every physical source."""
+    for base_name, periods in available_periods.items():
+        if base_name not in dataset or "time_period" not in dataset[base_name].dims:
+            continue
+        if not set(dataset.time_period.values).issubset(periods):
+            dataset[base_name] = dataset[base_name].where(
+                dataset.time_period.isin(list(periods)), 0
+            )
+    return dataset
+
+
+def _zero_fill_sparse_parquet(dataset, parquet_sources, ignore):
+    """Match the legacy loader's zero fill for absent sparse OD pairs."""
+    if isinstance(ignore, str):
+        ignore = [ignore]
+
+    # Sharrow's duplicate-column behavior is last-file-wins, so use the same
+    # source when selecting the OD pairs that are physically present.
+    matrix_sources = {}
+    for _, parquet_file in parquet_sources:
+        for matrix_name in parquet_file.data_cols:
+            if ignore and any(re.match(pattern, matrix_name) for pattern in ignore):
+                continue
+            matrix_sources[matrix_name] = parquet_file
+
+    presence_by_file = {}
+    for matrix_name, parquet_file in matrix_sources.items():
+        if parquet_file.layout != SPARSE:
+            continue
+
+        presence = presence_by_file.get(parquet_file)
+        if presence is None:
+            presence = np.zeros(parquet_file.shape, dtype=bool)
+            presence[parquet_file._orig_idx, parquet_file._dest_idx] = True
+            presence = xr.DataArray(
+                presence,
+                dims=("otaz", "dtaz"),
+                coords={
+                    "otaz": parquet_file.zone_ids,
+                    "dtaz": parquet_file.zone_ids,
+                },
+            )
+            presence_by_file[parquet_file] = presence
+
+        base_name, separator, period_name = matrix_name.partition("__")
+        if separator:
+            if base_name not in dataset:
+                continue
+            # Preserve explicit NaNs at present OD pairs and other periods;
+            # only combinations absent from this physical page become zero.
+            keep_value = presence | (dataset.time_period != period_name)
+            dataset[base_name] = dataset[base_name].where(keep_value, 0)
+        elif matrix_name in dataset:
+            dataset[matrix_name] = dataset[matrix_name].where(presence, 0)
+    return dataset
+
+
+def _load_skim_dataset_from_sources(
+    skim_file_paths,
+    *,
+    time_periods,
+    max_float_precision,
+    ignore,
+    parquet_file_metadata=None,
+):
+    """
+    Load OMX and/or Parquet skim files into one Sharrow-compatible Dataset.
+
+    Parquet index columns are identified from the first two columns in each
+    file, consistent with the legacy skim reader. Files with different index
+    column names are loaded in separate groups and aligned by their zone labels.
+
+    Returns
+    -------
+    dataset : xarray.Dataset
+    omx_file_handles : list
+        Open OMX handles retained for the optimized shared-memory reload path.
+    """
+    omx_file_paths = [f for f in skim_file_paths if not is_parquet_file(f)]
+    parquet_file_paths = [f for f in skim_file_paths if is_parquet_file(f)]
+    omx_file_handles = []
+    datasets = []
+
+    try:
+        if omx_file_paths:
+            omx_file_handles = [
+                openmatrix.open_file(f, mode="r") for f in omx_file_paths
+            ]
+            omx_dataset = sh.dataset.from_omx_3d(
+                omx_file_handles,
+                index_names=("otaz", "dtaz", "time_period"),
+                time_periods=time_periods,
+                max_float_precision=max_float_precision,
+                ignore=ignore,
+            )
+            omx_matrix_names = [
+                matrix_name
+                for handle in omx_file_handles
+                for matrix_name in handle.listMatrices()
+            ]
+            datasets.append((omx_dataset, omx_matrix_names))
+
+        if parquet_file_paths:
+            if not hasattr(sh.dataset, "from_parquet_3d"):
+                raise ImportError(
+                    "Parquet skims with Sharrow require Sharrow 2.16 or newer"
+                )
+            metadata_by_path = {
+                os.fspath(path): metadata
+                for path, metadata in (parquet_file_metadata or {}).items()
+            }
+            parquet_groups = {}
+            for file_path in parquet_file_paths:
+                parquet_file = metadata_by_path.get(os.fspath(file_path))
+                if parquet_file is None:
+                    parquet_file = ParquetSkimFile(file_path)
+                # Load sparse files independently. Sharrow derives each sparse
+                # axis only from labels present on that axis, so grouping a
+                # sparse file that omits an entire origin or destination can
+                # otherwise discard valid rows from another file in the group.
+                sparse_source = file_path if parquet_file.layout == SPARSE else None
+                group_key = (
+                    parquet_file.orig_col,
+                    parquet_file.dest_col,
+                    sparse_source,
+                )
+                parquet_groups.setdefault(group_key, []).append(
+                    (file_path, parquet_file)
+                )
+
+            for (orig_col, dest_col, _), parquet_sources in parquet_groups.items():
+                file_paths = [source[0] for source in parquet_sources]
+                parquet_dataset = sh.dataset.from_parquet_3d(
+                    file_paths,
+                    index_names=(orig_col, dest_col, "time_period"),
+                    time_periods=time_periods,
+                    max_float_precision=max_float_precision,
+                    ignore=ignore,
+                )
+
+                # Rename through temporary names so even swapped source names
+                # (e.g. dtaz/otaz) cannot collide during the rename.
+                parquet_dataset = parquet_dataset.rename(
+                    {
+                        orig_col: "__activitysim_parquet_origin__",
+                        dest_col: "__activitysim_parquet_destination__",
+                    }
+                ).rename(
+                    {
+                        "__activitysim_parquet_origin__": "otaz",
+                        "__activitysim_parquet_destination__": "dtaz",
+                    }
+                )
+                # Sparse xarray expansion can omit an entire coordinate when no
+                # row uses it. Normalize both dimensions to the full zone set,
+                # which also gives dense nonascending inputs legacy-compatible
+                # canonical ordering.
+                parquet_zone_ids = parquet_sources[0][1].zone_ids
+                parquet_dataset = parquet_dataset.reindex(
+                    otaz=parquet_zone_ids, dtaz=parquet_zone_ids
+                )
+                parquet_dataset = _zero_fill_sparse_parquet(
+                    parquet_dataset, parquet_sources, ignore
+                )
+                parquet_matrix_names = [
+                    matrix_name
+                    for _, parquet_file in parquet_sources
+                    for matrix_name in parquet_file.data_cols
+                ]
+                datasets.append((parquet_dataset, parquet_matrix_names))
+
+        if not datasets:
+            raise ValueError("no OMX or Parquet skim files were provided")
+        if len(datasets) == 1:
+            dataset = datasets[0][0]
+        else:
+            all_available_periods = {}
+            masked_datasets = []
+            for source_dataset, matrix_names in datasets:
+                source_periods = _matrix_time_periods(matrix_names, ignore)
+                for base_name, periods in source_periods.items():
+                    all_available_periods.setdefault(base_name, set()).update(periods)
+                masked_datasets.append(
+                    _mask_synthetic_time_periods(source_dataset, source_periods)
+                )
+            dataset = xr.merge(masked_datasets, compat="no_conflicts", join="outer")
+            dataset = _restore_synthetic_time_periods(dataset, all_available_periods)
+
+        # SkimDataset expects this coordinate even when all source matrices are
+        # time-agnostic and therefore do not otherwise create the dimension.
+        if "time_period" not in dataset.coords:
+            dataset = dataset.assign_coords(time_period=time_periods)
+
+        return dataset, omx_file_handles
+    except Exception:
+        for handle in omx_file_handles:
+            handle.close()
+        raise
+
+
 def load_skim_dataset_to_shared_memory(state, skim_tag="taz") -> xr.Dataset:
     """
     Load skims from disk into shared memory.
@@ -718,11 +948,12 @@ def load_skim_dataset_to_shared_memory(state, skim_tag="taz") -> xr.Dataset:
     if network_los_preload is None:
         raise ValueError("missing network_los_preload")
 
-    # find which OMX files are to be used.
+    # Find the source skim files to use; formats may be mixed.
     omx_file_paths = state.filesystem.expand_input_file_list(
         network_los_preload.omx_file_names(skim_tag),
     )
     omx_file_handles = []
+    source_has_parquet = any(is_parquet_file(f) for f in omx_file_paths)
     zarr_file = network_los_preload.zarr_file_name(skim_tag)
 
     if state.settings.disable_zarr:
@@ -834,23 +1065,22 @@ def _should_ignore(ignore, x):
             d = sh.dataset.from_zarr_with_attr(zarr_file)
             zarr_write_time = d.attrs.get("ZARR_WRITE_TIME", 0)
             if zarr_write_time < latest_file_modification_time(omx_file_paths):
-                logger.warning("zarr skims older than omx, not using them")
+                logger.warning("zarr skims older than source skims, not using them")
                 do_not_save_zarr = True
                 d = None
             else:
                 d = d.max_float_precision(max_float_precision)
         if d is None:
             if zarr_file and not do_not_save_zarr:
-                logger.info("did not find zarr skims, loading omx")
-            omx_file_handles = [
-                openmatrix.open_file(f, mode="r") for f in omx_file_paths
-            ]
-            d = sh.dataset.from_omx_3d(
-                omx_file_handles,
-                index_names=("otaz", "dtaz", "time_period"),
+                logger.info("did not find zarr skims, loading source skim files")
+            d, omx_file_handles = _load_skim_dataset_from_sources(
+                omx_file_paths,
                 time_periods=time_periods,
                 max_float_precision=max_float_precision,
                 ignore=state.settings.omx_ignore_patterns,
+                parquet_file_metadata=network_los_preload.skims_info[
+                    skim_tag
+                ].parquet_files,
             )
 
             if zarr_file:
@@ -950,11 +1180,15 @@ def _should_ignore(ignore, x):
         logger.info(
             "store_skims_in_shm is False, keeping skims in process-local memory"
         )
+        for f in omx_file_handles:
+            f.close()
         return d
     else:
         logger.info("writing skims to shared memory")
-        if dask_required:
-            # setting `load` to True uses dask to load the data into memory
+        if dask_required or source_has_parquet:
+            # Parquet-backed datasets cannot use reload_from_omx_3d, so copy
+            # their already-loaded data into shared memory. The same path is
+            # required when coordinate realignment created a dask graph.
             d_shared_mem = d.shm.to_shared_memory(backing, mode="r", load=True)
         else:
             # setting `load` to false then calling `reload_from_omx_3d` avoids