Python path: step plans' column names grow with depth and slow DataFusion's planning #211

🌀 Claude: from the xarray-sql session, which profiles its ML demos (xqlsystems/xarray-sql#285).

Summary

From Python, every step of a ddx program is decoded from Substrait by datafusion-python's Consumer.from_substrait_plan. That consumer names each computed column after its whole expression. When a step computes from computed columns, as a softmax does, each name contains the names below it, so names grow with depth. Two stacked softmaxes over 1,024 rows give a 25-node step whose plan prints as 789,000 characters, and a 137-node gradient step that prints as 1.8 million. DataFusion formats, hashes and compares these names as it optimizes and plans. ad.py's module docstring already notes the effect for deep chains of row-wise maps; this issue measures what it costs a training step.

Reproduction

ddxdb at 3dc0e62 (#203's branch, built locally), datafusion-python 55.0.0, Intel i7-8700:

import time

import numpy as np
import pyarrow as pa
from datafusion import SessionContext
from datafusion.substrait import Consumer, Serde
from ddxdb import ad

ctx = SessionContext()
g, k = np.divmod(np.arange(1024), 32)
ctx.register_record_batches("s", [[pa.record_batch({"g": g, "k": k, "x": np.sin(np.arange(1024.0))})]])
# A softmax over each g, then a second softmax over its output, as two layers of a model stack them.
for out, src, col in (("p1", "s", "x"), ("p2", "p1", "y")):
    ctx.sql(f"""CREATE VIEW {out} AS SELECT g, k, e / SUM(e) OVER (PARTITION BY g) AS y
                FROM (SELECT g, k, exp({col} - MAX({col}) OVER (PARTITION BY g)) AS e FROM {src})""").collect()
step = "WITH loss AS (SELECT SUM(y * y) AS l FROM p2) SELECT g, k, grad(loss, x) AS dx FROM s ORDER BY g, k"

# Measure each step's plan as ddx runs it.
materialize = ad.DataFusionBackend.materialize


def measured(self, name, plan):
    lp = Consumer.from_substrait_plan(ctx, Serde.deserialize_bytes(plan))
    text = lp.display_indent_schema()
    t = time.perf_counter()
    ctx.create_dataframe_from_logical_plan(lp).execution_plan()
    planning = time.perf_counter() - t
    print(f"{name.split('_1_')[-1]:<12} {len(text.splitlines()):>3} nodes, prints as {len(text):>7} chars "
          f"(longest line {max(map(len, text.splitlines())):>6}), planning {planning * 1e3:5.1f} ms")
    return materialize(self, name, plan)


ad.DataFusionBackend.materialize = measured
ad.sql(ctx, step).collect()

Output (abridged):

saved_0       18 nodes, prints as   96287 chars (longest line  17169), planning   2.0 ms
window_3      27 nodes, prints as  789387 chars (longest line  87571), planning   7.4 ms
window_6      25 nodes, prints as  165944 chars (longest line  13285), planning  13.0 ms
grad_0_s     137 nodes, prints as 1836628 chars (longest line  87685), planning  39.7 ms

Planning these 14 steps takes about 120 ms in all, for 1,024 rows.

What it costs

The names aren't the only cost of planning. DataFusion also spends time on each operator. To separate the two, I built the same chain of window layers with datafusion-python's DataFrame API twice: once with unaliased columns, so names nest as the consumer nests them, and once with each column aliased to a short name. At six layers (54,000-character names), optimizing and physical planning take 6.9 ms with the long names and 3.7 ms with short ones. At three layers, the two are the same. In the xarray-sql GPT demo, the deepest step, window_0, has 12 operators and prints as 173,000 characters.

So for deep steps, the names roughly double planning time. On a prepared GPT step, which plans 67 steps on every call, that's a share of the about 250 ms spent optimizing and planning (ad.prepare, DataFusion 55.0.0).

Possible directions

• Alias each computed column in the steps ddx writes. Substrait carries no alias for a project's output, but ddx could end a step with a project whose output names are set by the plan's RelRoot.names, as each step's table already does (c0, c1, …). Within a step, though, the nesting stays.
• Split a step where its names would grow past a size, as ddx already splits a step deeper than protobuf's decoding limit.
• Upstream: have DataFusion's Substrait consumer give a computed column a short generated name rather than its expression's display. That would also help every other Substrait user.
• Once a prepared program reuses physical plans instead of planning each call (see the plan-reuse note on perf(ddx-ad, ddx-datafusion, ddxdb): prepare statements to run again #203), names matter only on the first call.