Databricks pyspark Nested to Flatten format

This code bit helps us to extract dataframe objects to a flattened format.

# --- find all array paths (any depth) ---
def _array_paths(schema: StructType, prefix=()):
    paths = []
    for f in schema.fields:
        path = prefix + (f.name,)
        dt = f.dataType
        if isinstance(dt, ArrayType):
            paths.append(path)
        elif isinstance(dt, StructType):
            paths.extend(_array_paths(dt, path))
    return paths

# --- explode arrays anywhere (iterative; no renaming policy besides joining path with "_") ---
def _explode_all_arrays(df):
    processed = set()  # set of tuple paths we've already exploded (to avoid infinite loops)

while True:
        # collect current array paths (including those under columns created in previous iterations)
        paths = [p for p in _array_paths(df.schema) if p not in processed]
        if not paths:
            break

# explode one path at a time; schema changes each time
        path = paths[0]
        dotted = ".".join(path)       # e.g., "a.b.c"
        alias  = "_".join(path)       # e.g., "a_b_c" (keeps original names, just joined)

# create a top-level column with the exploded elements
        df = df.withColumn(alias, explode_outer(col(dotted)))
        processed.add(path)
        # NOTE: we don't try to "drop nested field inside a struct" (which is awkward in Spark).
        # We'll simply NOT select it in the final flatten/select step.

return df

# --- gather all leaf (non-struct, non-array) paths for final flat select ---
def _leaf_paths(schema: StructType, prefix=()):
    leaves = []
    for f in schema.fields:
        path = prefix + (f.name,)
        dt = f.dataType
        if isinstance(dt, StructType):
            leaves.extend(_leaf_paths(dt, path))
        elif isinstance(dt, ArrayType):
            # arrays should be gone by now due to _explode_all_arrays, but we skip defensively
            continue
        else:
            leaves.append(path)
    return leaves

# --- public function: explode everything and flatten to root ---
def explode_and_flatten(df):
    """
    Returns a DataFrame where:
      1) all arrays (at any depth) are exploded with outer semantics
      2) all structs are flattened so NO nested fields remain
    Column names are the original path components joined with "_".
    """
    df1 = _explode_all_arrays(df)

# final flat select: pick every leaf as a top-level column
    select_exprs, used = [], set()
    for path in _leaf_paths(df1.schema):
        dotted = ".".join(path)
        alias  = "_".join(path)
        # avoid duplicate aliases if they happen to collide
        if alias in used:
            # de-dup by appending an index
            i = 2
            new_alias = f"{alias}_{i}"
            while new_alias in used:
                i += 1
                new_alias = f"{alias}_{i}"
            alias = new_alias
        select_exprs.append(col(dotted).alias(alias))
        used.add(alias)

return df1.select(*select_exprs)

df_spark_cursor_flat = explode_and_flatten(df_spark_cursor)