Module distributed_plan 

Distributed physical-plan fragments and the stage builder (ADR-0003).

Phase 52 replaces stringly sql: <query> task bodies with protobuf-encoded DataFusion physical-plan subtrees. The krishiv_plan::TypedTaskFragment envelope stays as the carrier; this module owns the dfplan: body kind — encoding on the coordinator (stage builder) and decoding on the executor.

Body format: dfplan:v1:<partspec>:<base64(plan proto bytes)> where <partspec> names the output partition(s) of the decoded plan this task executes. The stage builder emits one partition per task (dfplan:v1:3:<b64>); Phase 54 AQE rewrites extend the grammar:

• Coalescing: dfplan:v1:1,4,7:<b64> — the task executes each listed root partition and concatenates the streams. Correct for any plan shape: root partitions are independent (each is a complete hash group), so the union of a task group’s outputs equals the union the original one-task-per-partition layout would produce.
• Skew split: dfplan:v1:5/s0m2-4:<b64> — the task executes root partition 5 but, for upstream stage 0, reads only map tasks [2, 4). Splitting is only correct when nothing above the shuffle read blocks on seeing the whole partition (see dfplan_body_is_split_safe).

The v1 segment is independent of the envelope version so plan-proto evolution (e.g. a DataFusion upgrade that changes the proto) is detected explicitly instead of failing deep inside prost decoding.

Stage building

build_distributed_stages cuts an optimized physical plan at hash RepartitionExec boundaries (Ballista-style): the subtree below each cut becomes a ShuffleMap stage whose tasks hash-partition their output into the shuffle store; the cut point is replaced by a ShuffleReadExec leaf that streams those partitions back on the reduce side. Any shape the builder cannot prove correct returns None — the caller falls back to today’s single-task sql: path (capability honesty).

Structs

ClusterCapacity

The compute capacity a query is being planned against.

DfplanMapRange

Restriction of a task’s shuffle reads to a subrange of one upstream stage’s map tasks (Phase 54 skew split).

DfplanTaskSpec

Parsed partition assignment of a dfplan:v1: task body.

DistributedStage

One stage of a distributed batch plan.

DistributedStagePlan

A batch query cut into shuffle-connected stages (Result stage last).

KrishivPhysicalCodec

Krishiv physical extension codec: (de)serializes ShuffleReadExec.

ParquetTableSpec

A parquet table to plan against, and what is known to be true about it.

ShuffleReadExec

Leaf node that streams an upstream ShuffleMap stage’s output partitions.

StageShuffleOutput

Shuffle-output contract of a ShuffleMap stage.

Constants

BROADCAST_JOIN_BYTES_ENV

Build-side byte ceiling under which a join is broadcast rather than hash-shuffled, on the staged path only. See planning_session_context for why the distributed default differs from DataFusion’s.

DFPLAN_BODY_PREFIX

Task-fragment body prefix for proto-encoded physical-plan subtrees.

DIMENSION_REDUCTION_ENV

Environment switch for the broadcast-dimension reducer. Default off.

SHUFFLE_FETCH_BUFFER_ENV

Env var overriding how many map fragments ONE reduce partition fetches at once. See the buffered call in ShuffleReadExec::execute.

STAGE_REUSE_ENV

Env flag for cross-stage reuse of identical leaf stages. Default off.

STAGE_SPLIT_ENV

Env var that disables stage splitting entirely (off/0/false).

STAGE_TARGET_PARTITIONS_ENV

Env var overriding the target partition count used when planning a distributed batch query (bounds both scan parallelism and shuffle partition count). Unset, the count is derived from the cluster — see resolve_stage_target_partitions.

Traits

ShufflePartitionReader

Executor-side access to upstream shuffle partitions.

Functions

build_distributed_stages

Cut a physical plan into shuffle-connected stages.

build_distributed_stages_with_udf_directives

See build_distributed_stages; udf_directive_source supplies the query’s inline Python-UDF directives so the decode rehearsal can resolve them.

build_stages_for_parquet_query

Plan a query over parquet tables and cut it into stages (coordinator seam — keeps DataFusion types out of the scheduler crate).

build_stages_for_parquet_tables

As build_stages_for_parquet_query, but each table may declare what is known about it — currently a primary key, which the optimizer turns into a functional dependency and uses to shrink GROUP BY lists.

decode_dfplan_task

Decode a dfplan:v1: fragment body into (partition spec, plan).

derive_stage_target_partitions

The pure derivation behind resolve_stage_target_partitions, with the environment and the machine passed in so it is testable (the workspace forbids unsafe, so tests cannot set environment variables).

dfplan_body_is_split_safe

True when a dfplan body’s decoded plan may be split by map-task ranges (Phase 54 skew split) without changing results.

dfplan_body_partition_spec

Parse the partition spec of a body without decoding the plan payload (cheap coordinator-side inspection).

dfplan_body_with_spec

Rewrite an existing dfplan body to a new partition spec, preserving the encoded plan bytes verbatim (no proto decode — coordinator-side AQE rewrites reuse the b64 payload untouched).

dfplan_task_body

Assemble the per-task fragment body: dfplan:v1:<partition>:<b64>.

dfplan_task_body_for_spec

Assemble a fragment body executing several root partitions (coalescing).

encode_dfplan_bytes

Encode a physical plan (sub)tree to raw proto bytes.

execute_dfplan_body

fragment_decode_session_context

The session context a dfplan:v1: fragment is decoded on.

is_dfplan_body

True when a task-fragment body carries a proto-encoded physical plan.

parse_dfplan_body

Split a dfplan:v1: body into (partition spec, plan proto bytes).

planning_session_context

Session context used to plan a query for distributed stage execution.

planning_session_context_with_join_threshold

As planning_session_context, with the spillable-join build-side threshold supplied instead of derived from this process’s cgroup.

planning_session_context_with_options

As planning_session_context_with_join_threshold, with the broadcast (CollectLeft) build-side ceiling supplied instead of read from BROADCAST_JOIN_BYTES_ENV.

redistribute_unsplittable_broadcast_joins

Convert broadcast joins that cannot be split — or that were chosen on an estimate claiming their build side is empty, or on a row ceiling blind to how wide those rows are — into hash-partitioned joins (see [is_unsplittable_broadcast_join], [is_degenerate_broadcast_join] and [broadcast_build_is_too_wide]).

register_parquet_table

Register one parquet table, attaching its declared key as a DataFusion constraint so the optimizer’s functional-dependency machinery can see it.

register_python_udf_signatures_and_strip

Register a signature-only DataFusion scalar UDF for every inline /* krishiv-register-python-udf:name:in,…:out:pickle */ directive in query, so the coordinator can plan a staged query that references it, and return query with the directives stripped (clean SQL for the parser).

resolve_stage_target_partitions

Resolve the planning-time target partition count for distributed stages.

shuffle_stage_key

Shuffle-store sub-stage key for one map task’s output.

stage_reuse_enabled

Whether identical leaf stages are collapsed into one.

stage_split_enabled

True unless stage splitting is disabled via STAGE_SPLIT_ENV.

Type Aliases

ShuffleFragmentStream

One upstream fragment’s batches, in write order.