diff --git a/binder/filespec.ipynb b/binder/filespec.ipynb
index 9e5e20ec1..a52dbc883 100644
--- a/binder/filespec.ipynb
+++ b/binder/filespec.ipynb
@@ -15,7 +15,8 @@
"\n",
"The FileSpec system provides:\n",
"- **Type-safe data structures** with automatic validation\n",
- "- **Automatic format detection** for ROOT and Parquet files\n",
+ "- **Automatic format detection** for ROOT (TTree and RNTuple) and Parquet files\n",
+ "- **Switchable execution backends** (`dask`, `iterative`, `futures`) for preprocessing\n",
"- **Seamless integration** with existing Coffea functions\n",
"- **JSON serialization/deserialization** for data persistence\n",
"- **Automatic promotion** between optional and concrete specifications\n",
@@ -37,7 +38,14 @@
"cell_type": "code",
"execution_count": 1,
"id": "f6c75a47",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:31:42.715172Z",
+ "iopub.status.busy": "2026-07-02T23:31:42.714932Z",
+ "iopub.status.idle": "2026-07-02T23:32:07.577593Z",
+ "shell.execute_reply": "2026-07-02T23:32:07.577053Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -69,6 +77,9 @@
" \n",
" # Dataset manipulation functions\n",
" preprocess,\n",
+ " preprocess_parquet,\n",
+ " preprocess_rntuple,\n",
+ " FuturesBackend,\n",
" apply_to_fileset,\n",
" max_chunks,\n",
" max_chunks_per_file,\n",
@@ -109,7 +120,14 @@
"cell_type": "code",
"execution_count": 2,
"id": "2ecaf55b",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:07.579034Z",
+ "iopub.status.busy": "2026-07-02T23:32:07.578839Z",
+ "iopub.status.idle": "2026-07-02T23:32:07.586602Z",
+ "shell.execute_reply": "2026-07-02T23:32:07.586149Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -212,7 +230,14 @@
"cell_type": "code",
"execution_count": 3,
"id": "d59e690c",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:07.587784Z",
+ "iopub.status.busy": "2026-07-02T23:32:07.587711Z",
+ "iopub.status.idle": "2026-07-02T23:32:07.592336Z",
+ "shell.execute_reply": "2026-07-02T23:32:07.591877Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -314,7 +339,14 @@
"cell_type": "code",
"execution_count": 4,
"id": "85d98abf",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:07.593469Z",
+ "iopub.status.busy": "2026-07-02T23:32:07.593381Z",
+ "iopub.status.idle": "2026-07-02T23:32:07.599317Z",
+ "shell.execute_reply": "2026-07-02T23:32:07.598796Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -465,7 +497,14 @@
"cell_type": "code",
"execution_count": 5,
"id": "38ee2f96",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:07.600473Z",
+ "iopub.status.busy": "2026-07-02T23:32:07.600391Z",
+ "iopub.status.idle": "2026-07-02T23:32:07.605639Z",
+ "shell.execute_reply": "2026-07-02T23:32:07.605231Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -519,26 +558,26 @@
"3 validation errors for CoffeaROOTFileSpec\n",
"steps\n",
" Field required [type=missing, input_value={'object_path': 'Events'}, input_type=dict]\n",
- " For further information visit https://errors.pydantic.dev/2.11/v/missing\n",
+ " For further information visit https://errors.pydantic.dev/2.12/v/missing\n",
"num_entries\n",
" Field required [type=missing, input_value={'object_path': 'Events'}, input_type=dict]\n",
- " For further information visit https://errors.pydantic.dev/2.11/v/missing\n",
+ " For further information visit https://errors.pydantic.dev/2.12/v/missing\n",
"uuid\n",
" Field required [type=missing, input_value={'object_path': 'Events'}, input_type=dict]\n",
- " For further information visit https://errors.pydantic.dev/2.11/v/missing\n",
+ " For further information visit https://errors.pydantic.dev/2.12/v/missing\n",
"\n"
],
"text/plain": [
"\u001b[1;36m3\u001b[0m validation errors for CoffeaROOTFileSpec\n",
"steps\n",
" Field required \u001b[1m[\u001b[0m\u001b[33mtype\u001b[0m=\u001b[35mmissing\u001b[0m, \u001b[33minput_value\u001b[0m=\u001b[1m{\u001b[0m\u001b[32m'object_path'\u001b[0m: \u001b[32m'Events'\u001b[0m\u001b[1m}\u001b[0m, \u001b[33minput_type\u001b[0m=\u001b[35mdict\u001b[0m\u001b[1m]\u001b[0m\n",
- " For further information visit \u001b[4;94mhttps://errors.pydantic.dev/2.11/v/missing\u001b[0m\n",
+ " For further information visit \u001b[4;94mhttps://errors.pydantic.dev/2.12/v/missing\u001b[0m\n",
"num_entries\n",
" Field required \u001b[1m[\u001b[0m\u001b[33mtype\u001b[0m=\u001b[35mmissing\u001b[0m, \u001b[33minput_value\u001b[0m=\u001b[1m{\u001b[0m\u001b[32m'object_path'\u001b[0m: \u001b[32m'Events'\u001b[0m\u001b[1m}\u001b[0m, \u001b[33minput_type\u001b[0m=\u001b[35mdict\u001b[0m\u001b[1m]\u001b[0m\n",
- " For further information visit \u001b[4;94mhttps://errors.pydantic.dev/2.11/v/missing\u001b[0m\n",
+ " For further information visit \u001b[4;94mhttps://errors.pydantic.dev/2.12/v/missing\u001b[0m\n",
"uuid\n",
" Field required \u001b[1m[\u001b[0m\u001b[33mtype\u001b[0m=\u001b[35mmissing\u001b[0m, \u001b[33minput_value\u001b[0m=\u001b[1m{\u001b[0m\u001b[32m'object_path'\u001b[0m: \u001b[32m'Events'\u001b[0m\u001b[1m}\u001b[0m, \u001b[33minput_type\u001b[0m=\u001b[35mdict\u001b[0m\u001b[1m]\u001b[0m\n",
- " For further information visit \u001b[4;94mhttps://errors.pydantic.dev/2.11/v/missing\u001b[0m\n"
+ " For further information visit \u001b[4;94mhttps://errors.pydantic.dev/2.12/v/missing\u001b[0m\n"
]
},
"metadata": {},
@@ -575,7 +614,14 @@
"cell_type": "code",
"execution_count": 6,
"id": "694daed3",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:07.606724Z",
+ "iopub.status.busy": "2026-07-02T23:32:07.606660Z",
+ "iopub.status.idle": "2026-07-02T23:32:07.611162Z",
+ "shell.execute_reply": "2026-07-02T23:32:07.610766Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -699,7 +745,14 @@
"cell_type": "code",
"execution_count": 7,
"id": "3cfd0d6a",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:07.612312Z",
+ "iopub.status.busy": "2026-07-02T23:32:07.612256Z",
+ "iopub.status.idle": "2026-07-02T23:32:07.621771Z",
+ "shell.execute_reply": "2026-07-02T23:32:07.621231Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -968,7 +1021,14 @@
"cell_type": "code",
"execution_count": 8,
"id": "b2db2dcd",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:07.622914Z",
+ "iopub.status.busy": "2026-07-02T23:32:07.622852Z",
+ "iopub.status.idle": "2026-07-02T23:32:07.627539Z",
+ "shell.execute_reply": "2026-07-02T23:32:07.627196Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -1107,7 +1167,14 @@
"cell_type": "code",
"execution_count": 9,
"id": "b31c13f5",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:07.628785Z",
+ "iopub.status.busy": "2026-07-02T23:32:07.628726Z",
+ "iopub.status.idle": "2026-07-02T23:32:07.633925Z",
+ "shell.execute_reply": "2026-07-02T23:32:07.633595Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -1243,7 +1310,14 @@
"cell_type": "code",
"execution_count": 10,
"id": "76aa5992",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:07.635068Z",
+ "iopub.status.busy": "2026-07-02T23:32:07.635001Z",
+ "iopub.status.idle": "2026-07-02T23:32:07.640202Z",
+ "shell.execute_reply": "2026-07-02T23:32:07.639734Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -1257,7 +1331,7 @@
"data": {
"text/html": [
"DatasetSpec(\n",
- " files=InputFiles(\n",
+ " files=PreprocessedFiles(\n",
" root={\n",
" 'processed_data_1.root': CoffeaROOTFileSpec(\n",
" object_path='Events',\n",
@@ -1290,7 +1364,7 @@
],
"text/plain": [
"\u001b[1;35mDatasetSpec\u001b[0m\u001b[1m(\u001b[0m\n",
- " \u001b[33mfiles\u001b[0m=\u001b[1;35mInputFiles\u001b[0m\u001b[1m(\u001b[0m\n",
+ " \u001b[33mfiles\u001b[0m=\u001b[1;35mPreprocessedFiles\u001b[0m\u001b[1m(\u001b[0m\n",
" \u001b[33mroot\u001b[0m=\u001b[1m{\u001b[0m\n",
" \u001b[32m'processed_data_1.root'\u001b[0m: \u001b[1;35mCoffeaROOTFileSpec\u001b[0m\u001b[1m(\u001b[0m\n",
" \u001b[33mobject_path\u001b[0m=\u001b[32m'Events'\u001b[0m,\n",
@@ -1373,7 +1447,14 @@
"cell_type": "code",
"execution_count": 11,
"id": "6af14880",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:07.641291Z",
+ "iopub.status.busy": "2026-07-02T23:32:07.641225Z",
+ "iopub.status.idle": "2026-07-02T23:32:07.643911Z",
+ "shell.execute_reply": "2026-07-02T23:32:07.643466Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -1382,7 +1463,7 @@
"=== Mixed Format Datasets ===\n",
"Validation error for mixed format dataset: 1 validation error for DatasetSpec\n",
" Value error, format: format must be one of {'root', 'parquet'} [type=value_error, input_value={'files': {'data.root': C...selected_entries=2000)}}, input_type=dict]\n",
- " For further information visit https://errors.pydantic.dev/2.11/v/value_error\n"
+ " For further information visit https://errors.pydantic.dev/2.12/v/value_error\n"
]
}
],
@@ -1421,7 +1502,14 @@
"cell_type": "code",
"execution_count": 12,
"id": "98564dea",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:07.644963Z",
+ "iopub.status.busy": "2026-07-02T23:32:07.644901Z",
+ "iopub.status.idle": "2026-07-02T23:32:07.650382Z",
+ "shell.execute_reply": "2026-07-02T23:32:07.650042Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -1589,7 +1677,14 @@
"cell_type": "code",
"execution_count": 13,
"id": "7109b611",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:07.651575Z",
+ "iopub.status.busy": "2026-07-02T23:32:07.651517Z",
+ "iopub.status.idle": "2026-07-02T23:32:07.660381Z",
+ "shell.execute_reply": "2026-07-02T23:32:07.659975Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -1848,7 +1943,14 @@
"cell_type": "code",
"execution_count": 14,
"id": "376a6765",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:07.661499Z",
+ "iopub.status.busy": "2026-07-02T23:32:07.661424Z",
+ "iopub.status.idle": "2026-07-02T23:32:07.664025Z",
+ "shell.execute_reply": "2026-07-02T23:32:07.663682Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -1894,7 +1996,14 @@
"cell_type": "code",
"execution_count": 15,
"id": "fc3b6e90",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:07.665153Z",
+ "iopub.status.busy": "2026-07-02T23:32:07.665075Z",
+ "iopub.status.idle": "2026-07-02T23:32:07.667908Z",
+ "shell.execute_reply": "2026-07-02T23:32:07.667560Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -1964,7 +2073,14 @@
"cell_type": "code",
"execution_count": 16,
"id": "8d360bf8",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:07.669063Z",
+ "iopub.status.busy": "2026-07-02T23:32:07.668989Z",
+ "iopub.status.idle": "2026-07-02T23:32:09.782645Z",
+ "shell.execute_reply": "2026-07-02T23:32:09.782076Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -2100,7 +2216,13 @@
"name": "stdout",
"output_type": "stream",
"text": [
- "Preprocessing the fileset...\n",
+ "Preprocessing the fileset...\n"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
"Fileset after preprocessing (excluding compressed_form string):\n"
]
},
@@ -2329,6 +2451,279 @@
"#rich.print({dname: {k: v for k, v in dataset_updated[dname].files.items() if k not in dataset_runnable[dname].files} for dname in dataset_updated})\n"
]
},
+ {
+ "cell_type": "markdown",
+ "id": "c6525825",
+ "metadata": {},
+ "source": [
+ "### 5.2 Switchable execution backends\n",
+ "\n",
+ "`preprocess` (and the format-specific `preprocess_root` / `preprocess_rntuple` / `preprocess_parquet`) accept a `backend=` argument that selects how the per-file work is executed:\n",
+ "\n",
+ "- `\"dask\"` (default): builds a dask-awkward task graph, which scales out to a cluster via the `scheduler` argument.\n",
+ "- `\"iterative\"`: immediate, synchronous, single-process execution. **No dask required.**\n",
+ "- `\"futures\"`: a `concurrent.futures` pool — threads by default (preprocessing is I/O-bound). Pass a `FuturesBackend(workers=N)` or `FuturesBackend(use_processes=True)` instance for finer control. **No dask required.**\n",
+ "\n",
+ "The dask-free backends produce results identical to the dask backend for both ROOT and Parquet inputs. TTree/RNTuple form extraction uses uproot's own (non-dask) form builder, so `save_form=True` also works without dask."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 17,
+ "id": "2e8b6aa5",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:09.784246Z",
+ "iopub.status.busy": "2026-07-02T23:32:09.784065Z",
+ "iopub.status.idle": "2026-07-02T23:32:10.556751Z",
+ "shell.execute_reply": "2026-07-02T23:32:10.556198Z"
+ }
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "iterative and futures backends agree: True\n"
+ ]
+ },
+ {
+ "data": {
+ "text/html": [
+ "{\n",
+ " 'ZJets': {\n",
+ " 'files': {\n",
+ " 'https://raw.githubusercontent.com/scikit-hep/coffea/master/tests/samples/nano_dy.root': {\n",
+ " 'object_path': 'Events',\n",
+ " 'steps': [[0, 20], [20, 40]],\n",
+ " 'num_entries': 40,\n",
+ " 'format': 'root',\n",
+ " 'lfn': None,\n",
+ " 'pfn': None,\n",
+ " 'uuid': 'a9490124-3648-11ea-89e9-f5b55c90beef',\n",
+ " 'num_selected_entries': 40\n",
+ " }\n",
+ " },\n",
+ " 'metadata': {},\n",
+ " 'format': 'root',\n",
+ " 'did': None\n",
+ " }\n",
+ "}\n",
+ "\n"
+ ],
+ "text/plain": [
+ "\u001b[1m{\u001b[0m\n",
+ " \u001b[32m'ZJets'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ " \u001b[32m'files'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ " \u001b[32m'https://raw.githubusercontent.com/scikit-hep/coffea/master/tests/samples/nano_dy.root'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ " \u001b[32m'object_path'\u001b[0m: \u001b[32m'Events'\u001b[0m,\n",
+ " \u001b[32m'steps'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m0\u001b[0m, \u001b[1;36m20\u001b[0m\u001b[1m]\u001b[0m, \u001b[1m[\u001b[0m\u001b[1;36m20\u001b[0m, \u001b[1;36m40\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m,\n",
+ " \u001b[32m'num_entries'\u001b[0m: \u001b[1;36m40\u001b[0m,\n",
+ " \u001b[32m'format'\u001b[0m: \u001b[32m'root'\u001b[0m,\n",
+ " \u001b[32m'lfn'\u001b[0m: \u001b[3;35mNone\u001b[0m,\n",
+ " \u001b[32m'pfn'\u001b[0m: \u001b[3;35mNone\u001b[0m,\n",
+ " \u001b[32m'uuid'\u001b[0m: \u001b[32m'a9490124-3648-11ea-89e9-f5b55c90beef'\u001b[0m,\n",
+ " \u001b[32m'num_selected_entries'\u001b[0m: \u001b[1;36m40\u001b[0m\n",
+ " \u001b[1m}\u001b[0m\n",
+ " \u001b[1m}\u001b[0m,\n",
+ " \u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m,\n",
+ " \u001b[32m'format'\u001b[0m: \u001b[32m'root'\u001b[0m,\n",
+ " \u001b[32m'did'\u001b[0m: \u001b[3;35mNone\u001b[0m\n",
+ " \u001b[1m}\u001b[0m\n",
+ "\u001b[1m}\u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "# 5.2 The same preprocessing, run with different backends\n",
+ "root_fileset = DataGroupSpec({\n",
+ " \"ZJets\": {\"files\": [\"https://raw.githubusercontent.com/scikit-hep/coffea/master/tests/samples/nano_dy.root:Events\"]},\n",
+ "})\n",
+ "\n",
+ "available_iterative, _ = preprocess(root_fileset, step_size=20, save_form=True, backend=\"iterative\")\n",
+ "available_futures, _ = preprocess(root_fileset, step_size=20, save_form=True, backend=FuturesBackend(workers=2))\n",
+ "\n",
+ "print(\"iterative and futures backends agree:\", available_iterative == available_futures)\n",
+ "rich.print({k: v.model_dump(exclude=\"compressed_form\") for k, v in available_iterative.items()})"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "b2071c3f",
+ "metadata": {},
+ "source": [
+ "### 5.3 Preprocessing Parquet datasets\n",
+ "\n",
+ "The file format is detected automatically from the filename, so `preprocess` handles Parquet datasets transparently; `preprocess_parquet` is the format-specific entry point. Unlike ROOT, Parquet files do not take an `object_path`, and steps can optionally follow the Parquet row groups with `use_row_groups=True`."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 18,
+ "id": "a60ff4e0",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:10.558360Z",
+ "iopub.status.busy": "2026-07-02T23:32:10.558268Z",
+ "iopub.status.idle": "2026-07-02T23:32:12.047819Z",
+ "shell.execute_reply": "2026-07-02T23:32:12.047180Z"
+ }
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Detected format: parquet\n"
+ ]
+ },
+ {
+ "data": {
+ "text/html": [
+ "{\n",
+ " 'ZJets': {\n",
+ " 'files': {\n",
+ " 'https://raw.githubusercontent.com/scikit-hep/coffea/master/tests/samples/nano_dy.parquet': {\n",
+ " 'object_path': None,\n",
+ " 'steps': [[0, 40]],\n",
+ " 'num_entries': 40,\n",
+ " 'format': 'parquet',\n",
+ " 'lfn': None,\n",
+ " 'pfn': None,\n",
+ " 'uuid': '24be1d49dc0e1cc4ae4a0ee765e3d17c27662d075bf923aeaade5fbec7e29fca',\n",
+ " 'num_selected_entries': 40\n",
+ " }\n",
+ " },\n",
+ " 'metadata': {},\n",
+ " 'format': 'parquet',\n",
+ " 'did': None\n",
+ " }\n",
+ "}\n",
+ "\n"
+ ],
+ "text/plain": [
+ "\u001b[1m{\u001b[0m\n",
+ " \u001b[32m'ZJets'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ " \u001b[32m'files'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ " \u001b[32m'https://raw.githubusercontent.com/scikit-hep/coffea/master/tests/samples/nano_dy.parquet'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ " \u001b[32m'object_path'\u001b[0m: \u001b[3;35mNone\u001b[0m,\n",
+ " \u001b[32m'steps'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m0\u001b[0m, \u001b[1;36m40\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m,\n",
+ " \u001b[32m'num_entries'\u001b[0m: \u001b[1;36m40\u001b[0m,\n",
+ " \u001b[32m'format'\u001b[0m: \u001b[32m'parquet'\u001b[0m,\n",
+ " \u001b[32m'lfn'\u001b[0m: \u001b[3;35mNone\u001b[0m,\n",
+ " \u001b[32m'pfn'\u001b[0m: \u001b[3;35mNone\u001b[0m,\n",
+ " \u001b[32m'uuid'\u001b[0m: \u001b[32m'24be1d49dc0e1cc4ae4a0ee765e3d17c27662d075bf923aeaade5fbec7e29fca'\u001b[0m,\n",
+ " \u001b[32m'num_selected_entries'\u001b[0m: \u001b[1;36m40\u001b[0m\n",
+ " \u001b[1m}\u001b[0m\n",
+ " \u001b[1m}\u001b[0m,\n",
+ " \u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m,\n",
+ " \u001b[32m'format'\u001b[0m: \u001b[32m'parquet'\u001b[0m,\n",
+ " \u001b[32m'did'\u001b[0m: \u001b[3;35mNone\u001b[0m\n",
+ " \u001b[1m}\u001b[0m\n",
+ "\u001b[1m}\u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "# 5.3 Preprocessing a Parquet dataset (note: parquet files have no object_path)\n",
+ "parquet_fileset = DataGroupSpec({\n",
+ " \"ZJets\": {\"files\": [\"https://raw.githubusercontent.com/scikit-hep/coffea/master/tests/samples/nano_dy.parquet\"]},\n",
+ "})\n",
+ "\n",
+ "parquet_available, _ = preprocess_parquet(parquet_fileset, step_size=20000, save_form=True, backend=\"iterative\")\n",
+ "print(\"Detected format:\", parquet_available[\"ZJets\"].format)\n",
+ "rich.print({k: v.model_dump(exclude=\"compressed_form\") for k, v in parquet_available.items()})"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "86c1f41a",
+ "metadata": {},
+ "source": [
+ "### 5.4 Preprocessing RNTuple datasets\n",
+ "\n",
+ "RNTuple-backed ROOT files are auto-detected: `preprocess` and `preprocess_root` transparently handle both TTrees and RNTuples. Use `preprocess_rntuple` when you want to *assert* that every object is an RNTuple — it raises a `ValueError` if it encounters a TTree."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 19,
+ "id": "c54655a5",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:12.049879Z",
+ "iopub.status.busy": "2026-07-02T23:32:12.049784Z",
+ "iopub.status.idle": "2026-07-02T23:32:13.451399Z",
+ "shell.execute_reply": "2026-07-02T23:32:13.450909Z"
+ }
+ },
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "{\n",
+ " 'ZJets': {\n",
+ " 'files': {\n",
+ " 'https://raw.githubusercontent.com/scikit-hep/coffea/master/tests/samples/nano_dy_rntuple.root': {\n",
+ " 'object_path': 'Events',\n",
+ " 'steps': [[0, 20], [20, 40]],\n",
+ " 'num_entries': 40,\n",
+ " 'format': 'root',\n",
+ " 'lfn': None,\n",
+ " 'pfn': None,\n",
+ " 'uuid': '5aa41eba-a9f5-11f0-ba63-0f00a8c0beef',\n",
+ " 'num_selected_entries': 40\n",
+ " }\n",
+ " },\n",
+ " 'metadata': {},\n",
+ " 'format': 'root',\n",
+ " 'did': None\n",
+ " }\n",
+ "}\n",
+ "\n"
+ ],
+ "text/plain": [
+ "\u001b[1m{\u001b[0m\n",
+ " \u001b[32m'ZJets'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ " \u001b[32m'files'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ " \u001b[32m'https://raw.githubusercontent.com/scikit-hep/coffea/master/tests/samples/nano_dy_rntuple.root'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ " \u001b[32m'object_path'\u001b[0m: \u001b[32m'Events'\u001b[0m,\n",
+ " \u001b[32m'steps'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m0\u001b[0m, \u001b[1;36m20\u001b[0m\u001b[1m]\u001b[0m, \u001b[1m[\u001b[0m\u001b[1;36m20\u001b[0m, \u001b[1;36m40\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m,\n",
+ " \u001b[32m'num_entries'\u001b[0m: \u001b[1;36m40\u001b[0m,\n",
+ " \u001b[32m'format'\u001b[0m: \u001b[32m'root'\u001b[0m,\n",
+ " \u001b[32m'lfn'\u001b[0m: \u001b[3;35mNone\u001b[0m,\n",
+ " \u001b[32m'pfn'\u001b[0m: \u001b[3;35mNone\u001b[0m,\n",
+ " \u001b[32m'uuid'\u001b[0m: \u001b[32m'5aa41eba-a9f5-11f0-ba63-0f00a8c0beef'\u001b[0m,\n",
+ " \u001b[32m'num_selected_entries'\u001b[0m: \u001b[1;36m40\u001b[0m\n",
+ " \u001b[1m}\u001b[0m\n",
+ " \u001b[1m}\u001b[0m,\n",
+ " \u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m,\n",
+ " \u001b[32m'format'\u001b[0m: \u001b[32m'root'\u001b[0m,\n",
+ " \u001b[32m'did'\u001b[0m: \u001b[3;35mNone\u001b[0m\n",
+ " \u001b[1m}\u001b[0m\n",
+ "\u001b[1m}\u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "# 5.4 Preprocessing an RNTuple dataset (auto-detected; preprocess_rntuple enforces RNTuple-only)\n",
+ "rntuple_fileset = DataGroupSpec({\n",
+ " \"ZJets\": {\"files\": [\"https://raw.githubusercontent.com/scikit-hep/coffea/master/tests/samples/nano_dy_rntuple.root:Events\"]},\n",
+ "})\n",
+ "\n",
+ "rntuple_available, _ = preprocess_rntuple(rntuple_fileset, step_size=20, save_form=True, backend=\"iterative\")\n",
+ "rich.print({k: v.model_dump(exclude=\"compressed_form\") for k, v in rntuple_available.items()})"
+ ]
+ },
{
"cell_type": "markdown",
"id": "6b974346",
@@ -2341,21 +2736,46 @@
},
{
"cell_type": "code",
- "execution_count": 17,
+ "execution_count": 20,
"id": "39a43fa8",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:13.453038Z",
+ "iopub.status.busy": "2026-07-02T23:32:13.452910Z",
+ "iopub.status.idle": "2026-07-02T23:32:21.747719Z",
+ "shell.execute_reply": "2026-07-02T23:32:21.747218Z"
+ }
+ },
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
- "/Users/nmangane/scikit-hep-dev-4/coffea/src/coffea/nanoevents/schemas/nanoaod.py:264: RuntimeWarning: Missing cross-reference index for LowPtElectron_electronIdx => Electron\n",
+ "/Users/nmangane/servicex_claude/coffea/src/coffea/nanoevents/schemas/nanoaod.py:297: RuntimeWarning: Missing cross-reference index for LowPtElectron_electronIdx => Electron\n",
+ " warnings.warn(\n",
+ "/Users/nmangane/servicex_claude/coffea/src/coffea/nanoevents/schemas/nanoaod.py:297: RuntimeWarning: Missing cross-reference index for LowPtElectron_genPartIdx => GenPart\n",
+ " warnings.warn(\n",
+ "/Users/nmangane/servicex_claude/coffea/src/coffea/nanoevents/schemas/nanoaod.py:297: RuntimeWarning: Missing cross-reference index for LowPtElectron_photonIdx => Photon\n",
" warnings.warn(\n",
- "/Users/nmangane/scikit-hep-dev-4/coffea/src/coffea/nanoevents/schemas/nanoaod.py:264: RuntimeWarning: Missing cross-reference index for LowPtElectron_genPartIdx => GenPart\n",
+ "/Users/nmangane/servicex_claude/coffea/src/coffea/nanoevents/schemas/nanoaod.py:297: RuntimeWarning: Missing cross-reference index for FatJet_genJetAK8Idx => GenJetAK8\n",
" warnings.warn(\n",
- "/Users/nmangane/scikit-hep-dev-4/coffea/src/coffea/nanoevents/schemas/nanoaod.py:264: RuntimeWarning: Missing cross-reference index for LowPtElectron_photonIdx => Photon\n",
+ "/Users/nmangane/servicex_claude/coffea/src/coffea/nanoevents/schemas/nanoaod.py:336: RuntimeWarning: Branch Photon_mass already exists but its values will be replaced with 0.0\n",
" warnings.warn(\n",
- "/Users/nmangane/scikit-hep-dev-4/coffea/src/coffea/nanoevents/schemas/nanoaod.py:264: RuntimeWarning: Missing cross-reference index for FatJet_genJetAK8Idx => GenJetAK8\n",
+ "/Users/nmangane/servicex_claude/coffea/src/coffea/nanoevents/schemas/nanoaod.py:336: RuntimeWarning: Branch Photon_charge already exists but its values will be replaced with 0.0\n",
+ " warnings.warn(\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/Users/nmangane/servicex_claude/coffea/src/coffea/nanoevents/schemas/nanoaod.py:297: RuntimeWarning: Missing cross-reference index for LowPtElectron_electronIdx => Electron\n",
+ " warnings.warn(\n",
+ "/Users/nmangane/servicex_claude/coffea/src/coffea/nanoevents/schemas/nanoaod.py:297: RuntimeWarning: Missing cross-reference index for LowPtElectron_photonIdx => Photon\n",
+ " warnings.warn(\n",
+ "/Users/nmangane/servicex_claude/coffea/src/coffea/nanoevents/schemas/nanoaod.py:336: RuntimeWarning: Branch Photon_mass already exists but its values will be replaced with 0.0\n",
+ " warnings.warn(\n",
+ "/Users/nmangane/servicex_claude/coffea/src/coffea/nanoevents/schemas/nanoaod.py:336: RuntimeWarning: Branch Photon_charge already exists but its values will be replaced with 0.0\n",
" warnings.warn(\n"
]
},
@@ -2453,9 +2873,16 @@
},
{
"cell_type": "code",
- "execution_count": 18,
+ "execution_count": 21,
"id": "e7919395",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:21.749318Z",
+ "iopub.status.busy": "2026-07-02T23:32:21.749215Z",
+ "iopub.status.idle": "2026-07-02T23:32:21.753528Z",
+ "shell.execute_reply": "2026-07-02T23:32:21.753161Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -2509,9 +2936,16 @@
},
{
"cell_type": "code",
- "execution_count": 19,
+ "execution_count": 22,
"id": "d860f60a",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:21.754954Z",
+ "iopub.status.busy": "2026-07-02T23:32:21.754878Z",
+ "iopub.status.idle": "2026-07-02T23:32:21.758814Z",
+ "shell.execute_reply": "2026-07-02T23:32:21.758442Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -2554,9 +2988,16 @@
},
{
"cell_type": "code",
- "execution_count": 20,
+ "execution_count": 23,
"id": "01044461",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:21.760066Z",
+ "iopub.status.busy": "2026-07-02T23:32:21.759987Z",
+ "iopub.status.idle": "2026-07-02T23:32:21.763465Z",
+ "shell.execute_reply": "2026-07-02T23:32:21.762990Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -2590,9 +3031,16 @@
},
{
"cell_type": "code",
- "execution_count": 21,
+ "execution_count": 24,
"id": "ee931c3e",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:21.764558Z",
+ "iopub.status.busy": "2026-07-02T23:32:21.764499Z",
+ "iopub.status.idle": "2026-07-02T23:32:21.767574Z",
+ "shell.execute_reply": "2026-07-02T23:32:21.767030Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -2655,9 +3103,16 @@
},
{
"cell_type": "code",
- "execution_count": 22,
+ "execution_count": 25,
"id": "bc8ca29a",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:21.768739Z",
+ "iopub.status.busy": "2026-07-02T23:32:21.768668Z",
+ "iopub.status.idle": "2026-07-02T23:32:21.775571Z",
+ "shell.execute_reply": "2026-07-02T23:32:21.775079Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -2753,9 +3208,16 @@
},
{
"cell_type": "code",
- "execution_count": 23,
+ "execution_count": 26,
"id": "a6237760",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:21.776638Z",
+ "iopub.status.busy": "2026-07-02T23:32:21.776576Z",
+ "iopub.status.idle": "2026-07-02T23:32:21.782356Z",
+ "shell.execute_reply": "2026-07-02T23:32:21.781926Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -2811,9 +3273,16 @@
},
{
"cell_type": "code",
- "execution_count": 24,
+ "execution_count": 27,
"id": "60f065ea",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:21.783482Z",
+ "iopub.status.busy": "2026-07-02T23:32:21.783408Z",
+ "iopub.status.idle": "2026-07-02T23:32:21.786933Z",
+ "shell.execute_reply": "2026-07-02T23:32:21.786436Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -2873,9 +3342,16 @@
},
{
"cell_type": "code",
- "execution_count": 25,
+ "execution_count": 28,
"id": "f22002f4",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:21.788092Z",
+ "iopub.status.busy": "2026-07-02T23:32:21.788023Z",
+ "iopub.status.idle": "2026-07-02T23:32:21.795655Z",
+ "shell.execute_reply": "2026-07-02T23:32:21.795224Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -2893,7 +3369,7 @@
"DataGroupSpec(\n",
" root={\n",
" 'higgs_m125': DatasetSpec(\n",
- " files=InputFiles(\n",
+ " files=PreprocessedFiles(\n",
" root={\n",
" 'higgs_m125_part0.root': CoffeaROOTFileSpec(\n",
" object_path='Events',\n",
@@ -2913,7 +3389,7 @@
" did=None\n",
" ),\n",
" 'higgs_m200': DatasetSpec(\n",
- " files=InputFiles(\n",
+ " files=PreprocessedFiles(\n",
" root={\n",
" 'higgs_m200_part0.root': CoffeaROOTFileSpec(\n",
" object_path='Events',\n",
@@ -2940,7 +3416,7 @@
"\u001b[1;35mDataGroupSpec\u001b[0m\u001b[1m(\u001b[0m\n",
" \u001b[33mroot\u001b[0m=\u001b[1m{\u001b[0m\n",
" \u001b[32m'higgs_m125'\u001b[0m: \u001b[1;35mDatasetSpec\u001b[0m\u001b[1m(\u001b[0m\n",
- " \u001b[33mfiles\u001b[0m=\u001b[1;35mInputFiles\u001b[0m\u001b[1m(\u001b[0m\n",
+ " \u001b[33mfiles\u001b[0m=\u001b[1;35mPreprocessedFiles\u001b[0m\u001b[1m(\u001b[0m\n",
" \u001b[33mroot\u001b[0m=\u001b[1m{\u001b[0m\n",
" \u001b[32m'higgs_m125_part0.root'\u001b[0m: \u001b[1;35mCoffeaROOTFileSpec\u001b[0m\u001b[1m(\u001b[0m\n",
" \u001b[33mobject_path\u001b[0m=\u001b[32m'Events'\u001b[0m,\n",
@@ -2960,7 +3436,7 @@
" \u001b[33mdid\u001b[0m=\u001b[3;35mNone\u001b[0m\n",
" \u001b[1m)\u001b[0m,\n",
" \u001b[32m'higgs_m200'\u001b[0m: \u001b[1;35mDatasetSpec\u001b[0m\u001b[1m(\u001b[0m\n",
- " \u001b[33mfiles\u001b[0m=\u001b[1;35mInputFiles\u001b[0m\u001b[1m(\u001b[0m\n",
+ " \u001b[33mfiles\u001b[0m=\u001b[1;35mPreprocessedFiles\u001b[0m\u001b[1m(\u001b[0m\n",
" \u001b[33mroot\u001b[0m=\u001b[1m{\u001b[0m\n",
" \u001b[32m'higgs_m200_part0.root'\u001b[0m: \u001b[1;35mCoffeaROOTFileSpec\u001b[0m\u001b[1m(\u001b[0m\n",
" \u001b[33mobject_path\u001b[0m=\u001b[32m'Events'\u001b[0m,\n",
@@ -3051,9 +3527,16 @@
},
{
"cell_type": "code",
- "execution_count": 26,
+ "execution_count": 29,
"id": "b5c82c85",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:21.796764Z",
+ "iopub.status.busy": "2026-07-02T23:32:21.796704Z",
+ "iopub.status.idle": "2026-07-02T23:32:21.799170Z",
+ "shell.execute_reply": "2026-07-02T23:32:21.798738Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -3115,9 +3598,16 @@
},
{
"cell_type": "code",
- "execution_count": 27,
+ "execution_count": 30,
"id": "f6f4842c",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:21.800305Z",
+ "iopub.status.busy": "2026-07-02T23:32:21.800246Z",
+ "iopub.status.idle": "2026-07-02T23:32:21.827349Z",
+ "shell.execute_reply": "2026-07-02T23:32:21.826827Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -4096,9 +4586,16 @@
},
{
"cell_type": "code",
- "execution_count": 28,
+ "execution_count": 31,
"id": "aeb145d5",
- "metadata": {},
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-02T23:32:21.828792Z",
+ "iopub.status.busy": "2026-07-02T23:32:21.828705Z",
+ "iopub.status.idle": "2026-07-02T23:32:21.834941Z",
+ "shell.execute_reply": "2026-07-02T23:32:21.834371Z"
+ }
+ },
"outputs": [
{
"name": "stdout",
@@ -4250,7 +4747,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
- "version": "3.14.0"
+ "version": "3.13.14"
}
},
"nbformat": 4,
diff --git a/binder/processing.ipynb b/binder/processing.ipynb
index 1cdc6b8b6..29720ce86 100644
--- a/binder/processing.ipynb
+++ b/binder/processing.ipynb
@@ -1563,7 +1563,9 @@
"metadata": {},
"source": [
"### Preprocessing\n",
- "There are dataset discovery tools inside of coffea to help construct such datasets. Those will not be demonstrated here. For now, we'll take the above `initial_fileset` and preprocess it."
+ "There are dataset discovery tools inside of coffea to help construct such datasets. Those will not be demonstrated here. For now, we'll take the above `initial_fileset` and preprocess it.\n",
+ "\n",
+ "> **New:** `preprocess` accepts a `backend=` argument — `\"dask\"` (default), or the dask-free `\"iterative\"` / `\"futures\"` backends — and handles Parquet and RNTuple inputs (auto-detected, or via the `preprocess_parquet` / `preprocess_rntuple` helpers). See the [filespec notebook](filespec.ipynb) for a focused walkthrough."
]
},
{
diff --git a/src/coffea/dataset_tools/__init__.py b/src/coffea/dataset_tools/__init__.py
index 0808e87b2..7992ad221 100644
--- a/src/coffea/dataset_tools/__init__.py
+++ b/src/coffea/dataset_tools/__init__.py
@@ -26,14 +26,23 @@
preprocess,
preprocess_legacy,
preprocess_parquet,
+ preprocess_rntuple,
preprocess_root,
)
+from coffea.dataset_tools.preprocess_backends import (
+ DaskBackend,
+ FuturesBackend,
+ IterativeBackend,
+ PreprocessBackend,
+ PreprocessJob,
+)
from coffea.dataset_tools.splitting import hash_fileset, split_fileset
__all__ = [
"preprocess",
"preprocess_legacy",
"preprocess_parquet",
+ "preprocess_rntuple",
"preprocess_root",
"split_fileset",
"hash_fileset",
@@ -58,4 +67,9 @@
"DatasetSpec",
"DataGroupSpec",
"ModelFactory",
+ "PreprocessBackend",
+ "PreprocessJob",
+ "DaskBackend",
+ "FuturesBackend",
+ "IterativeBackend",
]
diff --git a/src/coffea/dataset_tools/filespec.py b/src/coffea/dataset_tools/filespec.py
index 6f455a748..537e0750b 100644
--- a/src/coffea/dataset_tools/filespec.py
+++ b/src/coffea/dataset_tools/filespec.py
@@ -33,6 +33,21 @@ class GenericFileSpec(BaseModel):
format: str | None = None
lfn: str | None = None
pfn: str | None = None
+ # Per-file user metadata, e.g. filled by preprocess(metadata_extractor=...). Must be
+ # JSON-serializable. Carried by the pydantic models only; legacy dict conversions drop it.
+ metadata: dict[str, Any] | None = None
+ # Experimental: hex bitset over the owning dataset's union-form top-level fields
+ # (bit i set = field i is present in this file); meaningful only relative to that
+ # dataset's saved form field order. Subject to change; do not rely on the encoding.
+ experimental_field_bitset: str | None = None
+
+ def __eq__(self, other: Any) -> bool:
+ # experimental fields do not participate in equality: two specs describing the
+ # same file compare equal regardless of experimental annotations
+ if self.__class__ is not other.__class__:
+ return NotImplemented
+ excluded = {"experimental_field_bitset"}
+ return self.model_dump(exclude=excluded) == other.model_dump(exclude=excluded)
def __add__(self, other: GenericFileSpec) -> GenericFileSpec:
if not isinstance(other, GenericFileSpec):
@@ -60,6 +75,21 @@ def __add__(self, other: GenericFileSpec) -> GenericFileSpec:
new_spec["num_entries"] = self.num_entries
else:
new_spec["num_entries"] = max(self.num_entries, other.num_entries)
+ if self.metadata is None:
+ new_spec["metadata"] = other.metadata
+ elif other.metadata is not None:
+ # merge per-file metadata dictionaries, with other taking precedence
+ merged_meta = dict(self.metadata)
+ merged_meta.update(other.metadata)
+ new_spec["metadata"] = merged_meta
+ if self.experimental_field_bitset is None:
+ new_spec["experimental_field_bitset"] = other.experimental_field_bitset
+ elif other.experimental_field_bitset not in (
+ None,
+ self.experimental_field_bitset,
+ ):
+ # disagreeing bitsets cannot be reconciled without the owning forms
+ new_spec["experimental_field_bitset"] = None
return type(self)(**new_spec)
def __sub__(self, other: GenericFileSpec) -> GenericFileSpec:
@@ -466,6 +496,22 @@ def __eq__(self, other: Any) -> bool:
return self.form == other.form
def __add__(self, other: DatasetSpec) -> DatasetSpec:
+ return self.union_with(other)
+
+ def union_with(self, other: DatasetSpec, sort_fields: bool = False) -> DatasetSpec:
+ """Merge two DatasetSpecs, computing the union of their saved forms.
+
+ Files merge like ``+`` on the file collections and metadata merges with ``other``
+ taking precedence. When both operands carry a saved form, the result's form is
+ their union (every field appearing in either form), and per-file experimental
+ field bitsets are remapped to the union field order. When neither operand has a
+ form the result has none. Adding a form-bearing spec to a form-less one raises a
+ ValueError, since the union form could not describe the form-less operand's files.
+ ``sort_fields=True`` sorts record fields recursively so the serialized union form
+ is byte-stable regardless of operand order.
+
+ ``__add__`` delegates here with default options.
+ """
if not isinstance(other, DatasetSpec):
raise TypeError(
f"Can only add DatasetSpec to DatasetSpec, got {type(other)}"
@@ -475,20 +521,156 @@ def __add__(self, other: DatasetSpec) -> DatasetSpec:
raise ValueError(
f"Cannot add DatasetSpec with different dids: {self.did} and {other.did}"
)
+ if (self.compressed_form is None) != (other.compressed_form is None):
+ raise ValueError(
+ "Cannot add a DatasetSpec with a saved form to one without: the union "
+ "form could not describe the files of the form-less operand. Preprocess "
+ "it with save_form=True (or clear the other form) first."
+ )
new_spec = self.model_dump()
- new_spec["files"] = self.files + other.files
+ merged_files = self.files + other.files
+ new_spec["files"] = merged_files.model_dump()
# merge metadata dictionaries, with other taking precedence
new_metadata = copy.deepcopy(self.metadata)
new_metadata.update(other.metadata)
new_spec["metadata"] = new_metadata
# format will be re-evaluated in post validation
new_spec["format"] = None
- # compressed_form is not merged, set to None
- new_spec["compressed_form"] = None
# did is not merged, set to None
new_spec["did"] = self.did if self.did is not None else other.did
+
+ if self.compressed_form is None and other.compressed_form is None:
+ new_spec["compressed_form"] = None
+ else:
+ import awkward
+
+ from coffea.dataset_tools.forms import (
+ decode_field_bitset,
+ encode_field_bitset,
+ sort_form_fields,
+ union_form_jsonstr,
+ )
+ from coffea.util import compress_form
+
+ self_form = self.form
+ other_form = other.form
+ if list(self_form.fields) == list(other_form.fields) and (
+ self_form == other_form
+ ):
+ union_form = sort_form_fields(self_form) if sort_fields else self_form
+ union_jsonstr = union_form.to_json()
+ else:
+ union_jsonstr = union_form_jsonstr(
+ [self_form, other_form], sort_fields=sort_fields
+ )
+ union_form = awkward.forms.from_json(union_jsonstr)
+ new_spec["compressed_form"] = compress_form(union_jsonstr)
+
+ # remap per-file field bitsets from each operand's field order to the union's
+ union_fields = list(union_form.fields)
+
+ def _decoded(files, form):
+ fields = list(form.fields)
+ return {
+ fname: (
+ decode_field_bitset(fs.experimental_field_bitset, fields)
+ if fs.experimental_field_bitset is not None
+ else None
+ )
+ for fname, fs in files.items()
+ }
+
+ fields_self = _decoded(self.files, self_form)
+ fields_other = _decoded(other.files, other_form)
+ for fname, spec_dict in new_spec["files"].items():
+ present_self = fields_self.get(fname)
+ present_other = fields_other.get(fname)
+ if present_self is None:
+ present = present_other
+ elif present_other is None or present_self == present_other:
+ present = present_self
+ else:
+ # the same file reports different field sets; unknowable which is right
+ present = None
+ spec_dict["experimental_field_bitset"] = (
+ encode_field_bitset(present, union_fields)
+ if present is not None
+ else None
+ )
return type(self)(**new_spec)
+ def canonicalize_form(self) -> Self:
+ """Return a copy whose saved form has recursively sorted record fields.
+
+ Sorting makes the serialized form (and anything hashed from it) independent of the
+ union/merge history that produced it; per-file experimental field bitsets are
+ remapped to the sorted field order. A spec without a saved form is returned as an
+ unmodified copy.
+ """
+ spec = self.model_dump()
+ if self.compressed_form is not None:
+ from coffea.dataset_tools.forms import (
+ decode_field_bitset,
+ encode_field_bitset,
+ sort_form_fields,
+ )
+ from coffea.util import compress_form
+
+ old_form = self.form
+ old_fields = list(old_form.fields)
+ sorted_form = sort_form_fields(old_form)
+ new_fields = list(sorted_form.fields)
+ spec["compressed_form"] = compress_form(sorted_form.to_json())
+ for fname, spec_dict in spec["files"].items():
+ bitset = spec_dict.get("experimental_field_bitset")
+ if bitset is not None:
+ spec_dict["experimental_field_bitset"] = encode_field_bitset(
+ decode_field_bitset(bitset, old_fields), new_fields
+ )
+ return type(self)(**spec)
+
+ def _prune_form_for_files(self, spec: dict) -> dict:
+ """Prune the saved union form in a dumped ``spec`` to the fields its files carry.
+
+ Pruning applies only when every remaining file has an experimental field bitset
+ (otherwise the field content of some file is unknown and the form is kept as a
+ superset). Bitsets are remapped to the pruned field order.
+ """
+ if self.compressed_form is None or not spec["files"]:
+ return spec
+ bitsets = {
+ fname: file_spec.get("experimental_field_bitset")
+ for fname, file_spec in spec["files"].items()
+ }
+ if any(bitset is None for bitset in bitsets.values()):
+ return spec
+
+ from coffea.dataset_tools.forms import (
+ decode_field_bitset,
+ encode_field_bitset,
+ prune_form_fields,
+ )
+ from coffea.util import compress_form
+
+ form = self.form
+ old_fields = list(form.fields)
+ keep_fields = set()
+ present_by_file = {}
+ for fname, bitset in bitsets.items():
+ present = decode_field_bitset(bitset, old_fields)
+ present_by_file[fname] = present
+ keep_fields |= present
+ if keep_fields == set(old_fields):
+ return spec
+ pruned_form = prune_form_fields(form, keep_fields)
+ new_fields = list(pruned_form.fields)
+ spec["compressed_form"] = compress_form(pruned_form.to_json())
+ for fname, file_spec in spec["files"].items():
+ file_spec["experimental_field_bitset"] = encode_field_bitset(
+ present_by_file[fname], new_fields
+ )
+ return spec
+
def __sub__(self, other: DatasetSpec) -> DatasetSpec:
if not isinstance(other, DatasetSpec):
raise TypeError(
@@ -699,9 +881,14 @@ def limit_steps(self, max_steps: int | slice, per_file: bool = False) -> Self:
return type(self)(**spec)
def limit_files(self, max_files: int | slice | None) -> Self:
- """Limit the number of files."""
+ """Limit the number of files.
+
+ When every remaining file carries an experimental field bitset, the saved union
+ form is pruned to the fields those files carry; otherwise it is kept as a superset.
+ """
spec = self.model_dump()
- spec["files"] = self.files.limit_files(max_files)
+ spec["files"] = self.files.limit_files(max_files).model_dump()
+ spec = self._prune_form_for_files(spec)
return type(self)(**spec)
def filter_files(
@@ -709,11 +896,16 @@ def filter_files(
filter_name: str | None = None,
filter_callable: Callable[[FileSpecUnion], bool] | None = None,
) -> Self:
- """Filter files by a regex pattern on the file names(filter_name) or callable applied to Filespecs (filter_callable)."""
+ """Filter files by a regex pattern on the file names(filter_name) or callable applied to Filespecs (filter_callable).
+
+ When every remaining file carries an experimental field bitset, the saved union
+ form is pruned to the fields those files carry; otherwise it is kept as a superset.
+ """
spec = self.model_dump()
spec["files"] = self.files.filter_files(
filter_name=filter_name, filter_callable=filter_callable
- )
+ ).model_dump()
+ spec = self._prune_form_for_files(spec)
return type(self)(**spec)
@@ -937,7 +1129,8 @@ def filespec_to_dict(
raise TypeError(
f"{cls.__name__}.filespec_to_dict expects a Coffea(Parquet)FileSpec(Optional), got {type(input)} instead: {input}"
)
- return input.model_dump()
+ # the legacy dict format does not carry per-file metadata or experimental fields
+ return input.model_dump(exclude={"metadata", "experimental_field_bitset"})
@classmethod
def dict_to_datasetspec(cls, input: dict[str, Any], verbose=False) -> DatasetSpec:
@@ -953,6 +1146,11 @@ def datasetspec_to_dict(
input, DatasetSpec
), f"{cls.__name__}.datasetspec_to_dict expects a DatasetSpec, got {type(input)} instead: {input}"
if coerce_filespec_to_dict:
- return input.model_dump()
+ # the legacy dict format does not carry per-file metadata or experimental fields
+ return input.model_dump(
+ exclude={
+ "files": {"__all__": {"metadata", "experimental_field_bitset"}}
+ }
+ )
else:
return dict(input)
diff --git a/src/coffea/dataset_tools/forms.py b/src/coffea/dataset_tools/forms.py
new file mode 100644
index 000000000..f88150296
--- /dev/null
+++ b/src/coffea/dataset_tools/forms.py
@@ -0,0 +1,147 @@
+"""Helpers for combining awkward forms across files and datasets.
+
+A dataset's saved form is the union of its files' forms: NanoAOD-style inputs routinely have
+per-file field differences (e.g. ``HLT_*`` trigger bits or ``GenModel`` points present in only
+some files), so the dataset form must cover every field that appears anywhere while remaining
+readable for files that lack some of them. This module is import-light (awkward only) so both
+:mod:`coffea.dataset_tools.preprocess` and :mod:`coffea.dataset_tools.filespec` can use it.
+"""
+
+from __future__ import annotations
+
+import awkward
+
+__all__ = [
+ "union_form_jsonstr",
+ "sort_form_fields",
+ "prune_form_fields",
+ "encode_field_bitset",
+ "decode_field_bitset",
+]
+
+
+def union_form_jsonstr(forms: list, sort_fields: bool = False) -> str | None:
+ """Compute the union form (as a JSON string) over a list of awkward forms.
+
+ The input list is consumed. Returns None if the list is empty. Mirrors the merging of
+ flat-tuple-like schemas used when building a dataset's union form across files. Fields
+ appear in merge order by default; ``sort_fields=True`` recursively sorts record fields so
+ the serialized form is byte-stable regardless of merge order (awkward form equality is
+ field-order-insensitive either way).
+ """
+ union_array = None
+ while len(forms):
+ new_array = awkward.Array(forms.pop().length_zero_array())
+ if union_array is None:
+ union_array = new_array
+ else:
+ union_array = awkward.to_packed(
+ awkward.merge_union_of_records(
+ awkward.concatenate([union_array, new_array]), axis=0
+ )
+ )
+ union_array.layout.parameters.update(new_array.layout.parameters)
+ if union_array is None:
+ return None
+
+ union_form = union_array.layout.form
+ for icontent, content in enumerate(union_form.contents):
+ if isinstance(content, awkward.forms.IndexedOptionForm):
+ if (
+ not isinstance(content.content, awkward.forms.NumpyForm)
+ or content.content.primitive != "bool"
+ ):
+ raise ValueError(
+ "IndexedOptionArrays can only contain NumpyArrays of "
+ "bools in mergers of flat-tuple-like schemas!"
+ )
+ parameters = (
+ content.content.parameters.copy()
+ if content.content.parameters is not None
+ else {}
+ )
+ # re-create IndexOptionForm with parameters of lower level array
+ union_form.contents[icontent] = awkward.forms.IndexedOptionForm(
+ content.index,
+ content.content,
+ parameters=parameters,
+ form_key=content.form_key,
+ )
+ if sort_fields:
+ union_form = sort_form_fields(union_form)
+ return union_form.to_json()
+
+
+def _sort_record_nodes(node) -> None:
+ if isinstance(node, dict):
+ if (
+ node.get("class") == "RecordArray"
+ and isinstance(node.get("fields"), list)
+ and isinstance(node.get("contents"), list)
+ ):
+ pairs = sorted(
+ zip(node["fields"], node["contents"]), key=lambda pair: pair[0]
+ )
+ node["fields"] = [field for field, _ in pairs]
+ node["contents"] = [content for _, content in pairs]
+ for value in node.values():
+ _sort_record_nodes(value)
+ elif isinstance(node, list):
+ for value in node:
+ _sort_record_nodes(value)
+
+
+def sort_form_fields(form: awkward.forms.Form) -> awkward.forms.Form:
+ """Return a copy of ``form`` with record fields recursively sorted by name.
+
+ Tuple-like records (no field names) are left untouched. Sorting canonicalizes the
+ serialized byte order of forms whose field order depends on union/merge history.
+ """
+ form_dict = form.to_dict(verbose=True)
+ _sort_record_nodes(form_dict)
+ return awkward.forms.from_dict(form_dict)
+
+
+def prune_form_fields(
+ form: awkward.forms.Form, keep_fields: set[str]
+) -> awkward.forms.Form:
+ """Return a copy of ``form`` with only the top-level record fields in ``keep_fields``.
+
+ Only the outermost record is pruned: dataset union forms merge file forms at the top
+ level, so that is the level at which fields can disappear when files are removed.
+ """
+ form_dict = form.to_dict(verbose=True)
+ if not (
+ isinstance(form_dict.get("fields"), list)
+ and isinstance(form_dict.get("contents"), list)
+ ):
+ return form
+ pairs = [
+ (field, content)
+ for field, content in zip(form_dict["fields"], form_dict["contents"])
+ if field in keep_fields
+ ]
+ form_dict["fields"] = [field for field, _ in pairs]
+ form_dict["contents"] = [content for _, content in pairs]
+ return awkward.forms.from_dict(form_dict)
+
+
+def encode_field_bitset(present_fields, union_fields: list[str]) -> str:
+ """Encode which of ``union_fields`` a file carries as a compact hex bitset string.
+
+ Bit ``i`` (value ``1 << i``) corresponds to ``union_fields[i]``; a set bit means the
+ field is present in the file. Fields outside ``union_fields`` are ignored. The encoding
+ is only meaningful relative to the field order of the owning dataset's saved form.
+ """
+ present = set(present_fields)
+ bits = 0
+ for index, field in enumerate(union_fields):
+ if field in present:
+ bits |= 1 << index
+ return format(bits, "x")
+
+
+def decode_field_bitset(bitset: str, union_fields: list[str]) -> set[str]:
+ """Decode a hex bitset string (see :func:`encode_field_bitset`) into a set of field names."""
+ bits = int(bitset, 16)
+ return {field for index, field in enumerate(union_fields) if (bits >> index) & 1}
diff --git a/src/coffea/dataset_tools/mutable_steps.py b/src/coffea/dataset_tools/mutable_steps.py
new file mode 100644
index 000000000..e00d93ca5
--- /dev/null
+++ b/src/coffea/dataset_tools/mutable_steps.py
@@ -0,0 +1,257 @@
+"""Prototype for mutable (resizable) steps: generator-driven step iteration where the
+consumer may renegotiate the step size mid-stream.
+
+Stored ``steps`` on a file spec are a static tiling chosen at preprocess time. Execution
+backends with runtime resource information (e.g. a resource-monitoring scheduler that sees a
+worker near memory exhaustion) can do better by shrinking or growing subsequent chunks while
+a file is being processed. This module provides that as a ``send``-channel generator
+protocol: iterating yields ``[start, stop]`` steps, and ``generator.send(new_size)`` requests
+that subsequent steps be *at most* ``new_size`` entries, with the remainder of the current
+region re-tiled as evenly as possible::
+
+ n = ceil(remaining / new_size)
+ actual = ceil(remaining / n)
+
+so a resize triggered by resource exhaustion can never produce a larger-than-requested step,
+and the remaining steps stay near-uniform. The same channel shape is used by
+``Computable.gen_steps`` in the ``coffea.compute`` protocol, so consumers written against
+this prototype translate directly.
+
+Everything here is a prototype: APIs may change, and nothing is re-exported from
+``coffea.dataset_tools``.
+"""
+
+from __future__ import annotations
+
+import math
+import time
+from collections.abc import Callable, Generator
+from dataclasses import dataclass, field
+
+__all__ = [
+ "resizable_steps",
+ "coverage_regions",
+ "iter_file_steps",
+ "iter_dataset_steps",
+ "remaining_regions",
+ "completed_spec",
+ "WallTimeStepPolicy",
+ "run_adaptive_steps",
+]
+
+
+def _validate_size(size: int) -> None:
+ if size is None or size < 1:
+ raise ValueError(f"step size must be a positive integer (>= 1), got {size!r}.")
+
+
+def resizable_steps(
+ start: int, stop: int, step_size: int
+) -> Generator[list[int], int | None, None]:
+ """Yield ``[begin, end]`` steps tiling ``[start, stop)``, honoring resize requests.
+
+ Steps are at most the current target size and as even as possible: each iteration
+ re-tiles the remaining entries as ``ceil(remaining / ceil(remaining / size))``, which is
+ stable while the size is unchanged. Sending a positive integer sets the target size for
+ all subsequent steps; sending ``None`` (plain iteration) keeps the current size.
+ """
+ _validate_size(step_size)
+ current = step_size
+ pos = start
+ while pos < stop:
+ remaining = stop - pos
+ n_steps = math.ceil(remaining / current)
+ actual = math.ceil(remaining / n_steps)
+ end = min(pos + actual, stop)
+ sent = yield [pos, end]
+ pos = end
+ if sent is not None:
+ _validate_size(sent)
+ current = sent
+
+
+def coverage_regions(filespec) -> list[list[int]]:
+ """The contiguous entry regions a file spec covers, as ``[begin, end]`` pairs.
+
+ Adjacent steps (``stop == next start``) merge into one region. A spec without steps
+ covers ``[0, num_entries)``; a spec with neither steps nor num_entries cannot be tiled
+ and raises a ValueError.
+ """
+ if filespec.steps is not None:
+ regions: list[list[int]] = []
+ for begin, end in filespec.steps:
+ if regions and regions[-1][1] == begin:
+ regions[-1][1] = end
+ else:
+ regions.append([begin, end])
+ return regions
+ if filespec.num_entries is not None:
+ return [[0, filespec.num_entries]] if filespec.num_entries > 0 else []
+ raise ValueError(
+ "Cannot iterate steps for a file spec with neither steps nor num_entries; "
+ "preprocess the file first."
+ )
+
+
+def iter_file_steps(
+ filespec, step_size: int | None = None
+) -> Generator[list[int], int | None, None]:
+ """Yield resizable ``[start, stop]`` steps over a file spec's covered regions.
+
+ ``step_size`` is the initial target size (default: one step per contiguous region).
+ ``send(new_size)`` re-tiles from the next step onward, carrying across regions.
+ """
+ current = step_size
+ for begin, end in coverage_regions(filespec):
+ size = current if current is not None else end - begin
+ gen = resizable_steps(begin, end, size)
+ try:
+ step = next(gen)
+ while True:
+ sent = yield step
+ if sent is not None:
+ current = sent
+ step = gen.send(sent)
+ except StopIteration:
+ continue
+
+
+def iter_dataset_steps(
+ dataset, step_size: int | None = None
+) -> Generator[tuple[str, list[int]], int | None, None]:
+ """Yield resizable ``(filename, [start, stop])`` steps over every file of a DatasetSpec.
+
+ A resize request applies from the next step onward and carries across file boundaries,
+ so a consumer that shrinks the step size mid-file keeps the smaller size for later files
+ until it requests otherwise.
+ """
+ current = step_size
+ for fname, filespec in dataset.files.items():
+ gen = iter_file_steps(filespec, current)
+ try:
+ step = next(gen)
+ while True:
+ sent = yield (fname, step)
+ if sent is not None:
+ current = sent
+ step = gen.send(sent)
+ except StopIteration:
+ continue
+
+
+def remaining_regions(filespec, completed: list[list[int]]) -> list[list[int]]:
+ """The parts of a file spec's coverage not contained in ``completed`` ranges.
+
+ ``completed`` are arbitrary ``[start, stop]`` pairs (they need not align with the spec's
+ stored steps -- resized steps generally do not). Returns ``[begin, end]`` regions suitable
+ for resuming iteration via :func:`resizable_steps`.
+ """
+ merged: list[list[int]] = []
+ for begin, end in sorted([list(pair) for pair in completed]):
+ if merged and begin <= merged[-1][1]:
+ merged[-1][1] = max(merged[-1][1], end)
+ else:
+ merged.append([begin, end])
+ remaining: list[list[int]] = []
+ for begin, end in coverage_regions(filespec):
+ pos = begin
+ for done_begin, done_end in merged:
+ if done_end <= pos or done_begin >= end:
+ continue
+ if done_begin > pos:
+ remaining.append([pos, done_begin])
+ pos = max(pos, done_end)
+ if pos >= end:
+ break
+ if pos < end:
+ remaining.append([pos, end])
+ return remaining
+
+
+def completed_spec(filespec, completed: list[list[int]]):
+ """A copy of ``filespec`` whose steps are the ``completed`` ranges.
+
+ Useful for persisting progress: the result round-trips like any spec, and adding it to a
+ spec built from later completions accumulates coverage via the ordinary step arithmetic.
+ Returns ``None`` when nothing has completed.
+ """
+ if not completed:
+ return None
+ spec = filespec.model_dump()
+ spec["steps"] = sorted([list(pair) for pair in completed])
+ return type(filespec)(**spec)
+
+
+@dataclass
+class WallTimeStepPolicy:
+ """Toy resize policy targeting a fixed wall time per step.
+
+ After each step, the next target size is ``current * target_seconds / elapsed``, clamped
+ to ``[min_step_size, max_step_size]`` and damped by ``max_growth`` per adjustment so a
+ single fast outlier does not balloon the step size.
+ """
+
+ target_seconds: float
+ min_step_size: int = 1
+ max_step_size: int | None = None
+ max_growth: float = 2.0
+
+ def propose(self, current_size: int, elapsed_seconds: float) -> int | None:
+ """The new target size, or None to keep the current one."""
+ if elapsed_seconds <= 0.0:
+ scaled = current_size * self.max_growth
+ else:
+ scaled = current_size * self.target_seconds / elapsed_seconds
+ scaled = min(scaled, current_size * self.max_growth)
+ new_size = max(self.min_step_size, int(scaled))
+ if self.max_step_size is not None:
+ new_size = min(new_size, self.max_step_size)
+ return None if new_size == current_size else new_size
+
+
+@dataclass
+class AdaptiveRun:
+ """Result of :func:`run_adaptive_steps`."""
+
+ results: list
+ completed: dict[str, list[list[int]]]
+ step_sizes: list[int] = field(default_factory=list)
+
+
+def run_adaptive_steps(
+ dataset,
+ work: Callable[[str, list[int]], object],
+ step_size: int,
+ policy: WallTimeStepPolicy,
+ clock: Callable[[], float] = time.monotonic,
+) -> AdaptiveRun:
+ """Toy driver: process a DatasetSpec step by step, resizing steps from measured wall time.
+
+ ``work(filename, [start, stop])`` is called for each step; its wall time (measured with
+ ``clock``) feeds ``policy.propose``, and any proposed size is sent into the step
+ generator, re-tiling the remainder. This is the single-consumer analogue of what a
+ resource-monitoring scheduler does with worker feedback; ``clock`` is injectable so the
+ control loop is testable without real waiting.
+ """
+ _validate_size(step_size)
+ gen = iter_dataset_steps(dataset, step_size)
+ run = AdaptiveRun(results=[], completed={})
+ current = step_size
+ try:
+ item = next(gen)
+ while True:
+ fname, step = item
+ begin = clock()
+ run.results.append(work(fname, step))
+ elapsed = clock() - begin
+ run.completed.setdefault(fname, []).append(step)
+ run.step_sizes.append(step[1] - step[0])
+ proposed = policy.propose(current, elapsed)
+ if proposed is not None:
+ current = proposed
+ item = gen.send(proposed)
+ else:
+ item = next(gen)
+ except StopIteration:
+ pass
+ return run
diff --git a/src/coffea/dataset_tools/preprocess.py b/src/coffea/dataset_tools/preprocess.py
index 4439eea1e..b9cae8c82 100644
--- a/src/coffea/dataset_tools/preprocess.py
+++ b/src/coffea/dataset_tools/preprocess.py
@@ -2,6 +2,7 @@
import copy
import hashlib
+import json
import math
import warnings
from collections.abc import Callable
@@ -16,11 +17,31 @@
import uproot
from uproot._util import no_filter
+try:
+ # Private uproot helper that builds a TTree's awkward form without dask. It is exactly what
+ # uproot.dask() uses internally, so the result is byte-identical to
+ # uproot.dask(tree).layout.form. Guarded so a uproot release that moves or renames it
+ # degrades gracefully to the dask-based path below.
+ from uproot._dask import _get_ttree_form as _uproot_get_ttree_form
+except Exception: # pragma: no cover - depends on uproot internals
+ _uproot_get_ttree_form = None
+
from coffea.dataset_tools.filespec import (
DataGroupSpec,
DatasetSpec,
ModelFactory,
)
+from coffea.dataset_tools.forms import (
+ encode_field_bitset,
+)
+from coffea.dataset_tools.forms import union_form_jsonstr as _union_form_jsonstr
+from coffea.dataset_tools.preprocess_backends import (
+ DaskBackend,
+ PreprocessBackend,
+ PreprocessJob,
+ print_dask_backend_fallback_hint,
+ resolve_backend,
+)
from coffea.util import (
_import_dask,
_import_dask_awkward,
@@ -30,6 +51,179 @@
)
+def _validate_step_size(step_size: int | None) -> None:
+ """Reject a non-positive ``step_size`` at the public entrypoints.
+
+ ``step_size`` is the number of entries per step; ``None`` means "one step per file". A value
+ < 1 is meaningless and would otherwise surface as a bare ``ZeroDivisionError`` from
+ ``_even_steps`` deep inside a worker (or dask graph).
+ """
+ if step_size is not None and step_size < 1:
+ raise ValueError(
+ f"step_size must be a positive integer (>= 1) or None, got {step_size!r}."
+ )
+
+
+def _even_steps(num_entries: int, target_step_size: int) -> numpy.ndarray:
+ """Split ``num_entries`` into as-even-as-possible ``[start, stop]`` steps of ~target size."""
+ n_steps_target = max(round(num_entries / target_step_size), 1)
+ actual_step_size = math.ceil(num_entries / n_steps_target)
+ return numpy.array(
+ [
+ [i * actual_step_size, min((i + 1) * actual_step_size, num_entries)]
+ for i in range(n_steps_target)
+ ],
+ dtype="int64",
+ )
+
+
+def _aligned_steps(
+ boundaries,
+ target_step_size: int,
+ step_size_safety_factor: float,
+ file_label: str,
+ mode_label: str,
+) -> numpy.ndarray:
+ """Build ``[start, stop]`` steps that snap to natural boundaries (TTree clusters, RNTuple
+ cluster summaries, or parquet row groups).
+
+ ``boundaries`` is the increasing sequence of absolute entry offsets at which a step is
+ allowed to end, with the final element equal to ``num_entries``. Steps accumulate
+ boundaries until at least ``target_step_size`` entries are covered. ``mode_label`` is the
+ name of the user-facing option (``align_clusters`` or ``use_row_groups``) used in the
+ over-size warning.
+ """
+ out = [0]
+ for c in boundaries:
+ if c >= out[-1] + target_step_size:
+ out.append(c)
+ if boundaries[-1] != out[-1]:
+ out.append(boundaries[-1])
+ out = numpy.array(out, dtype="int64")
+ out = numpy.stack((out[:-1], out[1:]), axis=1)
+
+ step_mask = out[:, 1] - out[:, 0] > (1 + step_size_safety_factor) * target_step_size
+ if numpy.any(step_mask):
+ warnings.warn(
+ f"In file {file_label}, steps: {out[step_mask]} with {mode_label}=True are "
+ f"{step_size_safety_factor*100:.0f}% larger than target "
+ f"step size: {target_step_size}!"
+ )
+ return out
+
+
+def _rntuple_cluster_boundaries(rntuple, num_entries: int) -> list[int]:
+ """Absolute entry offsets at RNTuple cluster boundaries, terminating at ``num_entries``."""
+ boundaries = [cluster.num_first_entry for cluster in rntuple.cluster_summaries]
+ boundaries.append(num_entries)
+ return boundaries
+
+
+def _serialize_user_metadata(user_meta, filename) -> str | None:
+ """Validate and JSON-encode a metadata_extractor result for transport in the worker record."""
+ if user_meta is None:
+ return None
+ if not isinstance(user_meta, dict):
+ raise ValueError(
+ f"metadata_extractor must return a dict, got {type(user_meta).__name__} "
+ f"for file {filename!r}."
+ )
+ try:
+ return json.dumps(user_meta)
+ except TypeError as err:
+ raise ValueError(
+ f"metadata_extractor result for file {filename!r} is not JSON-serializable: {err}"
+ ) from err
+
+
+_FORM_AK_ADD_DOC = {"__doc__": "title", "typename": "typename"}
+
+
+def _form_json_with_null_keys(form) -> str:
+ """Serialize ``form`` to JSON with every ``form_key`` set to ``None``.
+
+ ``uproot.dask`` exposes its *meta* (typetracer) form, whose form keys are null, whereas
+ ``_get_ttree_form`` carries RNTuple column keys (``"column-N"``). Nulling the keys makes the
+ serialized form byte-identical to ``uproot.dask(...).layout.form.to_json()`` and is a no-op
+ for TTree forms (already keyless). The keys are stripped on the ``to_dict`` representation
+ and serialized directly, matching ``Form.to_json`` (``json.dumps`` of the verbose dict)
+ without rebuilding a ``Form`` tree in between.
+ """
+ form_dict = form.to_dict(verbose=True)
+
+ def _strip(node):
+ if isinstance(node, dict):
+ if "form_key" in node:
+ node["form_key"] = None
+ for value in node.values():
+ _strip(value)
+ elif isinstance(node, list):
+ for value in node:
+ _strip(value)
+
+ _strip(form_dict)
+ return json.dumps(form_dict)
+
+
+def _awkward_form_json(tree, is_rntuple: bool) -> str:
+ """Build a TTree or RNTuple awkward form JSON without importing dask.
+
+ Reuses uproot's own form builder (``uproot._dask._get_ttree_form``, which handles both TTree
+ branches and RNTuple fields) and nulls the resulting form keys, producing output
+ byte-identical to ``uproot.dask(...).layout.form.to_json()``. This is what lets the
+ ``iterative``/``futures`` backends extract forms in a dask-free environment for both formats.
+ """
+ # RNTuples filter on fields and use full field paths; TTrees filter on branches.
+ filter_kwarg = "filter_field" if is_rntuple else "filter_branch"
+ common_keys = tree.keys(
+ recursive=True,
+ filter_name=no_filter,
+ filter_typename=no_filter,
+ full_paths=is_rntuple,
+ ignore_duplicates=True,
+ **{filter_kwarg: partial(_is_interpretable, emit_warning=False)},
+ )
+ base_form = _uproot_get_ttree_form(awkward, tree, common_keys, _FORM_AK_ADD_DOC)
+ return _form_json_with_null_keys(base_form)
+
+
+def _dask_form_json(uproot_target, is_rntuple: bool, uproot_options: dict = {}) -> str:
+ """Build a form JSON via ``uproot.dask`` (requires dask).
+
+ Fallback used only when uproot's internal form builder (``_get_ttree_form``) is unavailable
+ (an unexpected uproot version); the dask-free :func:`_awkward_form_json` is preferred for both
+ TTree and RNTuple. ``uproot_target`` is an already-open TTree object, or a
+ ``{file: object_path}`` mapping (required for RNTuples, which cannot build a form from an
+ already-open object via ``uproot.dask``). ``uproot_options`` is forwarded to the file open
+ when ``uproot_target`` is a mapping.
+ """
+ if is_rntuple:
+ form_dask = uproot.dask(
+ uproot_target,
+ open_files=False,
+ full_paths=True,
+ ak_add_doc=_FORM_AK_ADD_DOC,
+ filter_name=no_filter,
+ filter_typename=no_filter,
+ filter_branch=partial(_is_interpretable, emit_warning=False),
+ **uproot_options,
+ )
+ else:
+ form_dask = uproot.dask(
+ uproot_target,
+ ak_add_doc=_FORM_AK_ADD_DOC,
+ filter_name=no_filter,
+ filter_typename=no_filter,
+ filter_branch=partial(_is_interpretable, emit_warning=False),
+ )
+ form_str = form_dask.layout.form.to_json()
+ # the function cache needs to be popped if present to prevent memory growth
+ dask = _import_dask()
+ if getattr(dask.base, "function_cache", None):
+ dask.base.function_cache.popitem()
+ return form_str
+
+
def get_steps(
normed_files: awkward.Array | dask_awkward.Array,
step_size: int | None = None,
@@ -41,6 +235,8 @@ def get_steps(
step_size_safety_factor: float = 0.5,
uproot_options: dict = {},
legacy_form_key: bool = True,
+ require_rntuple: bool = False,
+ metadata_extractor: Callable | None = None,
) -> awkward.Array | dask_awkward.Array:
"""
Given a list of normalized file and object paths (defined in uproot), determine the steps for each file according to the supplied processing options.
@@ -69,12 +265,22 @@ def get_steps(
legacy_form_key : bool, default True
Use "form" for the compressed form key in the output for backwards compatibility.
Set to False to use "compressed_form" instead.
+ require_rntuple : bool, default False
+ If True, require every object to be an RNTuple and raise a ValueError otherwise.
+ If False, TTree and RNTuple objects are auto-detected and handled transparently.
+ metadata_extractor : Callable or None, default None
+ User function called once per file with the open uproot file handle (the
+ ReadOnlyDirectory, not the tree); must return a JSON-serializable dict, which is
+ stored as that file's metadata. Runs inside the per-file error handling, so an
+ extraction failure participates in skip_bad_files/file_exceptions. Must be
+ picklable to run under process pools or distributed schedulers.
Returns
-------
array : awkward.Array or dask_awkward.Array
The normalized file descriptions, appended with the calculated steps for those files.
"""
+ _validate_step_size(step_size)
nf_backend = awkward.backend(normed_files)
lz_or_nf = awkward.typetracer.length_zero_if_typetracer(normed_files)
output_form_key = "form" if legacy_form_key else "compressed_form"
@@ -84,6 +290,18 @@ def get_steps(
try:
the_file = uproot.open({arg.file: None}, **uproot_options)
tree = the_file[arg.object_path]
+ is_rntuple = isinstance(tree, uproot.behaviors.RNTuple.HasFields)
+ if require_rntuple and not is_rntuple:
+ # raised inside this block so skip_bad_files/file_exceptions can skip the file
+ raise ValueError(
+ f"require_rntuple=True but {arg.object_path!r} in {arg.file!r} is a "
+ f"{type(tree).__name__}, not an RNTuple."
+ )
+ # run the user extractor on the open file handle, inside this block so a file
+ # whose extraction fails participates in skip_bad_files/file_exceptions
+ user_meta = (
+ metadata_extractor(the_file) if metadata_extractor is not None else None
+ )
except file_exceptions as e:
if skip_bad_files:
array.append(None)
@@ -91,22 +309,25 @@ def get_steps(
else:
raise e
+ user_meta_json = _serialize_user_metadata(user_meta, arg.file)
+
num_entries = tree.num_entries
form_json = None
form_hash = None
if save_form:
- form_str = uproot.dask(
- tree,
- ak_add_doc={"__doc__": "title", "typename": "typename"},
- filter_name=no_filter,
- filter_typename=no_filter,
- filter_branch=partial(_is_interpretable, emit_warning=False),
- ).layout.form.to_json()
- # the function cache needs to be popped if present to prevent memory growth
- dask = _import_dask()
- if hasattr(dask.base, "function_cache"):
- dask.base.function_cache.popitem()
+ if _uproot_get_ttree_form is not None:
+ # dask-free form extraction (TTree and RNTuple); byte-identical to uproot.dask
+ form_str = _awkward_form_json(tree, is_rntuple)
+ elif is_rntuple:
+ # uproot without the private form builder: fall back to the dask-based path
+ form_str = _dask_form_json(
+ {arg.file: arg.object_path},
+ is_rntuple=True,
+ uproot_options=uproot_options,
+ )
+ else:
+ form_str = _dask_form_json(tree, is_rntuple=False)
form_hash = hashlib.md5(form_str.encode("utf-8")).hexdigest()
form_json = compress_form(form_str)
@@ -128,46 +349,26 @@ def get_steps(
"uuid": file_uuid,
output_form_key: form_json,
"form_hash_md5": form_hash,
+ "user_metadata_json": user_meta_json,
}
)
continue
if out_uuid != file_uuid or recalculate_steps:
if align_clusters:
- clusters = tree.common_entry_offsets()
- out = [0]
- for c in clusters:
- if c >= out[-1] + target_step_size:
- out.append(c)
- if clusters[-1] != out[-1]:
- out.append(clusters[-1])
- out = numpy.array(out, dtype="int64")
- out = numpy.stack((out[:-1], out[1:]), axis=1)
-
- step_mask = (
- out[:, 1] - out[:, 0]
- > (1 + step_size_safety_factor) * target_step_size
+ if is_rntuple:
+ boundaries = _rntuple_cluster_boundaries(tree, num_entries)
+ else:
+ boundaries = tree.common_entry_offsets()
+ out = _aligned_steps(
+ boundaries,
+ target_step_size,
+ step_size_safety_factor,
+ arg.file,
+ "align_clusters",
)
- if numpy.any(step_mask):
- warnings.warn(
- f"In file {arg.file}, steps: {out[step_mask]} with align_clusters=True are "
- f"{step_size_safety_factor*100:.0f}% larger than target "
- f"step size: {target_step_size}!"
- )
-
else:
- n_steps_target = max(round(num_entries / target_step_size), 1)
- actual_step_size = math.ceil(num_entries / n_steps_target)
- out = numpy.array(
- [
- [
- i * actual_step_size,
- min((i + 1) * actual_step_size, num_entries),
- ]
- for i in range(n_steps_target)
- ],
- dtype="int64",
- )
+ out = _even_steps(num_entries, target_step_size)
out_uuid = file_uuid
out_steps = out.tolist()
@@ -184,6 +385,7 @@ def get_steps(
"uuid": out_uuid,
output_form_key: form_json,
"form_hash_md5": form_hash,
+ "user_metadata_json": user_meta_json,
}
)
@@ -198,6 +400,7 @@ def get_steps(
"uuid": "junk",
output_form_key: "junk",
"form_hash_md5": "junk",
+ "user_metadata_json": "junk",
},
None,
]
@@ -305,6 +508,7 @@ def preprocess_legacy(
out_updated : dict
The original set of datasets including files that were not accessible, updated to include the result of preprocessing where available.
"""
+ _validate_step_size(step_size)
dask = _import_dask()
dask_awkward = _import_dask_awkward()
@@ -413,7 +617,7 @@ def preprocess_legacy(
):
# skip trivially filled or empty files
form = awkward.forms.from_json(decompress_form(formstr))
- if num_entries >= 0 and set(form.fields) != _trivial_file_fields:
+ if set(form.fields) != _trivial_file_fields:
dataset_forms.append(form)
else:
warnings.warn(
@@ -425,46 +629,7 @@ def preprocess_legacy(
", by default, removes empty files each dataset in a fileset."
)
- union_array = None
- union_form_jsonstr = None
- while len(dataset_forms):
- new_array = awkward.Array(dataset_forms.pop().length_zero_array())
- if union_array is None:
- union_array = new_array
- else:
- union_array = awkward.to_packed(
- awkward.merge_union_of_records(
- awkward.concatenate([union_array, new_array]), axis=0
- )
- )
- union_array.layout.parameters.update(new_array.layout.parameters)
- if union_array is not None:
- union_form = union_array.layout.form
-
- for icontent, content in enumerate(union_form.contents):
- if isinstance(content, awkward.forms.IndexedOptionForm):
- if (
- not isinstance(content.content, awkward.forms.NumpyForm)
- or content.content.primitive != "bool"
- ):
- raise ValueError(
- "IndexedOptionArrays can only contain NumpyArrays of "
- "bools in mergers of flat-tuple-like schemas!"
- )
- parameters = (
- content.content.parameters.copy()
- if content.content.parameters is not None
- else {}
- )
- # re-create IndexOptionForm with parameters of lower level array
- union_form.contents[icontent] = awkward.forms.IndexedOptionForm(
- content.index,
- content.content,
- parameters=parameters,
- form_key=content.form_key,
- )
-
- union_form_jsonstr = union_form.to_json()
+ union_form_jsonstr = _union_form_jsonstr(dataset_forms)
files_available = {
item["file"]: {
@@ -547,6 +712,7 @@ def get_parquet_form_uuid_steps(
save_form: bool = False,
step_size_safety_factor: float = 0.5,
parquet_options: dict = {},
+ metadata_extractor: Callable | None = None,
) -> awkward.Array | dask_awkward.Array:
"""
Given a list of normalized file and object paths, determine the form, steps, uuid for each file according to the supplied processing options.
@@ -567,16 +733,23 @@ def get_parquet_form_uuid_steps(
file_exceptions : Exception | Warning | tuple[Exception | Warning], default (OSError,)
What exceptions to catch when skipping bad files.
save_form : bool, default False
- Extract the form of the parquet file so we can skip opening files later.
+ Extract the form from the parquet metadata so we can skip opening files later.
step_size_safety_factor : float, default 0.5
When using use_row_groups, if a resulting step is larger than step_size by this factor
warn the user that the resulting steps may be highly irregular.
+ metadata_extractor : Callable or None, default None
+ User function called once per file with the parquet metadata mapping returned by
+ ``awkward.metadata_from_parquet``; must return a JSON-serializable dict, which is
+ stored as that file's metadata. Runs inside the per-file error handling, so an
+ extraction failure participates in skip_bad_files/file_exceptions. Must be
+ picklable to run under process pools or distributed schedulers.
Returns
-------
array : awkward.Array | dask_awkward.Array
The normalized file descriptions, appended with the calculated steps for those files.
"""
+ _validate_step_size(step_size)
nf_backend = awkward.backend(normed_files)
lz_or_nf = awkward.typetracer.length_zero_if_typetracer(normed_files)
@@ -584,6 +757,11 @@ def get_parquet_form_uuid_steps(
for arg in lz_or_nf:
try:
the_file = awkward.metadata_from_parquet(arg.file, **parquet_options)
+ # run the user extractor on the parquet metadata, inside this block so a file
+ # whose extraction fails participates in skip_bad_files/file_exceptions
+ user_meta = (
+ metadata_extractor(the_file) if metadata_extractor is not None else None
+ )
except file_exceptions as e:
if skip_bad_files:
array.append(None)
@@ -591,13 +769,18 @@ def get_parquet_form_uuid_steps(
else:
raise e
+ user_meta_json = _serialize_user_metadata(user_meta, arg.file)
+
num_entries = the_file["num_rows"]
form_json = None
form_hash = None
if save_form:
+ # parquet metadata already carries the form; reading it builds no dask graph,
+ # so (unlike the TTree/RNTuple path) there is no function cache to pop here.
form = the_file["form"]
form_str = form.to_json()
+
form_hash = hashlib.md5(form_str.encode("utf-8")).hexdigest()
form_json = compress_form(form_str)
@@ -605,9 +788,8 @@ def get_parquet_form_uuid_steps(
file_uuid = the_file.get("uuid", None)
- out_uuid = arg.uuid
- out_steps = arg.steps
-
+ # A 0-row file gets a single trivial step: it has no row groups to align to and no
+ # entries to split into steps.
if num_entries == 0:
array.append(
{
@@ -618,47 +800,27 @@ def get_parquet_form_uuid_steps(
"uuid": file_uuid,
"compressed_form": form_json,
"form_hash_md5": form_hash,
+ "user_metadata_json": user_meta_json,
}
)
continue
+ out_uuid = arg.uuid
+ out_steps = arg.steps
+
if out_uuid != file_uuid or recalculate_steps:
if use_row_groups:
- row_group_entries = the_file["col_counts"]
- out = [0]
- this_offset = 0
- for c in row_group_entries:
- this_offset += c
- if this_offset >= out[-1] + target_step_size:
- out.append(this_offset)
- if this_offset != out[-1]:
- out.append(this_offset)
- out = numpy.array(out, dtype="int64")
- out = numpy.stack((out[:-1], out[1:]), axis=1)
-
- step_mask = (
- out[:, 1] - out[:, 0]
- > (1 + step_size_safety_factor) * target_step_size
+ # cumulative row counts give the absolute offset at each row-group boundary
+ boundaries = numpy.cumsum(the_file["col_counts"]).tolist()
+ out = _aligned_steps(
+ boundaries,
+ target_step_size,
+ step_size_safety_factor,
+ arg.file,
+ "use_row_groups",
)
- if numpy.any(step_mask):
- warnings.warn(
- f"In file {arg.file}, steps: {out[step_mask]} with use_row_groups=True are "
- f"{step_size_safety_factor*100:.0f}% larger than target "
- f"step size: {target_step_size}!"
- )
else:
- n_steps_target = max(round(num_entries / target_step_size), 1)
- actual_step_size = math.ceil(num_entries / n_steps_target)
- out = numpy.array(
- [
- [
- i * actual_step_size,
- min((i + 1) * actual_step_size, num_entries),
- ]
- for i in range(n_steps_target)
- ],
- dtype="int64",
- )
+ out = _even_steps(num_entries, target_step_size)
out_uuid = file_uuid
out_steps = out.tolist()
@@ -675,6 +837,7 @@ def get_parquet_form_uuid_steps(
"uuid": out_uuid,
"compressed_form": form_json,
"form_hash_md5": form_hash,
+ "user_metadata_json": user_meta_json,
}
)
@@ -689,6 +852,7 @@ def get_parquet_form_uuid_steps(
"uuid": "junk",
"compressed_form": "junk",
"form_hash_md5": "junk",
+ "user_metadata_json": "junk",
},
None,
]
@@ -718,10 +882,17 @@ def preprocess_root(
uproot_options: dict = {},
step_size_safety_factor: float = 0.5,
allow_empty_datasets: bool = False,
+ backend: str | PreprocessBackend = "dask",
+ require_rntuple: bool = False,
+ metadata_extractor: Callable | None = None,
+ metadata_reducer: Callable | None = None,
) -> tuple[DataGroupSpec, DataGroupSpec]:
"""
Given a list of normalized file and object paths (defined in uproot), determine the steps for each file according to the supplied processing options.
+ Both TTree and RNTuple objects are auto-detected and handled; pass ``require_rntuple=True``
+ (or use the :func:`preprocess_rntuple` alias) to require that every object is an RNTuple.
+
Parameters
----------
datagroupspec : DataGroupSpec
@@ -753,6 +924,24 @@ def preprocess_root(
allow_empty_datasets : bool, default False
When a dataset query comes back completely empty, this is normally considered a processing error.
Toggle this argument to True to change this to warnings and allow incomplete returned filesets.
+ backend : str or PreprocessBackend, default "dask"
+ Execution backend for preprocessing: "dask" (default), "iterative" (immediate,
+ synchronous, dask-free), "futures" (dask-free concurrent.futures thread pool), or a
+ PreprocessBackend instance. The ``scheduler`` argument only affects the dask backend.
+ require_rntuple : bool, default False
+ If True, require every dataset to be ROOT-format and every object to be an RNTuple.
+ A parquet dataset raises a ValueError; a TTree object raises a ValueError inside the
+ worker, subject to ``skip_bad_files``/``file_exceptions`` like any other per-file
+ error.
+ metadata_extractor : Callable or None, default None
+ User function called once per file with the open uproot file handle; must return a
+ JSON-serializable dict, stored as that file's ``metadata`` on its file spec. Runs
+ inside the per-file error handling (participates in skip_bad_files/file_exceptions)
+ and must be picklable.
+ metadata_reducer : Callable or None, default None
+ User function called once per dataset with ``{filename: extracted_dict}`` for the
+ available files; must return a dict, merged into the dataset's ``metadata`` (reducer
+ output takes precedence over existing keys) on both returned filesets.
Returns
-------
out_available : DataGroupSpec
@@ -773,9 +962,31 @@ def preprocess_root(
filetype_options=uproot_options,
step_size_safety_factor=step_size_safety_factor,
allow_empty_datasets=allow_empty_datasets,
+ backend=backend,
+ require_rntuple=require_rntuple,
+ metadata_extractor=metadata_extractor,
+ metadata_reducer=metadata_reducer,
)
+def preprocess_rntuple(
+ datagroupspec: DataGroupSpec,
+ **kwargs,
+) -> tuple[DataGroupSpec, DataGroupSpec]:
+ """
+ Preprocess datasets of ROOT files containing RNTuples, determining the steps for each file.
+
+ Alias for :func:`preprocess_root` with ``require_rntuple=True``: every dataset must be
+ ROOT-format (parquet datasets raise a ValueError) and every object must be an RNTuple
+ (a TTree raises a ValueError inside the worker, subject to ``skip_bad_files``/
+ ``file_exceptions``). ``preprocess`` and ``preprocess_root`` already auto-detect and handle
+ RNTuples transparently, so this function is for RNTuple-only workflows that want the
+ stricter contract. All other arguments are forwarded to :func:`preprocess_root`.
+ """
+ kwargs["require_rntuple"] = True
+ return preprocess_root(datagroupspec, **kwargs)
+
+
def preprocess_parquet(
datagroupspec: DataGroupSpec,
step_size: None | int = None,
@@ -789,6 +1000,9 @@ def preprocess_parquet(
parquet_options: dict = {},
step_size_safety_factor: float = 0.5,
allow_empty_datasets: bool = False,
+ backend: str | PreprocessBackend = "dask",
+ metadata_extractor: Callable | None = None,
+ metadata_reducer: Callable | None = None,
) -> tuple[DataGroupSpec, DataGroupSpec]:
"""
Given a list of normalized files, determine the form, steps, and add the metadata for each file according to the supplied processing options.
@@ -820,6 +1034,19 @@ def preprocess_parquet(
allow_empty_datasets : bool, default False
When a dataset query comes back completely empty, this is normally considered a processing error.
Toggle this argument to True to change this to warnings and allow incomplete returned filesets.
+ backend : str or PreprocessBackend, default "dask"
+ Execution backend for preprocessing: "dask" (default), "iterative" (immediate,
+ synchronous, dask-free), "futures" (dask-free concurrent.futures thread pool), or a
+ PreprocessBackend instance. The ``scheduler`` argument only affects the dask backend.
+ metadata_extractor : Callable or None, default None
+ User function called once per file with the parquet metadata mapping from
+ ``awkward.metadata_from_parquet``; must return a JSON-serializable dict, stored as
+ that file's ``metadata`` on its file spec. Runs inside the per-file error handling
+ (participates in skip_bad_files/file_exceptions) and must be picklable.
+ metadata_reducer : Callable or None, default None
+ User function called once per dataset with ``{filename: extracted_dict}`` for the
+ available files; must return a dict, merged into the dataset's ``metadata`` (reducer
+ output takes precedence over existing keys) on both returned filesets.
Returns
-------
out_available : DataGroupSpec
@@ -840,6 +1067,9 @@ def preprocess_parquet(
filetype_options=parquet_options,
step_size_safety_factor=step_size_safety_factor,
allow_empty_datasets=allow_empty_datasets,
+ backend=backend,
+ metadata_extractor=metadata_extractor,
+ metadata_reducer=metadata_reducer,
)
@@ -856,6 +1086,10 @@ def _preprocess_pydantic(
filetype_options: dict = {},
step_size_safety_factor: float = 0.5,
allow_empty_datasets: bool = False,
+ require_rntuple: bool = False,
+ backend: str | PreprocessBackend = "dask",
+ metadata_extractor: Callable | None = None,
+ metadata_reducer: Callable | None = None,
) -> tuple[DataGroupSpec, DataGroupSpec]:
"""
Internal function to preprocess either ROOT or parquet DatasetSpecs in a DataGroupSpec.
@@ -890,6 +1124,11 @@ def _preprocess_pydantic(
Warn if aligned steps exceed target by this factor.
allow_empty_datasets : bool, default False
If True, warn instead of raising when a dataset has no accessible files.
+ backend : str or PreprocessBackend, default "dask"
+ Execution backend for the per-dataset map-reduce. One of "dask" (default),
+ "iterative" (immediate, synchronous, dask-free), "futures" (a dask-free
+ concurrent.futures thread pool), or a PreprocessBackend instance for full control.
+ ``scheduler`` only affects the dask backend.
Returns
-------
@@ -909,15 +1148,18 @@ def _preprocess_pydantic(
raise ValueError(
f"_preprocess_pydantic expects a DataGroupSpec, got {type(datagroupspec)}"
)
-
- dask = _import_dask()
- dask_awkward = _import_dask_awkward()
+ _validate_step_size(step_size)
+ if len(datagroupspec) == 0:
+ return DataGroupSpec({}), DataGroupSpec({})
out_updated = datagroupspec.model_dump()
out_available = datagroupspec.model_dump()
+ # Build one map-reduce job per dataset. The map worker (get_steps /
+ # get_parquet_form_uuid_steps) and the concatenating reduce are backend-agnostic; only the
+ # execution strategy (dask graph vs. futures vs. synchronous) is selected via `backend`.
all_ak_norm_files = {}
- files_to_preprocess = {}
+ jobs = {}
for name, info in datagroupspec.items():
norm_files = _normalize_pydantic_file_info(info)
fields = ["file", "object_path", "steps", "num_entries", "uuid"]
@@ -927,26 +1169,9 @@ def _preprocess_pydantic(
)
all_ak_norm_files[name] = ak_norm_files
- dak_norm_files = dask_awkward.from_awkward(
- ak_norm_files, math.ceil(len(ak_norm_files) / files_per_batch)
- )
-
- concat_fn = partial(
- awkward.concatenate,
- axis=0,
- )
-
- split_every = 8
-
- files_trl_label = f"preprocess-{name}"
- files_trl_token = dask.base.tokenize(dak_norm_files, concat_fn, split_every)
- files_trl_name = f"{files_trl_label}-{files_trl_token}"
- files_trl_tree_node_name = f"{files_trl_label}-tree-node-{files_trl_token}"
-
if info.format == "root":
- files_part = dask_awkward.map_partitions(
+ map_fn = partial(
get_steps,
- dak_norm_files,
step_size=step_size,
align_clusters=use_alignment_boundaries,
recalculate_steps=recalculate_steps,
@@ -956,12 +1181,17 @@ def _preprocess_pydantic(
step_size_safety_factor=step_size_safety_factor,
legacy_form_key=False, # in the pydantic preprocess function, the output form key is always "compressed_form", "form" is a method to extract the uncompressed form
uproot_options=filetype_options,
- meta=dask_awkward.lib.core.empty_typetracer(),
+ require_rntuple=require_rntuple,
+ metadata_extractor=metadata_extractor,
)
elif info.format == "parquet":
- files_part = dask_awkward.map_partitions(
+ if require_rntuple:
+ raise ValueError(
+ f"require_rntuple=True but dataset {name!r} is parquet-format; "
+ "only ROOT files can contain RNTuples."
+ )
+ map_fn = partial(
get_parquet_form_uuid_steps,
- dak_norm_files,
step_size=step_size,
use_row_groups=use_alignment_boundaries,
recalculate_steps=recalculate_steps,
@@ -970,36 +1200,29 @@ def _preprocess_pydantic(
save_form=save_form,
step_size_safety_factor=step_size_safety_factor,
parquet_options=filetype_options,
- meta=dask_awkward.lib.core.empty_typetracer(),
+ metadata_extractor=metadata_extractor,
)
else:
raise ValueError(
f"Dataset {name} has unsupported format {info.format}, supported formats are 'root' and 'parquet'."
)
- files_trl = dask_awkward.layers.layers.AwkwardTreeReductionLayer(
- name=files_trl_name,
- name_input=files_part.name,
- npartitions_input=files_part.npartitions,
- concat_func=concat_fn,
- tree_node_func=lambda x: x,
- finalize_func=lambda x: x,
- split_every=split_every,
- tree_node_name=files_trl_tree_node_name,
- )
-
- files_graph = dask.highlevelgraph.HighLevelGraph.from_collections(
- files_trl_name, files_trl, dependencies=[files_part]
- )
-
- files_to_preprocess[name] = dask_awkward.lib.core.new_array_object(
- files_graph,
- files_trl_name,
- meta=dask_awkward.lib.core.empty_typetracer(),
- npartitions=len(files_trl.output_partitions),
+ jobs[name] = PreprocessJob(
+ array=ak_norm_files, map_fn=map_fn, files_per_batch=files_per_batch
)
- (all_processed_files,) = dask.compute(files_to_preprocess, scheduler=scheduler)
+ backend_obj = resolve_backend(backend, scheduler)
+ # Only submit() imports dask (for the dask backend); a ModuleNotFoundError here means the
+ # dask stack itself is missing, so we can point the user at the dask-free backends. Worker
+ # ImportErrors (e.g. a missing codec) surface later in result() and must NOT trigger the
+ # dask hint, so result() is called outside this guard.
+ try:
+ preprocess_task = backend_obj.submit(jobs)
+ except ModuleNotFoundError:
+ if isinstance(backend_obj, DaskBackend):
+ print_dask_backend_fallback_hint()
+ raise
+ all_processed_files = preprocess_task.result()
for name, processed_files in all_processed_files.items():
@@ -1019,7 +1242,14 @@ def _preprocess_pydantic(
continue
processed_files_without_forms = processed_files[
- ["file", "object_path", "steps", "num_entries", "uuid"]
+ [
+ "file",
+ "object_path",
+ "steps",
+ "num_entries",
+ "uuid",
+ "user_metadata_json",
+ ]
]
compressed_forms = processed_files[
@@ -1031,17 +1261,18 @@ def _preprocess_pydantic(
)
dataset_forms = []
+ fields_by_hash = {}
unique_forms = compressed_forms[unique_forms_idx]
- for thefile, formstr, num_entries in zip(
- unique_forms.file, unique_forms.compressed_form, unique_forms.num_entries
- ):
+ for item in unique_forms.to_list():
+ form = awkward.forms.from_json(decompress_form(item["compressed_form"]))
+ fields_by_hash[item["form_hash_md5"]] = set(form.fields)
# skip trivially filled or empty files
- form = awkward.forms.from_json(decompress_form(formstr))
- if num_entries >= 0 and set(form.fields) != _trivial_file_fields:
+ if set(form.fields) != _trivial_file_fields:
dataset_forms.append(form)
else:
warnings.warn(
- f"{thefile} has fields {form.fields} and num_entries={num_entries} "
+ f"{item['file']} has fields {form.fields} and "
+ f"num_entries={item['num_entries']} "
"and has been skipped during form-union determination. You will need "
"to skip this file when processing. You can either manually remove it "
"or, if it is an empty file, dynamically remove it with the function "
@@ -1049,72 +1280,87 @@ def _preprocess_pydantic(
", by default, removes empty files each dataset in a fileset."
)
- union_array = None
- union_form_jsonstr = None
- while len(dataset_forms):
- new_array = awkward.Array(dataset_forms.pop().length_zero_array())
- if union_array is None:
- union_array = new_array
- else:
- union_array = awkward.to_packed(
- awkward.merge_union_of_records(
- awkward.concatenate([union_array, new_array]), axis=0
- )
- )
- union_array.layout.parameters.update(new_array.layout.parameters)
- if union_array is not None:
- union_form = union_array.layout.form
-
- for icontent, content in enumerate(union_form.contents):
- if isinstance(content, awkward.forms.IndexedOptionForm):
- if (
- not isinstance(content.content, awkward.forms.NumpyForm)
- or content.content.primitive != "bool"
- ):
- raise ValueError(
- "IndexedOptionArrays can only contain NumpyArrays of "
- "bools in mergers of flat-tuple-like schemas!"
- )
- parameters = (
- content.content.parameters.copy()
- if content.content.parameters is not None
- else {}
- )
- # re-create IndexOptionForm with parameters of lower level array
- union_form.contents[icontent] = awkward.forms.IndexedOptionForm(
- content.index,
- content.content,
- parameters=parameters,
- form_key=content.form_key,
- )
+ union_form_jsonstr = _union_form_jsonstr(dataset_forms)
- union_form_jsonstr = union_form.to_json()
+ # Per-file experimental field bitsets: which top-level union-form fields each file
+ # carries, encoded against the union field order. These enable offline pruning of the
+ # union form when files are filtered out and per-file branch-set comparisons.
+ bitset_by_file = {}
+ if union_form_jsonstr is not None:
+ union_fields = awkward.forms.from_json(union_form_jsonstr).fields
+ for item in compressed_forms[["file", "form_hash_md5"]].to_list():
+ file_fields = fields_by_hash.get(item["form_hash_md5"])
+ if file_fields is not None:
+ bitset_by_file[item["file"]] = encode_field_bitset(
+ file_fields, union_fields
+ )
- files_available = {
+ # Index successfully-processed files by filename. Skipped/bad files were dropped as
+ # None by the worker and are simply absent here.
+ available_by_file = {
item["file"]: {
"object_path": item["object_path"],
"steps": item["steps"],
"num_entries": item["num_entries"],
"uuid": item["uuid"],
+ "experimental_field_bitset": bitset_by_file.get(item["file"]),
+ "metadata": (
+ json.loads(item["user_metadata_json"])
+ if item["user_metadata_json"] is not None
+ else None
+ ),
}
for item in awkward.drop_none(processed_files_without_forms).to_list()
}
+ # Assemble both outputs by filename in the original input order. Matching by key keeps
+ # correctness independent of the order in which the backend reduced/concatenated the
+ # per-file results -- each processed record carries its filename.
+ orig_items = all_ak_norm_files[name].to_list()
+
+ files_available = {
+ item["file"]: available_by_file[item["file"]]
+ for item in orig_items
+ if item["file"] in available_by_file
+ }
+
files_out = {}
- for proc_item, orig_item in zip(
- processed_files_without_forms.to_list(), all_ak_norm_files[name].to_list()
- ):
- item = orig_item if proc_item is None else proc_item
- files_out[item["file"]] = {
- "object_path": item["object_path"],
- "steps": item["steps"],
- "num_entries": item["num_entries"],
- "uuid": item["uuid"],
- }
+ for orig_item in orig_items:
+ filename = orig_item["file"]
+ # processed info when available, else fall back to the original input info
+ files_out[filename] = available_by_file.get(
+ filename,
+ {
+ "object_path": orig_item["object_path"],
+ "steps": orig_item["steps"],
+ "num_entries": orig_item["num_entries"],
+ "uuid": orig_item["uuid"],
+ },
+ )
out_updated[name]["files"] = files_out
out_available[name]["files"] = files_available
+ # Reduce per-file extracted metadata into the dataset-level metadata (e.g. summing
+ # per-file sums-of-weights); the reducer output takes precedence over existing keys.
+ if metadata_reducer is not None:
+ per_file_meta = {
+ fname: info["metadata"]
+ for fname, info in files_available.items()
+ if info.get("metadata") is not None
+ }
+ if per_file_meta:
+ reduced = metadata_reducer(per_file_meta)
+ if not isinstance(reduced, dict):
+ raise ValueError(
+ f"metadata_reducer must return a dict, got {type(reduced).__name__} "
+ f"for dataset {name!r}."
+ )
+ for out in (out_updated, out_available):
+ merged = dict(out[name].get("metadata") or {})
+ merged.update(reduced)
+ out[name]["metadata"] = merged
+
compressed_union_form = (
compress_form(union_form_jsonstr) if union_form_jsonstr else None
)
@@ -1141,8 +1387,16 @@ def _advertise_datagroupspec() -> None:
def _datagroupspec_to_dict(datagroupspec: DataGroupSpec) -> dict:
- """Convert a DataGroupSpec back to a plain (JSON-serializable) dict fileset."""
- return datagroupspec.model_dump()
+ """Convert a DataGroupSpec back to a plain (JSON-serializable) dict fileset.
+
+ The legacy dict format does not carry per-file metadata or experimental fields
+ (dataset-level metadata is kept).
+ """
+ return datagroupspec.model_dump(
+ exclude={
+ "__all__": {"files": {"__all__": {"metadata", "experimental_field_bitset"}}}
+ }
+ )
def preprocess(
@@ -1161,6 +1415,9 @@ def preprocess(
preprocess_legacy_root: bool = False,
use_row_groups: bool = False,
parquet_options: dict = {},
+ backend: str | PreprocessBackend = "dask",
+ metadata_extractor: Callable | None = None,
+ metadata_reducer: Callable | None = None,
) -> tuple[DataGroupSpec, DataGroupSpec] | tuple[dict, dict]:
"""
Given a list of normalized file and object paths (defined in uproot), determine the steps for each file according to the supplied processing options.
@@ -1208,6 +1465,24 @@ def preprocess(
Calculate steps according to the row_groups in the parquet files (only applies to DataGroupSpec datasets with parquet files).
parquet_options : dict, default {}
Options to pass to get_parquet_form_uuid_steps for opening parquet files (only applies to DataGroupSpec datasets with parquet files).
+ backend : str or PreprocessBackend, default "dask"
+ Execution backend for preprocessing: "dask" (default), "iterative" (immediate,
+ synchronous, dask-free), "futures" (dask-free concurrent.futures thread pool), or a
+ PreprocessBackend instance. The ``scheduler`` argument only affects the dask backend.
+ Ignored when ``preprocess_legacy_root=True`` (the legacy path is always dask-based).
+ metadata_extractor : Callable or None, default None
+ User function called once per file with the open file handle (uproot
+ ReadOnlyDirectory for ROOT files, ``awkward.metadata_from_parquet`` mapping for
+ parquet); must return a JSON-serializable dict, stored as that file's ``metadata``
+ on its file spec (pydantic output only; the legacy dict output does not carry
+ per-file metadata). Runs inside the per-file error handling and must be picklable.
+ Not supported with ``preprocess_legacy_root=True``.
+ metadata_reducer : Callable or None, default None
+ User function called once per dataset with ``{filename: extracted_dict}`` for the
+ available files; must return a dict, merged into the dataset's ``metadata``
+ (reducer output takes precedence over existing keys) on both returned filesets.
+ The reduced dataset-level metadata survives dict output as well. Not supported
+ with ``preprocess_legacy_root=True``.
Returns
-------
out_available : DataGroupSpec | dict
@@ -1217,6 +1492,14 @@ def preprocess(
"""
input_is_dict = not isinstance(fileset, DataGroupSpec)
+ if preprocess_legacy_root and (
+ metadata_extractor is not None or metadata_reducer is not None
+ ):
+ raise ValueError(
+ "metadata_extractor/metadata_reducer are not supported with "
+ "preprocess_legacy_root=True; use the pydantic preprocessing path."
+ )
+
if preprocess_legacy_root:
# use the legacy root TTree preprocessing function if requested;
# the legacy path historically defaulted to save_form=False
@@ -1269,6 +1552,10 @@ def preprocess(
"Entries assigned via item assignment are not validated; rebuild or re-validate the "
"DataGroupSpec (e.g. DataGroupSpec.model_validate(...)) so every entry is a DatasetSpec."
)
+ # Resolve the backend once so warnings (e.g. an ignored scheduler) are emitted a single
+ # time even for mixed root+parquet filesets, which dispatch to two sub-calls below.
+ backend = resolve_backend(backend, scheduler)
+ scheduler = None
# split datasetspecs into uproot and parquet files, keeping track of original order
original_order = list(datasetspecs.keys())
formats = [dss.format for dss in datasetspecs.values()]
@@ -1291,6 +1578,9 @@ def preprocess(
uproot_options=uproot_options,
step_size_safety_factor=step_size_safety_factor,
allow_empty_datasets=allow_empty_datasets,
+ backend=backend,
+ metadata_extractor=metadata_extractor,
+ metadata_reducer=metadata_reducer,
)
out_available_parquet, out_updated_parquet = preprocess_parquet(
datasetspecs.filter_datasets(
@@ -1307,6 +1597,9 @@ def preprocess(
parquet_options=parquet_options,
step_size_safety_factor=step_size_safety_factor,
allow_empty_datasets=allow_empty_datasets,
+ backend=backend,
+ metadata_extractor=metadata_extractor,
+ metadata_reducer=metadata_reducer,
)
# recombine outputs in original order, skipping datasets removed due to allow_empty_datasets.
# The sub-results are already-validated DatasetSpec instances, so use model_construct to
diff --git a/src/coffea/dataset_tools/preprocess_backends.py b/src/coffea/dataset_tools/preprocess_backends.py
new file mode 100644
index 000000000..45a09282a
--- /dev/null
+++ b/src/coffea/dataset_tools/preprocess_backends.py
@@ -0,0 +1,452 @@
+"""Switchable execution backends for pydantic preprocessing.
+
+Preprocessing a :class:`~coffea.dataset_tools.filespec.DataGroupSpec` is, per dataset, an
+ordered map-reduce: a worker (``get_steps`` / ``get_parquet_form_uuid_steps``) is mapped over
+batches of normalized file records and the resulting awkward arrays are concatenated. This
+module factors that map-reduce out from :func:`coffea.dataset_tools.preprocess._preprocess_pydantic`
+behind a small backend interface so the same worker can run on:
+
+- :class:`DaskBackend` -- the default: ``dask_awkward.map_partitions`` + an
+ ``AwkwardTreeReductionLayer`` + ``dask.compute``. This backend builds and executes a dask
+ task graph to orchestrate the map-reduce.
+- :class:`IterativeBackend` -- immediate (synchronous, single process) execution, mirroring
+ coffea's ``IterativeExecutor``.
+- :class:`FuturesBackend` -- a :mod:`concurrent.futures` pool, threads by default (preprocessing
+ is IO-bound) with an opt-in process pool (mirrors ``FuturesExecutor``).
+
+The iterative and futures backends do not build a dask *task graph* for orchestration, and they
+are dask-free for parquet input and for ROOT input of either flavor. ROOT form extraction
+(``save_form=True``) uses uproot's own non-dask form builder for both TTree and RNTuple, producing
+a form byte-identical to the ``uproot.dask`` form; it only falls back to ``uproot.dask`` (and thus
+dask) if that internal uproot helper is unavailable on an unexpected uproot version.
+
+A ``Backend`` turns a "computable" into a future-like ``Task`` (``result`` / ``partial_result``
+/ ``wait``) whose ``result()`` blocks. Here the computable is a mapping
+``{dataset_name: PreprocessJob}`` and the result is ``{dataset_name: awkward.Array}``.
+
+Ordering is not load-bearing for correctness: the reduction concatenates per-batch results, and
+the concrete backends happen to preserve submission order, but the downstream post-processing
+assembles its outputs *by filename* (see ``_preprocess_pydantic``), so a backend is free to
+reduce in any order.
+"""
+
+from __future__ import annotations
+
+import math
+import warnings
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterable, Mapping
+from concurrent.futures import Executor, Future, ProcessPoolExecutor, ThreadPoolExecutor
+from concurrent.futures import wait as futures_wait
+from dataclasses import dataclass, replace
+from functools import partial
+from typing import Protocol, runtime_checkable
+
+import awkward
+
+from coffea.util import _import_dask, _import_dask_awkward, coffea_console
+
+__all__ = [
+ "PreprocessJob",
+ "PreprocessTask",
+ "PreprocessBackend",
+ "IterativeBackend",
+ "FuturesBackend",
+ "DaskBackend",
+ "ordered_concat",
+ "resolve_backend",
+]
+
+
+def ordered_concat(parts: Iterable[awkward.Array | None]) -> awkward.Array | None:
+ """Concatenate mapped outputs, preserving input order.
+
+ ``None`` parts are dropped. Returns ``None`` if nothing remains. This is the reduction step
+ for preprocessing. Order is preserved for tidiness/determinism, but is not required for
+ correctness: the downstream consumer assembles its outputs by filename, so the reduce may run
+ in any order.
+ """
+ parts = [part for part in parts if part is not None]
+ if len(parts) == 0:
+ return None
+ if len(parts) == 1:
+ return parts[0]
+ return awkward.concatenate(parts, axis=0)
+
+
+def _iter_batches(array: awkward.Array, files_per_batch: int):
+ """Yield contiguous slices of ``array`` of at most ``files_per_batch`` rows, in order.
+
+ An empty array yields itself once so the worker still produces a correctly-typed empty
+ result (matching the single-partition dask behaviour).
+ """
+ n = len(array)
+ if n == 0:
+ yield array
+ return
+ step = max(1, files_per_batch)
+ for start in range(0, n, step):
+ yield array[start : start + step]
+
+
+@dataclass
+class PreprocessJob:
+ """One dataset's preprocessing work: map ``map_fn`` over batches of ``array`` and concat.
+
+ Parameters
+ ----------
+ array : awkward.Array
+ Normalized file records (fields ``file``, ``object_path``, ``steps``,
+ ``num_entries``, ``uuid``) for a single dataset.
+ map_fn : Callable[[awkward.Array], awkward.Array]
+ The per-batch worker, e.g. ``functools.partial(get_steps, **options)``. Must be a
+ top-level (picklable) callable so it can run under a process pool.
+ files_per_batch : int, default 1
+ Number of files handled per unit of work. Larger values mean fewer, heavier tasks.
+ """
+
+ array: awkward.Array
+ map_fn: Callable[[awkward.Array], awkward.Array]
+ files_per_batch: int = 1
+
+
+@runtime_checkable
+class PreprocessTask(Protocol):
+ """Future-like handle to a submitted preprocessing computable."""
+
+ def result(self) -> dict[str, awkward.Array]:
+ """Block until done and return ``{dataset_name: concatenated awkward.Array}``."""
+
+ def partial_result(self) -> dict[str, awkward.Array]:
+ """Return results for the work finished so far, without blocking on the rest.
+
+ Datasets/batches still running are omitted; a dataset appears only once at least one of
+ its batches has completed. Intended for resumable/progress use. Backends that cannot
+ produce a cheap partial (e.g. a fused dask graph) may return the full :meth:`result`.
+ """
+
+ def wait(self) -> None:
+ """Block until the computation has finished, without returning the result."""
+
+
+class PreprocessBackend(ABC):
+ """Base class for preprocessing execution backends."""
+
+ @abstractmethod
+ def submit(self, jobs: Mapping[str, PreprocessJob]) -> PreprocessTask:
+ """Start executing ``jobs`` and return a non-blocking :class:`PreprocessTask`."""
+
+ def compute(self, jobs: Mapping[str, PreprocessJob]) -> dict[str, awkward.Array]:
+ """Convenience blocking call: ``submit(jobs).result()``."""
+ return self.submit(jobs).result()
+
+
+@dataclass
+class _CompletedTask:
+ """A :class:`PreprocessTask` whose work is already done (used by the iterative backend)."""
+
+ _results: dict[str, awkward.Array]
+
+ def result(self) -> dict[str, awkward.Array]:
+ return self._results
+
+ def partial_result(self) -> dict[str, awkward.Array]:
+ # work already finished in submit(), so the partial result is the full result
+ return self._results
+
+ def wait(self) -> None:
+ return None
+
+
+@dataclass
+class IterativeBackend(PreprocessBackend):
+ """Execute preprocessing immediately -- synchronously, in the current thread.
+
+ Named to mirror coffea's ``IterativeExecutor``; the execution is "immediate" in that each
+ batch is mapped and reduced inline with no deferral or task graph. No dask required.
+ """
+
+ def submit(self, jobs: Mapping[str, PreprocessJob]) -> PreprocessTask:
+ results: dict[str, awkward.Array] = {}
+ for name, job in jobs.items():
+ parts = [
+ job.map_fn(batch)
+ for batch in _iter_batches(job.array, job.files_per_batch)
+ ]
+ results[name] = ordered_concat(parts)
+ return _CompletedTask(results)
+
+
+class _FuturesTask:
+ """A :class:`PreprocessTask` backed by :mod:`concurrent.futures` futures.
+
+ Futures are kept grouped and ordered per dataset so results reassemble in the original file
+ order regardless of completion order. A pool created by the backend is shut down once the
+ result has been gathered; an externally-supplied pool is left open.
+ """
+
+ def __init__(
+ self,
+ name_to_futures: dict[str, list[Future]],
+ pool: Executor,
+ owns_pool: bool,
+ ):
+ self._name_to_futures = name_to_futures
+ self._pool = pool
+ self._owns_pool = owns_pool
+
+ def _all_futures(self) -> list[Future]:
+ return [fut for futs in self._name_to_futures.values() for fut in futs]
+
+ def wait(self) -> None:
+ futures_wait(self._all_futures())
+
+ def result(self) -> dict[str, awkward.Array]:
+ try:
+ results: dict[str, awkward.Array] = {}
+ for name, futs in self._name_to_futures.items():
+ # .result() re-raises worker exceptions here; ordering follows submission
+ parts = [fut.result() for fut in futs]
+ results[name] = ordered_concat(parts)
+ except BaseException:
+ # fail fast: cancel batches that have not started so a fatal error does not wait
+ # for the rest of the fileset to be processed
+ for fut in self._all_futures():
+ fut.cancel()
+ if self._owns_pool:
+ self._pool.shutdown(wait=False, cancel_futures=True)
+ raise
+ if self._owns_pool:
+ self._pool.shutdown()
+ return results
+
+ def partial_result(self) -> dict[str, awkward.Array]:
+ # Gather only the batches that have completed successfully, per dataset, without blocking.
+ # The owned pool is not shut down here; unfinished work may still be running.
+ results: dict[str, awkward.Array] = {}
+ for name, futs in self._name_to_futures.items():
+ done_parts = [
+ fut.result()
+ for fut in futs
+ if fut.done() and not fut.cancelled() and fut.exception() is None
+ ]
+ merged = ordered_concat(done_parts)
+ if merged is not None:
+ results[name] = merged
+ return results
+
+ def __del__(self):
+ # Safety net: an owned pool is normally shut down in result(); if the caller only
+ # wait()s or drops the task without calling result(), avoid leaking the pool.
+ # Executor.shutdown is idempotent, so a double call after result() is harmless.
+ if getattr(self, "_owns_pool", False):
+ self._pool.shutdown(wait=False)
+
+
+@dataclass
+class FuturesBackend(PreprocessBackend):
+ """Execute preprocessing over a :mod:`concurrent.futures` pool. No dask required.
+
+ Batches from every dataset are submitted to a single shared pool for good cross-dataset
+ load balancing, then reassembled per dataset in order.
+
+ Parameters
+ ----------
+ workers : int or None, default None
+ Number of workers when this backend creates the pool. ``None`` uses the executor's
+ default sizing (for :class:`~concurrent.futures.ThreadPoolExecutor`,
+ ``min(32, os.cpu_count() + 4)``; for
+ :class:`~concurrent.futures.ProcessPoolExecutor`, ``os.cpu_count()``).
+ use_processes : bool, default False
+ Create a :class:`~concurrent.futures.ProcessPoolExecutor` instead of the default
+ :class:`~concurrent.futures.ThreadPoolExecutor`. Threads are preferred because
+ preprocessing is IO-bound (opening files / reading metadata) and avoids pickling
+ the awkward arrays; use processes when form extraction is CPU-heavy.
+ pool : concurrent.futures.Executor or Callable, optional
+ An existing executor instance to reuse (left open by this backend), or a callable
+ (e.g. an ``Executor`` subclass) invoked as ``pool(max_workers=workers)``. Overrides
+ ``use_processes`` when given.
+ """
+
+ workers: int | None = None
+ use_processes: bool = False
+ pool: Executor | Callable[..., Executor] | None = None
+
+ def _make_pool(self) -> tuple[Executor, bool]:
+ if isinstance(self.pool, Executor):
+ return self.pool, False
+ if self.pool is not None:
+ return self.pool(max_workers=self.workers), True
+ if self.use_processes:
+ return ProcessPoolExecutor(max_workers=self.workers), True
+ return ThreadPoolExecutor(max_workers=self.workers), True
+
+ def submit(self, jobs: Mapping[str, PreprocessJob]) -> PreprocessTask:
+ pool, owns_pool = self._make_pool()
+ try:
+ name_to_futures: dict[str, list[Future]] = {}
+ for name, job in jobs.items():
+ name_to_futures[name] = [
+ pool.submit(job.map_fn, batch)
+ for batch in _iter_batches(job.array, job.files_per_batch)
+ ]
+ except BaseException:
+ if owns_pool:
+ pool.shutdown()
+ raise
+ return _FuturesTask(name_to_futures, pool, owns_pool)
+
+
+class _DaskTask:
+ """A :class:`PreprocessTask` wrapping unmaterialized dask-awkward collections."""
+
+ def __init__(self, collections: dict, scheduler):
+ self._collections = collections
+ self._scheduler = scheduler
+ self._computed: dict[str, awkward.Array] | None = None
+
+ def wait(self) -> None:
+ self.result()
+
+ def result(self) -> dict[str, awkward.Array]:
+ if self._computed is None:
+ dask = _import_dask()
+ (self._computed,) = dask.compute(
+ self._collections, scheduler=self._scheduler
+ )
+ return self._computed
+
+ def partial_result(self) -> dict[str, awkward.Array]:
+ # A fused dask graph has no cheap partial: fall back to the full (blocking) result.
+ return self.result()
+
+
+@dataclass
+class DaskBackend(PreprocessBackend):
+ """Execute preprocessing as a dask-awkward task graph (the default).
+
+ Per dataset this builds ``from_awkward`` -> ``map_partitions`` -> ``AwkwardTreeReductionLayer``
+ and lets a single ``dask.compute`` materialize all datasets together.
+
+ Parameters
+ ----------
+ scheduler : None or Callable or str, default None
+ Passed through to ``dask.compute``.
+ split_every : int, default 8
+ Fan-in of the tree reduction that concatenates per-batch results.
+ """
+
+ scheduler: None | Callable | str = None
+ split_every: int = 8
+
+ def _build_collection(self, name: str, job: PreprocessJob, dask, dask_awkward):
+ ak_norm_files = job.array
+ # guard files_per_batch<1 so the dask backend matches the iterative/futures batching
+ # (which clamp via max(1, files_per_batch)) instead of dividing by zero
+ files_per_batch = max(1, job.files_per_batch)
+ dak_norm_files = dask_awkward.from_awkward(
+ ak_norm_files, math.ceil(len(ak_norm_files) / files_per_batch)
+ )
+
+ concat_fn = partial(awkward.concatenate, axis=0)
+ split_every = self.split_every
+
+ files_trl_label = f"preprocess-{name}"
+ files_trl_token = dask.base.tokenize(dak_norm_files, concat_fn, split_every)
+ files_trl_name = f"{files_trl_label}-{files_trl_token}"
+ files_trl_tree_node_name = f"{files_trl_label}-tree-node-{files_trl_token}"
+
+ files_part = dask_awkward.map_partitions(
+ job.map_fn,
+ dak_norm_files,
+ meta=dask_awkward.lib.core.empty_typetracer(),
+ )
+
+ files_trl = dask_awkward.layers.layers.AwkwardTreeReductionLayer(
+ name=files_trl_name,
+ name_input=files_part.name,
+ npartitions_input=files_part.npartitions,
+ concat_func=concat_fn,
+ tree_node_func=lambda x: x,
+ finalize_func=lambda x: x,
+ split_every=split_every,
+ tree_node_name=files_trl_tree_node_name,
+ )
+
+ files_graph = dask.highlevelgraph.HighLevelGraph.from_collections(
+ files_trl_name, files_trl, dependencies=[files_part]
+ )
+
+ return dask_awkward.lib.core.new_array_object(
+ files_graph,
+ files_trl_name,
+ meta=dask_awkward.lib.core.empty_typetracer(),
+ npartitions=len(files_trl.output_partitions),
+ )
+
+ def submit(self, jobs: Mapping[str, PreprocessJob]) -> PreprocessTask:
+ dask = _import_dask()
+ dask_awkward = _import_dask_awkward()
+ collections = {
+ name: self._build_collection(name, job, dask, dask_awkward)
+ for name, job in jobs.items()
+ }
+ return _DaskTask(collections, self.scheduler)
+
+
+def resolve_backend(
+ backend: str | PreprocessBackend | None,
+ scheduler: None | Callable | str = None,
+) -> PreprocessBackend:
+ """Turn a ``backend`` selector into a :class:`PreprocessBackend` instance.
+
+ ``backend`` may be an existing :class:`PreprocessBackend`, or one of the strings ``"dask"``
+ (default), ``"iterative"``, or ``"futures"``. ``scheduler`` is forwarded to a
+ default-constructed :class:`DaskBackend`, and is injected into a passed :class:`DaskBackend`
+ instance whose own ``scheduler`` is unset (a copy is returned; the instance is not mutated).
+ If ``scheduler`` is set while a non-dask backend is selected, or the :class:`DaskBackend`
+ instance already carries its own scheduler, a warning is issued because the argument has no
+ effect there.
+ """
+ if isinstance(backend, PreprocessBackend):
+ if scheduler is not None:
+ if isinstance(backend, DaskBackend) and backend.scheduler is None:
+ return replace(backend, scheduler=scheduler)
+ warnings.warn(
+ "The 'scheduler' argument is ignored when a PreprocessBackend instance "
+ "carries its own execution configuration; set it on the instance, e.g. "
+ "DaskBackend(scheduler=...).",
+ stacklevel=2,
+ )
+ return backend
+
+ if backend is None or backend == "dask":
+ return DaskBackend(scheduler=scheduler)
+
+ if scheduler is not None:
+ warnings.warn(
+ "The 'scheduler' argument only affects the dask backend and is ignored here.",
+ stacklevel=2,
+ )
+ if backend == "iterative":
+ return IterativeBackend()
+ if backend == "futures":
+ return FuturesBackend()
+ raise ValueError(
+ f"Unknown preprocessing backend {backend!r}; expected 'dask', 'iterative', "
+ "'futures', or a PreprocessBackend instance."
+ )
+
+
+def print_dask_backend_fallback_hint() -> None:
+ """Print (via ``coffea_console``) a hint that dask-free backends exist.
+
+ Called when the default dask backend fails because dask / dask-awkward are not importable.
+ """
+ coffea_console.print(
+ "[bold red]The dask preprocessing backend is unavailable because dask / dask-awkward "
+ "could not be imported.[/bold red]\n"
+ "Preprocessing can run without dask for parquet and ROOT input: pass "
+ "[bold]backend='iterative'[/bold] (single process) or [bold]backend='futures'[/bold] "
+ "(thread pool) to preprocess() and the format-specific preprocess_* functions."
+ )
diff --git a/tests/samples/nano_genmodel_with20_700_1_with0_1100_200_with0_950_400.parquet b/tests/samples/nano_genmodel_with20_700_1_with0_1100_200_with0_950_400.parquet
new file mode 100644
index 000000000..e33c44d69
Binary files /dev/null and b/tests/samples/nano_genmodel_with20_700_1_with0_1100_200_with0_950_400.parquet differ
diff --git a/tests/samples/nano_genmodel_with20_700_1_with0_1100_200_with0_950_400.root b/tests/samples/nano_genmodel_with20_700_1_with0_1100_200_with0_950_400.root
new file mode 100644
index 000000000..add9bee5e
Binary files /dev/null and b/tests/samples/nano_genmodel_with20_700_1_with0_1100_200_with0_950_400.root differ
diff --git a/tests/samples/nano_genmodel_without_700_1_with0_1100_200_with20_950_400.parquet b/tests/samples/nano_genmodel_without_700_1_with0_1100_200_with20_950_400.parquet
new file mode 100644
index 000000000..e80c61d65
Binary files /dev/null and b/tests/samples/nano_genmodel_without_700_1_with0_1100_200_with20_950_400.parquet differ
diff --git a/tests/samples/nano_genmodel_without_700_1_with0_1100_200_with20_950_400.root b/tests/samples/nano_genmodel_without_700_1_with0_1100_200_with20_950_400.root
new file mode 100644
index 000000000..8d40d7a71
Binary files /dev/null and b/tests/samples/nano_genmodel_without_700_1_with0_1100_200_with20_950_400.root differ
diff --git a/tests/test_dataset_tools.py b/tests/test_dataset_tools.py
index 77c785d05..924fb5686 100644
--- a/tests/test_dataset_tools.py
+++ b/tests/test_dataset_tools.py
@@ -1,11 +1,13 @@
import contextlib
import copy
import json
+from functools import partial
from pathlib import Path
import awkward
import pytest
import uproot
+from uproot._util import no_filter
from uproot.exceptions import KeyInFileError
from coffea.dataset_tools import (
@@ -17,16 +19,19 @@
max_chunks_per_file,
max_files,
preprocess,
+ preprocess_rntuple,
+ preprocess_root,
slice_chunks,
slice_files,
split_fileset,
)
from coffea.dataset_tools.filespec import (
DataGroupSpec,
+ PreprocessedFiles,
)
from coffea.nanoevents import BaseSchema, NanoAODSchema
from coffea.processor.test_items import NanoEventsProcessor, NanoTestProcessor
-from coffea.util import decompress_form
+from coffea.util import _is_interpretable, decompress_form
dask_awkward = pytest.importorskip("dask_awkward")
@@ -749,6 +754,58 @@ def test_preprocess_calculate_form(dask_client, preprocess_legacy_root):
)
+def test_preprocess_rntuple():
+ """RNTuple .root files are preprocessed via preprocess_rntuple, auto-detected by
+ preprocess_root, and preprocess_rntuple rejects TTree files.
+
+ Uses the synchronous scheduler (RNTuple preprocessing does not require a distributed
+ client and this keeps the test fast and deterministic).
+ """
+ rntuple_file = "tests/samples/nano_dy_rntuple.root"
+ rntuple_fileset = DataGroupSpec({"ZJets": {"files": {rntuple_file: "Events"}}})
+
+ available, _updated = preprocess_rntuple(
+ rntuple_fileset, step_size=15, save_form=True, scheduler="synchronous"
+ )
+
+ dataset = available["ZJets"]
+ assert isinstance(dataset.files, PreprocessedFiles)
+ spec = dataset.files[rntuple_file]
+ assert spec.num_entries == 40
+ assert spec.steps == [[0, 14], [14, 28], [28, 40]]
+
+ # the stored union form matches the form uproot produces under the same interpretable-branch
+ # filter that preprocessing applies (which, for RNTuples, drops _collection* subfields)
+ raw_form = uproot.dask(
+ {rntuple_file: "Events"},
+ open_files=False,
+ ak_add_doc={"__doc__": "title", "typename": "typename"},
+ filter_name=no_filter,
+ filter_typename=no_filter,
+ filter_branch=partial(_is_interpretable, emit_warning=False),
+ ).layout.form.to_json()
+ assert decompress_form(dataset.compressed_form) == raw_form
+
+ # preprocess_root auto-detects the RNTuple and yields the same steps
+ available_auto, _ = preprocess_root(
+ DataGroupSpec({"ZJets": {"files": {rntuple_file: "Events"}}}),
+ step_size=15,
+ save_form=False,
+ scheduler="synchronous",
+ )
+ auto_spec = available_auto["ZJets"].files[rntuple_file]
+ assert auto_spec.steps == [[0, 14], [14, 28], [28, 40]]
+
+ # preprocess_rntuple must reject a TTree file
+ ttree_fileset = DataGroupSpec(
+ {"ZJets": {"files": {"tests/samples/nano_dy.root": "Events"}}}
+ )
+ with pytest.raises(ValueError, match="require_rntuple"):
+ preprocess_rntuple(
+ ttree_fileset, step_size=15, save_form=False, scheduler="synchronous"
+ )
+
+
@pytest.mark.dask_client
def test_preprocess_failed_file(dask_client):
with dask_client.as_current() as _, pytest.raises(FileNotFoundError):
@@ -950,6 +1007,22 @@ def test_filter_files(the_fileset):
assert filtered_files == target
+def test_filter_files_returns_preprocessed_files():
+ """Filtering a preprocessed DataGroupSpec keeps each dataset's files as a
+ PreprocessedFiles (all surviving files are concrete), and empty files are removed.
+ """
+ filtered = filter_files(DataGroupSpec(_updated_result))
+ assert isinstance(filtered, DataGroupSpec)
+ for name, dataset in filtered.items():
+ assert isinstance(
+ dataset.files, PreprocessedFiles
+ ), f"{name} files should be PreprocessedFiles, got {type(dataset.files).__name__}"
+ # the empty file present in _updated_result must have been filtered out
+ assert all(
+ spec.num_entries and spec.num_entries > 0 for spec in dataset.files.values()
+ )
+
+
@pytest.mark.parametrize(
"the_fileset", [_updated_result, DataGroupSpec(_updated_result)]
)
diff --git a/tests/test_dataset_tools_backends.py b/tests/test_dataset_tools_backends.py
new file mode 100644
index 000000000..42a55dfc5
--- /dev/null
+++ b/tests/test_dataset_tools_backends.py
@@ -0,0 +1,464 @@
+"""Tests for the switchable preprocessing backends in coffea.dataset_tools.preprocess_backends."""
+
+import shutil
+
+import awkward
+import pytest
+
+from coffea.dataset_tools import (
+ DaskBackend,
+ DataGroupSpec,
+ FuturesBackend,
+ IterativeBackend,
+ PreprocessBackend,
+ PreprocessJob,
+ preprocess,
+)
+from coffea.dataset_tools.preprocess_backends import (
+ _iter_batches,
+ ordered_concat,
+ resolve_backend,
+)
+
+
+def _multi_file_fileset(tmp_path):
+ """A single dataset with two identical-form ROOT files (so form union is trivial)."""
+ src = "tests/samples/nano_dy.root"
+ copy = tmp_path / "nano_dy_copy.root"
+ shutil.copy(src, copy)
+ return DataGroupSpec(
+ {
+ "ZJets": {"files": {src: "Events", str(copy): "Events"}},
+ "ZJets2": {"files": {src: "Events"}},
+ }
+ )
+
+
+# --------------------------------------------------------------------------------------
+# ordered_concat / _iter_batches unit tests
+# --------------------------------------------------------------------------------------
+
+
+def test_ordered_concat_preserves_order():
+ a = awkward.Array([1, 2])
+ b = awkward.Array([3, 4])
+ c = awkward.Array([5])
+ assert ordered_concat([a, b, c]).to_list() == [1, 2, 3, 4, 5]
+
+
+def test_ordered_concat_drops_none_and_handles_empty():
+ a = awkward.Array([1, 2])
+ assert ordered_concat([None, a, None]).to_list() == [1, 2]
+ # single element is returned as-is
+ assert ordered_concat([a]) is a
+ # nothing to concatenate
+ assert ordered_concat([]) is None
+ assert ordered_concat([None, None]) is None
+
+
+def test_iter_batches_slices_in_order():
+ arr = awkward.Array([{"x": i} for i in range(5)])
+ batches = list(_iter_batches(arr, files_per_batch=2))
+ assert [b.x.to_list() for b in batches] == [[0, 1], [2, 3], [4]]
+
+
+def test_iter_batches_empty_yields_once():
+ arr = awkward.Array([{"x": 1}])[0:0]
+ batches = list(_iter_batches(arr, files_per_batch=1))
+ assert len(batches) == 1
+ assert len(batches[0]) == 0
+
+
+# --------------------------------------------------------------------------------------
+# resolve_backend
+# --------------------------------------------------------------------------------------
+
+
+def test_resolve_backend_strings():
+ assert isinstance(resolve_backend("dask"), DaskBackend)
+ assert isinstance(resolve_backend(None), DaskBackend)
+ assert isinstance(resolve_backend("iterative"), IterativeBackend)
+ assert isinstance(resolve_backend("futures"), FuturesBackend)
+
+
+def test_resolve_backend_instance_passthrough():
+ inst = IterativeBackend()
+ assert resolve_backend(inst) is inst
+
+
+def test_resolve_backend_unknown_raises():
+ with pytest.raises(ValueError, match="Unknown preprocessing backend"):
+ resolve_backend("nonsense")
+
+
+def test_resolve_backend_scheduler_warns_for_non_dask_string():
+ with pytest.warns(UserWarning, match="only affects the dask backend"):
+ resolve_backend("iterative", scheduler="synchronous")
+
+
+def test_resolve_backend_scheduler_injected_into_dask_instance(recwarn):
+ """A DaskBackend instance with no scheduler of its own receives the scheduler argument
+ (as a copy; the original instance is not mutated), without warning."""
+ inst = DaskBackend(split_every=4)
+ resolved = resolve_backend(inst, scheduler="synchronous")
+ assert resolved.scheduler == "synchronous"
+ assert resolved.split_every == 4
+ assert inst.scheduler is None
+ assert not any("ignored" in str(w.message) for w in recwarn)
+
+
+def test_resolve_backend_scheduler_warns_for_instance():
+ # scheduler cannot be injected into a non-dask instance or a DaskBackend that already
+ # carries its own scheduler, so it warns instead of silently dropping the argument
+ with pytest.warns(UserWarning, match="ignored when a PreprocessBackend instance"):
+ resolve_backend(FuturesBackend(), scheduler="synchronous")
+ with pytest.warns(UserWarning, match="ignored when a PreprocessBackend instance"):
+ inst = DaskBackend(scheduler="threads")
+ assert resolve_backend(inst, scheduler="synchronous") is inst
+
+
+def test_resolve_backend_scheduler_no_warn_for_dask_string(recwarn):
+ resolve_backend("dask", scheduler="synchronous")
+ assert not any("ignored" in str(w.message) for w in recwarn)
+
+
+# --------------------------------------------------------------------------------------
+# Backend equivalence (dask-free backends do not import dask)
+# --------------------------------------------------------------------------------------
+
+
+def test_iterative_and_threads_agree(tmp_path):
+ dgs = _multi_file_fileset(tmp_path)
+ a_iter, u_iter = preprocess(dgs, step_size=7, save_form=True, backend="iterative")
+ a_thr, u_thr = preprocess(
+ dgs, step_size=7, save_form=True, backend=FuturesBackend(workers=2)
+ )
+ # DatasetSpec.__eq__ compares decoded forms (ignoring non-deterministic compressed bytes)
+ assert a_iter == a_thr
+ assert u_iter == u_thr
+ # order of files within a dataset is preserved by the ordered reduce
+ src = "tests/samples/nano_dy.root"
+ assert list(a_iter["ZJets"].files)[0] == src
+
+
+def test_iterative_preserves_steps_and_form(tmp_path):
+ dgs = _multi_file_fileset(tmp_path)
+ available, _ = preprocess(dgs, step_size=7, save_form=True, backend="iterative")
+ ds = available["ZJets"]
+ assert len(ds.files) == 2
+ for fs in ds.files.values():
+ assert fs.num_entries == 40
+ assert fs.uuid is not None
+ assert fs.steps[0][0] == 0
+ assert ds.compressed_form is not None
+
+
+def test_dask_matches_dask_free(tmp_path):
+ pytest.importorskip("dask")
+ pytest.importorskip("dask_awkward")
+ dgs = _multi_file_fileset(tmp_path)
+ a_dask, u_dask = preprocess(
+ dgs, step_size=7, save_form=True, backend="dask", scheduler="synchronous"
+ )
+ a_iter, u_iter = preprocess(dgs, step_size=7, save_form=True, backend="iterative")
+ assert a_dask == a_iter
+ assert u_dask == u_iter
+
+
+# --------------------------------------------------------------------------------------
+# dask-unavailable fallback hint
+# --------------------------------------------------------------------------------------
+
+
+def test_dask_import_failure_prints_hint(tmp_path, monkeypatch):
+ """When the default dask backend can't import dask, a hint about the dask-free
+ backends is emitted before the ModuleNotFoundError propagates."""
+ import importlib
+
+ backends_mod = importlib.import_module("coffea.dataset_tools.preprocess_backends")
+ # the package re-exports the `preprocess` function, shadowing the submodule attribute,
+ # so fetch the module object explicitly rather than via attribute access
+ preprocess_mod = importlib.import_module("coffea.dataset_tools.preprocess")
+
+ def _boom():
+ raise ModuleNotFoundError("no dask here")
+
+ monkeypatch.setattr(backends_mod, "_import_dask", _boom)
+
+ called = {"hint": False}
+ orig_hint = preprocess_mod.print_dask_backend_fallback_hint
+
+ def _hint():
+ called["hint"] = True
+ orig_hint()
+
+ monkeypatch.setattr(preprocess_mod, "print_dask_backend_fallback_hint", _hint)
+
+ dgs = _multi_file_fileset(tmp_path)
+ with pytest.raises(ModuleNotFoundError):
+ preprocess(dgs, step_size=7, save_form=True, backend="dask")
+ assert called["hint"] is True
+
+
+def test_worker_modulenotfound_does_not_trigger_dask_hint(tmp_path, monkeypatch):
+ """A ModuleNotFoundError raised by a worker during compute (surfaced in Task.result())
+ must NOT print the 'dask could not be imported' hint -- only a failure to import dask in
+ submit() should."""
+ import importlib
+
+ preprocess_mod = importlib.import_module("coffea.dataset_tools.preprocess")
+
+ class _RaisingTask:
+ def result(self):
+ raise ModuleNotFoundError("missing xrootd inside a worker")
+
+ class _FakeDaskBackend(DaskBackend):
+ def submit(self, jobs): # submit succeeds; the error is deferred to result()
+ return _RaisingTask()
+
+ called = {"hint": False}
+ monkeypatch.setattr(
+ preprocess_mod,
+ "print_dask_backend_fallback_hint",
+ lambda: called.__setitem__("hint", True),
+ )
+ monkeypatch.setattr(
+ preprocess_mod, "resolve_backend", lambda b, s=None: _FakeDaskBackend()
+ )
+
+ dgs = _multi_file_fileset(tmp_path)
+ with pytest.raises(ModuleNotFoundError, match="missing xrootd"):
+ preprocess(dgs, step_size=7, save_form=True, backend="dask")
+ assert called["hint"] is False
+
+
+def test_empty_parquet_file_does_not_crash(tmp_path):
+ """A 0-row parquet file must yield steps=[[0, 0]] instead of a ZeroDivisionError,
+ mirroring the ROOT get_steps num_entries==0 guard."""
+ from coffea.dataset_tools import preprocess_parquet
+
+ path = tmp_path / "empty.parquet"
+ awkward.to_parquet(awkward.Array([{"x": 1.0, "y": 2}])[0:0], str(path))
+ dgs = DataGroupSpec({"E": {"files": {str(path): None}}})
+
+ # recalculate_steps=True forces the step-computation branch
+ _available, updated = preprocess_parquet(
+ dgs, recalculate_steps=True, save_form=True, backend="iterative"
+ )
+ fs = next(iter(updated["E"].files.values()))
+ assert fs.num_entries == 0
+ assert fs.steps == [[0, 0]]
+
+
+def test_fallback_hint_mentions_backends(capsys):
+ from coffea.dataset_tools.preprocess_backends import (
+ print_dask_backend_fallback_hint,
+ )
+
+ print_dask_backend_fallback_hint()
+ out = capsys.readouterr().out
+ assert "iterative" in out
+ assert "futures" in out
+
+
+@pytest.mark.parametrize(
+ "sample, is_rntuple",
+ [
+ ("tests/samples/nano_dy.root", False),
+ ("tests/samples/nano_dy_rntuple.root", True),
+ ],
+)
+def test_awkward_form_json_matches_uproot_dask(sample, is_rntuple):
+ """The dask-free form builder must stay byte-identical to uproot.dask's form for both TTree
+ and RNTuple, otherwise the iterative/futures backends would store forms incompatible with the
+ dask read path used at analysis time."""
+ pytest.importorskip("dask")
+ pytest.importorskip("dask_awkward")
+ from functools import partial
+
+ import uproot
+ from uproot._util import no_filter
+
+ from coffea.dataset_tools.preprocess import _FORM_AK_ADD_DOC, _awkward_form_json
+ from coffea.util import _is_interpretable
+
+ tree = uproot.open({sample: None})["Events"]
+ filt = partial(_is_interpretable, emit_warning=False)
+ if is_rntuple:
+ # RNTuples cannot build a form from an already-open object via uproot.dask; pass the spec
+ dask_form = uproot.dask(
+ {sample: "Events"},
+ open_files=False,
+ full_paths=True,
+ ak_add_doc=_FORM_AK_ADD_DOC,
+ filter_name=no_filter,
+ filter_typename=no_filter,
+ filter_branch=filt,
+ ).layout.form.to_json()
+ else:
+ dask_form = uproot.dask(
+ tree,
+ ak_add_doc=_FORM_AK_ADD_DOC,
+ filter_name=no_filter,
+ filter_typename=no_filter,
+ filter_branch=filt,
+ ).layout.form.to_json()
+ assert _awkward_form_json(tree, is_rntuple) == dask_form
+
+
+def test_step_size_zero_or_negative_raises(tmp_path):
+ """step_size < 1 is rejected with a clear ValueError at the entrypoint rather than a bare
+ ZeroDivisionError from deep inside a worker."""
+ dgs = _multi_file_fileset(tmp_path)
+ for bad in (0, -5):
+ with pytest.raises(ValueError, match="step_size must be a positive integer"):
+ preprocess(dgs, step_size=bad, save_form=False, backend="iterative")
+ # None and >= 1 are accepted (smoke: no raise)
+ preprocess(dgs, step_size=1, save_form=False, backend="iterative")
+
+
+def test_step_size_zero_raises_in_legacy_path():
+ """The legacy path validates step_size too, before importing dask or opening files."""
+ from coffea.dataset_tools import preprocess_legacy
+
+ with pytest.raises(ValueError, match="step_size must be a positive integer"):
+ preprocess_legacy(
+ {"ds": {"files": {"nonexistent.root": "Events"}}}, step_size=0
+ )
+
+
+def test_backends_are_preprocessbackend_instances():
+ assert isinstance(DaskBackend(), PreprocessBackend)
+ assert isinstance(IterativeBackend(), PreprocessBackend)
+ assert isinstance(FuturesBackend(), PreprocessBackend)
+
+
+def test_skipped_bad_file_assembled_by_filename(tmp_path):
+ """A skipped bad file is absent from `available` but retained in `updated` with its original
+ input info, assembled by filename in the original input order (no positional zip).
+ """
+ good = "tests/samples/nano_dy.root"
+ bad = str(tmp_path / "does_not_exist.root")
+ dgs = DataGroupSpec({"ZJets": {"files": {bad: "Events", good: "Events"}}})
+ available, updated = preprocess(
+ dgs, step_size=20, save_form=True, skip_bad_files=True, backend="iterative"
+ )
+ # updated keeps both files, in the original input order; available drops the bad one
+ assert list(updated["ZJets"].files) == [bad, good]
+ assert bad not in available["ZJets"].files
+ assert good in available["ZJets"].files
+
+
+def test_partial_result_equals_result_when_complete():
+ """partial_result() agrees with result() once all work has finished, for the eager and
+ futures backends."""
+ arr = awkward.Array([{"x": i} for i in range(4)])
+ jobs = {
+ "d": PreprocessJob(array=arr, map_fn=lambda batch: batch, files_per_batch=1)
+ }
+ for backend in (IterativeBackend(), FuturesBackend(workers=2)):
+ task = backend.submit(jobs)
+ task.wait()
+ partial = task.partial_result()
+ result = task.result()
+ assert result["d"].to_list() == arr.to_list()
+ assert partial["d"].to_list() == result["d"].to_list()
+
+
+def test_futures_default_workers_uses_executor_default():
+ """FuturesBackend defaults workers to None, deferring pool sizing to the executor
+ (parallel by default), including via the string selector."""
+ assert FuturesBackend().workers is None
+ assert resolve_backend("futures").workers is None
+
+
+def test_futures_string_backend_matches_iterative(tmp_path):
+ dgs = _multi_file_fileset(tmp_path)
+ a_fut, u_fut = preprocess(dgs, step_size=7, save_form=True, backend="futures")
+ a_iter, u_iter = preprocess(dgs, step_size=7, save_form=True, backend="iterative")
+ assert a_fut == a_iter
+ assert u_fut == u_iter
+
+
+def test_futures_result_failure_cancels_pending():
+ """When a batch fails, result() re-raises after cancelling batches that have not started,
+ so a fatal error does not wait for the rest of the fileset to be processed."""
+ import threading
+ from concurrent.futures import Future, ThreadPoolExecutor
+
+ from coffea.dataset_tools.preprocess_backends import _FuturesTask
+
+ release = threading.Event()
+ ran = threading.Event()
+ pool = ThreadPoolExecutor(max_workers=1)
+ try:
+ pool.submit(release.wait) # occupies the only worker
+ pending = pool.submit(ran.set) # queued behind the blocker
+
+ failed = Future()
+ failed.set_exception(ValueError("boom"))
+
+ task = _FuturesTask({"d": [failed, pending]}, pool, owns_pool=True)
+ with pytest.raises(ValueError, match="boom"):
+ task.result()
+ assert pending.cancelled()
+ assert not ran.is_set()
+ finally:
+ release.set()
+
+
+def test_preprocess_rntuple_skips_ttree_file_when_requested():
+ """With require_rntuple, a TTree object raises a ValueError that participates in
+ skip_bad_files/file_exceptions like any other per-file error."""
+ from coffea.dataset_tools import preprocess_rntuple
+
+ rnt = "tests/samples/nano_dy_rntuple.root"
+ ttree = "tests/samples/nano_dy.root"
+ dgs = DataGroupSpec({"D": {"files": {rnt: "Events", ttree: "Events"}}})
+
+ with pytest.raises(ValueError, match="not an RNTuple"):
+ preprocess_rntuple(dgs, save_form=False, backend="iterative")
+
+ available, updated = preprocess_rntuple(
+ dgs,
+ save_form=False,
+ backend="iterative",
+ skip_bad_files=True,
+ file_exceptions=(OSError, ValueError),
+ )
+ assert list(available["D"].files) == [rnt]
+ assert list(updated["D"].files) == [rnt, ttree]
+
+
+def test_preprocess_rntuple_rejects_parquet(tmp_path):
+ """require_rntuple rejects parquet-format datasets with a clear error."""
+ from coffea.dataset_tools import preprocess_rntuple
+
+ path = tmp_path / "d.parquet"
+ awkward.to_parquet(awkward.Array([{"x": 1.0}]), str(path))
+ dgs = DataGroupSpec({"P": {"files": {str(path): None}}})
+ with pytest.raises(ValueError, match="parquet-format"):
+ preprocess_rntuple(dgs, save_form=False, backend="iterative")
+
+
+def test_workers_validate_step_size():
+ """get_steps and get_parquet_form_uuid_steps reject step_size < 1 directly, including
+ negative values (which would otherwise silently produce a single step per file)."""
+ from coffea.dataset_tools.preprocess import get_parquet_form_uuid_steps, get_steps
+
+ normed = awkward.Array(
+ [
+ {
+ "file": "tests/samples/nano_dy.root",
+ "object_path": "Events",
+ "steps": None,
+ "num_entries": None,
+ "uuid": None,
+ }
+ ]
+ )
+ for bad in (0, -5):
+ with pytest.raises(ValueError, match="step_size must be a positive integer"):
+ get_steps(normed, step_size=bad)
+ with pytest.raises(ValueError, match="step_size must be a positive integer"):
+ get_parquet_form_uuid_steps(normed, step_size=bad)
diff --git a/tests/test_dataset_tools_forms.py b/tests/test_dataset_tools_forms.py
new file mode 100644
index 000000000..6ebe9b487
--- /dev/null
+++ b/tests/test_dataset_tools_forms.py
@@ -0,0 +1,454 @@
+"""Tests for union-form generation via DatasetSpec addition and the forms helpers."""
+
+import awkward
+import numpy
+import pytest
+import uproot
+
+from coffea.dataset_tools import DataGroupSpec, preprocess
+from coffea.dataset_tools.forms import (
+ decode_field_bitset,
+ encode_field_bitset,
+ prune_form_fields,
+ sort_form_fields,
+ union_form_jsonstr,
+)
+
+_DY = "tests/samples/nano_dy.root"
+
+# Two CMS NanoAOD files carrying disjoint subsets of GenModel_TChiZH_* model-point flags
+# (the GenModel case #1478 targets): file A has GenModel_TChiZH_700_1, file B does not.
+_GENMODEL_A = "tests/samples/nano_genmodel_with20_700_1_with0_1100_200_with0_950_400"
+_GENMODEL_B = "tests/samples/nano_genmodel_without_700_1_with0_1100_200_with20_950_400"
+
+
+def _record_form(**fields):
+ return awkward.Array([fields]).layout.form
+
+
+@pytest.fixture
+def hlt_files(tmp_path):
+ """Two small TTree files sharing a flat branch but with disjoint HLT-style bool branches
+ (the per-file field variation the union form exists to cover)."""
+ file_a = str(tmp_path / "hlt_a.root")
+ file_b = str(tmp_path / "hlt_b.root")
+ with uproot.recreate(file_a) as f:
+ f["Events"] = {
+ "x": numpy.arange(10, dtype="f8"),
+ "HLT_a": numpy.ones(10, dtype=bool),
+ }
+ with uproot.recreate(file_b) as f:
+ f["Events"] = {
+ "x": numpy.arange(10, dtype="f8"),
+ "HLT_b": numpy.zeros(10, dtype=bool),
+ }
+ return file_a, file_b
+
+
+def _preprocessed(files, name="D"):
+ dgs = DataGroupSpec({name: {"files": files}})
+ available, _ = preprocess(dgs, save_form=True, backend="iterative")
+ return available[name]
+
+
+# --------------------------------------------------------------------------------------
+# forms helpers
+# --------------------------------------------------------------------------------------
+
+
+def test_union_form_jsonstr_first_seen_and_sorted():
+ fa = _record_form(x=1.0, HLT_a=True)
+ fb = _record_form(x=1.0, HLT_b=False)
+ ab = union_form_jsonstr([fa, fb])
+ ba = union_form_jsonstr([fb, fa])
+ # merge order perturbs the byte-level field order but not form equality
+ assert ab != ba
+ assert awkward.forms.from_json(ab) == awkward.forms.from_json(ba)
+ # sort_fields canonicalizes the bytes independent of merge order
+ fa2 = _record_form(x=1.0, HLT_a=True)
+ fb2 = _record_form(x=1.0, HLT_b=False)
+ ab_sorted = union_form_jsonstr([fa2, fb2], sort_fields=True)
+ ba_sorted = union_form_jsonstr([fb2, fa2], sort_fields=True)
+ assert ab_sorted == ba_sorted
+ assert awkward.forms.from_json(ab_sorted).fields == sorted(
+ awkward.forms.from_json(ab).fields
+ )
+
+
+def test_union_form_jsonstr_empty_returns_none():
+ assert union_form_jsonstr([]) is None
+
+
+def test_sort_form_fields_recursive_and_equal():
+ form = _record_form(b=1.0, a={"y": 1, "x": 2.0})
+ sorted_form = sort_form_fields(form)
+ assert sorted_form.fields == ["a", "b"]
+ assert sorted_form.contents[0].fields == ["x", "y"]
+ assert sorted_form == form
+
+
+def test_prune_form_fields_top_level():
+ form = _record_form(x=1.0, y=2, z=True)
+ pruned = prune_form_fields(form, {"x", "z"})
+ assert pruned.fields == ["x", "z"]
+
+
+def test_field_bitset_roundtrip():
+ union_fields = ["a", "b", "c", "d"]
+ present = {"a", "c"}
+ bitset = encode_field_bitset(present, union_fields)
+ assert decode_field_bitset(bitset, union_fields) == present
+ # fields outside the union are ignored on encode
+ assert encode_field_bitset({"a", "c", "zz"}, union_fields) == bitset
+
+
+# --------------------------------------------------------------------------------------
+# DatasetSpec.__add__ / union_with
+# --------------------------------------------------------------------------------------
+
+
+def test_add_unions_forms_matches_joint_preprocess(hlt_files):
+ """Adding two separately preprocessed DatasetSpecs yields the same union form as
+ preprocessing all files together, without re-opening any file."""
+ file_a, file_b = hlt_files
+ ds_a = _preprocessed({file_a: "Events"})
+ ds_b = _preprocessed({file_b: "Events"})
+ joint = _preprocessed({file_a: "Events", file_b: "Events"})
+
+ combined = ds_a + ds_b
+ assert combined.compressed_form is not None
+ assert combined.form == joint.form
+ assert set(combined.form.fields) == {"x", "HLT_a", "HLT_b"}
+ # per-file bitsets decode to each file's own field set
+ union_fields = list(combined.form.fields)
+ assert decode_field_bitset(
+ combined.files[file_a].experimental_field_bitset, union_fields
+ ) == {"x", "HLT_a"}
+ assert decode_field_bitset(
+ combined.files[file_b].experimental_field_bitset, union_fields
+ ) == {"x", "HLT_b"}
+
+
+def test_add_equal_forms_short_circuits(tmp_path):
+ import shutil
+
+ copy = str(tmp_path / "nano_dy_copy.root")
+ shutil.copy(_DY, copy)
+ ds_a = _preprocessed({_DY: "Events"})
+ ds_b = _preprocessed({copy: "Events"})
+ combined = ds_a + ds_b
+ assert combined.form == ds_a.form
+ assert list(combined.form.fields) == list(ds_a.form.fields)
+
+
+def test_add_one_sided_form_raises(hlt_files):
+ file_a, file_b = hlt_files
+ with_form = _preprocessed({file_a: "Events"})
+ without_form = DataGroupSpec({"D": {"files": {file_b: "Events"}}})["D"]
+ with pytest.raises(ValueError, match="saved form"):
+ with_form + without_form
+ with pytest.raises(ValueError, match="saved form"):
+ without_form + with_form
+
+
+def test_add_no_forms_stays_none(hlt_files):
+ file_a, file_b = hlt_files
+ ds_a = DataGroupSpec({"D": {"files": {file_a: "Events"}}})["D"]
+ ds_b = DataGroupSpec({"D": {"files": {file_b: "Events"}}})["D"]
+ assert (ds_a + ds_b).compressed_form is None
+
+
+def test_union_with_sort_fields_is_order_independent(hlt_files):
+ from coffea.util import decompress_form
+
+ file_a, file_b = hlt_files
+ ds_a = _preprocessed({file_a: "Events"})
+ ds_b = _preprocessed({file_b: "Events"})
+ ab = ds_a.union_with(ds_b, sort_fields=True)
+ ba = ds_b.union_with(ds_a, sort_fields=True)
+ assert decompress_form(ab.compressed_form) == decompress_form(ba.compressed_form)
+ assert list(ab.form.fields) == sorted(ab.form.fields)
+
+
+def test_datagroupspec_add_unions_same_name_dataset(hlt_files):
+ file_a, file_b = hlt_files
+ group_a = DataGroupSpec({"D": _preprocessed({file_a: "Events"}).model_dump()})
+ group_b = DataGroupSpec({"D": _preprocessed({file_b: "Events"}).model_dump()})
+ combined = group_a + group_b
+ assert combined["D"].compressed_form is not None
+ assert set(combined["D"].form.fields) == {"x", "HLT_a", "HLT_b"}
+
+
+# --------------------------------------------------------------------------------------
+# canonicalize_form / filter pruning
+# --------------------------------------------------------------------------------------
+
+
+def test_canonicalize_form_sorts_and_remaps_bitsets(hlt_files):
+ file_a, file_b = hlt_files
+ ds = _preprocessed({file_a: "Events", file_b: "Events"})
+ canonical = ds.canonicalize_form()
+ assert list(canonical.form.fields) == sorted(ds.form.fields)
+ assert canonical.form == ds.form
+ for fname in ds.files:
+ before = decode_field_bitset(
+ ds.files[fname].experimental_field_bitset, list(ds.form.fields)
+ )
+ after = decode_field_bitset(
+ canonical.files[fname].experimental_field_bitset,
+ list(canonical.form.fields),
+ )
+ assert before == after
+
+
+def test_canonicalize_form_without_form_is_copy():
+ ds = DataGroupSpec({"D": {"files": {_DY: "Events"}}})["D"]
+ assert ds.canonicalize_form() == ds
+
+
+def test_filter_files_prunes_union_form(hlt_files):
+ file_a, file_b = hlt_files
+ ds = _preprocessed({file_a: "Events", file_b: "Events"})
+
+ filtered = ds.filter_files(filter_name=".*hlt_a.root")
+ assert list(filtered.files) == [file_a]
+ # the union form shrinks to exactly the fields the remaining file carries
+ assert set(filtered.form.fields) == {"x", "HLT_a"}
+ # the remaining file's bitset covers the whole pruned form
+ assert decode_field_bitset(
+ filtered.files[file_a].experimental_field_bitset, list(filtered.form.fields)
+ ) == set(filtered.form.fields)
+
+
+def test_filter_files_without_bitsets_keeps_superset(hlt_files):
+ file_a, file_b = hlt_files
+ ds = _preprocessed({file_a: "Events", file_b: "Events"})
+ spec = ds.model_dump()
+ for file_spec in spec["files"].values():
+ file_spec["experimental_field_bitset"] = None
+ ds_nobits = type(ds)(**spec)
+
+ filtered = ds_nobits.filter_files(filter_name=".*hlt_a.root")
+ assert filtered.form == ds.form
+
+
+def test_limit_files_prunes_union_form(hlt_files):
+ file_a, file_b = hlt_files
+ ds = _preprocessed({file_a: "Events", file_b: "Events"})
+ limited = ds.limit_files(1)
+ remaining = next(iter(limited.files))
+ remaining_fields = decode_field_bitset(
+ limited.files[remaining].experimental_field_bitset, list(limited.form.fields)
+ )
+ assert remaining_fields == set(limited.form.fields)
+ assert len(limited.form.fields) < len(ds.form.fields)
+
+
+def test_preprocess_populates_bitsets(hlt_files):
+ file_a, file_b = hlt_files
+ ds = _preprocessed({file_a: "Events", file_b: "Events"})
+ union_fields = list(ds.form.fields)
+ all_fields = set()
+ for fname, fs in ds.files.items():
+ assert fs.experimental_field_bitset is not None
+ present = decode_field_bitset(fs.experimental_field_bitset, union_fields)
+ assert present <= set(union_fields)
+ assert len(present) > 0
+ all_fields |= present
+ # the union of the per-file field sets is the full union form
+ assert all_fields == set(union_fields)
+
+
+def test_bitset_roundtrips_through_json(hlt_files):
+ file_a, file_b = hlt_files
+ ds = _preprocessed({file_a: "Events", file_b: "Events"})
+ from coffea.dataset_tools.filespec import DatasetSpec
+
+ restored = DatasetSpec.model_validate_json(ds.model_dump_json())
+ assert restored == ds
+ for fname in ds.files:
+ assert (
+ restored.files[fname].experimental_field_bitset
+ == ds.files[fname].experimental_field_bitset
+ )
+
+
+# --------------------------------------------------------------------------------------
+# GenModel: real CMS NanoAOD files with disjoint model-point (GenModel) branch subsets
+# --------------------------------------------------------------------------------------
+
+_GM_700 = "GenModel_TChiZH_700_1"
+
+
+@pytest.mark.parametrize(
+ "ext, object_path",
+ [(".root", "Events"), (".parquet", None)],
+)
+def test_genmodel_union_makes_absent_fields_optional(ext, object_path):
+ """Adding two DatasetSpecs whose files carry disjoint GenModel model-point flags unions
+ their forms: the result is a superset of both, and a flag present in only one file becomes
+ an option type so it stays readable (as None) for the file that lacks it -- the GenModel
+ behavior #1478 targets, for both ROOT and parquet inputs."""
+ files_a = {_GENMODEL_A + ext: object_path}
+ files_b = {_GENMODEL_B + ext: object_path}
+ da, _ = preprocess(
+ DataGroupSpec({"genmodel": {"files": files_a}}),
+ save_form=True,
+ backend="iterative",
+ )
+ db, _ = preprocess(
+ DataGroupSpec({"genmodel": {"files": files_b}}),
+ save_form=True,
+ backend="iterative",
+ )
+ fields_a = set(da["genmodel"].form.fields)
+ fields_b = set(db["genmodel"].form.fields)
+ # the fixtures differ in their GenModel subset, and only file A has GenModel_TChiZH_700_1
+ assert fields_a != fields_b
+ assert _GM_700 in fields_a and _GM_700 not in fields_b
+
+ combined = da["genmodel"] + db["genmodel"]
+ union = combined.form
+ assert set(union.fields) == fields_a | fields_b
+
+ # a flag present in only one file is an IndexedOptionArray(bool) in the union
+ only_in_a = sorted(fields_a - fields_b)
+ content = union.contents[union.fields.index(only_in_a[0])]
+ assert isinstance(content, awkward.forms.IndexedOptionForm)
+ assert isinstance(content.content, awkward.forms.NumpyForm)
+ assert content.content.primitive == "bool"
+ # GenModel_TChiZH_700_1 (present only in A) is optional in the union
+ assert isinstance(
+ union.contents[union.fields.index(_GM_700)], awkward.forms.IndexedOptionForm
+ )
+
+
+def test_genmodel_matches_joint_preprocess():
+ """The union built by adding two separately-preprocessed GenModel DatasetSpecs equals the
+ form built by preprocessing both files together in one dataset."""
+ da, _ = preprocess(
+ DataGroupSpec({"genmodel": {"files": {_GENMODEL_A + ".root": "Events"}}}),
+ save_form=True,
+ backend="iterative",
+ )
+ db, _ = preprocess(
+ DataGroupSpec({"genmodel": {"files": {_GENMODEL_B + ".root": "Events"}}}),
+ save_form=True,
+ backend="iterative",
+ )
+ joint, _ = preprocess(
+ DataGroupSpec(
+ {
+ "genmodel": {
+ "files": {
+ _GENMODEL_A + ".root": "Events",
+ _GENMODEL_B + ".root": "Events",
+ }
+ }
+ }
+ ),
+ save_form=True,
+ backend="iterative",
+ )
+ assert (da["genmodel"] + db["genmodel"]).form == joint["genmodel"].form
+
+
+def test_genmodel_bitsets_and_filter_prune():
+ """Each GenModel file's experimental bitset decodes to its own branch set, and filtering the
+ combined dataset back to one file prunes the union form to that file's branches."""
+ combined, _ = preprocess(
+ DataGroupSpec(
+ {
+ "genmodel": {
+ "files": {
+ _GENMODEL_A + ".root": "Events",
+ _GENMODEL_B + ".root": "Events",
+ }
+ }
+ }
+ ),
+ save_form=True,
+ backend="iterative",
+ )
+ ds = combined["genmodel"]
+ union_fields = list(ds.form.fields)
+ present_a = decode_field_bitset(
+ ds.files[_GENMODEL_A + ".root"].experimental_field_bitset, union_fields
+ )
+ assert _GM_700 in present_a
+
+ only_a = ds.filter_files(filter_name=".*with20_700_1.*")
+ assert list(only_a.files) == [_GENMODEL_A + ".root"]
+ # the pruned form is exactly file A's fields, and still carries GenModel_TChiZH_700_1
+ assert set(only_a.form.fields) == present_a
+ assert _GM_700 in only_a.form.fields
+
+
+def _genmodel_events(files, form, mode="dask"):
+ from coffea.nanoevents import NanoAODSchema, NanoEventsFactory
+
+ NanoAODSchema.warn_missing_crossrefs = False
+ return NanoEventsFactory.from_root(
+ files, schemaclass=NanoAODSchema, known_base_form=form, mode=mode
+ ).events()
+
+
+def _genmodel_form(files):
+ out, _ = preprocess(
+ DataGroupSpec({"g": {"files": files}}), save_form=True, backend="iterative"
+ )
+ return out["g"].form
+
+
+def test_genmodel_union_form_preserves_masked_events():
+ """Events passing a GenModel model-point mask survive the union form unchanged. File A's
+ GenModel_TChiZH_700_1 selection is the same read with A's own saved form or the dataset union
+ form, and for each model point the combined (A+B) count equals the per-file sum and the known
+ per-file truth (700_1: 20 from A only; 950_400: 20 from B only; 1100_200: 0). Uses the dask
+ read path, where the union form injects the option-typed flag physically absent from a file.
+ """
+ pytest.importorskip("dask")
+ pytest.importorskip("dask_awkward")
+
+ a_root = _GENMODEL_A + ".root"
+ b_root = _GENMODEL_B + ".root"
+ form_a = _genmodel_form({a_root: "Events"})
+ union = _genmodel_form({a_root: "Events", b_root: "Events"})
+
+ def selected(files, form, point):
+ events = _genmodel_events(files, form)
+ mask = awkward.fill_none(events.GenModel[point], False)
+ return int(awkward.count_nonzero(mask).compute())
+
+ # the precise regression: applying the union form does not falsify A's True flags
+ assert selected({a_root: "Events"}, form_a, "TChiZH_700_1") == 20
+ assert selected({a_root: "Events"}, union, "TChiZH_700_1") == 20
+
+ for point, truth in [
+ ("TChiZH_700_1", 20),
+ ("TChiZH_1100_200", 0),
+ ("TChiZH_950_400", 20),
+ ]:
+ a = selected({a_root: "Events"}, union, point)
+ b = selected({b_root: "Events"}, union, point)
+ together = selected({a_root: "Events", b_root: "Events"}, union, point)
+ assert a + b == together == truth, point
+
+
+def test_genmodel_union_flag_requires_fill_none_for_masking():
+ """A model-point flag absent from a file is injected as None by the union form, so selecting
+ events with fill_none(flag, False) keeps only the True rows (20, all from file A), while the
+ raw option-type mask keeps the None rows too and over-selects every event (40)."""
+ pytest.importorskip("dask")
+ pytest.importorskip("dask_awkward")
+
+ a_root = _GENMODEL_A + ".root"
+ b_root = _GENMODEL_B + ".root"
+ union = _genmodel_form({a_root: "Events", b_root: "Events"})
+ events = _genmodel_events({a_root: "Events", b_root: "Events"}, union)
+ flag = events.GenModel.TChiZH_700_1
+
+ n_fill = int(awkward.num(events[awkward.fill_none(flag, False)], axis=0).compute())
+ n_raw = int(awkward.num(events[flag], axis=0).compute())
+ assert n_fill == 20
+ assert n_raw == 40
diff --git a/tests/test_dataset_tools_metadata.py b/tests/test_dataset_tools_metadata.py
new file mode 100644
index 000000000..776b67748
--- /dev/null
+++ b/tests/test_dataset_tools_metadata.py
@@ -0,0 +1,231 @@
+"""Tests for user-supplied metadata extraction during pydantic preprocessing."""
+
+import awkward
+import pytest
+
+from coffea.dataset_tools import DataGroupSpec, preprocess
+
+_DY = "tests/samples/nano_dy.root"
+
+
+def _nbranches(file_handle):
+ """Extractor: number of branches/fields in the Events object of an open ROOT file."""
+ return {"nbranches": len(file_handle["Events"].keys())}
+
+
+def _sumw_style(file_handle):
+ """Extractor shaped like the ATLAS CutBookkeepers sum-of-weights pattern."""
+ return {"sumw": float(file_handle["Events"].num_entries)}
+
+
+def _sum_sumw(per_file):
+ """Reducer: dataset-level sum of the per-file sumw values."""
+ return {"sumw_dataset": sum(meta["sumw"] for meta in per_file.values())}
+
+
+def _parquet_rows(parquet_metadata):
+ return {"rows": int(parquet_metadata["num_rows"])}
+
+
+def _returns_non_dict(file_handle):
+ return 42
+
+
+def _returns_unserializable(file_handle):
+ return {"handle": object()}
+
+
+def _raises_oserror(file_handle):
+ raise OSError("broken metadata source")
+
+
+@pytest.fixture
+def two_files(tmp_path):
+ import shutil
+
+ copy = str(tmp_path / "nano_dy_copy.root")
+ shutil.copy(_DY, copy)
+ return _DY, copy
+
+
+def test_extractor_fills_file_metadata(two_files):
+ file_a, file_b = two_files
+ dgs = DataGroupSpec({"D": {"files": {file_a: "Events", file_b: "Events"}}})
+ available, updated = preprocess(
+ dgs, save_form=False, backend="iterative", metadata_extractor=_nbranches
+ )
+ for out in (available, updated):
+ for fs in out["D"].files.values():
+ assert fs.metadata is not None
+ assert fs.metadata["nbranches"] > 0
+
+
+def test_reducer_fills_dataset_metadata(two_files):
+ file_a, file_b = two_files
+ dgs = DataGroupSpec(
+ {"D": {"files": {file_a: "Events", file_b: "Events"}, "metadata": {"xs": 1.5}}}
+ )
+ available, updated = preprocess(
+ dgs,
+ save_form=False,
+ backend="iterative",
+ metadata_extractor=_sumw_style,
+ metadata_reducer=_sum_sumw,
+ )
+ for out in (available, updated):
+ # both files have 40 entries; existing dataset metadata is preserved
+ assert out["D"].metadata["sumw_dataset"] == 80.0
+ assert out["D"].metadata["xs"] == 1.5
+ for fs in out["D"].files.values():
+ assert fs.metadata == {"sumw": 40.0}
+
+
+def test_extractor_parquet(tmp_path):
+ path = str(tmp_path / "d.parquet")
+ awkward.to_parquet(awkward.Array([{"x": 1.0}, {"x": 2.0}]), path)
+ dgs = DataGroupSpec({"P": {"files": {path: None}}})
+ available, _ = preprocess(
+ dgs, save_form=False, backend="iterative", metadata_extractor=_parquet_rows
+ )
+ assert available["P"].files[path].metadata == {"rows": 2}
+
+
+def test_extractor_backend_equivalence(two_files):
+ file_a, file_b = two_files
+ dgs = DataGroupSpec({"D": {"files": {file_a: "Events", file_b: "Events"}}})
+ kwargs = dict(
+ save_form=True,
+ metadata_extractor=_sumw_style,
+ metadata_reducer=_sum_sumw,
+ )
+ a_iter, u_iter = preprocess(dgs, backend="iterative", **kwargs)
+ a_fut, u_fut = preprocess(dgs, backend="futures", **kwargs)
+ assert a_iter == a_fut
+ assert u_iter == u_fut
+ assert a_iter["D"].metadata["sumw_dataset"] == 80.0
+
+
+def test_extractor_dask_backend(two_files):
+ pytest.importorskip("dask")
+ pytest.importorskip("dask_awkward")
+ file_a, file_b = two_files
+ dgs = DataGroupSpec({"D": {"files": {file_a: "Events", file_b: "Events"}}})
+ kwargs = dict(
+ save_form=True,
+ metadata_extractor=_sumw_style,
+ metadata_reducer=_sum_sumw,
+ )
+ a_dask, u_dask = preprocess(dgs, backend="dask", scheduler="synchronous", **kwargs)
+ a_iter, u_iter = preprocess(dgs, backend="iterative", **kwargs)
+ assert a_dask == a_iter
+ assert u_dask == u_iter
+
+
+def test_extractor_failure_participates_in_skip_bad_files(two_files, tmp_path):
+ file_a, file_b = two_files
+ dgs = DataGroupSpec({"D": {"files": {file_a: "Events", file_b: "Events"}}})
+ with pytest.raises(OSError, match="broken metadata source"):
+ preprocess(
+ dgs,
+ save_form=False,
+ backend="iterative",
+ metadata_extractor=_raises_oserror,
+ )
+ # a failing extractor is a per-file failure: skip_bad_files drops the files
+ available, updated = preprocess(
+ dgs,
+ save_form=False,
+ backend="iterative",
+ skip_bad_files=True,
+ allow_empty_datasets=True,
+ metadata_extractor=_raises_oserror,
+ )
+ assert "D" not in available or len(available["D"].files) == 0
+
+
+def test_extractor_non_dict_raises(two_files):
+ file_a, _ = two_files
+ dgs = DataGroupSpec({"D": {"files": {file_a: "Events"}}})
+ with pytest.raises(ValueError, match="must return a dict"):
+ preprocess(
+ dgs,
+ save_form=False,
+ backend="iterative",
+ metadata_extractor=_returns_non_dict,
+ )
+
+
+def test_extractor_unserializable_raises(two_files):
+ file_a, _ = two_files
+ dgs = DataGroupSpec({"D": {"files": {file_a: "Events"}}})
+ with pytest.raises(ValueError, match="not JSON-serializable"):
+ preprocess(
+ dgs,
+ save_form=False,
+ backend="iterative",
+ metadata_extractor=_returns_unserializable,
+ )
+
+
+def test_extractor_rejected_on_legacy_path(two_files):
+ file_a, _ = two_files
+ with pytest.raises(ValueError, match="not supported"):
+ preprocess(
+ {"D": {"files": {file_a: "Events"}}},
+ preprocess_legacy_root=True,
+ metadata_extractor=_nbranches,
+ )
+
+
+def test_file_metadata_survives_json_roundtrip(two_files):
+ file_a, file_b = two_files
+ dgs = DataGroupSpec({"D": {"files": {file_a: "Events", file_b: "Events"}}})
+ available, _ = preprocess(
+ dgs, save_form=True, backend="iterative", metadata_extractor=_nbranches
+ )
+ restored = DataGroupSpec.model_validate_json(available.model_dump_json())
+ assert restored == available
+ for fname, fs in available["D"].files.items():
+ assert restored["D"].files[fname].metadata == fs.metadata
+
+
+def test_file_metadata_excluded_from_legacy_dict_output(two_files):
+ file_a, file_b = two_files
+ available, _ = preprocess(
+ {"D": {"files": {file_a: "Events", file_b: "Events"}}},
+ save_form=False,
+ backend="iterative",
+ metadata_extractor=_sumw_style,
+ metadata_reducer=_sum_sumw,
+ )
+ # dict-in/dict-out: per-file metadata is not part of the legacy format, but the
+ # reduced dataset-level metadata is
+ for file_info in available["D"]["files"].values():
+ assert "metadata" not in file_info
+ assert available["D"]["metadata"]["sumw_dataset"] == 80.0
+
+
+def test_filespec_metadata_add_merges(two_files):
+ file_a, _ = two_files
+ dgs = DataGroupSpec({"D": {"files": {file_a: "Events"}}})
+ available, _ = preprocess(
+ dgs,
+ save_form=False,
+ step_size=20,
+ backend="iterative",
+ metadata_extractor=_nbranches,
+ )
+ fs = available["D"].files[file_a]
+ assert len(fs.steps) == 2
+ merged = fs.limit_steps(slice(0, 1)) + fs.limit_steps(slice(1, 2))
+ assert merged.metadata == fs.metadata
+
+
+def test_extractor_with_rntuple(tmp_path):
+ """The extractor receives the open file for RNTuple inputs too."""
+ rnt = "tests/samples/nano_dy_rntuple.root"
+ dgs = DataGroupSpec({"D": {"files": {rnt: "Events"}}})
+ available, _ = preprocess(
+ dgs, save_form=False, backend="iterative", metadata_extractor=_sumw_style
+ )
+ assert available["D"].files[rnt].metadata == {"sumw": 40.0}
diff --git a/tests/test_dataset_tools_mutable_steps.py b/tests/test_dataset_tools_mutable_steps.py
new file mode 100644
index 000000000..9dd207340
--- /dev/null
+++ b/tests/test_dataset_tools_mutable_steps.py
@@ -0,0 +1,218 @@
+"""Tests for the mutable (resizable) steps prototype."""
+
+import shutil
+
+import pytest
+
+from coffea.dataset_tools import DataGroupSpec, preprocess
+from coffea.dataset_tools.mutable_steps import (
+ WallTimeStepPolicy,
+ completed_spec,
+ coverage_regions,
+ iter_dataset_steps,
+ iter_file_steps,
+ remaining_regions,
+ resizable_steps,
+ run_adaptive_steps,
+)
+
+_DY = "tests/samples/nano_dy.root"
+
+
+def _contiguous(steps, begin, end):
+ """Steps tile [begin, end) exactly, in order, with no gaps or overlaps."""
+ assert steps[0][0] == begin
+ assert steps[-1][1] == end
+ for previous, current in zip(steps, steps[1:]):
+ assert previous[1] == current[0]
+
+
+@pytest.fixture
+def dataset(tmp_path):
+ copy = str(tmp_path / "nano_dy_copy.root")
+ shutil.copy(_DY, copy)
+ dgs = DataGroupSpec({"D": {"files": {_DY: "Events", copy: "Events"}}})
+ available, _ = preprocess(dgs, save_form=False, backend="iterative")
+ return available["D"]
+
+
+# --------------------------------------------------------------------------------------
+# resizable_steps
+# --------------------------------------------------------------------------------------
+
+
+def test_resizable_steps_even_tiling():
+ # 100 entries at target 30: 4 steps of 25 (as even as possible, at most the target)
+ steps = list(resizable_steps(0, 100, 30))
+ assert steps == [[0, 25], [25, 50], [50, 75], [75, 100]]
+
+
+def test_resizable_steps_shrink_retiles_remainder():
+ gen = resizable_steps(0, 100, 30)
+ steps = [next(gen)]
+ # request at most 10 for the remaining 75 entries
+ step = gen.send(10)
+ while True:
+ steps.append(step)
+ assert step[1] - step[0] <= 10
+ try:
+ step = next(gen)
+ except StopIteration:
+ break
+ _contiguous(steps, 0, 100)
+
+
+def test_resizable_steps_grow_is_at_most_requested():
+ gen = resizable_steps(0, 100, 25)
+ next(gen)
+ # remaining 75 at target 50: 2 steps of at most 38 = ceil(75/2)
+ step = gen.send(50)
+ assert step == [25, 63]
+
+
+def test_resizable_steps_invalid_sizes_raise():
+ with pytest.raises(ValueError, match="positive integer"):
+ next(resizable_steps(0, 10, 0))
+ gen = resizable_steps(0, 10, 5)
+ next(gen)
+ with pytest.raises(ValueError, match="positive integer"):
+ gen.send(-1)
+
+
+# --------------------------------------------------------------------------------------
+# coverage_regions / iter_file_steps
+# --------------------------------------------------------------------------------------
+
+
+def test_coverage_regions_merges_adjacent(dataset):
+ fs = next(iter(dataset.files.values()))
+ spec = fs.model_dump()
+ spec["steps"] = [[0, 20], [20, 40]]
+ merged = type(fs)(**spec)
+ assert coverage_regions(merged) == [[0, 40]]
+
+ spec["steps"] = [[0, 10], [20, 30]]
+ disjoint = type(fs)(**spec)
+ assert coverage_regions(disjoint) == [[0, 10], [20, 30]]
+
+
+def test_coverage_regions_without_steps_uses_num_entries(dataset):
+ fs = next(iter(dataset.files.values()))
+ spec = fs.model_dump()
+ spec["steps"] = None
+ from coffea.dataset_tools.filespec import CoffeaROOTFileSpecOptional
+
+ nosteps = CoffeaROOTFileSpecOptional(**spec)
+ assert coverage_regions(nosteps) == [[0, fs.num_entries]]
+
+ spec["num_entries"] = None
+ unknown = CoffeaROOTFileSpecOptional(**spec)
+ with pytest.raises(ValueError, match="neither steps nor num_entries"):
+ coverage_regions(unknown)
+
+
+def test_iter_file_steps_respects_disjoint_regions(dataset):
+ fs = next(iter(dataset.files.values()))
+ spec = fs.model_dump()
+ spec["steps"] = [[0, 10], [20, 30]]
+ disjoint = type(fs)(**spec)
+ steps = list(iter_file_steps(disjoint, 5))
+ assert steps == [[0, 5], [5, 10], [20, 25], [25, 30]]
+
+
+def test_iter_dataset_steps_resize_carries_across_files(dataset):
+ gen = iter_dataset_steps(dataset, 40)
+ fname_first, step = next(gen)
+ assert step == [0, 40]
+ # shrink after the first file's single step: the second file re-tiles at 10
+ seen = []
+ try:
+ item = gen.send(10)
+ while True:
+ seen.append(item)
+ item = next(gen)
+ except StopIteration:
+ pass
+ fnames = {fname for fname, _ in seen}
+ assert fnames == set(dataset.files) - {fname_first} or fnames == set(dataset.files)
+ for _, step in seen:
+ assert step[1] - step[0] <= 10
+ second_file_steps = [step for fname, step in seen if fname != fname_first]
+ _contiguous(second_file_steps, 0, 40)
+
+
+# --------------------------------------------------------------------------------------
+# resumption helpers
+# --------------------------------------------------------------------------------------
+
+
+def test_remaining_regions_interval_subtraction(dataset):
+ fs = next(iter(dataset.files.values()))
+ # completed ranges need not align with stored steps
+ assert remaining_regions(fs, [[0, 25]]) == [[25, 40]]
+ assert remaining_regions(fs, [[10, 15], [15, 20]]) == [[0, 10], [20, 40]]
+ assert remaining_regions(fs, [[0, 40]]) == []
+ assert remaining_regions(fs, []) == coverage_regions(fs)
+
+
+def test_completed_spec_roundtrip_and_accumulation(dataset):
+ fs = next(iter(dataset.files.values()))
+ first = completed_spec(fs, [[0, 15]])
+ second = completed_spec(fs, [[15, 40]])
+ assert completed_spec(fs, []) is None
+ merged = first + second
+ assert merged.steps == [[0, 15], [15, 40]]
+ assert remaining_regions(fs, merged.steps) == []
+
+
+# --------------------------------------------------------------------------------------
+# toy adaptive driver
+# --------------------------------------------------------------------------------------
+
+
+class _FakeClock:
+ """Deterministic clock advanced by the fake work function."""
+
+ def __init__(self):
+ self.now = 0.0
+
+ def __call__(self):
+ return self.now
+
+
+def test_run_adaptive_steps_converges_to_target(dataset):
+ clock = _FakeClock()
+ per_entry_seconds = 0.01
+
+ def work(fname, step):
+ clock.now += (step[1] - step[0]) * per_entry_seconds
+ return (fname, tuple(step))
+
+ # equilibrium size: target 0.1 s at 0.01 s/entry -> 10 entries per step
+ policy = WallTimeStepPolicy(target_seconds=0.1)
+ run = run_adaptive_steps(dataset, work, step_size=40, policy=policy, clock=clock)
+
+ # the first (full-file) step triggers a shrink; later steps sit at the equilibrium
+ assert run.step_sizes[0] == 40
+ assert set(run.step_sizes[1:]) == {10}
+ # every file is fully covered exactly once
+ assert set(run.completed) == set(dataset.files)
+ for fname, steps in run.completed.items():
+ _contiguous(steps, 0, dataset.files[fname].num_entries)
+ assert len(run.results) == len(run.step_sizes)
+
+
+def test_run_adaptive_steps_growth_is_damped(dataset):
+ clock = _FakeClock()
+
+ def instant_work(fname, step):
+ return None
+
+ policy = WallTimeStepPolicy(target_seconds=1.0, max_step_size=64, max_growth=2.0)
+ run = run_adaptive_steps(
+ dataset, instant_work, step_size=4, policy=policy, clock=clock
+ )
+ # growth per adjustment is bounded by max_growth and capped at max_step_size
+ for previous, current in zip(run.step_sizes, run.step_sizes[1:]):
+ assert current <= max(previous * 2, 1) or current <= 64
+ assert max(run.step_sizes) <= 64