From 9f2ee60169387f2252685acdd84cf6ca0d8ecbc0 Mon Sep 17 00:00:00 2001 From: Thomas Roehr Date: Tue, 6 Oct 2026 13:07:18 +0200 Subject: [PATCH 1/2] core: generalize tracking of applying pipeline elements --- src/damast/core/dataprocessing.py | 122 +++++++++--- src/damast/core/tracking.py | 215 +++++++++++++++++++++ src/damast/integrations/mlflow_tracker.py | 13 +- src/damast/integrations/tracking.py | 120 +++++++++++- tests/damast/core/test_dataprocessing.py | 116 ++++++++++- tests/damast/core/test_tracking.py | 165 ++++++++++++++++ tests/damast/integrations/test_tracking.py | 118 ++++++++++- 7 files changed, 824 insertions(+), 45 deletions(-) create mode 100644 src/damast/core/tracking.py create mode 100644 tests/damast/core/test_tracking.py diff --git a/src/damast/core/dataprocessing.py b/src/damast/core/dataprocessing.py index 5988617..eb08ff3 100644 --- a/src/damast/core/dataprocessing.py +++ b/src/damast/core/dataprocessing.py @@ -9,7 +9,6 @@ import tempfile import traceback as tc from collections import OrderedDict -from datetime import datetime, timezone from logging import Logger, getLogger from pathlib import Path from typing import Any @@ -26,6 +25,7 @@ from .dataframe import AnnotatedDataFrame from .metadata import DataSpecification, MetaData from .pydantic_export import PydanticExporter +from .tracking import PipelineElementTracker, RowCountTracker, TimingTracker __all__ = [ "DataProcessingPipeline", @@ -38,6 +38,9 @@ """Suffix of :class:`DataProcessingPipeline` files created in :func:`DataProcessingPipeline.save` and used by :func:`DataProcessingPipeline.load`""" +DAMAST_STATS_SUFFIX: str = ".stats.yaml" +"""Suffix of the per-run report created by :func:`DataProcessingPipeline.save_stats`""" + VAEX_STATE_SUFFIX: str = ".vaex-state.json" """Suffix of :class:`DataProcessingPipeline` files created in :func:`DataProcessingPipeline.save_state` and used by :func:`DataProcessingPipeline.load_state`""" @@ -53,6 +56,12 @@ class DataProcessingPipeline(PipelineElement): :param inplace_transformation: If true, the input :class:`damast.core.dataframe.AnnotatedDataFrame` is not copied when calling :func:`transform`. Else input data-frame is untouched :param name_mappings: Name mappings that should be applied to individual transformations. + :param trackers: Observers of what each step does to the data, see + :class:`damast.core.tracking.PipelineElementTracker`. What they report ends up in + :func:`processing_stats`. Defaults to a :class:`damast.core.tracking.TimingTracker` and a + :class:`damast.core.tracking.RowCountTracker`; pass a list to add to or replace those, e.g. + a :class:`damast.core.tracking.KeyTracker` to also count distinct entities. An empty list + collects nothing. :raises ValueError: If any of the transformer names are `None` :raises AttributeError: If the transformer is missing the :func:`transform` function @@ -80,6 +89,11 @@ class DataProcessingPipeline(PipelineElement): _name_mappings: dict[str, dict[str, str]] _processing_stats: dict[str, dict[str, Any]] + #: Observers of what each step does to the data + _trackers: list[PipelineElementTracker] + #: uuids of the steps between their start and their end, to keep a tracker from seeing one twice + _steps_in_flight: set[Any] + _meta: dict[str, str] def __init__(self, *, @@ -90,6 +104,7 @@ def __init__(self, *, inplace_transformation: bool = False, name_mappings: dict[str, dict[str, str]] = { DAMAST_DEFAULT_DATASOURCE: {}}, meta: dict[str, str] | None = None, + trackers: list[PipelineElementTracker] | None = None, ): super().__init__() @@ -103,6 +118,9 @@ def __init__(self, *, self._name_mappings = name_mappings self._processing_stats = {} + self._trackers = list(trackers) if trackers is not None else [TimingTracker(), RowCountTracker()] + self._steps_in_flight = set() + self.processing_graph = ProcessingGraph() if processing_graph: if isinstance(processing_graph, ProcessingGraph): @@ -159,15 +177,30 @@ def base_dir(self, value): @property def processing_stats(self) -> dict[str, dict[str, Any]]: """ - Per-step timing and row-count statistics collected by the last :func:`transform` run. + What the :attr:`trackers` reported about each step of the last :func:`transform` run. - Keyed by step name, e.g. ``{"my_step": {"processing_time_in_s": 0.1, - "input_dataframe_length": {...}, "output_dataframe_length": 123, ...}}``. Populated - incrementally by :func:`on_transform_start`/:func:`on_transform_end` as steps run, so - it also reflects partial progress if a pipeline fails mid-run. + Keyed by step name, holding the merged result of every tracker, e.g. + ``{"my_step": {"processing_time_in_s": 0.1, "input_dataframe_length": {...}, + "output_dataframe_length": 123, "rows_removed": 7, ...}}``. Filled in as steps run, so it + also reflects partial progress if a pipeline fails mid-run. """ return self._processing_stats + @property + def trackers(self) -> list[PipelineElementTracker]: + """ + The observers of what each step does to the data, see + :class:`damast.core.tracking.PipelineElementTracker`. + + Settable, so that a pipeline from :func:`load` - which is built from the saved file alone, + and whose trackers are therefore the defaults - can be given different ones. + """ + return self._trackers + + @trackers.setter + def trackers(self, value: list[PipelineElementTracker]): + self._trackers = list(value) + @property def output_specs(self): if not self.is_ready: @@ -512,6 +545,24 @@ def save(self, dir: str | Path) -> Path: yaml.dump(dict(self), f) return filename + def save_stats(self, dir: str | Path | None = None) -> Path: + """ + Save the statistics of the last :func:`transform` run - see :func:`processing_stats`. + + Written automatically at the end of :func:`transform`, so that a run leaves an audit trail + of what each step did even without an experiment tracker. + + :param dir: Directory to write to, :attr:`base_dir` if not given + :return: The path of the written report + """ + base_dir = Path(dir) if dir is not None else self.base_dir + base_dir.mkdir(parents=True, exist_ok=True) + filename = base_dir / f"{self.name}{DAMAST_STATS_SUFFIX}" + + with open(filename, "w") as f: + yaml.dump({"name": self.name, "steps": self._processing_stats}, f) + return filename + def __iter__(self): yield "name", self.name yield "description", self.description @@ -675,6 +726,10 @@ def transform(self, :param df: The input dataframe :returns: The transformed dataframe """ + # the statistics describe this run, so anything a previous one left behind has to go + self._processing_stats = {} + self._steps_in_flight = set() + dataframes = { DAMAST_DEFAULT_DATASOURCE: df } for x in self.processing_graph.get_joins(): if x.name not in kwargs: @@ -707,6 +762,8 @@ def transform(self, adf = pipeline._run(in_dataframes, verbose=verbose) assert isinstance(adf, AnnotatedDataFrame) adf.validate_metadata(validation_mode=damast.core.ValidationMode.UPDATE_METADATA) + + self.save_stats() return adf def _run(self, @@ -760,47 +817,49 @@ def on_transform_start(self, step: PipelineElement, dataframes: dict[str, AnnotatedDataFrame]): """ - Default implementation of the on_transform_start callback. + Let every tracker see a step's input before it runs - see :attr:`trackers`. - This output is a message in INFO log level + This also outputs a message in INFO log level. """ - if hasattr(step, "transform_start"): + # a step is announced once per datasource, so without this a tracker would see a join twice + if step.uuid in self._steps_in_flight: return + self._steps_in_flight.add(step.uuid) logger.info(f"[transform] start: {step.__class__.__name__} - {step.name_mappings}") - start_time = datetime.now(timezone.utc) - step.transform_start = start_time step_name = self.processing_graph[step.uuid].name - self._processing_stats[step_name] = { - "input_dataframe_length": {x: y.shape[0] for x,y in dataframes.items()}, - "start_time": start_time - } + self._processing_stats[step_name] = {} + for tracker in self._trackers: + tracker.on_step_start(step, dataframes) def on_transform_end(self, step: PipelineElement, adf: AnnotatedDataFrame): """ - Default implementation of the on_transform_end callback. + Collect what every tracker has to report about a step that just ran - see :attr:`trackers`. - This outputs a message in INFO log level + The trackers' results are merged into :func:`processing_stats` under the step's name. + This also outputs a message in INFO log level. """ - if not hasattr(step, "transform_start"): + if step.uuid not in self._steps_in_flight: return - - start = step.transform_start - end_time = datetime.now(timezone.utc) - delta = (end_time - start).total_seconds() - delattr(step, "transform_start") - logger.info(f"[transform] end: {step.__class__.__name__} - {step.name_mappings}: " - f"{delta} seconds, {adf.shape[0]} remaining rows)") + self._steps_in_flight.discard(step.uuid) step_name = self.processing_graph[step.uuid].name - self._processing_stats[step_name].update({ - "processing_time_in_s": delta, - "output_dataframe_length": adf.shape[0], - "end_time": end_time - }) + statistics = self._processing_stats.setdefault(step_name, {}) + for tracker in self._trackers: + statistics.update(tracker.on_step_end(step, {DAMAST_DEFAULT_DATASOURCE: adf})) + + # Report whatever the trackers happened to measure, so that a custom set of them degrades + # the message instead of breaking it + details = [] + if "processing_time_in_s" in statistics: + details.append(f"{statistics['processing_time_in_s']} seconds") + if "output_dataframe_length" in statistics: + details.append(f"{statistics['output_dataframe_length']} remaining rows") + logger.info(f"[transform] end: {step.__class__.__name__} - {step.name_mappings}" + f"{': ' + ', '.join(details) if details else ''}") def __repr__(self) -> str: """ @@ -820,6 +879,9 @@ def __deepcopy__(self, memo) -> DataProcessingPipeline: pipeline._processing_stats = copy.deepcopy(self._processing_stats) pipeline.processing_graph = copy.deepcopy(self.processing_graph) pipeline._output_specs = [copy.deepcopy(x) for x in self._output_specs] + # '_trackers' stays shared on purpose: 'prepare' copies the pipeline, but the transformers + # keep calling the original's hooks, so the caller's own tracker instances are the ones + # that run and hold state return pipeline def describe(self) -> str: diff --git a/src/damast/core/tracking.py b/src/damast/core/tracking.py new file mode 100644 index 0000000..cd3449c --- /dev/null +++ b/src/damast/core/tracking.py @@ -0,0 +1,215 @@ +""" +Trackers that observe what each step of a :class:`damast.core.dataprocessing.DataProcessingPipeline` +does to the data. + +A pipeline is given a list of :class:`PipelineElementTracker` instances and calls each one before +and after every step. What a tracker returns at the end of a step is merged into the pipeline's +:func:`damast.core.dataprocessing.DataProcessingPipeline.processing_stats`, so a new kind of +measurement is a new class here rather than a change to the pipeline. + +Not to be confused with :class:`damast.integrations.tracking.ExperimentTracker`, which reports a +finished run to an experiment tracker such as MLflow. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from datetime import datetime, timezone +from typing import Any + +from damast.core.dataframe import AnnotatedDataFrame +from damast.core.transformations import PipelineElement + +__all__ = [ + "KeyTracker", + "PipelineElementTracker", + "RowCountTracker", + "TimingTracker", +] + + +class PipelineElementTracker(ABC): + """ + Observes a single :class:`damast.core.transformations.PipelineElement` as a pipeline runs it. + + A tracker is called twice per step - :func:`on_step_start` with the step's inputs, and + :func:`on_step_end` with its output - and returns the statistics for that step from + :func:`on_step_end`. The pipeline merges the returned dictionaries of all its trackers, so two + trackers must not report under the same key. + + A tracker that needs to compare the two ends correlates them by ``step.uuid`` and keeps its own + state; it is never told the step's name, since the pipeline keys the merged result by that. + """ + + def on_step_start(self, step: PipelineElement, dataframes: dict[str, AnnotatedDataFrame]) -> None: + """ + Called before a step runs. Does nothing by default, so a stateless tracker can ignore it. + + :param step: The step that is about to run + :param dataframes: The step's input, keyed by datasource - several entries for a join + """ + + @abstractmethod + def on_step_end(self, step: PipelineElement, dataframes: dict[str, AnnotatedDataFrame]) -> dict[str, Any]: + """ + Called after a step ran. + + :param step: The step that just ran + :param dataframes: The step's output, keyed by datasource + :return: The statistics for this step, merged into the pipeline's processing stats + """ + + +class TimingTracker(PipelineElementTracker): + """ + How long each step took. + + Reports ``start_time``, ``end_time`` and ``processing_time_in_s``. + """ + + def __init__(self): + #: step uuid -> when it started, held only while that step runs + self._started: dict[Any, datetime] = {} + + def on_step_start(self, step: PipelineElement, dataframes: dict[str, AnnotatedDataFrame]) -> None: + self._started[step.uuid] = datetime.now(timezone.utc) + + def on_step_end(self, step: PipelineElement, dataframes: dict[str, AnnotatedDataFrame]) -> dict[str, Any]: + end_time = datetime.now(timezone.utc) + start_time = self._started.pop(step.uuid, None) + if start_time is None: + return {"end_time": end_time} + + return { + "start_time": start_time, + "end_time": end_time, + "processing_time_in_s": (end_time - start_time).total_seconds(), + } + + +class RowCountTracker(PipelineElementTracker): + """ + How many rows each step consumed and produced. + + Reports ``input_dataframe_length`` (per datasource), ``output_dataframe_length`` and + ``rows_removed`` - the latter against the total of all inputs, so for a join it states what the + step reduced rather than comparing against one side. + + The row count materialises the frame, which the pipeline does anyway in order to run the next + step. + """ + + def __init__(self): + #: step uuid -> total input rows, held only while that step runs + self._input_lengths: dict[Any, int] = {} + + def on_step_start(self, step: PipelineElement, dataframes: dict[str, AnnotatedDataFrame]) -> None: + lengths = {label: adf.shape[0] for label, adf in dataframes.items()} + self._input_lengths[step.uuid] = lengths + + def on_step_end(self, step: PipelineElement, dataframes: dict[str, AnnotatedDataFrame]) -> dict[str, Any]: + output_length = sum(adf.shape[0] for adf in dataframes.values()) + lengths = self._input_lengths.pop(step.uuid, None) + if lengths is None: + return {"output_dataframe_length": output_length} + + return { + "input_dataframe_length": lengths, + "output_dataframe_length": output_length, + "rows_removed": sum(lengths.values()) - output_length, + } + + +class KeyTracker(PipelineElementTracker): + """ + How many distinct entities each step consumed and produced. + + For every tracked column this reports, under ``keys``, the distinct counts ``unique_in``, + ``unique_out`` and ``unique_removed``, plus ``removed_sample`` - at most + :attr:`max_removed_keys` of the removed values, with ``removed_truncated`` saying whether the + sample is partial. The removed set itself is unbounded, hence the sample. + + Example: + + .. highlight:: python + .. code-block:: python + + pipeline = DataProcessingPipeline( + name="prepare", base_dir=out, + trackers=[TimingTracker(), RowCountTracker(), KeyTracker(["mmsi"])]) + + :param keys: Columns identifying an entity, e.g. ``["mmsi"]`` + :param max_removed_keys: How many removed values to keep per step and column; 0 keeps the + counts but no values + :raise ValueError: If ``max_removed_keys`` is negative + + .. note:: + The distinct values of each tracked column are held in memory while a step runs, so a + very high-cardinality key is not free. The aggregation itself runs over a frame the + pipeline has already materialised to count its rows. + """ + + def __init__(self, keys: list[str], max_removed_keys: int = 10): + if max_removed_keys < 0: + raise ValueError(f"{self.__class__.__name__}.__init__: max_removed_keys must not be" + f" negative, got {max_removed_keys}") + + self.keys = list(keys) + self.max_removed_keys = max_removed_keys + #: step uuid -> column -> distinct input values, held only while that step runs + self._values_in: dict[Any, dict[str, set[Any]]] = {} + + @staticmethod + def _distinct_values(adf: AnnotatedDataFrame, column: str) -> set[Any]: + """ + The distinct values of a column of an already collected dataframe. + + :param adf: The dataframe holding the column + :param column: Name of the column + :return: The distinct values, with null left out - it identifies no entity + """ + return set(adf.lazyframe.compat.collected().get_column(column).drop_nulls().unique().to_list()) + + def _collect(self, dataframes: dict[str, AnnotatedDataFrame]) -> dict[str, set[Any]]: + """ + The distinct values of every tracked column present in the given dataframes. + + :param dataframes: The dataframes to scan, keyed by datasource + :return: Column name to its distinct values, union-ed over the dataframes that carry it. + A column no dataframe carries is left out. + """ + values = {} + for key in self.keys: + for adf in dataframes.values(): + if key in adf.column_names: + values.setdefault(key, set()).update(self._distinct_values(adf, key)) + return values + + def on_step_start(self, step: PipelineElement, dataframes: dict[str, AnnotatedDataFrame]) -> None: + # Materialise now: with 'inplace_transformation' the input frame is the same object the + # step mutates, so a lazy 'before' would resolve to the state afterwards. + self._values_in[step.uuid] = self._collect(dataframes) + + def on_step_end(self, step: PipelineElement, dataframes: dict[str, AnnotatedDataFrame]) -> dict[str, Any]: + values_in = self._values_in.pop(step.uuid, {}) + if not values_in: + return {} + + values_out = self._collect(dataframes) + statistics = {} + for key, before in values_in.items(): + if key not in values_out: + # the step is allowed to drop the column - that is not 'every entity was removed' + continue + + removed = before - values_out[key] + sample = sorted(removed)[:self.max_removed_keys] + statistics[key] = { + "unique_in": len(before), + "unique_out": len(values_out[key]), + "unique_removed": len(removed), + "removed_sample": sample, + "removed_truncated": len(sample) < len(removed), + } + + return {"keys": statistics} if statistics else {} diff --git a/src/damast/integrations/mlflow_tracker.py b/src/damast/integrations/mlflow_tracker.py index bb8931e..e0cf6db 100644 --- a/src/damast/integrations/mlflow_tracker.py +++ b/src/damast/integrations/mlflow_tracker.py @@ -23,7 +23,7 @@ from damast.utils import ensure_packages -from .tracking import ExperimentTracker, flatten_step_stats +from .tracking import ExperimentTracker, flatten_step_params, flatten_step_stats if TYPE_CHECKING: from damast.core.dataprocessing import DataProcessingPipeline @@ -71,10 +71,10 @@ def track_pipeline( Report a `DataProcessingPipeline` run to MLflow. Starts an MLflow run, yields a `MLflowTracker` to log the result against (call - `tracker.log_result(adf)` on the transformed dataframe), and on exit always logs the - pipeline's per-step timing/row-count stats (`DataProcessingPipeline.processing_stats`, - partial if the pipeline raised) and ends the run - `FAILED` if the `with` block raised, - `FINISHED` otherwise. + `tracker.log_result(adf)` on the transformed dataframe), and on exit always logs whatever + the pipeline's trackers reported about each step (`DataProcessingPipeline.processing_stats`, + partial if the pipeline raised) - the measurements as metrics, the rest as parameters - and + ends the run: `FAILED` if the `with` block raised, `FINISHED` otherwise. Example: @@ -116,4 +116,7 @@ def track_pipeline( step_metrics = flatten_step_stats(pipeline.processing_stats) if step_metrics: tracker.log_metrics(step_metrics) + step_params = flatten_step_params(pipeline.processing_stats) + if step_params: + tracker.log_params(step_params) tracker.end_run(status=status) diff --git a/src/damast/integrations/tracking.py b/src/damast/integrations/tracking.py index 995b954..77f70cf 100644 --- a/src/damast/integrations/tracking.py +++ b/src/damast/integrations/tracking.py @@ -7,8 +7,10 @@ from __future__ import annotations import json +import re import tempfile from abc import ABC, abstractmethod +from collections.abc import Iterator, Mapping from pathlib import Path from typing import TYPE_CHECKING, Any @@ -16,7 +18,71 @@ from damast.core.dataframe import AnnotatedDataFrame from damast.core.metadata import MetaData -__all__ = ["ExperimentTracker", "flatten_metadata", "flatten_step_stats"] +__all__ = [ + "ExperimentTracker", + "flatten_metadata", + "flatten_step_params", + "flatten_step_stats", +] + +#: Characters MLflow rejects in a metric or parameter name - it permits alphanumerics, +#: underscore, dash, period, space and slash +_UNSAFE_NAME = re.compile(r"[^0-9a-zA-Z_\-./ ]") + + +def _safe_name(path: str) -> str: + """ + Make a dotted path usable as a metric or parameter name. + + A tracker names its own statistics, and a name can reach here from data - a tracked column, + for instance, which damast itself may label `latitude (deg)`. Sanitising here keeps such a + name from failing the run when it is logged. + + Args: + path: The dotted path to sanitise. + + Returns: + The path with every character MLflow rejects replaced by an underscore. + """ + return _UNSAFE_NAME.sub("_", path) + + +def _leaves(value: Any, prefix: str = "") -> Iterator[tuple[str, Any]]: + """ + Every non-mapping leaf of a nested mapping, with its dotted path. + + A list is a leaf rather than something to descend into, so that e.g. a sample of removed + values stays a single entry. + + Args: + value: The value to walk. + prefix: The path accumulated so far. + + Yields: + A `(dotted path, leaf value)` pair per leaf. + """ + if isinstance(value, Mapping): + for key, child in value.items(): + yield from _leaves(child, f"{prefix}.{key}" if prefix else str(key)) + else: + yield prefix, value + + +def _is_metric(value: Any) -> bool: + """ + Whether a leaf is a measurement rather than something to record as a parameter. + + `bool` is excluded although `isinstance(True, int)` holds in Python: a flag such as "the + sample was truncated" is not a quantity to plot. :func:`flatten_step_params` takes exactly + the leaves this rejects, so nothing a tracker reports is dropped. + + Args: + value: The leaf value to classify. + + Returns: + `True` if the value belongs in metrics. + """ + return isinstance(value, (int, float)) and not isinstance(value, bool) def flatten_metadata( @@ -71,7 +137,20 @@ def flatten_metadata( def flatten_step_stats(processing_stats: dict[str, dict[str, Any]]) -> dict[str, float]: """ - Flatten `DataProcessingPipeline.processing_stats` into tracker-agnostic metrics. + Flatten the numeric part of `DataProcessingPipeline.processing_stats` into metrics. + + Every numeric leaf becomes `step..`, whatever + `damast.core.tracking.PipelineElementTracker` produced it - so a new tracker's measurements + are reported without a change here. Non-numeric leaves and flags go to + :func:`flatten_step_params`. + + Example: + + ```python + metrics = flatten_step_stats(pipeline.processing_stats) + # {"step.valid_mmsi.rows_removed": 4.0, + # "step.valid_mmsi.keys.mmsi.unique_removed": 4.0, ...} + ``` Args: processing_stats: Per-step stats, as returned by @@ -82,16 +161,41 @@ def flatten_step_stats(processing_stats: dict[str, dict[str, Any]]) -> dict[str, """ metrics: dict[str, float] = {} for step_name, stats in processing_stats.items(): - for key in ("processing_time_in_s", "output_dataframe_length"): - if key in stats: - metrics[f"step.{step_name}.{key}"] = float(stats[key]) - - for source, length in stats.get("input_dataframe_length", {}).items(): - metrics[f"step.{step_name}.input_dataframe_length.{source}"] = float(length) + for path, value in _leaves(stats): + if _is_metric(value): + metrics[_safe_name(f"step.{step_name}.{path}")] = float(value) return metrics +def flatten_step_params(processing_stats: dict[str, dict[str, Any]]) -> dict[str, Any]: + """ + Flatten the non-numeric part of `DataProcessingPipeline.processing_stats` into parameters. + + The counterpart of :func:`flatten_step_stats`: it takes exactly the leaves that are not + measurements - timestamps, flags such as whether a sample was truncated, and the sampled + values themselves - so that nothing a tracker reports is lost. Containers are JSON-encoded, + as in :func:`flatten_metadata`. + + Args: + processing_stats: Per-step stats, as returned by + `DataProcessingPipeline.processing_stats`. + + Returns: + A flat mapping of parameter name to value. + """ + params: dict[str, Any] = {} + for step_name, stats in processing_stats.items(): + for path, value in _leaves(stats): + if _is_metric(value): + continue + if isinstance(value, (dict, list)): + value = json.dumps(value, default=str) + params[_safe_name(f"step.{step_name}.{path}")] = value + + return params + + class ExperimentTracker(ABC): """ Minimal interface a pipeline run can be reported to. diff --git a/tests/damast/core/test_dataprocessing.py b/tests/damast/core/test_dataprocessing.py index 810898a..16f08ff 100644 --- a/tests/damast/core/test_dataprocessing.py +++ b/tests/damast/core/test_dataprocessing.py @@ -10,6 +10,7 @@ from damast.core.data_description import CyclicMinMax, MinMax from damast.core.dataframe import AnnotatedDataFrame from damast.core.dataprocessing import DataProcessingPipeline, PipelineElement +from damast.core.tracking import KeyTracker, RowCountTracker from damast.core.decorators import ( DAMAST_DEFAULT_DATASOURCE, DECORATED_INPUT_SPECS, @@ -1124,4 +1125,117 @@ def test_exclusive_output_join(tmp_path): assert {c.name for c in pipeline.output_metadata().columns} == {"key"} adf = pipeline.transform(df=_annotated(key=[1], x=[2]), other=_annotated(key=[3], y=[4])) - assert adf.column_names == ["key"] \ No newline at end of file + assert adf.column_names == ["key"] + + +# --- per-step effect statistics ------------------------------------------------------------- + +def _keyed_adf(mmsis: list[int]) -> AnnotatedDataFrame: + return AnnotatedDataFrame( + polars.LazyFrame({"mmsi": mmsis}), + metadata=MetaData([DataSpecification(name="mmsi")]), + validation_mode=damast.core.ValidationMode.IGNORE, + ) + + +def _keep_valid_mmsi(tmp_path, **kwargs) -> DataProcessingPipeline: + from damast.data_handling.transformers.filters import Filter + + return DataProcessingPipeline(name="prepare", base_dir=tmp_path, **kwargs).add( + "valid_mmsi", Filter(operator=">=", value=999999), name_mappings={"x": "mmsi"}) + + +def test_default_trackers_report_timing_and_rows(tmp_path): + """The out-of-the-box tracker set covers what a pipeline reported before trackers existed.""" + pipeline = _keep_valid_mmsi(tmp_path) + pipeline.transform(df=_keyed_adf([0, 1, 2, 3, 999999, 1000000, 1000000, 227991212])) + + stats = pipeline.processing_stats["valid_mmsi"] + assert set(stats) == {"start_time", "end_time", "processing_time_in_s", + "input_dataframe_length", "output_dataframe_length", "rows_removed"} + assert stats["input_dataframe_length"] == {"df": 8} + assert stats["output_dataframe_length"] == 4 + assert stats["rows_removed"] == 4 + + +def test_trackers_can_be_replaced(tmp_path): + """A given list replaces the defaults rather than adding to them.""" + pipeline = _keep_valid_mmsi(tmp_path, trackers=[RowCountTracker()]) + pipeline.transform(df=_keyed_adf([0, 999999])) + + stats = pipeline.processing_stats["valid_mmsi"] + assert "rows_removed" in stats + assert "processing_time_in_s" not in stats + + +def test_without_trackers_nothing_is_collected(tmp_path): + pipeline = _keep_valid_mmsi(tmp_path, trackers=[]) + pipeline.transform(df=_keyed_adf([0, 999999])) + + assert pipeline.processing_stats["valid_mmsi"] == {} + + +def test_key_tracker_reports_removed_entities_through_the_pipeline(tmp_path): + pipeline = _keep_valid_mmsi( + tmp_path, trackers=[RowCountTracker(), KeyTracker(["mmsi"], max_removed_keys=3)]) + pipeline.transform(df=_keyed_adf([0, 1, 2, 3, 999999, 1000000, 1000000, 227991212])) + + keys = pipeline.processing_stats["valid_mmsi"]["keys"]["mmsi"] + assert keys["unique_in"] == 7 + assert keys["unique_out"] == 3 + assert keys["unique_removed"] == 4 + assert keys["removed_sample"] == [0, 1, 2] + assert keys["removed_truncated"] is True + + +def test_the_callers_own_tracker_instance_is_the_one_that_runs(tmp_path): + """ + 'prepare' deep-copies the pipeline, so a copied tracker would silently collect instead - + leaving the instance the caller holds empty. + """ + tracker = KeyTracker(["mmsi"]) + pipeline = _keep_valid_mmsi(tmp_path, trackers=[tracker]) + pipeline.transform(df=_keyed_adf([0, 999999])) + + assert pipeline.trackers[0] is tracker + + +def test_trackers_can_be_set_on_a_loaded_pipeline(tmp_path): + """A pipeline from load() is built from the file alone, so its trackers are the defaults.""" + _keep_valid_mmsi(tmp_path).save(tmp_path) + + pipeline = DataProcessingPipeline.load(tmp_path / "prepare.damast.ppl") + pipeline.base_dir = tmp_path + pipeline.trackers = [KeyTracker(["mmsi"])] + pipeline.transform(df=_keyed_adf([0, 1, 999999])) + + assert pipeline.processing_stats["valid_mmsi"]["keys"]["mmsi"]["unique_removed"] == 2 + + +def test_processing_stats_are_reset_per_run(tmp_path): + """The stats describe the last run, as the property documents.""" + pipeline = _keep_valid_mmsi(tmp_path, trackers=[RowCountTracker(), KeyTracker(["mmsi"])]) + + pipeline.transform(df=_keyed_adf([0, 1, 999999])) + assert pipeline.processing_stats["valid_mmsi"]["rows_removed"] == 2 + + pipeline.transform(df=_keyed_adf([999999, 1000000])) + assert pipeline.processing_stats["valid_mmsi"]["rows_removed"] == 0 + # reported even when nothing was removed, so the counts are a series across runs + assert pipeline.processing_stats["valid_mmsi"]["keys"]["mmsi"]["unique_removed"] == 0 + + +def test_transform_writes_a_statistics_report(tmp_path): + """A run leaves an audit trail on disk, without needing an experiment tracker.""" + import yaml + + pipeline = _keep_valid_mmsi(tmp_path, trackers=[RowCountTracker(), KeyTracker(["mmsi"])]) + pipeline.transform(df=_keyed_adf([0, 1, 999999])) + + report_file = tmp_path / "prepare.stats.yaml" + assert report_file.exists() + + report = yaml.safe_load(report_file.read_text()) + assert report["name"] == "prepare" + assert report["steps"]["valid_mmsi"]["rows_removed"] == 2 + assert report["steps"]["valid_mmsi"]["keys"]["mmsi"]["removed_sample"] == [0, 1] diff --git a/tests/damast/core/test_tracking.py b/tests/damast/core/test_tracking.py new file mode 100644 index 0000000..7c9f3f2 --- /dev/null +++ b/tests/damast/core/test_tracking.py @@ -0,0 +1,165 @@ +import polars +import pytest + +import damast.core +from damast.core.constants import DAMAST_DEFAULT_DATASOURCE +from damast.core.dataframe import AnnotatedDataFrame +from damast.core.metadata import DataSpecification, MetaData +from damast.core.tracking import KeyTracker, RowCountTracker, TimingTracker + + +def adf_of(**columns) -> AnnotatedDataFrame: + return AnnotatedDataFrame( + polars.LazyFrame(columns), + metadata=MetaData([DataSpecification(name=name) for name in columns]), + validation_mode=damast.core.ValidationMode.IGNORE, + ) + + +class FakeStep: + """Stands in for a PipelineElement - a tracker only ever reads its uuid.""" + + def __init__(self, uuid: str = "step-1"): + self.uuid = uuid + + +def test_timing_tracker_reports_a_duration(): + step = FakeStep() + tracker = TimingTracker() + + tracker.on_step_start(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(x=[1])}) + stats = tracker.on_step_end(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(x=[1])}) + + assert stats["end_time"] >= stats["start_time"] + assert stats["processing_time_in_s"] >= 0.0 + assert stats["processing_time_in_s"] == pytest.approx( + (stats["end_time"] - stats["start_time"]).total_seconds()) + + +def test_timing_tracker_without_a_start_still_reports_the_end(): + """A step whose start was never seen must not make a tracker raise.""" + stats = TimingTracker().on_step_end(FakeStep(), {DAMAST_DEFAULT_DATASOURCE: adf_of(x=[1])}) + + assert "end_time" in stats + assert "processing_time_in_s" not in stats + + +def test_row_count_tracker_reports_removed_rows(): + step = FakeStep() + tracker = RowCountTracker() + + tracker.on_step_start(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(x=[1, 2, 3, 4])}) + stats = tracker.on_step_end(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(x=[1, 2])}) + + assert stats["input_dataframe_length"] == {DAMAST_DEFAULT_DATASOURCE: 4} + assert stats["output_dataframe_length"] == 2 + assert stats["rows_removed"] == 2 + + +def test_row_count_tracker_totals_the_inputs_of_a_join(): + """A join consumes several frames, so 'removed' is against their total.""" + step = FakeStep() + tracker = RowCountTracker() + + tracker.on_step_start(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(x=[1, 2, 3]), + "other": adf_of(x=[4, 5])}) + stats = tracker.on_step_end(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(x=[1, 2, 3, 4, 5])}) + + assert stats["input_dataframe_length"] == {DAMAST_DEFAULT_DATASOURCE: 3, "other": 2} + assert stats["output_dataframe_length"] == 5 + assert stats["rows_removed"] == 0 + + +def test_key_tracker_counts_distinct_entities(): + """1000000 appears twice but is one entity, so rows and keys differ.""" + step = FakeStep() + tracker = KeyTracker(["mmsi"]) + + tracker.on_step_start(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(mmsi=[1, 2, 3, 1000000, 1000000])}) + stats = tracker.on_step_end(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(mmsi=[1000000, 1000000])}) + + assert stats["keys"]["mmsi"] == { + "unique_in": 4, + "unique_out": 1, + "unique_removed": 3, + "removed_sample": [1, 2, 3], + "removed_truncated": False, + } + + +def test_key_tracker_bounds_the_removed_sample(): + step = FakeStep() + tracker = KeyTracker(["mmsi"], max_removed_keys=2) + + tracker.on_step_start(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(mmsi=[1, 2, 3, 4, 9])}) + stats = tracker.on_step_end(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(mmsi=[9])}) + + assert stats["keys"]["mmsi"]["unique_removed"] == 4 + assert stats["keys"]["mmsi"]["removed_sample"] == [1, 2] + assert stats["keys"]["mmsi"]["removed_truncated"] is True + + +def test_key_tracker_ignores_nulls(): + """A null identifies no entity, so it is neither counted nor reported as removed.""" + step = FakeStep() + tracker = KeyTracker(["mmsi"]) + + tracker.on_step_start(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(mmsi=[1, None, 2])}) + stats = tracker.on_step_end(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(mmsi=[1, None])}) + + assert stats["keys"]["mmsi"]["unique_in"] == 2 + assert stats["keys"]["mmsi"]["unique_out"] == 1 + assert stats["keys"]["mmsi"]["removed_sample"] == [2] + + +def test_key_tracker_skips_a_key_the_step_drops(): + """A step may drop the column - that is not 'every entity was removed'.""" + step = FakeStep() + tracker = KeyTracker(["mmsi"]) + + tracker.on_step_start(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(mmsi=[1, 2], value=[3, 4])}) + stats = tracker.on_step_end(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(value=[3, 4])}) + + assert stats == {} + + +def test_key_tracker_skips_a_key_that_was_never_there(): + """Tracking a column the data does not have is harmless.""" + step = FakeStep() + tracker = KeyTracker(["mmsi"]) + + tracker.on_step_start(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(value=[1, 2])}) + stats = tracker.on_step_end(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(value=[1])}) + + assert stats == {} + + +def test_key_tracker_unions_the_keys_of_a_join(): + step = FakeStep() + tracker = KeyTracker(["mmsi"]) + + tracker.on_step_start(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(mmsi=[1, 2]), + "other": adf_of(mmsi=[2, 3])}) + stats = tracker.on_step_end(step, {DAMAST_DEFAULT_DATASOURCE: adf_of(mmsi=[1, 2, 3])}) + + assert stats["keys"]["mmsi"]["unique_in"] == 3 + assert stats["keys"]["mmsi"]["unique_removed"] == 0 + + +def test_key_tracker_tracks_each_step_separately(): + """The transient 'before' state is per step, so two steps cannot bleed into each other.""" + tracker = KeyTracker(["mmsi"]) + first, second = FakeStep("a"), FakeStep("b") + + tracker.on_step_start(first, {DAMAST_DEFAULT_DATASOURCE: adf_of(mmsi=[1, 2, 3])}) + tracker.on_step_start(second, {DAMAST_DEFAULT_DATASOURCE: adf_of(mmsi=[7, 8])}) + + assert tracker.on_step_end(second, {DAMAST_DEFAULT_DATASOURCE: adf_of(mmsi=[7])})[ + "keys"]["mmsi"]["removed_sample"] == [8] + assert tracker.on_step_end(first, {DAMAST_DEFAULT_DATASOURCE: adf_of(mmsi=[1])})[ + "keys"]["mmsi"]["removed_sample"] == [2, 3] + + +def test_key_tracker_rejects_a_negative_bound(): + with pytest.raises(ValueError, match="max_removed_keys must not be negative"): + KeyTracker(["mmsi"], max_removed_keys=-1) diff --git a/tests/damast/integrations/test_tracking.py b/tests/damast/integrations/test_tracking.py index 19ea3b0..09b99be 100644 --- a/tests/damast/integrations/test_tracking.py +++ b/tests/damast/integrations/test_tracking.py @@ -1,9 +1,17 @@ +import datetime +import json +import re + from astropy import units from damast.core.annotations import Annotation from damast.core.data_description import MinMax, NumericValueStats from damast.core.metadata import DataSpecification, MetaData -from damast.integrations.tracking import flatten_metadata, flatten_step_stats +from damast.integrations.tracking import ( + flatten_metadata, + flatten_step_params, + flatten_step_stats, +) def test_flatten_metadata_splits_contract_from_stats(): @@ -59,5 +67,113 @@ def test_flatten_step_stats(): assert metrics["step.step_a.input_dataframe_length.df"] == 50.0 +def test_flatten_step_stats_includes_removed_rows_and_keys(): + processing_stats = { + "step_a": { + "processing_time_in_s": 1.5, + "output_dataframe_length": 42, + "input_dataframe_length": {"df": 50}, + "rows_removed": 8, + "keys": {"mmsi": {"unique_in": 7, "unique_out": 3, "unique_removed": 4, + "removed_sample": [1, 2, 3], "removed_truncated": True}}, + } + } + + metrics = flatten_step_stats(processing_stats) + + assert metrics["step.step_a.rows_removed"] == 8.0 + assert metrics["step.step_a.keys.mmsi.unique_in"] == 7.0 + assert metrics["step.step_a.keys.mmsi.unique_out"] == 3.0 + assert metrics["step.step_a.keys.mmsi.unique_removed"] == 4.0 + + # a list and a bool are not metrics - the contract is dict[str, float] + assert not [name for name in metrics if "removed_sample" in name or "removed_truncated" in name] + assert all(isinstance(value, float) for value in metrics.values()) + + def test_flatten_step_stats_empty(): assert flatten_step_stats({}) == {} + + +# --- flattening is generic over whatever a tracker reports ----------------------------------- + +#: what MLflow accepts in a metric or parameter name +MLFLOW_SAFE_NAME = re.compile(r"^[0-9a-zA-Z_\-./ ]+$") + +STATS_WITH_AN_UNKNOWN_TRACKER = { + "step_a": { + "rows_removed": 8, + "my_block": {"my_count": 3, "nested": {"deeper": 1.5}}, + } +} + + +def test_flatten_step_stats_reports_an_unknown_trackers_measurements(): + """The point of the generic walk: a tracker this module knows nothing about is reported.""" + metrics = flatten_step_stats(STATS_WITH_AN_UNKNOWN_TRACKER) + + assert metrics["step.step_a.my_block.my_count"] == 3.0 + # the path carries every level, so a nested block cannot collide with a top-level one + assert metrics["step.step_a.my_block.nested.deeper"] == 1.5 + assert "step.step_a.nested.deeper" not in metrics + assert metrics["step.step_a.rows_removed"] == 8.0 + + +def test_flatten_step_params_takes_what_is_not_a_measurement(): + stats = { + "step_a": { + "rows_removed": 8, + "start_time": datetime.datetime(2026, 6, 1, tzinfo=datetime.timezone.utc), + "keys": {"mmsi": {"unique_removed": 4, + "removed_sample": [1, 2, 3], + "removed_truncated": True}}, + } + } + + metrics = flatten_step_stats(stats) + params = flatten_step_params(stats) + + # a flag is not a quantity, even though isinstance(True, int) holds + assert "step.step_a.keys.mmsi.removed_truncated" not in metrics + assert params["step.step_a.keys.mmsi.removed_truncated"] is True + assert params["step.step_a.start_time"] == stats["step_a"]["start_time"] + # a container is JSON-encoded, as flatten_metadata does + assert json.loads(params["step.step_a.keys.mmsi.removed_sample"]) == [1, 2, 3] + assert metrics["step.step_a.keys.mmsi.unique_removed"] == 4.0 + + +def test_metrics_and_params_partition_the_stats(): + """Exact complements, so nothing a tracker reports is silently dropped.""" + stats = { + "step_a": { + "rows_removed": 8, + "processing_time_in_s": 1.5, + "start_time": datetime.datetime(2026, 6, 1, tzinfo=datetime.timezone.utc), + "input_dataframe_length": {"df": 50, "other": 10}, + "keys": {"mmsi": {"unique_in": 7, "removed_sample": [], "removed_truncated": False}}, + } + } + + metrics = flatten_step_stats(stats) + params = flatten_step_params(stats) + + assert not set(metrics) & set(params) + # 8 leaves in total: 2 scalars + 1 datetime + 2 lengths + 3 key entries + assert len(metrics) + len(params) == 8 + + +def test_names_are_sanitised_for_mlflow(): + """A tracked column can carry characters MLflow rejects - damast writes 'latitude (deg)'.""" + stats = {"step:a": {"keys": {"latitude (deg)": {"unique_in": 3}, + "a:b": {"unique_in": 1}}}} + + metrics = flatten_step_stats(stats) + + assert all(MLFLOW_SAFE_NAME.match(name) for name in metrics), metrics + # parentheses are not in MLflow's set either, so they go too - a space is fine + assert metrics["step.step_a.keys.latitude _deg_.unique_in"] == 3.0 + assert metrics["step.step_a.keys.a_b.unique_in"] == 1.0 + + +def test_flatten_step_params_empty(): + assert flatten_step_params({}) == {} From ccbaf171d7eaf8e56262a2174878f9153bcbee5d Mon Sep 17 00:00:00 2001 From: Thomas Roehr Date: Tue, 6 Oct 2026 17:24:53 +0200 Subject: [PATCH 2/2] cli: add viz subcommand and stats visualization --- src/damast/cli/data_visualize.py | 96 ++++++++ src/damast/cli/main.py | 5 + src/damast/core/dataprocessing.py | 4 +- src/damast/viz/__init__.py | 2 + src/damast/viz/stats_export.py | 281 +++++++++++++++++++++++ tests/damast/cli/test_viz_cli.py | 45 ++++ tests/damast/core/test_dataprocessing.py | 43 ++++ tests/damast/viz/test_stats_export.py | 133 +++++++++++ 8 files changed, 608 insertions(+), 1 deletion(-) create mode 100644 src/damast/cli/data_visualize.py create mode 100644 src/damast/viz/stats_export.py create mode 100644 tests/damast/cli/test_viz_cli.py create mode 100644 tests/damast/viz/test_stats_export.py diff --git a/src/damast/cli/data_visualize.py b/src/damast/cli/data_visualize.py new file mode 100644 index 0000000..8560976 --- /dev/null +++ b/src/damast/cli/data_visualize.py @@ -0,0 +1,96 @@ +""" +Visualize a pipeline - either its structure, or what a run of it did to the data. + +Which of the two is drawn follows from the input file: a saved pipeline (``*.damast.ppl``) +has a structure to render, a run's statistics report (``*.stats.yaml``, written by +`damast.core.dataprocessing.DataProcessingPipeline.save_stats`) has per-step numbers to plot. +""" + +from argparse import ArgumentParser +from pathlib import Path + +from damast.cli.base import BaseParser +from damast.core.dataprocessing import ( + DAMAST_PIPELINE_SUFFIX, + DAMAST_STATS_SUFFIX, + DataProcessingPipeline, +) +from damast.viz.mermaid_export import MermaidExporter +from damast.viz.stats_export import StatsExporter +from damast.viz.svg_export import SvgExporter + +#: Default file extension per output format +_EXTENSIONS = {"svg": ".svg", "html": ".html", "mermaid": ".mmd"} + +#: Formats each kind of input can be drawn as, the first being its default +_FORMATS = { + "stats": ["svg"], + "pipeline": ["svg", "html", "mermaid"], +} + + +def resolve_kind(filename: Path) -> str: + """ + Decide what the given file holds, from its name. + + :param filename: The file to visualize + :return: ``"stats"`` or ``"pipeline"`` + :raise ValueError: If the suffix matches neither + """ + name = filename.name + if name.endswith(DAMAST_STATS_SUFFIX): + return "stats" + if name.endswith(DAMAST_PIPELINE_SUFFIX): + return "pipeline" + + raise ValueError(f"Cannot visualize '{filename}': expected a pipeline ('{DAMAST_PIPELINE_SUFFIX}')" + f" or a run's statistics ('{DAMAST_STATS_SUFFIX}')") + + +class DataVisualizeParser(BaseParser): + """Argparser for the 'viz' subcommand.""" + + def __init__(self, parser: ArgumentParser): + super().__init__(parser=parser) + + parser.description = "damast viz - visualize a pipeline or a run of it" + + parser.add_argument("-f", "--filename", type=str, required=True, + help=f"A saved pipeline ('{DAMAST_PIPELINE_SUFFIX}') to draw the structure" + f" of, or a run's statistics ('{DAMAST_STATS_SUFFIX}') to plot") + parser.add_argument("-o", "--output-file", type=str, default=None, + help="Where to write the result, by default next to the input file") + parser.add_argument("--format", type=str, default=None, + choices=sorted({f for formats in _FORMATS.values() for f in formats}), + help="Output format - 'svg' for statistics; 'svg', 'html' or 'mermaid'" + " for a pipeline (default: svg)") + + def execute(self, args): + super().execute(args) + + filename = Path(args.filename) + if not filename.exists(): + raise FileNotFoundError(f"File '{filename}' does not exist") + + kind = resolve_kind(filename) + permitted = _FORMATS[kind] + output_format = args.format or permitted[0] + if output_format not in permitted: + raise ValueError(f"Cannot write a {kind} visualization as '{output_format}' -" + f" available: {', '.join(permitted)}") + + output_file = Path(args.output_file) if args.output_file else \ + filename.parent / f"{filename.name.split('.')[0]}{_EXTENSIONS[output_format]}" + + if kind == "stats": + written = StatsExporter.from_report(filename).export_svg(output_file) + else: + pipeline = DataProcessingPipeline.load(filename) + if output_format == "svg": + written = SvgExporter(pipeline).export_svg(output_file) + elif output_format == "html": + written = MermaidExporter(pipeline).export_html(output_file) + else: + written = MermaidExporter(pipeline).export_mermaid(output_file) + + print(f"Written: {written}") diff --git a/src/damast/cli/main.py b/src/damast/cli/main.py index 9f6e6ff..50b9c80 100644 --- a/src/damast/cli/main.py +++ b/src/damast/cli/main.py @@ -13,6 +13,7 @@ from damast.cli.data_converter import DataConvertParser from damast.cli.data_inspect import DataInspectParser from damast.cli.data_processing import DataProcessingParser +from damast.cli.data_visualize import DataVisualizeParser from damast.cli.data_watch import DataWatchParser from damast.cli.experiment import ExperimentParser from damast.cli.plugins import PluginsParser @@ -89,6 +90,10 @@ def run(): help="Process data by running a predefined pipeline", parser_klass=DataProcessingParser) + main_parser.attach_subcommand_parser(subcommand="viz", + help="Visualize a pipeline, or what a run of it did to the data", + parser_klass=DataVisualizeParser) + main_parser.attach_subcommand_parser(subcommand="plugins", help="List transformer plugins registered by installed packages", parser_klass=PluginsParser) diff --git a/src/damast/core/dataprocessing.py b/src/damast/core/dataprocessing.py index eb08ff3..4f5ad36 100644 --- a/src/damast/core/dataprocessing.py +++ b/src/damast/core/dataprocessing.py @@ -560,7 +560,9 @@ def save_stats(self, dir: str | Path | None = None) -> Path: filename = base_dir / f"{self.name}{DAMAST_STATS_SUFFIX}" with open(filename, "w") as f: - yaml.dump({"name": self.name, "steps": self._processing_stats}, f) + # sort_keys=False: the steps are a sequence of events, so alphabetising them would + # lose the execution order the statistics only make sense in + yaml.dump({"name": self.name, "steps": self._processing_stats}, f, sort_keys=False) return filename def __iter__(self): diff --git a/src/damast/viz/__init__.py b/src/damast/viz/__init__.py index 54529ae..ab5d72d 100644 --- a/src/damast/viz/__init__.py +++ b/src/damast/viz/__init__.py @@ -4,9 +4,11 @@ """ from .mermaid_export import MermaidExporter +from .stats_export import StatsExporter from .svg_export import SvgExporter DAMAST_VIZ_EXPORTER = { MermaidExporter, + StatsExporter, SvgExporter } diff --git a/src/damast/viz/stats_export.py b/src/damast/viz/stats_export.py new file mode 100644 index 0000000..990a954 --- /dev/null +++ b/src/damast/viz/stats_export.py @@ -0,0 +1,281 @@ +""" +Render what a `DataProcessingPipeline` run did to the data, step by step. + +Takes the statistics the pipeline's trackers collected - see +`damast.core.dataprocessing.DataProcessingPipeline.processing_stats` and the report +`save_stats` writes - and draws them as panels sharing one x axis of pipeline steps, so the +point at which rows or entities disappear is visible against where the time went. + +Each measure gets its own panel rather than a second y axis: rows, distinct keys and seconds +do not share a scale, and neither do a per-step duration and its running total. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import yaml + +__all__ = ["StatsExporter"] + +#: Chart surface and ink, one shade off each other so grid and axes stay recessive +_SURFACE = "#fcfcfb" +_TEXT_PRIMARY = "#0b0b0b" +_TEXT_SECONDARY = "#52514e" +_MUTED = "#898781" +_AXIS = "#c3c2b7" + +#: Categorical slots 1 and 2. Every panel uses them with one fixed meaning - slot 1 is what +#: entered a step, slot 2 what left it - so the two are never a different pair panel to panel. +_SERIES_IN = "#2a78d6" +_SERIES_OUT = "#eb6834" + +#: Width of a bar within its step slot, leaving a surface gap between neighbours +_BAR_WIDTH = 0.38 + + +@dataclass +class _Panel: + """One measure over the steps: its title, y label and the series to draw.""" + title: str + y_label: str + #: label -> (values, colour); a value may be None where a step did not report it + series: dict[str, tuple[list[float | None], str]] = field(default_factory=dict) + #: draw as bars per step, else as a line + bars: bool = True + + +class StatsExporter: + """ + Draw a pipeline run's per-step statistics as an SVG. + + Example: + + .. highlight:: python + .. code-block:: python + + pipeline.transform(df=adf) + StatsExporter(pipeline.processing_stats).export_svg("run.svg") + + # or from the report a run left behind + StatsExporter.from_report("results/prepare.stats.yaml").export_svg("run.svg") + + :param processing_stats: Per-step statistics, keyed by step name in execution order + :param name: Name of the pipeline, used in the figure title + """ + + def __init__(self, processing_stats: dict[str, dict[str, Any]], name: str | None = None): + self.processing_stats = processing_stats + self.name = name + + @classmethod + def from_report(cls, path: str | Path) -> StatsExporter: + """ + Load the report written by + :func:`damast.core.dataprocessing.DataProcessingPipeline.save_stats`. + + :param path: Path of the ``*.stats.yaml`` report + :return: An exporter for the statistics it holds + :raise ValueError: If the file holds no per-step statistics + """ + report = yaml.safe_load(Path(path).read_text()) + if not isinstance(report, dict) or "steps" not in report: + raise ValueError(f"{cls.__name__}.from_report: '{path}' is not a pipeline statistics" + f" report - expected a mapping with a 'steps' entry") + + return cls(processing_stats=report["steps"] or {}, name=report.get("name")) + + @property + def steps(self) -> list[str]: + """ + The step names, in the order the pipeline ran them. + + Insertion order already is execution order, but a report written before ``save_stats`` + stopped alphabetising its keys has lost that - so where every step recorded a + ``start_time``, that is what the order is taken from. + """ + starts = {name: stats.get("start_time") for name, stats in self.processing_stats.items()} + if starts and all(start is not None for start in starts.values()): + return sorted(starts, key=lambda name: starts[name]) + + return list(self.processing_stats) + + def _values(self, *path: str) -> list[float | None]: + """ + One value per step, reaching into the nested statistics by key path. + + :param path: Keys to follow per step, e.g. ``("keys", "mmsi", "unique_in")`` + :return: The value per step, None where that step did not report it + """ + values: list[float | None] = [] + for step in self.steps: + value: Any = self.processing_stats[step] + for key in path: + value = value.get(key) if isinstance(value, dict) else None + values.append(value if isinstance(value, (int, float)) and not isinstance(value, bool) + else None) + return values + + def _tracked_keys(self) -> list[str]: + """The key columns any step reported on, in first-seen order.""" + columns: list[str] = [] + for step in self.steps: + for column in self.processing_stats[step].get("keys", {}): + if column not in columns: + columns.append(column) + return columns + + def panels(self) -> list[_Panel]: + """ + The panels to draw, leaving out every measure this run did not record - a pipeline + without a `damast.core.tracking.KeyTracker` simply has no key panel. + + :return: The panels, in drawing order + """ + panels = [] + + # 'input_dataframe_length' is per datasource; a join has several, so total them + rows_in = [sum(self.processing_stats[step].get("input_dataframe_length", {}).values()) or None + for step in self.steps] + rows_out = self._values("output_dataframe_length") + if any(v is not None for v in rows_in + rows_out): + panels.append(_Panel( + title="Rows", y_label="rows", + series={"entering": (rows_in, _SERIES_IN), "remaining": (rows_out, _SERIES_OUT)})) + + for column in self._tracked_keys(): + unique_in = self._values("keys", column, "unique_in") + unique_out = self._values("keys", column, "unique_out") + panels.append(_Panel( + title=f"Distinct '{column}'", y_label=f"distinct {column}", + series={"entering": (unique_in, _SERIES_IN), + "remaining": (unique_out, _SERIES_OUT)})) + + durations = self._values("processing_time_in_s") + if any(v is not None for v in durations): + panels.append(_Panel( + title="Duration per step", y_label="seconds", + series={"duration": (durations, _SERIES_IN)})) + + running = 0.0 + cumulative: list[float | None] = [] + for value in durations: + running += value or 0.0 + cumulative.append(running) + panels.append(_Panel( + title="Cumulative duration", y_label="seconds", + series={"cumulative": (cumulative, _SERIES_OUT)}, bars=False)) + + return panels + + def export_svg(self, path: str | Path) -> Path: + """ + Draw the panels and write them as an SVG. + + :param path: Where to write the figure + :return: The path written + :raise ValueError: If the statistics hold nothing to draw + """ + import matplotlib.pyplot as plt + + figure = self.build_figure() + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(path, format=path.suffix.lstrip(".") or "svg", facecolor=_SURFACE) + plt.close(figure) + return path + + def build_figure(self): + """ + Build the figure, so that a caller can embed or rasterise it instead of writing an SVG. + + :return: The :class:`matplotlib.figure.Figure` holding one panel per recorded measure + :raise ValueError: If the statistics hold nothing to draw + """ + # imported here, so that importing this module does not pull in a backend + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + from matplotlib.ticker import FuncFormatter + + panels = self.panels() + if not panels: + raise ValueError(f"{self.__class__.__name__}.build_figure: no statistics to draw -" + f" the run recorded nothing for its {len(self.steps)} step(s)") + + steps = self.steps + positions = range(len(steps)) + # the x band has to fit the rotated step names, so grow the figure with the step count + figure, axes = plt.subplots( + nrows=len(panels), ncols=1, sharex=True, squeeze=False, + figsize=(max(7.0, 0.62 * len(steps)), 1.85 * len(panels) + 1.6), + ) + figure.patch.set_facecolor(_SURFACE) + + for panel, axis in zip(panels, (a for row in axes for a in row)): + axis.set_facecolor(_SURFACE) + offsets = self._bar_offsets(len(panel.series)) if panel.bars else None + + for index, (label, (values, colour)) in enumerate(panel.series.items()): + if panel.bars: + axis.bar([p + offsets[index] for p in positions], + [v if v is not None else 0 for v in values], + width=_BAR_WIDTH, label=label, color=colour, linewidth=0) + else: + axis.plot(list(positions), values, label=label, color=colour, + linewidth=2.0, marker="o", markersize=4, + markeredgecolor=_SURFACE, markeredgewidth=1.5) + + axis.set_title(panel.title, loc="left", fontsize=10, color=_TEXT_PRIMARY, pad=6) + axis.set_ylabel(panel.y_label, fontsize=8, color=_TEXT_SECONDARY) + axis.yaxis.set_major_formatter(FuncFormatter(self._format_tick)) + self._style_axis(axis) + + # a legend names the series only where there is more than one to tell apart + if len(panel.series) > 1: + legend = axis.legend(loc="upper right", fontsize=8, frameon=False, + labelcolor=_TEXT_SECONDARY, handlelength=1.2) + legend.set_zorder(5) + + bottom = axes[-1][0] + bottom.set_xticks(list(positions)) + bottom.set_xticklabels(steps, rotation=40, ha="right", fontsize=8, color=_TEXT_SECONDARY) + bottom.set_xlabel("pipeline step", fontsize=9, color=_TEXT_SECONDARY, labelpad=6) + + title = f"{self.name} - per-step statistics" if self.name else "Per-step statistics" + figure.suptitle(title, x=0.01, ha="left", fontsize=12, color=_TEXT_PRIMARY) + figure.tight_layout(rect=(0, 0, 1, 0.97)) + return figure + + @staticmethod + def _bar_offsets(count: int) -> list[float]: + """ + Where each series' bar sits within a step's slot, with a gap between neighbours. + + :param count: Number of series in the panel + :return: The offset per series + """ + span = _BAR_WIDTH * count + return [-span / 2 + _BAR_WIDTH * (index + 0.5) for index in range(count)] + + @staticmethod + def _format_tick(value: float, _position: int) -> str: + """Abbreviate a tick, so that a row count of 30 million does not crowd the axis.""" + for threshold, suffix in ((1e9, "B"), (1e6, "M"), (1e3, "k")): + if abs(value) >= threshold: + return f"{value / threshold:.3g}{suffix}" + return f"{value:.3g}" + + @staticmethod + def _style_axis(axis) -> None: + """Hairline, solid, recessive chrome - a dashed grid reads as a threshold.""" + axis.grid(axis="y", color=_AXIS, linewidth=0.5, linestyle="-", alpha=0.6) + axis.set_axisbelow(True) + for side in ("top", "right"): + axis.spines[side].set_visible(False) + for side in ("left", "bottom"): + axis.spines[side].set_color(_AXIS) + axis.spines[side].set_linewidth(0.6) + axis.tick_params(colors=_MUTED, labelsize=8, length=3, width=0.6) diff --git a/tests/damast/cli/test_viz_cli.py b/tests/damast/cli/test_viz_cli.py new file mode 100644 index 0000000..b69dc22 --- /dev/null +++ b/tests/damast/cli/test_viz_cli.py @@ -0,0 +1,45 @@ +from pathlib import Path + +import pytest +import yaml + +from damast.cli.data_visualize import resolve_kind + + +@pytest.mark.parametrize("name,expected", [ + ("prepare.stats.yaml", "stats"), + ("prepare.damast.ppl", "pipeline"), + ("/some/dir/ais_join_sources.damast.ppl", "pipeline"), +]) +def test_resolve_kind_follows_the_suffix(name: str, expected: str): + assert resolve_kind(Path(name)) == expected + + +@pytest.mark.parametrize("name", ["README.md", "data.parquet", "prepare.yaml"]) +def test_resolve_kind_rejects_anything_else(name: str): + with pytest.raises(ValueError, match="expected a pipeline"): + resolve_kind(Path(name)) + + +def test_viz_renders_a_statistics_report(script_runner, tmp_path): + report = tmp_path / "prepare.stats.yaml" + report.write_text(yaml.dump({"name": "prepare", "steps": { + "df": {"processing_time_in_s": 1.0, + "input_dataframe_length": {"df": 10}, "output_dataframe_length": 10}, + }}, sort_keys=False)) + output = tmp_path / "run.svg" + + result = script_runner.run(["damast", "viz", "-f", str(report), "-o", str(output)]) + + assert result.returncode == 0, result.stderr + assert output.exists() + + +def test_viz_rejects_a_format_the_input_cannot_be_drawn_as(script_runner, tmp_path): + report = tmp_path / "prepare.stats.yaml" + report.write_text(yaml.dump({"name": "p", "steps": {}}, sort_keys=False)) + + result = script_runner.run(["damast", "viz", "-f", str(report), "--format", "mermaid"]) + + assert result.returncode != 0 + assert "Cannot write a stats visualization as 'mermaid'" in result.stdout + result.stderr diff --git a/tests/damast/core/test_dataprocessing.py b/tests/damast/core/test_dataprocessing.py index 16f08ff..79179bb 100644 --- a/tests/damast/core/test_dataprocessing.py +++ b/tests/damast/core/test_dataprocessing.py @@ -1239,3 +1239,46 @@ def test_transform_writes_a_statistics_report(tmp_path): assert report["name"] == "prepare" assert report["steps"]["valid_mmsi"]["rows_removed"] == 2 assert report["steps"]["valid_mmsi"]["keys"]["mmsi"]["removed_sample"] == [0, 1] + + +def _two_step_pipeline(tmp_path, name: str) -> DataProcessingPipeline: + """A pipeline whose steps are named so that alphabetical order is not execution order.""" + from damast.data_handling.transformers.filters import Filter + + pipeline = DataProcessingPipeline(name=name, base_dir=tmp_path) + pipeline.add("zzz_first", Filter(operator=">=", value=1), name_mappings={"x": "mmsi"}) + pipeline.add("aaa_second", Filter(operator=">=", value=999999), name_mappings={"x": "mmsi"}) + return pipeline + + +def test_statistics_report_keeps_the_steps_in_execution_order(tmp_path): + """ + The steps are a sequence of events, and 'yaml.dump' alphabetises keys unless told not to. + Sorted, anything read back cumulatively - a running duration, say - adds up in an order the + pipeline never ran in. + """ + import yaml + + pipeline = _two_step_pipeline(tmp_path, name="ordered") + pipeline.transform(df=_keyed_adf([0, 1, 999999])) + + report = yaml.safe_load((tmp_path / "ordered.stats.yaml").read_text()) + + assert list(report["steps"]) == ["df", "zzz_first", "aaa_second"] + assert list(report["steps"]) == list(pipeline.processing_stats) + # guards the test itself: these names have to discriminate, or it would pass either way + assert list(report["steps"]) != sorted(report["steps"]) + + +def test_save_stats_writes_to_a_given_directory(tmp_path): + """Without an argument the report lands in base_dir, which transform relies on.""" + import yaml + + pipeline = _two_step_pipeline(tmp_path, name="elsewhere") + pipeline.transform(df=_keyed_adf([0, 1, 999999])) + + target = tmp_path / "reports" / "nested" + written = pipeline.save_stats(target) + + assert written == target / "elsewhere.stats.yaml" + assert list(yaml.safe_load(written.read_text())["steps"]) == list(pipeline.processing_stats) diff --git a/tests/damast/viz/test_stats_export.py b/tests/damast/viz/test_stats_export.py new file mode 100644 index 0000000..9c056dc --- /dev/null +++ b/tests/damast/viz/test_stats_export.py @@ -0,0 +1,133 @@ +import datetime + +import pytest +import yaml + +from damast.viz.stats_export import StatsExporter + + +def _stats(**overrides) -> dict: + """Two steps' worth of statistics, as the default trackers record them.""" + start = datetime.datetime(2026, 6, 1, 12, 0, tzinfo=datetime.timezone.utc) + stats = { + "df": { + "start_time": start, + "end_time": start + datetime.timedelta(seconds=2), + "processing_time_in_s": 2.0, + "input_dataframe_length": {"df": 100}, + "output_dataframe_length": 100, + "rows_removed": 0, + }, + "filter": { + "start_time": start + datetime.timedelta(seconds=2), + "end_time": start + datetime.timedelta(seconds=5), + "processing_time_in_s": 3.0, + "input_dataframe_length": {"df": 100}, + "output_dataframe_length": 40, + "rows_removed": 60, + }, + } + stats.update(overrides) + return stats + + +def _with_keys(stats: dict) -> dict: + stats["df"]["keys"] = {"mmsi": {"unique_in": 9, "unique_out": 9, "unique_removed": 0}} + stats["filter"]["keys"] = {"mmsi": {"unique_in": 9, "unique_out": 3, "unique_removed": 6}} + return stats + + +def test_panels_cover_the_recorded_measures(): + panels = StatsExporter(_with_keys(_stats())).panels() + + assert [panel.title for panel in panels] == [ + "Rows", "Distinct 'mmsi'", "Duration per step", "Cumulative duration"] + + +def test_key_panel_is_absent_without_a_key_tracker(): + """The default tracker set records no keys, so there is nothing to draw for them.""" + titles = [panel.title for panel in StatsExporter(_stats()).panels()] + + assert "Distinct 'mmsi'" not in titles + assert "Rows" in titles + + +def test_cumulative_duration_accumulates_in_execution_order(): + panels = {panel.title: panel for panel in StatsExporter(_stats()).panels()} + + assert panels["Duration per step"].series["duration"][0] == [2.0, 3.0] + assert panels["Cumulative duration"].series["cumulative"][0] == [2.0, 5.0] + + +def test_steps_follow_execution_order_not_the_reports_key_order(): + """ + A report written before save_stats stopped alphabetising its keys has lost the order the + statistics only make sense in - start_time recovers it. + """ + alphabetised = dict(sorted(_stats().items())) + assert list(alphabetised) == ["df", "filter"] + + renamed = {"zzz_first": _stats()["df"], "aaa_second": _stats()["filter"]} + exporter = StatsExporter(dict(sorted(renamed.items()))) + + assert list(exporter.processing_stats) == ["aaa_second", "zzz_first"] + assert exporter.steps == ["zzz_first", "aaa_second"] + # the series have to follow the steps, or they would be plotted against the wrong labels + panels = {panel.title: panel for panel in exporter.panels()} + assert panels["Duration per step"].series["duration"][0] == [2.0, 3.0] + + +def test_steps_keep_insertion_order_without_timestamps(): + stats = _stats() + for step in stats.values(): + del step["start_time"] + + assert StatsExporter(stats).steps == ["df", "filter"] + + +def test_a_join_totals_the_rows_of_its_inputs(): + stats = _stats() + stats["filter"]["input_dataframe_length"] = {"df": 100, "other": 25} + + panels = {panel.title: panel for panel in StatsExporter(stats).panels()} + assert panels["Rows"].series["entering"][0] == [100, 125] + + +def test_export_svg_writes_a_figure(tmp_path): + written = StatsExporter(_with_keys(_stats()), name="prepare").export_svg(tmp_path / "run.svg") + + assert written.exists() + content = written.read_text() + assert content.lstrip().startswith("