Capture query plan in pds-h output by TomAugspurger · Pull Request #21315 · NVIDIA/cudf

TomAugspurger · 2026-02-03T19:33:22Z

Description

Currently, pds-h benchmarks will optionally print out the query plan if --explain is passed. With this PR, we'll also persist the query plan in the pdsh_results.jsonl file, under RunConfig.plans, a dict mapping query ID to (serialized) query plan. This will facilitate downstream analysis.

To make it easier for downstream tools to work with, we persist a structured version of the query plan, rather than the tree-like textual representation.

Here's an example:

$ python -m cudf_polars.experimental.benchmarks.pdsh --path /datasets/toaugspurger/tpch-rs/scale-10/ --suffix "" --iterations 1 1 --explain --runtime rapidsmpf
$ tail -n1 pdsh_results.jsonl | jq '.plans."1"'

giving

Details

{
  "roots": [
    "1598104009"
  ],
  "nodes": {
    "1598104009": {
      "id": "1598104009",
      "children": [
        "2141244145"
      ],
      "schema": {
        "l_returnflag": "STRING",
        "l_linestatus": "STRING",
        "sum_qty": "FLOAT64",
        "sum_base_price": "FLOAT64",
        "sum_disc_price": "FLOAT64",
        "sum_charge": "FLOAT64",
        "avg_qty": "FLOAT64",
        "avg_price": "FLOAT64",
        "avg_disc": "FLOAT64",
        "count_order": "UINT32"
      },
      "properties": {
        "by": [
          "l_returnflag",
          "l_linestatus"
        ],
        "order": [
          "ASCENDING",
          "ASCENDING"
        ]
      },
      "type": "Sort"
    },
    "2141244145": {
      "id": "2141244145",
      "children": [
        "3011781837"
      ],
      "schema": {
        "l_returnflag": "STRING",
        "l_linestatus": "STRING",
        "sum_qty": "FLOAT64",
        "sum_base_price": "FLOAT64",
        "sum_disc_price": "FLOAT64",
        "sum_charge": "FLOAT64",
        "avg_qty": "FLOAT64",
        "avg_price": "FLOAT64",
        "avg_disc": "FLOAT64",
        "count_order": "UINT32"
      },
      "properties": {
        "columns": [
          "l_returnflag",
          "l_linestatus",
          "sum_qty",
          "sum_base_price",
          "sum_disc_price",
          "sum_charge",
          "avg_qty",
          "avg_price",
          "avg_disc",
          "count_order"
        ]
      },
      "type": "Select"
    },
    "3011781837": {
      "id": "3011781837",
      "children": [
        "3056697853"
      ],
      "schema": {
        "l_returnflag": "STRING",
        "l_linestatus": "STRING",
        "sum_qty": "FLOAT64",
        "sum_base_price": "FLOAT64",
        "sum_disc_price": "FLOAT64",
        "sum_charge": "FLOAT64",
        "avg_qty": "FLOAT64",
        "avg_price": "FLOAT64",
        "avg_disc": "FLOAT64",
        "______________11": "UINT32"
      },
      "properties": {
        "columns": [
          "l_returnflag",
          "l_linestatus",
          "sum_qty",
          "sum_base_price",
          "sum_disc_price",
          "sum_charge",
          "avg_qty",
          "avg_price",
          "avg_disc",
          "______________11"
        ]
      },
      "type": "Select"
    },
    "3056697853": {
      "id": "3056697853",
      "children": [
        "3958692417"
      ],
      "schema": {
        "l_returnflag": "STRING",
        "l_linestatus": "STRING",
        "sum_qty": "FLOAT64",
        "sum_base_price": "FLOAT64",
        "sum_disc_price": "FLOAT64",
        "sum_charge": "FLOAT64",
        "________________0__mean_sum": "FLOAT64",
        "________________1__mean_count": "INT32",
        "________________2__mean_sum": "FLOAT64",
        "________________3__mean_count": "INT32",
        "________________4__mean_sum": "FLOAT64",
        "________________5__mean_count": "INT32",
        "______________11": "UINT32"
      },
      "properties": {
        "keys": [
          "l_returnflag",
          "l_linestatus"
        ]
      },
      "type": "GroupBy"
    },
    "3958692417": {
      "id": "3958692417",
      "children": [
        "2166696742"
      ],
      "schema": {
        "l_returnflag": "STRING",
        "l_linestatus": "STRING",
        "sum_qty": "FLOAT64",
        "sum_base_price": "FLOAT64",
        "sum_disc_price": "FLOAT64",
        "sum_charge": "FLOAT64",
        "________________0__mean_sum": "FLOAT64",
        "________________1__mean_count": "INT32",
        "________________2__mean_sum": "FLOAT64",
        "________________3__mean_count": "INT32",
        "________________4__mean_sum": "FLOAT64",
        "________________5__mean_count": "INT32",
        "______________11": "UINT32"
      },
      "properties": {},
      "type": "Repartition"
    },
    "2166696742": {
      "id": "2166696742",
      "children": [
        "101741285"
      ],
      "schema": {
        "l_returnflag": "STRING",
        "l_linestatus": "STRING",
        "sum_qty": "FLOAT64",
        "sum_base_price": "FLOAT64",
        "sum_disc_price": "FLOAT64",
        "sum_charge": "FLOAT64",
        "________________0__mean_sum": "FLOAT64",
        "________________1__mean_count": "INT32",
        "________________2__mean_sum": "FLOAT64",
        "________________3__mean_count": "INT32",
        "________________4__mean_sum": "FLOAT64",
        "________________5__mean_count": "INT32",
        "______________11": "UINT32"
      },
      "properties": {
        "keys": [
          "l_returnflag",
          "l_linestatus"
        ]
      },
      "type": "GroupBy"
    },
    "101741285": {
      "id": "101741285",
      "children": [
        "3768794358"
      ],
      "schema": {
        "l_quantity": "FLOAT64",
        "l_extendedprice": "FLOAT64",
        "l_discount": "FLOAT64",
        "l_tax": "FLOAT64",
        "l_returnflag": "STRING",
        "l_linestatus": "STRING",
        "l_shipdate": "TIMESTAMP_MILLISECONDS",
        "__POLARS_CSER_0x2d17d531a95f17f3": "FLOAT64"
      },
      "properties": {
        "columns": [
          "__POLARS_CSER_0x2d17d531a95f17f3"
        ]
      },
      "type": "HStack"
    },
    "3768794358": {
      "id": "3768794358",
      "children": [],
      "schema": {
        "l_quantity": "FLOAT64",
        "l_extendedprice": "FLOAT64",
        "l_discount": "FLOAT64",
        "l_tax": "FLOAT64",
        "l_returnflag": "STRING",
        "l_linestatus": "STRING",
        "l_shipdate": "TIMESTAMP_MILLISECONDS"
      },
      "properties": {
        "typ": "parquet",
        "paths": [
          "/datasets/toaugspurger/tpch-rs/scale-10//lineitem/part.0.parquet"
        ]
      },
      "type": "Scan"
    }
  },
  "partition_info": {
    "1598104009": {
      "count": 1,
      "partitioned_on": [
        "l_returnflag",
        "l_linestatus"
      ]
    },
    "2141244145": {
      "count": 1,
      "partitioned_on": [
        "l_returnflag",
        "l_linestatus"
      ]
    },
    "3011781837": {
      "count": 1,
      "partitioned_on": [
        "l_returnflag",
        "l_linestatus"
      ]
    },
    "3056697853": {
      "count": 1,
      "partitioned_on": []
    },
    "3958692417": {
      "count": 1,
      "partitioned_on": []
    },
    "2166696742": {
      "count": 3,
      "partitioned_on": []
    },
    "101741285": {
      "count": 3,
      "partitioned_on": []
    },
    "3768794358": {
      "count": 3,
      "partitioned_on": []
    }
  }
}

