class_docstr: GroupByPivot operator.
  
  Let's say we have a lineage of artifacts (executions omitted for brevity):
  
  ```dot
  digraph {
      a1 -> b1 -> c1
      a2 -> b2 -> c2
      a2 -> b3 -> c3
      a3 -> b4 -> c4
  }
  ```
  
  Then `GroupByPivot` would group artifacts by each pivot artifact from
  the input artifacts where artifacts in the same group is reachable from the
  pivot.
  
  ```python
  inputs = {
      'a': [a1, a2, a3],
      'b': [b1, b2, b3, b4],
      'c': [c1, c2, c3, c4],
  }
  # 'c' is empty they are not adjacent provenance from 'a'.
  GroupByPivot(inputs, pivot_key='a') == [
      {'a': [a1], 'b': [b1], 'c': []},
      {'a': [a2], 'b': [b2, b3], 'c': []},
      {'a': [a3], 'b': [b4], 'c': []},
  ]
  # Both 'a' and 'c' is not empty as they are adjacent provenance from 'b'.
  GroupByPivot(inputs, pivot_key='b') == [
      {'a': [a1], 'b': [b1], 'c': [c1]},
      {'a': [a2], 'b': [b2], 'c': [c2]},
      {'a': [a2], 'b': [b3], 'c': [c3]},
      {'a': [a3], 'b': [b4], 'c': [c4]},
  ]
  ```
  
  The result of the operator is a list of dictionary, where each dictionary
  contains individual pivot artifact. Non-pivot artifacts could be included in
  multiple dictionaries if they are associated with multiple pivots, or some
  dictionary might have empty artifact list for non-pivot artifacts if adjacent
  provenances are not found.
  
  CAVEAT: Lineage is only searched for the 2-hop distances (i.e. artifact ->
  execution -> artifact) and the artifacts farther than 2 hops from the pivot
  artifacts would NOT be included in the result.
class_docstr_tok: ["GroupByPivot","operator",".","Let","'s","say","we","have","a","lineage","of","artifacts","(","executions","omitted","for","brevity",")",":","``","`","dot","digraph","{","a1","-",">","b1","-",">","c1","a2","-",">","b2","-",">","c2","a2","-",">","b3","-",">","c3","a3","-",">","b4","-",">","c4","}","``","`","Then","`","GroupByPivot","`","would","group","artifacts","by","each","pivot","artifact","from","the","input","artifacts","where","artifacts","in","the","same","group","is","reachable","from","the","pivot",".","``","`","python","inputs","=","{","'","a","'",":","[","a1",",","a2",",","a3","]",",","'","b","'",":","[","b1",",","b2",",","b3",",","b4","]",",","'","c","'",":","[","c1",",","c2",",","c3",",","c4","]",",","}","#","'","c","'","is","empty","they","are","not","adjacent","provenance","from","'","a","'",".","GroupByPivot","(","inputs",",","pivot_key=","'","a","'",")","==","[","{","'","a","'",":","[","a1","]",",","'","b","'",":","[","b1","]",",","'","c","'",":","[","]","}",",","{","'","a","'",":","[","a2","]",",","'","b","'",":","[","b2",",","b3","]",",","'","c","'",":","[","]","}",",","{","'","a","'",":","[","a3","]",",","'","b","'",":","[","b4","]",",","'","c","'",":","[","]","}",",","]","#","Both","'","a","'","and","'","c","'","is","not","empty","as","they","are","adjacent","provenance","from","'","b","'",".","GroupByPivot","(","inputs",",","pivot_key=","'","b","'",")","==","[","{","'","a","'",":","[","a1","]",",","'","b","'",":","[","b1","]",",","'","c","'",":","[","c1","]","}",",","{","'","a","'",":","[","a2","]",",","'","b","'",":","[","b2","]",",","'","c","'",":","[","c2","]","}",",","{","'","a","'",":","[","a2","]",",","'","b","'",":","[","b3","]",",","'","c","'",":","[","c3","]","}",",","{","'","a","'",":","[","a3","]",",","'","b","'",":","[","b4","]",",","'","c","'",":","[","c4","]","}",",","]","``","`","The","result","of","the","operator","is","a","list","of","dictionary",",","where","each","dictionary","contains","individual","pivot","artifact",".","Non-pivot","artifacts","could","be","included","in","multiple","dictionaries","if","they","are","associated","with","multiple","pivots",",","or","some","dictionary","might","have","empty","artifact","list","for","non-pivot","artifacts","if","adjacent","provenances","are","not","found",".","CAVEAT",":","Lineage","is","only","searched","for","the","2-hop","distances","(","i.e",".","artifact","-",">","execution","-",">","artifact",")","and","the","artifacts","farther","than","2","hops","from","the","pivot","artifacts","would","NOT","be","included","in","the","result","."]
code: def apply(
      self, artifact_map: typing_utils.ArtifactMultiMap
  ) -> List[typing_utils.ArtifactMultiDict]:
      if self.pivot_key not in artifact_map:
          raise exceptions.FailedPreconditionError(
              f'Pivot "{self.pivot_key}" does not exist in the artifact map. '
              f"Containing keys: {list(artifact_map.keys())}"
          )
      if not artifact_map[self.pivot_key]:
          return []
  
      artifacts_by_id = {}
      input_keys_by_id = collections.defaultdict(set)
      for input_key, artifacts in artifact_map.items():
          for a in artifacts:
              input_keys_by_id[a.id].add(input_key)
              artifacts_by_id[a.id] = a
  
      events = self.context.store.get_events_by_artifact_ids(
          artifact_ids=artifacts_by_id
      )
  
      neighbors = collections.defaultdict(set)
      for a, b in _get_neighbor_artifact_pairs(events):
          neighbors[a].add(b)
          neighbors[b].add(a)
  
      result = []
      # Preserve the initial order in artifact_map[pivot_key].
      for pivot in artifact_map[self.pivot_key]:
          result_item = {input_key: [] for input_key in artifact_map}
          result_item[self.pivot_key].append(pivot)
          # Sort for deterministic result.
          for artifact_id in sorted(neighbors[pivot.id]):
              for input_key in input_keys_by_id[artifact_id]:
                  result_item[input_key].append(
                      artifacts_by_id[artifact_id]
                  )
          if not self.require_all or all(result_item.values()):
              result.append(result_item)
      return result
  
file_path: tfx/tfx/dsl/input_resolution/ops/group_by_lineage_op.py
imports: import collections
  from typing import List, Iterable, Tuple
  from tfx.dsl.input_resolution import resolver_op
  from tfx.orchestration.portable.input_resolution import exceptions
  from tfx.utils import typing_utils
  from ml_metadata.proto import metadata_store_pb2
label_desc: Use image node_id 1 for calling the GroupByPivot obj's underlying member method code with example usage: obj.apply(artifact_map) and returns: result, list