feat!: switch over to dask-based processing idioms, improve dataset handling

.github/workflows/ci.yml

@@ -126,33 +126,33 @@ jobs:
       env:
         GH_PAT: ${{ secrets.GITHUB_OAUTH }}
 
-  testwq:
-    runs-on: ubuntu-latest
-    needs: pre-commit
-    strategy:
-      matrix:
-        python-version: ["3.11"]
-    name: test coffea-workqueue
-
-    steps:
-    - uses: actions/checkout@v4
-    - name: Set up Conda
-      uses: conda-incubator/setup-miniconda@v2
-      env:
-        ACTIONS_ALLOW_UNSECURE_COMMANDS: true
-      with:
-        auto-update-conda: true
-        python-version: ${{ matrix.python-version }}
-        channels: conda-forge
-    - name: Test work_queue
-      shell: bash -l {0}
-      run: |
-        conda create --yes --name coffea-env -c conda-forge python=${{ matrix.python-version }} ndcctools dill conda-pack conda
-        conda activate coffea-env
-        python -m pip install --ignore-installed .
-        cd tests
-        conda-pack --output coffea-env.tar.gz
-        python wq.py coffea-env.tar.gz
+#  testwq:
+#    runs-on: ubuntu-latest
+#    needs: pre-commit
+#    strategy:
+#      matrix:
+#        python-version: ["3.11"]
+#    name: test coffea-workqueue
+#
+#    steps:
+#    - uses: actions/checkout@v4
+#    - name: Set up Conda
+#      uses: conda-incubator/setup-miniconda@v2
+#      env:
+#        ACTIONS_ALLOW_UNSECURE_COMMANDS: true
+#      with:
+#        auto-update-conda: true
+#        python-version: ${{ matrix.python-version }}
+#        channels: conda-forge
+#    - name: Test work_queue
+#      shell: bash -l {0}
+#      run: |
+#        conda create --yes --name coffea-env -c conda-forge python=${{ matrix.python-version }} ndcctools dill conda-pack conda
+#        conda activate coffea-env
+#        python -m pip install --ignore-installed .
+#        cd tests
+#        conda-pack --output coffea-env.tar.gz
+#        python wq.py coffea-env.tar.gz
 
 #  testskyhookjob:
 #    runs-on: ubuntu-latest
@@ -178,7 +178,7 @@ jobs:
   release:
     if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')
     runs-on: ubuntu-latest
-    needs: [test, testwq]
+    needs: [test]
     strategy:
       matrix:
         python-version: ["3.11"]
@@ -201,7 +201,7 @@ jobs:
         password: ${{ secrets.PYPI_TOKEN }}
 
   pass:
-    needs: [test, testwq]
+    needs: [test]
     runs-on: ubuntu-latest
     steps:
       - run: echo "All jobs passed"

pyproject.toml

@@ -85,6 +85,11 @@ servicex = [
   "servicex>=2.5.3",
   "func-adl_servicex",
 ]
+rucio = [
+  "rucio-clients>=32;python_version>'3.8'",
+  "rucio-clients<32;python_version<'3.9'",
+  "cmd2",
+]
 dev = [
   "pre-commit",
   "flake8",

src/coffea/dataset_tools/__init__.py

@@ -0,0 +1,11 @@
+from coffea.dataset_tools.apply_processor import apply_to_fileset, apply_to_one_dataset
+from coffea.dataset_tools.manipulations import max_chunks, slice_chunks
+from coffea.dataset_tools.preprocess import preprocess
+
+__all__ = [
+    "preprocess",
+    "apply_to_one_dataset",
+    "apply_to_fileset",
+    "max_chunks",
+    "slice_chunks",
+]

src/coffea/dataset_tools/apply_processor.py

@@ -0,0 +1,38 @@
+import copy
+from typing import Callable, Union
+
+from coffea.nanoevents import NanoAODSchema, NanoEventsFactory
+from coffea.processor import ProcessorABC
+
+
+def apply_to_one_dataset(
+    data_manipulation: Union[ProcessorABC, Callable],
+    dataset,
+    schemaclass=NanoAODSchema,
+    metadata={},
+):
+    files = dataset["files"]
+    events = NanoEventsFactory.from_root(
+        files,
+        metadata=metadata,
+        schemaclass=schemaclass,
+    ).events()
+    if isinstance(data_manipulation, ProcessorABC):
+        return data_manipulation.process(events)
+    elif isinstance(data_manipulation, Callable):
+        return data_manipulation(events)
+    else:
+        raise ValueError("data_manipulation must either be a ProcessorABC or Callable")
+
+
+def apply_to_fileset(
+    data_manipulation: Union[ProcessorABC, Callable], fileset, schemaclass=NanoAODSchema
+):
+    out = {}
+    for name, dataset in fileset.items():
+        metadata = copy.deepcopy(dataset.get("metadata", {}))
+        metadata["dataset"] = name
+        out[name] = apply_to_one_dataset(
+            data_manipulation, dataset, schemaclass, metadata
+        )
+    return out