One design note: I've flattened the query plan for serialization. See the updated overview.md, but the serialized object has a nodes key that's a mapping from node ID to node for all the nodes in the query plan. That seemed easier to work with than having to traverse the tree. The roots key gives you the root node ID(s), and you can build the tree from there by inspecting the children on each node.

Currently, pds-h benchmarks will optionally print out the query plan if `--explain` is passed. With this PR, we'll also persist the query plan in the `pdsh_results.jsonl` file. This will facilitate downstream analysis. To make it easier for downstream tools to work with, we persist a structured version of the query plan, rather than the tree-like textual representation.

copy-pr-bot · 2026-02-03T19:33:26Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

TomAugspurger · 2026-02-03T23:00:09Z

+    cudf_polars.dsl.ir.IR node.
+    """
+
+    id: str


ir.get_stable_id() returns an integer. But everything in in this module uses string IDs, since they end up as keys inside DAG.nodes, and JSON requires string keys.

TomAugspurger · 2026-02-04T14:07:23Z

In case it's helpful to see where this is (potentially) going, https://github.com/TomAugspurger/pygdf/tree/tom/structured-plans%2Btracing adds some additional things to our structlog records (similar to #21316).

With the changes there, we can run Q1 and capture the plan:

❯ CUDF_POLARS_LOG_TRACES=1 CUDF_POLARS_LOG_TRACES_MEMORY=0 python -m cudf_polars.experimental.benchmarks.pdsh --path /datasets/toaugspurger/tpch-rs/scale-10/ --suffix "" --iterations 1 1 --no-print-results --no-summarize --runtime rapidsmpf --collect-traces --explain

Query 1 - Physical plan

SORT ('l_returnflag', 'l_linestatus') ('l_returnflag', 'l_linestatus', 'sum_qty', '...', 'avg_disc', 'count_order') [1]
  REPARTITION ('l_returnflag', 'l_linestatus', 'sum_qty', '...', 'avg_disc', 'count_order') [1]
    SELECT ('l_returnflag', 'l_linestatus', 'sum_qty', '...', 'avg_disc', 'count_order') [1]
      SELECT ('l_returnflag', 'l_linestatus', 'sum_qty', '...', 'avg_disc', '______________11') [1]
        GROUPBY ('l_returnflag', 'l_linestatus') ('l_returnflag', 'l_linestatus', 'sum_qty', '...', '________________5__mean_count', '______________11') [1]
          REPARTITION ('l_returnflag', 'l_linestatus', 'sum_qty', '...', '________________5__mean_count', '______________11') [1]
            GROUPBY ('l_returnflag', 'l_linestatus') ('l_returnflag', 'l_linestatus', 'sum_qty', '...', '________________5__mean_count', '______________11') [3]
              HSTACK ('l_quantity', 'l_extendedprice', 'l_discount', '...', 'l_shipdate', '__POLARS_CSER_0xd3b9042278ae2444') [3]
                SCAN PARQUET ('l_quantity', 'l_extendedprice', 'l_discount', '...', 'l_linestatus', 'l_shipdate') [3]

Query 1 - Iteration 0 finished in 0.5329s

And then reconstruct it (with the annotations, like number of chunks) from the serialized JSON:

❯ python python/cudf_polars/scripts/visualize_traces.py pdsh_results.jsonl
SORT ('l_returnflag', 'l_linestatus') ('l_returnflag', 'l_linestatus', 'sum_qty', '...', 'avg_disc', 'count_order') chunks=1
  REPARTITION ('l_returnflag', 'l_linestatus', 'sum_qty', '...', 'avg_disc', 'count_order') chunks=1
    SELECT ('l_returnflag', 'l_linestatus', 'sum_qty', '...', 'avg_disc', 'count_order') chunks=1
      SELECT ('l_returnflag', 'l_linestatus', 'sum_qty', '...', 'avg_disc', '______________11') chunks=1
        GROUPBY ('l_returnflag', 'l_linestatus') ('l_returnflag', 'l_linestatus', 'sum_qty', '...', '________________5__mean_count', '______________11') chunks=1
          REPARTITION ('l_returnflag', 'l_linestatus', 'sum_qty', '...', '________________5__mean_count', '______________11') chunks=1
            GROUPBY ('l_returnflag', 'l_linestatus') ('l_returnflag', 'l_linestatus', 'sum_qty', '...', '________________5__mean_count', '______________11') chunks=3
              HSTACK ('l_quantity', 'l_extendedprice', 'l_discount', '...', 'l_shipdate', '__POLARS_CSER_0xd3b9042278ae2444') chunks=3
                SCAN PARQUET ('l_quantity', 'l_extendedprice', 'l_discount', '...', 'l_linestatus', 'l_shipdate') chunks=3

none of that work is vetted yet, but I'm hopeful that it shows we can get some (and hopefully all) of the information we want out of these logs.

use this serialize when logging.

rjzamora · 2026-02-06T21:50:43Z

Just a note that I was planning to review this much earlier, but keep getting distracted. That said, I'm on board with the general changes - Nothin big is jumping out at me.

rjzamora

Thanks Tom! I left minor comments, but this should be good after the test conflict is resolved.

TomAugspurger · 2026-02-10T14:55:32Z

/merge

github-actions Bot assigned TomAugspurger Feb 3, 2026

github-actions Bot added Python Affects Python cuDF API. cudf-polars Issues specific to cudf-polars labels Feb 3, 2026

github-project-automation Bot added this to cuDF Python Feb 3, 2026

GPUtester moved this to In Progress in cuDF Python Feb 3, 2026

TomAugspurger commented Feb 3, 2026

View reviewed changes

Comment thread python/cudf_polars/cudf_polars/experimental/explain.py

Comment thread python/cudf_polars/cudf_polars/experimental/explain.py

Comment thread python/cudf_polars/cudf_polars/experimental/explain.py Outdated

TomAugspurger added non-breaking Non-breaking change improvement Improvement / enhancement to an existing function labels Feb 3, 2026

TomAugspurger added 2 commits February 3, 2026 14:57

use string IDs

aed6708

Merge remote-tracking branch 'upstream/main' into tom/structured-plans

1234b64

TomAugspurger commented Feb 3, 2026

View reviewed changes

TomAugspurger added 2 commits February 3, 2026 15:02

fix

66e3356

fixes

01a42ce

TomAugspurger marked this pull request as ready for review February 3, 2026 23:04

TomAugspurger requested a review from a team as a code owner February 3, 2026 23:04

TomAugspurger requested review from mroeschke and vyasr February 3, 2026 23:04

test debugging

fd6a52c

TomAugspurger added 5 commits February 4, 2026 07:04

polars compat

748731f

Merge remote-tracking branch 'upstream/main' into tom/structured-plans

0ebdbcc

cache stable id

53f30b1

test coverage

5924b82

Merge remote-tracking branch 'upstream/main' into tom/structured-plans

1813d88

TomAugspurger commented Feb 4, 2026

View reviewed changes

Comment thread python/cudf_polars/cudf_polars/experimental/explain.py Outdated

TomAugspurger added 2 commits February 4, 2026 14:18

testing

89f34dd

coverage

0429ef1

TomAugspurger mentioned this pull request Feb 4, 2026

Full structlog-based tracing #21316

Merged

3 tasks

TomAugspurger added 6 commits February 5, 2026 05:28

doc strings

fe6e794

Merge remote-tracking branch 'upstream/main' into tom/structured-plans

1333359

more docs

6295afe

polars 1.30 compat

bb57d6f

Merge remote-tracking branch 'upstream/main' into tom/structured-plans

71e7312

Move to get_stable_id

f3f908a

TomAugspurger commented Feb 5, 2026

View reviewed changes

Comment thread python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py Outdated

Updates for logging changes

a1ea778

use this serialize when logging.

rjzamora mentioned this pull request Feb 6, 2026

[FEA] Add mechanism to print annotated physical plan #21367

Open

Merge branch 'main' into tom/structured-plans

e60fc0a

TomAugspurger added 2 commits February 6, 2026 14:39

Merge remote-tracking branch 'upstream/main' into tom/structured-plans

a8868f8

dag -> plan

75027f5

rjzamora approved these changes Feb 8, 2026

View reviewed changes

Comment thread python/cudf_polars/cudf_polars/experimental/explain.py Outdated

Comment thread python/cudf_polars/cudf_polars/experimental/explain.py Outdated

Comment thread python/cudf_polars/cudf_polars/experimental/explain.py Outdated

TomAugspurger added 3 commits February 9, 2026 06:38

Merge remote-tracking branch 'upstream/main' into tom/structured-plans

a01d803

example

5728aa9

Merge branch 'main' into tom/structured-plans

a85e322

rapids-bot Bot merged commit d1ed501 into NVIDIA:main Feb 10, 2026
108 checks passed

github-project-automation Bot moved this from In Progress to Done in cuDF Python Feb 10, 2026

mroeschke added a commit to mroeschke/cudf that referenced this pull request Mar 6, 2026

Update test_log_query_plan after NVIDIA#21315

30294dd

Navigation Menu

Uh oh!

Capture query plan in pds-h output - #21315

Capture query plan in pds-h output#21315
rapids-bot[bot] merged 26 commits into
NVIDIA:mainfrom
TomAugspurger:tom/structured-plans

TomAugspurger commented Feb 3, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Feb 3, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

TomAugspurger Feb 3, 2026

Uh oh!

TomAugspurger commented Feb 4, 2026

Uh oh!

Uh oh!

Uh oh!

rjzamora commented Feb 6, 2026

Uh oh!

rjzamora left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

TomAugspurger commented Feb 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

TomAugspurger commented Feb 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Uh oh!

copy-pr-bot Bot commented Feb 3, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

TomAugspurger Feb 3, 2026

Choose a reason for hiding this comment

Uh oh!

TomAugspurger commented Feb 4, 2026

Uh oh!

Uh oh!

Uh oh!

rjzamora commented Feb 6, 2026

Uh oh!

rjzamora left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

TomAugspurger commented Feb 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

TomAugspurger commented Feb 3, 2026 •

edited

Loading