from collections import deque
[docs]
def propagate_sparse_columns(mapper, transform_id):
"""Update the columns to keep in sparse data around a transform.
Columns required in forward data (``required_sparse_columns_fwd``) are
propagated upstream, and columns required in adjoint data
(``required_sparse_columns_adj``) downstream. Propagation stops at
transforms producing the columns (``produced_sparse_columns_fwd`` and
``produced_sparse_columns_adj``) and at data that is neither sparse nor
sampled. Results are stored in ``keep_sparse_columns_fwd`` and
``keep_sparse_columns_adj``; columns no longer needed are removed.
Args:
mapper: Dictionary mapping each transform id to its inputs/outputs/
precursors/successors metadata, mutated in place.
transform_id: Id of the transform to start from.
"""
seeds = start_nodes(mapper, transform_id)
# Pipes can have loops, in which columns may support themselves.
# Thus, columns that lost their support are first removed from all
# dependent nodes, then re-derived from the remaining needs
# (delete and re-derive), which terminates and flushes loops.
# Lost support is detected from changes in what a node is computed from
removed = []
for node in seeds:
lost = lost_support(mapper, node) \
| (get_value(mapper, node) - (compute(mapper, node) or set()))
if lost:
removed.append((node, lost))
to_rederive = set()
while removed:
node, lost = removed.pop()
present = get_value(mapper, node) & lost
if not present:
continue
set_value(mapper, node, get_value(mapper, node) - present)
to_rederive.add(node)
removed.extend((dep, present) for dep in dependents(mapper, node))
# Re-derive values by growing only
queue = deque(seeds | to_rederive)
queued = set(queue)
while queue:
node = queue.popleft()
queued.discard(node)
columns = compute(mapper, node)
if columns is None:
set_value(mapper, node, None)
continue
old = get_value(mapper, node)
set_value(mapper, node, old | columns)
if columns <= old:
continue
for dep in dependents(mapper, node):
if dep not in queued:
queue.append(dep)
queued.add(dep)
[docs]
def strip_copied_sparse_columns(transf_mapper, mapper):
"""Remove sparse-column attributes copied from other transforms' mappers.
Transforms often build their mapper from their neighbours' ones, which
would also copy their column declarations. Declarations must thus be
new lists in ``ini_mapper``: lists shared with another transform's
mapper are considered copied and removed.
Args:
transf_mapper: Mapper of the new transform, mutated in place.
mapper: Mappers of the transforms already in the pipe.
"""
declarations = [
f"{what}_sparse_columns_{kind}"
for what in ["required", "produced"] for kind in ["fwd", "adj"]
]
copied = {
id(trid_dict[key])
for tr_mapper in mapper.values()
for side in ["inputs", "outputs"]
for trid_dict in tr_mapper.get(side, {}).values()
for key in declarations
if key in trid_dict
}
for side in ["inputs", "outputs"]:
for trid_dict in transf_mapper.get(side, {}).values():
for kind in ["fwd", "adj"]:
trid_dict.pop(f"keep_sparse_columns_{kind}", None)
trid_dict.pop(f"keep_sparse_columns_{kind}_base", None)
for key in declarations:
if id(trid_dict.get(key)) in copied:
del trid_dict[key]
# A node is (transform, side, trid, kind), kind being "fwd" or "adj"
[docs]
def start_nodes(mapper, transform_id):
# The transform itself and its direct neighbours,
# whose links may have changed when inserting the transform
transf_mapper = mapper[transform_id]
nodes = set()
for side, other_side, links in [
("inputs", "outputs", "precursors"),
("outputs", "inputs", "successors"),
]:
for trid in transf_mapper[side]:
for kind in ["fwd", "adj"]:
nodes.add((transform_id, side, trid, kind))
for tr in transf_mapper.get(links, {}).get(trid, []):
if trid in mapper.get(tr, {}).get(other_side, {}):
nodes.add((tr, other_side, trid, kind))
return nodes
[docs]
def is_sparse(trid_dict):
return bool(trid_dict.get("sparse_data", False)
or trid_dict.get("sampled", False))
[docs]
def is_tracked(mapper, node):
# Columns are only tracked for sparse or sampled data.
# Flags of linked neighbours are also checked, as some transforms
# (e.g. inserted by init_sparse) do not flag their sparse inputs
transform, side, trid, kind = node
if is_sparse(mapper[transform][side][trid]):
return True
links = "precursors" if side == "inputs" else "successors"
other_side = "outputs" if side == "inputs" else "inputs"
return any(
is_sparse(mapper[tr][other_side][trid])
for tr in mapper[transform].get(links, {}).get(trid, [])
if trid in mapper.get(tr, {}).get(other_side, {})
)
[docs]
def get_columns(trid_dict, key):
return set(trid_dict.get(key, []))
[docs]
def get_value(mapper, node):
transform, side, trid, kind = node
return get_columns(
mapper[transform][side][trid], f"keep_sparse_columns_{kind}")
[docs]
def set_value(mapper, node, columns):
# Not stored for data that is neither sparse nor sampled,
# so that nothing is dropped there
transform, side, trid, kind = node
trid_dict = mapper[transform][side][trid]
if columns is None or not is_tracked(mapper, node):
trid_dict.pop(f"keep_sparse_columns_{kind}", None)
else:
trid_dict[f"keep_sparse_columns_{kind}"] = sorted(columns)
[docs]
def get_base(mapper, node):
# Everything a node value is computed from, apart from other node values
transform, side, trid, kind = node
transf_mapper = mapper[transform]
trid_dict = transf_mapper[side][trid]
produced = get_columns(trid_dict, f"produced_sparse_columns_{kind}")
if is_across_link(side, kind):
links = "successors" if side == "outputs" else "precursors"
linked = [str(tr) for tr in transf_mapper.get(links, {}).get(trid, [])]
else:
other_side = "outputs" if side == "inputs" else "inputs"
mapping = transf_mapper.get(f"{side}2{other_side}", {})
linked = []
for other_trid in mapping.get(trid, [trid]):
if other_trid not in transf_mapper[other_side]:
continue
linked.append(str(other_trid))
produced |= get_columns(
transf_mapper[other_side][other_trid],
f"produced_sparse_columns_{kind}")
return {
"sparse": is_tracked(mapper, node),
"required": sorted(get_columns(
trid_dict, f"required_sparse_columns_{kind}")),
"produced": sorted(produced),
"linked": sorted(linked),
}
[docs]
def lost_support(mapper, node):
# Columns of the node value whose support may have been removed
# since last visit; the reference is updated for next visits
transform, side, trid, kind = node
trid_dict = mapper[transform][side][trid]
key = f"keep_sparse_columns_{kind}_base"
old, new = trid_dict.get(key), get_base(mapper, node)
trid_dict[key] = new
value = get_value(mapper, node)
if old is None or old == new:
return set()
if not new["sparse"] or set(old["linked"]) - set(new["linked"]):
return value
return value & (
(set(old["required"]) - set(new["required"]))
| (set(new["produced"]) - set(old["produced"]))
)
[docs]
def is_across_link(side, kind):
# Forward needs come from successors, adjoint needs from precursors
return (kind, side) in [("fwd", "outputs"), ("adj", "inputs")]
[docs]
def compute(mapper, node):
transform, side, trid, kind = node
transf_mapper = mapper[transform]
trid_dict = transf_mapper[side][trid]
if not is_tracked(mapper, node):
return None
columns = get_columns(trid_dict, f"required_sparse_columns_{kind}")
# Needs of neighbour transforms on the same trid
if is_across_link(side, kind):
links = "successors" if side == "outputs" else "precursors"
other_side = "inputs" if side == "outputs" else "outputs"
for tr in transf_mapper.get(links, {}).get(trid, []):
if trid in mapper.get(tr, {}).get(other_side, {}):
columns |= get_value(mapper, (tr, other_side, trid, kind))
return columns
# Needs from the other side of the transform, except produced columns
other_side = "outputs" if side == "inputs" else "inputs"
mapping = transf_mapper.get(f"{side}2{other_side}", {})
produced = f"produced_sparse_columns_{kind}"
passed = set()
for other_trid in mapping.get(trid, [trid]):
if other_trid not in transf_mapper[other_side]:
continue
other_dict = transf_mapper[other_side][other_trid]
passed |= (get_value(mapper, (transform, other_side, other_trid, kind))
- get_columns(other_dict, produced))
return columns | (passed - get_columns(trid_dict, produced))
[docs]
def dependents(mapper, node):
# Nodes whose computed value depends on this node
transform, side, trid, kind = node
transf_mapper = mapper[transform]
# Needs passed through the transform to the other side
if is_across_link(side, kind):
other_side = "inputs" if side == "outputs" else "outputs"
mapping = transf_mapper.get(f"{side}2{other_side}", {})
return [
(transform, other_side, other_trid, kind)
for other_trid in mapping.get(trid, [trid])
if other_trid in transf_mapper[other_side]
]
# Needs passed to neighbour transforms on the same trid
links = "precursors" if side == "inputs" else "successors"
other_side = "outputs" if side == "inputs" else "inputs"
return [
(tr, other_side, trid, kind)
for tr in transf_mapper.get(links, {}).get(trid, [])
if trid in mapper.get(tr, {}).get(other_side, {})
]
