PySpark column lineage

 def _node2column(node: str) -> str:
      """Inner function. Transform the node from plan to column name.
      Like: col_11#1234L -> col_11.
      """
      match_ = re.match(r"([\w\d]+)#[\w\d]+", node)
      if match_:
          return match_.groups()[0]

  def _get_aliases(col: str, line: str) -> tuple[list[str], str]:
      """Inner function. Returns all the aliases from the expr and expr itself."""
      alias_exp = _extract_alias_expressions(col, line)
      # Regexp to extract columns: each column has a pattern like col_name#1234
      return (re.findall(r"[\w\d]+#[\w\d]+", alias_exp), alias_exp)

  def _extract_alias_expressions(col: str, line: str) -> str:
      """Inner function. Extract expression before ... AS col from the line."""
      num_close_parentheses = 0 # our counter
      idx = line.index(f" AS {col}") # the end of the alias expression we need to parse
      alias_expr = [] # buffer to store what we are parsing

      if line[idx - 1] != ")":
          """It is possible that there is no expression.
          It is the case when we just make a rename of the column. In the plan
          it will look like `col#123 AS col#321`;
          """
          for j in range(idx - 1, 0, -1):
              alias_expr.append(line[j])
              if line[j - 1] == "[":
                  break
              if line[j - 1] == " ":
                  break
          return "".join(alias_expr)[::-1]

      """In all other cases there will be `(` at the end of the expr before AS.
      Our goal is to go symbol by symbol back until we balance all the parentheses.
      """
      for i in range(idx - 1, 0, -1):
          alias_expr.append(line[i])
          if line[i] == ")":
              # Add parenthesis
              num_close_parentheses += 1
          if line[i] == "(":
              if num_close_parentheses == 1:
                  # Parentheses are balanced
                  break
              # Remove parenthesis
              num_close_parentheses -= 1

      """After balancing parentheses we need to parse leading expression.
      It is always here because we checked single alias case separately."""
      for j in range(i, 0, -1):
          alias_expr.append(line[j])
          if line[j - 1] == "[":
              break
          if line[j - 1] == " ":
              break

      return "".join(alias_expr[::-1])

  _extract_alias_expressions(
      "id1_and_id2#2437",
      "Project [id1#1321, id2#1322, id4#1323, id5#1324, v2#1325, id1_renamed#2430, concat_ws(_, id1#1321, id2#1322) AS id1_and_id2#2437]"
  )

  def _add_aggr_or_not(expr: str, line: str) -> str:
      """If the expr is aggregation we should add agg keys to the beginning."""
      # We are checking for aggregation pattern
      match_ = re.match(r"^[\s\+\-:]*Aggregate\s\[([\w\d#,\s]+)\].*$", line)
      if match_:
          agg_expr = match_.groups()[0]
          return (
              "GroupBy: " + re.sub(r"([\w\d]+)#([\w\d]+)", r"\1", agg_expr) + f"\n{expr}"
          )

      # If not just return an original expr
      return expr

  def _get_graph(lines: list[str], node: str):
      nodes = []
      edges = []
      node_attrs = {}

      for i, l in enumerate(lines):  # noqa: E741
          """Iteration over lines of logical plan."""

          # We should use hash of line + node as a key in the graph.
          # It is not enough to use only hash of line because the same line
          # may be related to multiple nodes!
          # A good example is reading the CSV that is represented by one line!
          h = hash(l + node)

          # If the current node is not root we need to store hash of previous node.
          prev_h = None if not nodes else nodes[-1]

          if node not in l:
              continue
          if f"AS {node}" in l:
              """It is a hard case, when a node is an alias to some expression."""
              aliases, expr = _get_aliases(node, l)
              # For visualization we need to transform from nodes to columns
              expr = re.sub(r"([\w\d]+)#([\w\d]+)", r"\1", expr)

              # Append a new node
              nodes.append(h)
              # Append expr as an attribute of the node
              node_attrs[h] = _add_aggr_or_not(f"{expr} AS {_node2column(node)}", l)

              if len(aliases) == 1:
                  # It is the case of simple alis
                  # Like col1#123 AS col2#321
                  # In this case we just replace an old node by new one.
                  if prev_h:
                      edges.append([h, prev_h])
                  node = aliases[0]
              else:
                  # It is a case of complex expression.
                  # Here we recursively go through all the nodes from expr.
                  if prev_h:
                      edges.append([h, prev_h])
                  for aa in aliases:
                      # Get graph from sub-column
                      sub_nodes, sub_edges, sub_attrs = _get_graph(lines[i:], aa)

                      # Add everything to the current graph
                      nodes.extend(sub_nodes)
                      edges.extend(sub_edges)
                      node_attrs = {**node_attrs, **sub_attrs}

                      # Add connection between top subnode and node
                      edges.append([sub_nodes[0], h])
                  return (nodes, edges, node_attrs)
          else:
              # Continue of the simple alias or expr case
              # In the future that may be more cases, that is the reason of nested if instead of elif
              if "Relation" in l:
                  nodes.append(h)
                  if prev_h:
                      edges.append([h, prev_h])

                  # It is a pattern, related to data-sources (like CSV)
                  match_ = re.match(r"[\s\+\-:]*Relation\s\[.*\]\s(\w+)", l)
                  if match_:
                      s_ = "Read from {}: {}"
                      # Add data-source as a node
                      node_attrs[h] = s_.format(match_.groups()[0], _node2column(node))
                  else:
                      # We need it to avoid empty graphs and related runtime exceptions
                      print(l)
                      node_attrs[h] = f"Relation to: {_node2column(node)}"

              elif "Join" in l:
                  nodes.append(h)
                  if prev_h:
                      edges.append([h, prev_h])
                  match_ = re.match(r"[\s\+\-:]*Join\s(\w+),\s\((.*)\)", l)
                  if match_:
                      join_type = match_.groups()[0]
                      join_expr = match_.groups()[1]
                      join_expr_clr = re.sub(r"([\w\d]+)#([\w\d]+)", r"\1", join_expr)
                      node_attrs[h] = f"{join_type}: {join_expr_clr}"
              else:
                  continue

      if not nodes:
          # Just the case of empty return. We need to avoid it.
          # I'm not sure that line is reachable.
          nodes.append(h)
          node_attrs[h] = f"Select: {_node2column(node)}"

      return (nodes, edges, node_attrs)

  def get_column_lineage(df: DataFrame, column: str) -> ColumnLineageGraph:
      """Get data lineage on the level of the given column.

      Currently Union operation is not supported! API is unstable, no guarantee
      that custom spark operations or connectors won't break it!

      :param df: DataFrame
      :param column: column
      :returns: Struct with nodes, edges and attributes
      """
      lines = get_plan_from_df(df, PlanType.ANALYZED_LOGICAL_PLAN).split("\n")

      # Top line should contain plan-id of our column. We need it.
      # Regular pattern of node is column#12345L or [\w\d]+#[\w\d]+
      match_ = re.match(r".*(" + column + r"#[\w\d]+).*", lines[0])
      if match_:
          node = match_.groups()[0]
      else:
          err = f"There is no column {column} in the final schema of DF!"
          raise KeyError(err)

      nodes, edges, attrs = _get_graph(lines, node)

      return ColumnLineageGraph(nodes, edges, attrs)