diff --git a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/connector.py b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/connector.py index 7fab9e4ce..3a8961ae4 100644 --- a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/connector.py +++ b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/connector.py @@ -1,129 +1,78 @@ -"""Convert ARFBench questions, metric series, and charts to TimeF. +"""Convert the ARFBench questions and metric series to TimeF. The release holds one question table, one series file per metric and rollup interval, and one -chart per metric or cited metric pair. Each metric becomes one Record, each chart one Record, and -each question a series task and a chart task. The loop that walks the release is in ``convert``. -:mod:`~timenet_connectors.datasets.datadog.arfbench.series` reads the series files, and -:mod:`~timenet_connectors.datasets.datadog.arfbench.questions` reads the question table. +chart per metric or cited metric pair. This connector reads the table and the series. Each metric +becomes one Record, and each question becomes one series task. The loop over the release is in +``convert``. :mod:`~timenet_connectors.datasets.datadog.arfbench.series` reads the series files. +:mod:`~timenet_connectors.datasets.datadog.arfbench.questions` reads the question table. The +``datadog/arfbench-plots`` dataset uses this dataset as its parent and adds the charts. """ from __future__ import annotations -from dataclasses import dataclass -from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar -from timenet.connectors import BaseConnector -from timenet.dataset import Record, Source, TimeFDataset -from timenet.errors import TimeNetConnectorError, TimeNetDownloadError -from timenet_connectors.datasets.datadog.arfbench.questions import ARFBenchQuestion, read_questions +from timenet.dataset import TimeFDataset +from timenet.errors import TimeNetConnectorError +from timenet_connectors.bases.huggingface import BaseHuggingFaceConnector +from timenet_connectors.datasets.datadog.arfbench.questions import read_questions from timenet_connectors.datasets.datadog.arfbench.series import MetricSeriesFiles -from timenet_connectors.sources.huggingface_hub import hub_snapshot -from timenet_connectors.sources.images import image_signal if TYPE_CHECKING: - from collections.abc import Iterator + from pathlib import Path + + from timenet_connectors.bases.huggingface import HubFileRef REPO = "Datadog/ARFBench" REVISION = "1cc8ae54e9633f596b755f7c4bce54ccb0cb9f5a" QA_CSV = "arfbench-qa.csv" SERIES_DIR = "arfbench-ts-data" # {metric}_{interval}.parquet: one metric at one rollup interval, in seconds -CHART_DIR = "arfbench-images" # {metric}.png, and {metric}-{metric}.png for the pairs the questions cite -def _create_chart_record(path: Path) -> Record: - """Create the Record of a chart, with the image as its one signal. +def release_paths(raw_refs: list[HubFileRef], *parts: str) -> dict[str, Path]: + """Map each named part of the release to its local path. Returns: - The Record, not yet added to a dataset. - """ - return Record( - metadata={"source_record_id": f"arfbench-chart-{path.stem}"}, - sources=(Source(name="ARFBench chart", signals=(image_signal(path, name="chart"),)),), - ) - - -@dataclass(frozen=True) -class ARFBenchSource: - """Paths to the ARFBench question table, series files, and charts.""" - - qa_csv: Path - series_dir: Path - chart_dir: Path - - def read_questions(self) -> Iterator[ARFBenchQuestion]: - """Read the question table, in table order. - - Returns: - Each parsed question. - """ - return read_questions(self.qa_csv) - - def metric_series_files(self) -> list[MetricSeriesFiles]: - """List the series files, grouped by metric. - - Returns: - The metrics in ID order. - """ - return MetricSeriesFiles.from_series_dir(self.series_dir) + The path of each part: a file, or the directory of a pattern part. - def chart_paths(self) -> list[Path]: - """List the chart files. - - Returns: - The PNG files in name order. - """ - return sorted(self.chart_dir.glob("*.png")) - - -class ARFBenchConnector(BaseConnector[ARFBenchSource]): - """Connector for ARFBench (Hub repo ``Datadog/ARFBench``).""" + Raises: + TimeNetConnectorError: If a part is missing. + """ + paths = {ref.part: ref.path for ref in raw_refs} + missing = [part for part in parts if part not in paths] + if missing: + raise TimeNetConnectorError(f"the ARFBench release lacks {', '.join(missing)}") + return paths - values_backend = "zarr" # the chart tensors need it - async def download_async(self, cache_dir: Path) -> list[ARFBenchSource]: # ruff: ignore[no-self-use] - BaseConnector override - """Download the question table, series files, and charts from the pinned release. +class ARFBenchConnector(BaseHuggingFaceConnector): + """Connector for the ARFBench questions about the metric series (Hub repo ``Datadog/ARFBench``). - Args: - cache_dir: Directory for cached Hub files. - - Returns: - A list with the paths to the table and the two directories. + The values are plain series, so the dataset uses the default Parquet backend. + """ - Raises: - TimeNetDownloadError: If the download fails or the question table is missing. - """ - root = await hub_snapshot(REPO, REVISION, cache_dir, (QA_CSV, f"{SERIES_DIR}/*.parquet", f"{CHART_DIR}/*.png")) - if not (root / QA_CSV).is_file(): - raise TimeNetDownloadError(f"{REPO!r} at {REVISION!r} holds no {QA_CSV} under {root}") - return [ARFBenchSource(qa_csv=root / QA_CSV, series_dir=root / SERIES_DIR, chart_dir=root / CHART_DIR)] + HF_REPO = REPO + SOURCE_REVISION = REVISION + SOURCE_FILES: ClassVar = {"questions": QA_CSV, "series": f"{SERIES_DIR}/*.parquet"} - def convert(self, raw_refs: list[ARFBenchSource]) -> TimeFDataset: - """Create a Record for each metric and chart, and two tasks for each question. + def convert(self, raw_refs: list[HubFileRef]) -> TimeFDataset: + """Create a Record for each metric and a series task for each question. Args: - raw_refs: The one source from :meth:`download_async`. + raw_refs: The question table and the series directory from :meth:`download_async`. Returns: The populated dataset. - - Raises: - TimeNetConnectorError: If there is not exactly one source. """ - if len(raw_refs) != 1: - raise TimeNetConnectorError(f"expected one ARFBench source, got {len(raw_refs)}") - source = raw_refs[0] + paths = release_paths(raw_refs, "questions", "series") dataset = TimeFDataset(metadata=self.metadata()) metric_records = { - files.metric: dataset.add_record(record=files.create_record()) for files in source.metric_series_files() - } - chart_records = { - path.stem: dataset.add_record(record=_create_chart_record(path)) for path in source.chart_paths() + files.metric: dataset.add_record(record=files.create_record()) + for files in MetricSeriesFiles.from_series_dir(paths["series"]) } - for question in source.read_questions(): + for question in read_questions(paths["questions"]): dataset.add_task(task=question.create_series_task(metric_records)) - dataset.add_task(task=question.create_chart_task(chart_records)) return dataset diff --git a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/dataset.yaml b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/dataset.yaml index 5d43bfaca..da9710a71 100644 --- a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/dataset.yaml +++ b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/dataset.yaml @@ -1,38 +1,43 @@ -# yaml-language-server: $schema=https://docs.timenet.ai/schemas/dataset-card.schema.json +# yaml-language-server: $schema=https://docs.timenet.ai/schemas/dataset-card-v1.schema.json # # ARFBench asks 750 multiple-choice questions about 142 metrics from 63 Datadog incidents. A metric # is one monitored quantity, such as the request count of a service. Datadog splits a metric by -# tags, so a metric holds one series per tag group, for example one per host. +# tags. Thus a metric holds one series per tag group, for example one series per host. # # Datadog stores a metric at several rollups. A rollup is the metric on a fixed time grid, with one # value per interval of 10, 60, 300, 1800, 3600, or 86400 seconds. The rollups of a metric are -# separate exports, not resamplings of one series: their time ranges and tag groups differ. +# separate exports, not resamplings of one series. Their time ranges and tag groups differ. # # Each metric is one record, with one source per tag group and one signal per rollup. A consumer -# chooses the rollup. The paper and the charts use the finest one, which the annotation -# `finest_interval_s` names. +# chooses the rollup. The paper and the charts use the finest rollup. The annotation +# `finest_interval_s` names this rollup. # # A rollup file lists the tag groups at the steps of its grid. Where the metric reported nothing -# for a group, the step holds a null or is left out. Each series is cut to the window from its -# group's first value to its last, so a group with no value at a rollup has no series there. Inside -# the window, each such step is a missing timestep on the regular grid, with the spec's reason -# `not_recorded`. Values are float64, and the release names no units. The release stores naive -# epochs, which the dataset reads as UTC. The charts and the answer options use the same clock. +# for a group, the step holds a null or is not in the file. The dataset cuts each series to the +# window from the first value of its group to the last value. Thus a group with no value at a +# rollup has no series at that rollup. Inside the window, each such step is a missing timestep on +# the regular grid, with the reason `not_recorded` from the spec. Values are float64. The release +# names no units. The release stores naive epochs. The dataset reads them as UTC. The answer +# options use the same clock. # -# Each question becomes two test tasks: one over the cited metric records, one over the released -# chart images. Both carry the question's row as `question_index` in their metadata, so the test -# split holds 1,500 tasks. +# Each question becomes one test task over the cited metric records. Thus the test split holds 750 +# tasks. A task carries the row of the question as `question_index` in its metadata. The released +# charts are in `datadog/arfbench-plots`. The chart tasks of that dataset carry the same +# `question_index`. # -# Kept as released: `query_name`, an identifier that differs from the metric ID for 122 of the 142 -# metrics and that the release does not explain; one end-time answer (row 191, metric 36464_0) -# 140 seconds after its series ends; and 29 time options that lie outside every rollup. +# The dataset keeps three items as released: +# - `query_name`. This identifier differs from the metric ID for 122 of the 142 metrics. The +# release does not explain it. +# - One end-time answer (row 191, metric 36464_0). This answer is 140 seconds after the end of its +# series. +# - 29 time options that are outside every rollup. yaml_schema_version: 1 dataset_id: datadog/arfbench dataset_version: 1.0.0 name: ARFBench description: >- 750 multiple-choice questions about anomalies in 142 metrics from 63 Datadog incidents. - The dataset includes every published metric rollup and the released charts. + The dataset includes every published metric rollup. The charts are in datadog/arfbench-plots. license: Apache-2.0 source_url: https://huggingface.co/datasets/Datadog/ARFBench citation: >- diff --git a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/questions.py b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/questions.py index d3ffc273a..c96eeea60 100644 --- a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/questions.py +++ b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/questions.py @@ -1,4 +1,4 @@ -"""The ARFBench question table: 750 multiple-choice questions, each a series task and a chart task.""" +"""The ARFBench question table: 750 multiple-choice questions, each a series task or a chart task.""" from __future__ import annotations @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Self, TypeVar from timenet.errors import TimeNetConnectorError -from timenet.types import Annotation, AnswerTask, InputModality, Split +from timenet.types import Annotation, ClassificationTask, InputModality, Label, LabelSet, Split if TYPE_CHECKING: @@ -167,7 +167,7 @@ def chart_stems(self) -> tuple[str, ...]: return self.cited_metrics return ("-".join(self.cited_metrics), *self.cited_metrics) - def create_series_task(self, metric_records: Mapping[str, Record]) -> AnswerTask: + def create_series_task(self, metric_records: Mapping[str, Record]) -> ClassificationTask: """Create the task that asks the question about the cited metric records. Args: @@ -185,7 +185,7 @@ def create_series_task(self, metric_records: Mapping[str, Record]) -> AnswerTask inputs = tuple(metric_records[metric] for metric in self.cited_metrics) return self._create_task(inputs, InputModality.TIME_SERIES) - def create_chart_task(self, chart_records: Mapping[str, Record]) -> AnswerTask: + def create_chart_task(self, chart_records: Mapping[str, Record]) -> ClassificationTask: """Create the task that asks the question about the released charts. Args: @@ -203,25 +203,28 @@ def create_chart_task(self, chart_records: Mapping[str, Record]) -> AnswerTask: inputs = tuple(chart_records[stem] for stem in self.chart_stems()) return self._create_task(inputs, InputModality.IMAGE) - def _create_task(self, inputs: tuple[Record, ...], modality: InputModality) -> AnswerTask: - """Create a test task with the question as prompt and the correct answer as target. + def _create_task(self, inputs: tuple[Record, ...], modality: InputModality) -> ClassificationTask: + """Create a test task that asks for the correct answer from the options of the question. - Both tasks of a question carry its row index in their metadata, so a consumer can pair - them. The question's annotations describe it, and the interpolation flags describe what - an engineer saw. Nothing here interpolates. + The options are the label set of the task. The label set has the name of the question row, + so both tasks of a question share it. Both tasks also have the row index in their metadata, + so a consumer can pair them. The annotations describe the question. The interpolation flags + describe what an engineer saw. This method does not interpolate. Returns: The task, not yet added to a dataset. """ - task = AnswerTask( + task = ClassificationTask( inputs=inputs, prompt=self.question, targets=(self.correct_answer,), input_modalities=frozenset({InputModality.TEXT, modality}), split=Split.TEST, metadata={"question_index": self.index}, + label_set=LabelSet( + name=f"question_{self.index}", labels=tuple(Label(name=option) for option in self.answer_options) + ), ) - task.annotate(Annotation(key="answer_options", value=list(self.answer_options))) task.annotate(Annotation(key="task_category", value=self.task_category.value)) task.annotate(Annotation(key="difficulty", value=self.difficulty.value)) task.annotate(Annotation(key="interpolate_1", value=self.interpolate_1)) diff --git a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/requirements.txt b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/requirements.txt index 98c1a7241..47df73a21 100644 --- a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/requirements.txt +++ b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/requirements.txt @@ -1,7 +1,2 @@ -# huggingface_hub downloads the release from the HuggingFace Hub. -huggingface_hub>=0.24 -# pyarrow reads the series files. -pyarrow>=21 -# Pillow decodes chart PNGs; Zarr stores their tensor values. -pillow>=11 -zarr>=3 +# huggingface_hub downloads the release from the HuggingFace Hub. Version 1.1 adds tqdm_class. +huggingface_hub>=1.1 diff --git a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/tests/release.py b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/tests/release.py new file mode 100644 index 000000000..7c6308d86 --- /dev/null +++ b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/tests/release.py @@ -0,0 +1,128 @@ +"""Write a small ARFBench release for the connector tests: two metrics and the question table. + +The plots tests add the charts. This module does not import Pillow. Thus the environment of the +series connector can run it. +""" + +import csv +from datetime import datetime, timedelta +from pathlib import Path + +import pyarrow as pa +import pyarrow.parquet as pq + +from timenet.types import read_value +from timenet_connectors.bases.huggingface import HubFileRef +from timenet_connectors.datasets.datadog.arfbench.connector import QA_CSV, SERIES_DIR +from timenet_connectors.datasets.datadog.arfbench.questions import QUESTION_COLUMNS + + +START = datetime(2025, 3, 7, 9, 30) # naive, as the release stores its epochs. The connector reads it as UTC. +START_US = 1_741_339_800_000_000 + + +def write_series_file(path: Path, rows: list[tuple[int, str, float | None]]) -> None: + """Write (seconds after START, tag label, value) rows the way the release does, in arbitrary order.""" + pq.write_table( + pa.table( + { + "epoch": pa.array([START + timedelta(seconds=s) for s, _, _ in rows], type=pa.timestamp("ns")), + "group": [label for _, label, _ in rows], + "value": pa.array([value for _, _, value in rows], type=pa.float64()), + "num_groups": [1] * len(rows), + "query_name": ["35997_9"] * len(rows), + "__index_level_0__": list(range(len(rows))), + } + ), + path, + ) + + +def write_series_dir(root: Path) -> Path: + """Write the series files of metrics 35997_0 and 35997_1. + + Returns: + The series directory. + """ + series_dir = root / SERIES_DIR + series_dir.mkdir() + # Metric 35997_0 has two tag groups at 10 s, and dc:2 skips a step. At 60 s a third group appears, + # as in the wider windows of the release. The values of this group start at the second grid step. + # Its first step is grid padding. dc:1 has a null inside its window. dc:4 is on the grid without a value. + write_series_file( + series_dir / "35997_0_10.parquet", [(10, "dc:1", 2.0), (0, "dc:2", 5.0), (0, "dc:1", 1.0), (20, "dc:2", 6.0)] + ) + write_series_file( + series_dir / "35997_0_60.parquet", + [ + (0, "dc:3", None), + (60, "dc:3", 9.0), + (60, "dc:1", None), + (0, "dc:1", 1.5), + (120, "dc:1", 2.5), + (0, "dc:4", None), + (60, "dc:4", None), + (120, "dc:4", None), + ], + ) + # Metric 35997_1 has no tag grouping, so its label is empty. The metric has only a 10 s file. + write_series_file(series_dir / "35997_1_10.parquet", [(30, "", 0.5), (40, "", 0.6)]) + return series_dir + + +PRESENCE_ROW = { + "Unnamed: 0": "0", + "question": "Is there an anomaly?\nTime-series: A gauge.", + "task_category": "Anomaly Presence", + "difficulty": "Tier 1", + "options_str": '["Yes", "No"]', + "correct_answer": "Yes", + "query_group": "35997_1", + "options": '[{"value": "Yes"}, {"value": "No"}]', + "interpolate_1": "0", + "interpolate_2": "0", +} +CORRELATION_ROW = { + "Unnamed: 0": "1", + "question": "Are they correlated?\nTime-series 1: A gauge.\nTime-series 2: A count.", + "task_category": "Anomaly Correlation", + "difficulty": "Tier 3", + "options_str": '["Yes", "No"]', + "correct_answer": "No", + "query_group": "35997_0, 35997_1", + "options": '[{"value": "Yes"}, {"value": "No"}]', + "interpolate_1": "1", + "interpolate_2": "0", +} +ROWS = [PRESENCE_ROW, CORRELATION_ROW] + + +def write_qa_csv(root: Path, rows: list[dict[str, str]], *, columns: tuple[str, ...] = QUESTION_COLUMNS) -> Path: + """Write the question table. + + Returns: + The table path. + """ + qa_csv = root / QA_CSV + with qa_csv.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=columns, extrasaction="ignore") + writer.writeheader() + writer.writerows(rows) + return qa_csv + + +def facts(annotated) -> dict: + """Read the annotations of a record or task as a dict.""" + return {annotation.key: read_value(annotation.value) for annotation in annotated.annotations} + + +def write_release(root: Path) -> list[HubFileRef]: + """Write the question table and the series files. + + Returns: + The refs that the series connector downloads. + """ + return [ + HubFileRef(write_qa_csv(root, ROWS), "questions"), + HubFileRef(write_series_dir(root), "series"), + ] diff --git a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/tests/test_connector.py b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/tests/test_connector.py index b2caf88be..45beb3317 100644 --- a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/tests/test_connector.py +++ b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench/tests/test_connector.py @@ -1,142 +1,58 @@ -"""Check that a metric keeps every rollup per tag group, that a question gets a series task and a chart task, -and that a malformed source raises with its location.""" +"""Check four properties of the ARFBench series connector. + +A metric keeps every rollup per tag group. A question becomes a classification task over its options. The dataset +round-trips through Parquet with its missing timesteps. A malformed source raises an error with its location.""" -import csv -from datetime import datetime, timedelta from pathlib import Path -import numpy as np -from PIL import Image -import pyarrow as pa -import pyarrow.parquet as pq import pytest from timenet.dataset import RegularAxis from timenet.errors import TimeNetConnectorError from timenet.reader import TimeFReader from timenet.registry import open_local_version -from timenet.types import InputModality, MissingReason, Split, read_value, read_values +from timenet.types import ClassificationTask, InputModality, MissingReason, Split, read_value, read_values from timenet.writer import TimeFWriter -from timenet_connectors.datasets.datadog.arfbench.connector import ARFBenchConnector, ARFBenchSource +from timenet_connectors.bases.huggingface import HubFileRef +from timenet_connectors.datasets.datadog.arfbench.connector import ARFBenchConnector from timenet_connectors.datasets.datadog.arfbench.questions import QUESTION_COLUMNS +from timenet_connectors.datasets.datadog.arfbench.tests.release import ( + CORRELATION_ROW, + PRESENCE_ROW, + ROWS, + START_US, + facts, + write_qa_csv, + write_release, + write_series_file, +) -START = datetime(2025, 3, 7, 9, 30) # naive, as the release stores its epochs; the connector reads it as UTC -START_US = 1_741_339_800_000_000 - - -def _series_file(path: Path, rows: list[tuple[int, str, float | None]]) -> None: - """Write (seconds after START, tag label, value) rows the way the release does, in arbitrary order.""" - pq.write_table( - pa.table( - { - "epoch": pa.array([START + timedelta(seconds=s) for s, _, _ in rows], type=pa.timestamp("ns")), - "group": [label for _, label, _ in rows], - "value": pa.array([value for _, _, value in rows], type=pa.float64()), - "num_groups": [1] * len(rows), - "query_name": ["35997_9"] * len(rows), - "__index_level_0__": list(range(len(rows))), - } - ), - path, - ) - - -def _write_series_dir(directory: Path) -> Path: - series_dir = directory / "arfbench-ts-data" - series_dir.mkdir() - # Metric 35997_0: two tag groups at 10 s, where dc:2 skips a step. At 60 s a third group appears, - # as it does in the wider windows of the release, but only from the second grid step: its first - # step is grid padding. dc:1 has a null inside its window, and dc:4 is on the grid without a value. - _series_file( - series_dir / "35997_0_10.parquet", [(10, "dc:1", 2.0), (0, "dc:2", 5.0), (0, "dc:1", 1.0), (20, "dc:2", 6.0)] - ) - _series_file( - series_dir / "35997_0_60.parquet", - [ - (0, "dc:3", None), - (60, "dc:3", 9.0), - (60, "dc:1", None), - (0, "dc:1", 1.5), - (120, "dc:1", 2.5), - (0, "dc:4", None), - (60, "dc:4", None), - (120, "dc:4", None), - ], - ) - # Metric 35997_1 has no tag grouping, so its label is empty, and only a 10 s file. - _series_file(series_dir / "35997_1_10.parquet", [(30, "", 0.5), (40, "", 0.6)]) - return series_dir - - -def _write_chart_dir(directory: Path) -> Path: - chart_dir = directory / "arfbench-images" - chart_dir.mkdir() - for metric in ("35997_0", "35997_1"): - Image.new("RGBA", (3, 2), (10, 20, 30, 42)).save(chart_dir / f"{metric}.png") - Image.new("RGB", (6, 2), (255, 255, 255)).save(chart_dir / "35997_0-35997_1.png") - return chart_dir - - -_PRESENCE_ROW = { - "Unnamed: 0": "0", - "question": "Is there an anomaly?\nTime-series: A gauge.", - "task_category": "Anomaly Presence", - "difficulty": "Tier 1", - "options_str": '["Yes", "No"]', - "correct_answer": "Yes", - "query_group": "35997_1", - "options": '[{"value": "Yes"}, {"value": "No"}]', - "interpolate_1": "0", - "interpolate_2": "0", -} -_CORRELATION_ROW = { - "Unnamed: 0": "1", - "question": "Are they correlated?\nTime-series 1: A gauge.\nTime-series 2: A count.", - "task_category": "Anomaly Correlation", - "difficulty": "Tier 3", - "options_str": '["Yes", "No"]', - "correct_answer": "No", - "query_group": "35997_0, 35997_1", - "options": '[{"value": "Yes"}, {"value": "No"}]', - "interpolate_1": "1", - "interpolate_2": "0", -} - +@pytest.fixture +def release(tmp_path: Path) -> list[HubFileRef]: + """A release with two metrics and one question about each metric count.""" + return write_release(tmp_path) -def _write_qa_csv(directory: Path, rows: list[dict[str, str]], *, columns: tuple[str, ...] = QUESTION_COLUMNS) -> Path: - qa_csv = directory / "arfbench-qa.csv" - with qa_csv.open("w", newline="", encoding="utf-8") as handle: - writer = csv.DictWriter(handle, fieldnames=columns, extrasaction="ignore") - writer.writeheader() - writer.writerows(rows) - return qa_csv +def _path(refs: list[HubFileRef], part: str) -> Path: + return next(ref.path for ref in refs if ref.part == part) -@pytest.fixture -def release(tmp_path: Path) -> ARFBenchSource: - """A release with two metrics, their three charts, and one question about each metric count.""" - return ARFBenchSource( - qa_csv=_write_qa_csv(tmp_path, [_PRESENCE_ROW, _CORRELATION_ROW]), - series_dir=_write_series_dir(tmp_path), - chart_dir=_write_chart_dir(tmp_path), - ) +def test_the_connector_downloads_the_table_and_the_series_by_pattern(): + connector = ARFBenchConnector() -def _facts(annotated): - return {annotation.key: read_value(annotation.value) for annotation in annotated.annotations} + assert connector.metadata().dataset_id == "datadog/arfbench" + assert connector.values_backend == "parquet" + assert dict(connector.SOURCE_FILES) == {"questions": "arfbench-qa.csv", "series": "arfbench-ts-data/*.parquet"} def test_a_metric_record_keeps_every_rollup_per_tag_group(release): - dataset = ARFBenchConnector().convert([release]) - first, second = dataset.records[:2] + dataset = ARFBenchConnector().convert(release) + first, second = dataset.records assert [record.metadata["source_record_id"] for record in dataset.records] == [ "arfbench-35997_0", "arfbench-35997_1", - "arfbench-chart-35997_0-35997_1", - "arfbench-chart-35997_0", - "arfbench-chart-35997_1", ] # A source is a tag group, its signals that group's rollups, finest first. assert [source.name for source in first.sources] == ["dc:1", "dc:2", "dc:3"] @@ -162,7 +78,7 @@ def test_a_metric_record_keeps_every_rollup_per_tag_group(release): assert dc3_60s.to_arrow().to_pylist() == [9.0] assert isinstance(dc3_60s.time_axis, RegularAxis) and dc3_60s.time_axis.offset_us == 60 * 1_000_000 assert first.start_time.timestamp == START_US and second.start_time.timestamp == START_US + 30 * 1_000_000 - assert _facts(first) == { + assert facts(first) == { "metric_id": "35997_0", "incident_id": "35997", "query_name": "35997_9", @@ -170,32 +86,28 @@ def test_a_metric_record_keeps_every_rollup_per_tag_group(release): } -def test_each_question_gets_a_series_task_and_a_chart_task(release): - dataset = ARFBenchConnector().convert([release]) - tasks = dataset.get_all() +def test_each_question_is_a_classification_task_over_its_options(release): + dataset = ARFBenchConnector().convert(release) + presence, correlation = tasks = dataset.get_all() - # Two tasks per question, in the test split, paired by the question's row index. assert dataset.get_test() == tasks and dataset.get_train() == () - assert [task.metadata["question_index"] for task in tasks] == [0, 0, 1, 1] - assert [read_value(task.prompt) for task in tasks[2:]] == [_CORRELATION_ROW["question"]] * 2 - assert [read_values(task.targets) for task in tasks] == [("Yes",), ("Yes",), ("No",), ("No",)] - # The series task reads the cited metrics in citation order, the chart task the combined chart - # first and then each metric's own. - assert [record.metadata["source_record_id"] for record in tasks[2].inputs] == [ + assert [task.metadata["question_index"] for task in tasks] == [0, 1] + assert [read_value(task.prompt) for task in tasks] == [row["question"] for row in ROWS] + assert [read_values(task.targets) for task in tasks] == [("Yes",), ("No",)] + # The options are the task's label set, named after the question's row. + assert isinstance(presence, ClassificationTask) and isinstance(correlation, ClassificationTask) + assert [presence.label_set.name, correlation.label_set.name] == ["question_0", "question_1"] + assert [label.name for label in correlation.label_set.labels] == ["Yes", "No"] + assert not correlation.label_set.multi_label + # The task reads the cited metrics in citation order. + assert [record.metadata["source_record_id"] for record in correlation.inputs] == [ "arfbench-35997_0", "arfbench-35997_1", ] - assert [record.metadata["source_record_id"] for record in tasks[3].inputs] == [ - "arfbench-chart-35997_0-35997_1", - "arfbench-chart-35997_0", - "arfbench-chart-35997_1", - ] - assert [record.metadata["source_record_id"] for record in tasks[1].inputs] == ["arfbench-chart-35997_1"] - assert tasks[2].input_modalities == frozenset({InputModality.TEXT, InputModality.TIME_SERIES}) - assert tasks[3].input_modalities == frozenset({InputModality.TEXT, InputModality.IMAGE}) + assert [record.metadata["source_record_id"] for record in presence.inputs] == ["arfbench-35997_1"] + assert all(task.input_modalities == frozenset({InputModality.TEXT, InputModality.TIME_SERIES}) for task in tasks) assert all(task.split is Split.TEST for task in tasks) - assert _facts(tasks[3]) == { - "answer_options": ["Yes", "No"], + assert facts(correlation) == { "task_category": "Anomaly Correlation", "difficulty": "Tier 3", "interpolate_1": True, @@ -203,28 +115,32 @@ def test_each_question_gets_a_series_task_and_a_chart_task(release): } -def test_the_dataset_round_trips_through_storage(release, tmp_path): - dataset = ARFBenchConnector().convert([release]) +def test_the_dataset_round_trips_through_parquet_with_its_missing_timesteps(release, tmp_path): + connector = ARFBenchConnector() + dataset = connector.convert(release) dataset.derive_schema() - with TimeFWriter(tmp_path / "built", dataset, values_backend="zarr") as output: + with TimeFWriter(tmp_path / "built", dataset, values_backend=connector.values_backend) as output: output.write() with TimeFReader(open_local_version(tmp_path / "built/datadog/arfbench/1.0.0")) as reader: reader.verify() + assert reader.values_backend == "parquet" back = {record.metadata["source_record_id"]: record for record in reader.iter_records()} tasks = reader.read().get_all() - assert [source.name for source in back["arfbench-35997_0"].sources] == ["dc:1", "dc:2", "dc:3"] - assert back["arfbench-35997_0"].sources[0].signals[1].to_arrow().to_pylist() == [1.5, None, 2.5] - chart = back["arfbench-chart-35997_0"].signals[0] - assert chart.spec.spec_type == "rgba_image_2x3" - np.testing.assert_array_equal(chart.to_numpy()[0, 0, 0], [10, 20, 30, 42]) - assert [task.metadata["question_index"] for task in tasks] == [0, 0, 1, 1] + first = back["arfbench-35997_0"] + assert [source.name for source in first.sources] == ["dc:1", "dc:2", "dc:3"] + dc1_60s = first.sources[0].signals[1] + assert dc1_60s.to_arrow().to_pylist() == [1.5, None, 2.5] + assert dc1_60s.spec.missing is not None and dc1_60s.spec.missing.reason is MissingReason.NOT_RECORDED + assert first.sources[1].signals[0].to_arrow().to_pylist() == [5.0, None, 6.0] + assert [task.metadata["question_index"] for task in tasks] == [0, 1] + assert isinstance(tasks[1], ClassificationTask) + assert [label.name for label in tasks[1].label_set.labels] == ["Yes", "No"] @pytest.mark.parametrize( ("change", "message"), [ ({"query_group": "35997_2"}, r"question 2: no series file for cited metric 35997_2"), - ({"query_group": "35997_1, 35997_0"}, r"question 2: no chart for 35997_1-35997_0"), ({"correct_answer": "Maybe"}, r"data row 3: correct_answer 'Maybe' is not one of"), ({"options_str": "Yes"}, r"data row 3: options_str is not JSON"), ({"interpolate_1": "yes"}, r"data row 3: interpolate_1 must be 0 or 1"), @@ -233,7 +149,6 @@ def test_the_dataset_round_trips_through_storage(release, tmp_path): ], ids=[ "a cited metric without series", - "a pair without a combined chart", "an answer outside the options", "options that are not JSON", "a flag that is not 0 or 1", @@ -242,25 +157,31 @@ def test_the_dataset_round_trips_through_storage(release, tmp_path): ], ) def test_a_malformed_question_raises_with_its_location(release, change, message): - rows = [_PRESENCE_ROW, _CORRELATION_ROW, {**_CORRELATION_ROW, "Unnamed: 0": "2"} | change] - _write_qa_csv(release.qa_csv.parent, rows) + rows = [*ROWS, {**CORRELATION_ROW, "Unnamed: 0": "2"} | change] + write_qa_csv(_path(release, "questions").parent, rows) with pytest.raises(TimeNetConnectorError, match=rf"(arfbench-qa\.csv )?{message}"): - ARFBenchConnector().convert([release]) + ARFBenchConnector().convert(release) def test_an_unexpected_header_raises(release): - _write_qa_csv(release.qa_csv.parent, [_PRESENCE_ROW], columns=QUESTION_COLUMNS[1:]) + write_qa_csv(_path(release, "questions").parent, [PRESENCE_ROW], columns=QUESTION_COLUMNS[1:]) with pytest.raises(TimeNetConnectorError, match=r"arfbench-qa\.csv: expected columns"): - ARFBenchConnector().convert([release]) + ARFBenchConnector().convert(release) + + +def test_a_missing_part_raises(release): + with pytest.raises(TimeNetConnectorError, match=r"the ARFBench release lacks series"): + ARFBenchConnector().convert([ref for ref in release if ref.part == "questions"]) def test_a_row_off_the_rollup_grid_raises(release): - _series_file(release.series_dir / "35997_1_10.parquet", [(30, "", 0.5), (35, "", 0.6)]) + write_series_file(_path(release, "series") / "35997_1_10.parquet", [(30, "", 0.5), (35, "", 0.6)]) with pytest.raises(TimeNetConnectorError, match=r"35997_1_10\.parquet group '': rows are off the 10 s grid"): - ARFBenchConnector().convert([release]) + ARFBenchConnector().convert(release) def test_a_series_file_without_an_interval_raises(release): - (release.series_dir / "35997_1_10.parquet").rename(release.series_dir / "35997_1_10s.parquet") + series_dir = _path(release, "series") + (series_dir / "35997_1_10.parquet").rename(series_dir / "35997_1_10s.parquet") with pytest.raises(TimeNetConnectorError, match=r"35997_1_10s\.parquet is not named"): - ARFBenchConnector().convert([release]) + ARFBenchConnector().convert(release) diff --git a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/__init__.py b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/__init__.py new file mode 100644 index 000000000..17e3626e9 --- /dev/null +++ b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/__init__.py @@ -0,0 +1,6 @@ +"""The arfbench_plots connector package: its connector implementation and dataset card.""" + +from timenet_connectors.datasets.datadog.arfbench_plots.connector import ( + CONNECTOR as CONNECTOR, + ARFBenchPlotsConnector as ARFBenchPlotsConnector, +) diff --git a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/connector.py b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/connector.py new file mode 100644 index 000000000..f9354eaff --- /dev/null +++ b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/connector.py @@ -0,0 +1,76 @@ +"""Build the ARFBench charts as a child dataset of ``datadog/arfbench``. + +This dataset is a child of ``datadog/arfbench``. Its card declares that parent with the alias +``series``. It imports every series task of the parent, with the metric Records that the tasks cite. +It then adds one Record per chart in the release and one chart task per question. The imported +series continue to refer to the stored values of the parent, so this release stores only the +charts. The charts are image tensors, and only the Zarr backend can store them. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar + +from timenet.dataset import Record, Source, TimeFDataset +from timenet.types.reference_keys import ParentAlias +from timenet_connectors.bases.huggingface import BaseHuggingFaceConnector +from timenet_connectors.datasets.datadog.arfbench.connector import QA_CSV, REPO, REVISION, release_paths +from timenet_connectors.datasets.datadog.arfbench.questions import read_questions +from timenet_connectors.sources.images import image_signal + + +if TYPE_CHECKING: + from pathlib import Path + + from timenet.connectors import BuildContext + from timenet_connectors.bases.huggingface import HubFileRef + +CHART_DIR = "arfbench-images" # {metric}.png, and {metric}-{metric}.png for the pairs that the questions cite +SERIES_PARENT = ParentAlias("series") + + +def _create_chart_record(path: Path) -> Record: + """Create the Record of a chart, with the image as its one signal. + + Returns: + The Record, not yet added to a dataset. + """ + return Record( + metadata={"source_record_id": f"arfbench-chart-{path.stem}"}, + sources=(Source(name="ARFBench chart", signals=(image_signal(path, name="chart"),)),), + ) + + +class ARFBenchPlotsConnector(BaseHuggingFaceConnector): + """Connector for the ARFBench charts, as a child of the ``datadog/arfbench`` series.""" + + values_backend = "zarr" # the chart tensors need the Zarr backend + HF_REPO = REPO + SOURCE_REVISION = REVISION + SOURCE_FILES: ClassVar = {"questions": QA_CSV, "charts": f"{CHART_DIR}/*.png"} + + def compose(self, raw_refs: list[HubFileRef], context: BuildContext) -> TimeFDataset: + """Import the parent's series tasks, then add a Record for each chart and a chart task for each question. + + Args: + raw_refs: The question table and the chart directory from :meth:`download_async`. + context: The opened ``datadog/arfbench`` parent, under the alias ``series``. + + Returns: + The dataset: the series tasks first, then the chart tasks, paired by ``question_index``. + """ + self.check_refs(raw_refs) + paths = release_paths(raw_refs, "questions", "charts") + series = context.parents[SERIES_PARENT] + dataset = TimeFDataset(metadata=self.metadata()) + dataset.import_from(series, tasks=[task.id for task in series.iter_tasks()]) + chart_records = { + path.stem: dataset.add_record(record=_create_chart_record(path)) + for path in sorted(paths["charts"].glob("*.png")) + } + for question in read_questions(paths["questions"]): + dataset.add_task(task=question.create_chart_task(chart_records)) + return dataset + + +CONNECTOR = ARFBenchPlotsConnector diff --git a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/dataset.yaml b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/dataset.yaml new file mode 100644 index 000000000..d01d8f598 --- /dev/null +++ b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/dataset.yaml @@ -0,0 +1,65 @@ +# yaml-language-server: $schema=https://docs.timenet.ai/schemas/dataset-card-v1.schema.json +# +# ARFBench asks 750 multiple-choice questions about 142 metrics from 63 Datadog incidents. A metric +# is one monitored quantity, such as the request count of a service. Datadog splits a metric by +# tags. Thus a metric holds one series per tag group, for example one series per host. +# +# Datadog stores a metric at several rollups. A rollup is the metric on a fixed time grid, with one +# value per interval of 10, 60, 300, 1800, 3600, or 86400 seconds. The rollups of a metric are +# separate exports, not resamplings of one series. Their time ranges and tag groups differ. +# +# Each metric is one record, with one source per tag group and one signal per rollup. A consumer +# chooses the rollup. The paper and the charts use the finest rollup. The annotation +# `finest_interval_s` names this rollup. +# +# A rollup file lists the tag groups at the steps of its grid. Where the metric reported nothing +# for a group, the step holds a null or is not in the file. The dataset cuts each series to the +# window from the first value of its group to the last value. Thus a group with no value at a +# rollup has no series at that rollup. Inside the window, each such step is a missing timestep on +# the regular grid, with the reason `not_recorded` from the spec. Values are float64. The release +# names no units. The release stores naive epochs. The dataset reads them as UTC. The charts +# and the answer options use the same clock. +# +# This dataset is a child of `datadog/arfbench`. The name of this parent is `series`. The dataset +# imports the series tasks of the parent, together with the metric records that these tasks cite. +# The values of these records stay in the parent release. The dataset also adds the released +# charts. A chart is one record with the PNG as an RGBA image signal. There is one chart per +# metric, and one chart per metric pair that a question cites. The charts are tensors. Thus this +# dataset uses the Zarr values backend for the charts that it stores. +# +# Each question becomes two test tasks. One task is over the cited metric records. The other task +# is over the charts. For a metric pair, the combined chart comes first, then the chart of each +# metric. The series tasks come first, then the chart tasks. Thus the test split holds 1,500 +# tasks. Both tasks of a question carry its row as `question_index` in their metadata. +# +# The dataset keeps three items as released: +# - `query_name`. This identifier differs from the metric ID for 122 of the 142 metrics. The +# release does not explain it. +# - One end-time answer (row 191, metric 36464_0). This answer is 140 seconds after the end of its +# series. +# - 29 time options that are outside every rollup. +yaml_schema_version: 1 +dataset_id: datadog/arfbench-plots +dataset_version: 1.0.0 +name: ARFBench plots +description: >- + 750 multiple-choice questions about anomalies in 142 metrics from 63 Datadog incidents. + The dataset includes every published metric rollup and the released charts. +license: Apache-2.0 +source_url: https://huggingface.co/datasets/Datadog/ARFBench +citation: >- + Xie, S., Cohen, B., Goswami, M., Shen, J., Khwaja, E., Liu, C., Asker, D., Abou-Amal, O., and + Talwalkar, A. (2026). ARFBench: Benchmarking Time Series Question Answering Ability for Software + Incident Response. arXiv:2604.21199. https://arxiv.org/abs/2604.21199 +access: open +domains: + - observability +tags: + - observability + - anomaly-detection + - qa + - multiple-choice + - datadog + - huggingface +parents: + series: datadog/arfbench@1.0.0 diff --git a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/requirements.txt b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/requirements.txt new file mode 100644 index 000000000..7a9827be2 --- /dev/null +++ b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/requirements.txt @@ -0,0 +1,5 @@ +# huggingface_hub downloads the release from the HuggingFace Hub. Version 1.1 adds tqdm_class. +huggingface_hub>=1.1 +# Pillow decodes the chart PNGs. Zarr stores their tensor values. The connector base installs neither. +pillow>=11 +zarr>=3 diff --git a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/tests/__init__.py b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/tests/__init__.py new file mode 100644 index 000000000..9096a502e --- /dev/null +++ b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/tests/__init__.py @@ -0,0 +1 @@ +"""ARFBench plots connector tests.""" diff --git a/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/tests/test_connector.py b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/tests/test_connector.py new file mode 100644 index 000000000..58f14a353 --- /dev/null +++ b/packages/timenet-connectors/src/timenet_connectors/datasets/datadog/arfbench_plots/tests/test_connector.py @@ -0,0 +1,152 @@ +"""Check that the plots dataset composes on its ``datadog/arfbench`` parent. + +The dataset imports the series tasks. It adds a chart record per chart and a chart task per question. It stores +only the charts. It reads the series through the Parquet release of the parent. A question without its charts +raises an error.""" + +from collections.abc import Generator +from contextlib import contextmanager +from pathlib import Path + +import numpy as np +from PIL import Image +import pytest + +from timenet.dataset import TimeFDataset +from timenet.engine import open_parents, store_dataset +from timenet.errors import TimeNetConnectorError +from timenet.reader import TimeFReader +from timenet.registry import LocalRegistry, open_local_version +from timenet.types import ClassificationTask, InputModality, Split, read_values +from timenet_connectors.bases.huggingface import HubFileRef +from timenet_connectors.datasets.datadog.arfbench.connector import ARFBenchConnector +from timenet_connectors.datasets.datadog.arfbench.tests.release import ( + CORRELATION_ROW, + ROWS, + facts, + write_qa_csv, + write_release, +) +from timenet_connectors.datasets.datadog.arfbench_plots.connector import CHART_DIR, ARFBenchPlotsConnector + + +def _write_chart_dir(root: Path) -> Path: + chart_dir = root / CHART_DIR + chart_dir.mkdir() + for metric in ("35997_0", "35997_1"): + Image.new("RGBA", (3, 2), (10, 20, 30, 42)).save(chart_dir / f"{metric}.png") + Image.new("RGB", (6, 2), (255, 255, 255)).save(chart_dir / "35997_0-35997_1.png") + return chart_dir + + +@pytest.fixture +def registry(tmp_path: Path) -> Path: + """A local registry that holds the ``datadog/arfbench`` parent, built from the test release.""" + root = tmp_path / "registry" + release = tmp_path / "series" + release.mkdir() + parent = ARFBenchConnector() + store_dataset(parent.convert(write_release(release)), root, values_backend=parent.values_backend) + return root + + +@pytest.fixture +def charts(tmp_path: Path) -> list[HubFileRef]: + """The question table and the three charts of the test release.""" + root = tmp_path / "plots" + root.mkdir() + return [HubFileRef(write_qa_csv(root, ROWS), "questions"), HubFileRef(_write_chart_dir(root), "charts")] + + +@contextmanager +def _composed(registry: Path, refs: list[HubFileRef]) -> Generator[TimeFDataset]: + """Compose the plots dataset inside its parent context, as the engine does.""" + connector = ARFBenchPlotsConnector() + with open_parents(connector.metadata(), LocalRegistry(registry)) as context: + yield connector.compose(refs, context) + + +def test_the_connector_declares_the_series_dataset_as_its_parent(): + connector = ARFBenchPlotsConnector() + + assert connector.metadata().dataset_id == "datadog/arfbench-plots" + assert connector.values_backend == "zarr" + assert dict(connector.SOURCE_FILES) == {"questions": "arfbench-qa.csv", "charts": "arfbench-images/*.png"} + assert {str(alias): str(ref) for alias, ref in connector.metadata().parents.items()} == { + "series": "datadog/arfbench@1.0.0" + } + + +def test_the_child_imports_the_series_tasks_and_adds_a_chart_task_per_question(registry, charts): + with _composed(registry, charts) as dataset: + records = [record.metadata["source_record_id"] for record in dataset.records] + tasks = dataset.get_all() + + # The import brings the metric records in the order that the tasks first cite them. The charts come next. + assert records == [ + "arfbench-35997_1", + "arfbench-35997_0", + "arfbench-chart-35997_0-35997_1", + "arfbench-chart-35997_0", + "arfbench-chart-35997_1", + ] + # The series tasks come first, then the chart tasks. The row index of the question pairs the two tasks. + assert [task.metadata["question_index"] for task in tasks] == [0, 1, 0, 1] + assert [read_values(task.targets) for task in tasks] == [("Yes",), ("No",)] * 2 + series_correlation, chart_correlation = tasks[1], tasks[3] + assert series_correlation.input_modalities == frozenset({InputModality.TEXT, InputModality.TIME_SERIES}) + assert chart_correlation.input_modalities == frozenset({InputModality.TEXT, InputModality.IMAGE}) + # Both tasks of a question pick among its options, from the same label set. + assert isinstance(series_correlation, ClassificationTask) and isinstance(chart_correlation, ClassificationTask) + assert series_correlation.label_set == chart_correlation.label_set + assert series_correlation.label_set.name == "question_1" + # The chart task reads the combined chart first, then the chart of each metric. + assert [record.metadata["source_record_id"] for record in chart_correlation.inputs] == [ + "arfbench-chart-35997_0-35997_1", + "arfbench-chart-35997_0", + "arfbench-chart-35997_1", + ] + assert [record.metadata["source_record_id"] for record in tasks[2].inputs] == ["arfbench-chart-35997_1"] + assert all(task.split is Split.TEST for task in tasks) + assert facts(chart_correlation) == facts(series_correlation) + + +def test_the_child_stores_only_the_charts_and_reads_the_series_from_the_parent(registry, charts): + connector = ARFBenchPlotsConnector() + with _composed(registry, charts) as dataset: + dataset.derive_schema() + version_dir = store_dataset(dataset, registry, values_backend=connector.values_backend) + with TimeFReader(open_local_version(version_dir)) as reader: + reader.verify() + assert reader.values_backend == "zarr" + back = {record.metadata["source_record_id"]: record for record in reader.iter_records()} + tasks = reader.read().get_all() + # The series values stay in the Parquet release of the parent. The child names it as a payload source. + assert [str(source) for source in open_local_version(version_dir).manifest.payload_sources] == [ + "datadog/arfbench@1.0.0" + ] + # The child stores its control plane and the two chart shapes. It stores no series values of its own. + assert sorted(path.relative_to(version_dir).parts[0] for path in version_dir.iterdir()) == [ + "control.duckdb", + "manifest.json", + "time_series.zarr", + ] + assert sorted(path.name for path in (version_dir / "time_series.zarr").iterdir() if path.is_dir()) == [ + "rgb_image_2x6", + "rgba_image_2x3", + ] + assert back["arfbench-35997_0"].sources[0].signals[1].to_arrow().to_pylist() == [1.5, None, 2.5] + chart = back["arfbench-chart-35997_0"].signals[0] + assert chart.spec.spec_type == "rgba_image_2x3" + np.testing.assert_array_equal(chart.to_numpy()[0, 0, 0], [10, 20, 30, 42]) + assert [task.metadata["question_index"] for task in tasks] == [0, 1, 0, 1] + + +def test_a_pair_without_a_combined_chart_raises(registry, charts): + questions = next(ref.path for ref in charts if ref.part == "questions") + write_qa_csv(questions.parent, [*ROWS, {**CORRELATION_ROW, "Unnamed: 0": "2", "query_group": "35997_1, 35997_0"}]) + with ( + pytest.raises(TimeNetConnectorError, match=r"question 2: no chart for 35997_1-35997_0"), + _composed(registry, charts), + ): + pass