from __future__ import annotations
from attrs import define, field
import contextlib
import copy
import awkward as ak
from typing import Any, ClassVar
import logging
import enum
from analyzer.utils.structure_tools import freeze
[docs]
logger = logging.getLogger("analyzer.core")
[docs]
def coerceFields(data):
if isinstance(data, str):
return tuple(data.split("."))
elif isinstance(data, Column):
return data.fields
else:
return data
# Poorly handled right now.
[docs]
class EventBackend(str, enum.Enum):
[docs]
coffea_virtual = "coffea_virtual"
[docs]
coffea_dask = "coffea_dask"
[docs]
coffea_imm = "coffea_eager"
@define(frozen=True)
[docs]
class Column:
"""
Simple wrapper for a "path" in events.
Stores a single tuple representing the path to the desired field.
Automatically converts "dot" delimited paths.
"""
[docs]
fields: tuple[str, ...] = field(converter=coerceFields)
[docs]
def contains(self, other: Column | str | tuple[str,...]):
if not isinstance(other, Column):
other = Column(other)
if len(self.fields) > len(other.fields):
return False
return other.fields[: len(self.fields)] == self.fields
[docs]
def extract(self, events: ak.Array):
for f in self.fields:
events = events[f]
return events
[docs]
def __iter__(self):
return iter(self.fields)
[docs]
def __eq__(self, other):
return self.fields == other.fields
[docs]
def __len__(self):
return len(self.fields)
[docs]
def __getitem__(self, key: int):
return Column(self.fields.__getitem__(key))
[docs]
def __add__(self, other: Column):
return Column(self.fields + Column(other).fields)
[docs]
def __radd__(self, other):
return Column(Column(other).fields + self.fields)
[docs]
def __hash__(self):
return hash(self.fields)
[docs]
def __str__(self):
return ".".join(self.fields)
[docs]
def __repr__(self):
return ".".join(self.fields)
@property
[docs]
def adl_name(self):
return self.fields[-1] if self.fields else ""
@classmethod
def _structure(cls, data: str, conv):
if isinstance(data, str):
return Column(data)
[docs]
def parents(self):
return [Column(self.fields[: i + 1]) for i in range(len(self.fields) - 1)]
[docs]
def setColumn(events: ak.Array, column: Column, value: ak.Array)->ak.Array:
"""
Set a potentially column in an array.
Automatically creates intermediate fields if needed.
"""
if not isinstance(column, Column):
column = Column(column)
if len(column) == 1:
return ak.with_field(events, value, column.fields)
# If column has more than two parts, we must potentially create the intermediates.
head = column.fields[0]
rest = column.fields[1:]
if head not in events.fields:
# If the top level field is not present, we create a nested arraay with zip
for c in reversed(rest):
value = ak.zip({c: value}, depth_limit=1)
return ak.with_field(events, value, head)
else:
# If top level is present, recurse
return ak.with_field(events, setColumn(events[head], Column(rest), value), head)
@define
[docs]
class ColumnCollection:
[docs]
columns: set[Column] = field(converter=lambda z: set(Column(x) for x in z))
[docs]
def __iter__(self):
return iter(self.columns)
[docs]
def contains(self, other: Column):
"""
Check if any of the columns in the set is a parent of other.
"""
if not isinstance(other, Column):
other = Column(other)
f = other.fields
for i in range(1, len(f) + 1):
if Column(f[:i]) in self.columns:
return True
return False
[docs]
def intersect(self, other: ColumnCollection):
ret = {
x
for x in self.columns
if any((x.contains(o) or o.contains(x)) for o in other)
}
return ret
[docs]
def getAllColumns(events, cur_col=None, cur_depth=0, max_depth=None) -> set[Column]:
"""
Extract all columns from a potentially nested layout.
"""
if fields := getattr(events, "fields"):
ret = set()
for f in fields:
if cur_col is not None:
n = Column(cur_col.fields + (f,))
else:
n = Column(f)
if cur_depth == max_depth:
ret.add(n)
else:
ret |= getAllColumns(
events[f], n, max_depth=max_depth, cur_depth=cur_depth + 1
)
if cur_col is not None:
ret.add(cur_col)
return ret
else:
return {cur_col}
@define
[docs]
class TrackedColumns:
[docs]
INTERNAL_USE_COL: ClassVar[Column] = Column("INTERNAL_USE")
_events: Any
_column_provenance: dict[Column, int]
_current_provenance: int
_allowed_inputs: ColumnCollection | None = None
_allowed_outputs: ColumnCollection | None = None
_allow_filter: bool = True
[docs]
pipeline_data: dict[str, Any] = field(factory=dict)
# Temporary store to avoid excessive "syncing" to the main array
_lazy_columns: dict[Column, Any] = field(factory=dict)
@property
[docs]
def events(self):
self.flush()
return self._events
[docs]
def flush(self):
"""
Synchronize stored lazy columns with underlying events
"""
if not self._lazy_columns:
return
for col, val in self._lazy_columns.items():
self._events = setColumn(self._events, col, val)
self._lazy_columns.clear()
@property
[docs]
def fields(self):
s = set(self._events.fields)
for col in self._lazy_columns:
if len(col.fields) > 0:
s.add(col.fields[0])
return list(s)
[docs]
def updatedColumns(self, old, limit=None):
"""
Determine which in the current TrackedColumns have been updated from old by comparing their provenance.
"""
ans = []
for x, y in self._column_provenance.items():
if limit is not None and not limit.contains(x):
continue
if x not in old._column_provenance or y != old._column_provenance[x]:
ans.append(x)
return ans
[docs]
def copy(self):
return TrackedColumns(
# events=copy.copy(self._events),
events=self._events,
column_provenance=self._column_provenance.copy(),
current_provenance=self._current_provenance,
metadata=copy.copy(self.metadata),
pipeline_data=copy.deepcopy(self.pipeline_data),
backend=self.backend,
lazy_columns=self._lazy_columns.copy(),
)
@staticmethod
[docs]
def fromEvents(events, metadata, backend, provenance: int):
return TrackedColumns(
events=events,
column_provenance={
x: provenance
for x in getAllColumns(
events.layout, max_depth=events.layout.minmax_depth[1] - 2
)
}, # Do try to recurse into leaves when finding all columns
current_provenance=provenance,
backend=backend,
metadata=metadata,
)
[docs]
def getKeyForColumns(self, columns):
"""
Get an excecution key for the column.
Returns a hash dependent on the provenance of all the columns contains in the input.
"""
ret = []
for column in columns:
try:
ret.append((column, self._column_provenance[column]))
except KeyError:
continue
# for col, prov in self._column_provenance.items():
# if any(x.contains(col) for x in columns):
# ret.append((col, self._column_provenance))
# logger.debug(f"Relevant columns for {columns} are :\n {ret}")
return hash((freeze(self.metadata), freeze(self.pipeline_data), tuple(ret)))
[docs]
def getKeyForAll(self):
ret = []
for column in self._column_provenance:
ret.append((column, self._column_provenance[column]))
return hash((freeze(self.metadata), freeze(self.pipeline_data), tuple(ret)))
[docs]
def __setitem__(self, column, value):
"""
Set a column, potentially lazy.
"""
if not isinstance(column, Column):
column = Column(column)
if (
self._allowed_outputs is not None
and not TrackedColumns.INTERNAL_USE_COL.contains(column)
and not self._allowed_outputs.contains(column)
):
raise RuntimeError(
f"Column {column} is not in the list of outputs {self._allowed_outputs}"
)
if any(c.contains(column) for c in self._lazy_columns if c != column):
# If any of the lazy columns contains the columns we are trying to set,
# we first synchronize the lazy columns to the events
self.flush()
self._lazy_columns[column] = value
# Delete any lazy column that was contained in the one we just set, it is overriden
to_del = [k for k in self._lazy_columns if k != column and column.contains(k)]
for k in to_del:
del self._lazy_columns[k]
self._column_provenance[column] = self._current_provenance
# Get all the columns the column we are trying to set and update their provenance
if hasattr(value, "layout"):
all_columns = getAllColumns(value.layout, column)
else:
all_columns = {column}
for c in all_columns:
self._column_provenance[c] = self._current_provenance
# For every parent column, update it provenance to be based on its previous provenance and the provenance of the set child
for c in column.parents():
p = hash((self._column_provenance.get(c, None), self._current_provenance))
self._column_provenance[c] = p
logger.debug(
f"Updating parent {c} of {column} to events with provenance {p}"
)
[docs]
def __getitem__(self, column):
"""
Get a column. May use the lazy cache.
"""
if not isinstance(column, Column):
column = Column(column)
if self._allowed_inputs is not None and not self._allowed_inputs.contains(
column
):
raise RuntimeError(
f"Column {column} is not in the list of inputs {self._allowed_inputs}"
)
if column in self._lazy_columns:
return self._lazy_columns[column]
# If this column is the child of a lazy column, extract the relevant portion
for c in column.parents():
if c in self._lazy_columns:
remainder = Column(column.fields[len(c.fields) :])
return remainder.extract(self._lazy_columns[c])
# If we are trying to extract a column that contains a lazy column, first synchronize
if any(column.contains(c) for c in self._lazy_columns):
self.flush()
# Do the extraction with the synchronized columns
return column.extract(self._events)
[docs]
def addColumnsFrom(self, other, columns):
"""
Merge in columns from another TrackedColumns.
"""
for column in columns:
with self.useKey(other._column_provenance[column]):
self[column] = other[column]
# self._setIndividualColumnWithProvenance(
# column, other[column], other._column_provenance[column]
# )
[docs]
def filter(self, mask):
"""
Filter at the event level.
"""
if not self._allow_filter:
raise RuntimeError()
self._events = self._events[mask]
for c in list(self._lazy_columns.keys()):
self._lazy_columns[c] = self._lazy_columns[c][mask]
for c in self._column_provenance:
self._column_provenance[c] = hash(
(self._column_provenance[c], self._current_provenance)
)
return self
@contextlib.contextmanager
[docs]
def useKey(self, provenance):
old_provenance = self._current_provenance
self._current_provenance = provenance
yield
self._current_provenance = old_provenance
@contextlib.contextmanager
@contextlib.contextmanager
[docs]
def allowedOutputs(self, columns: list[Column] | ColumnCollection):
if not isinstance(columns, ColumnCollection):
columns = ColumnCollection(columns)
old_outputs = self._allowed_outputs
self._allowed_outputs = columns
yield
self._allowed_outputs = old_outputs
[docs]
def mergeColumns(column_views):
ret = copy.copy(column_views[0])
for other in column_views[1:]:
ret.addColumnsFrom(other, list(other._column_provenance.keys()))
return ret
[docs]
def addSelection(columns, name, data):
column = Column(("Selection", name))
columns[column] = data
if "Selections" not in columns.pipeline_data:
columns.pipeline_data["Selections"] = {}
columns.pipeline_data["Selections"][name] = False