pyspark-inspect 0.1.0__tar.gz

pyspark_inspect-0.1.0/PKG-INFO

@@ -0,0 +1,44 @@
+Metadata-Version: 2.4
+Name: pyspark-inspect
+Version: 0.1.0
+Summary: Inspect pyspark query plans
+License: MIT
+Author: Matthias Ossadnik
+Author-email: ossadnik.matthias@gmail.com
+Requires-Python: >=3.11,<4.0
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Requires-Dist: pyspark (==3.5.*)
+Project-URL: homepage, https://github.com/mossadnik/pyspark-inspect
+Description-Content-Type: text/markdown
+
+# pyspark-inspect
+
+pyspark-inspect allows to inspect the query plans of pyspark dataframes from Python.
+
+It converts analyzed Catalyst plans into Python data structures that can be queried programmatically.
+
+## Compatibility & Limitiations
+
+- Tested against Spark 3.5
+- Spark Connect is not supported
+- Coverage of Spark SQL operations is currently very limited
+
+## Basic Usage
+
+```python
+from pypspark.sql import functions as F
+from pyspark_inspect import inspect_dataframe
+
+df = (
+    spark.createDataFrame([[1, 2]], ['a', 'b'])
+    .withColumn('c', F.col('a') + F.col('b'))
+)
+plan = inspect_dataframe(df)
+```
+

pyspark_inspect-0.1.0/README.md

@@ -0,0 +1,24 @@
+# pyspark-inspect
+
+pyspark-inspect allows to inspect the query plans of pyspark dataframes from Python.
+
+It converts analyzed Catalyst plans into Python data structures that can be queried programmatically.
+
+## Compatibility & Limitiations
+
+- Tested against Spark 3.5
+- Spark Connect is not supported
+- Coverage of Spark SQL operations is currently very limited
+
+## Basic Usage
+
+```python
+from pypspark.sql import functions as F
+from pyspark_inspect import inspect_dataframe
+
+df = (
+    spark.createDataFrame([[1, 2]], ['a', 'b'])
+    .withColumn('c', F.col('a') + F.col('b'))
+)
+plan = inspect_dataframe(df)
+```

pyspark_inspect-0.1.0/pyproject.toml

@@ -0,0 +1,34 @@
+[tool.poetry]
+name = "pyspark-inspect"
+version = "0.1.0"
+description = "Inspect pyspark query plans"
+authors = [
+    "Matthias Ossadnik <ossadnik.matthias@gmail.com>"
+]
+classifiers = [
+    "Development Status :: 3 - Alpha"
+]
+readme = "README.md"
+license = "MIT"
+
+[tool.poetry.urls]
+homepage = "https://github.com/mossadnik/pyspark-inspect"
+
+
+[tool.poetry.dependencies]
+python = "^3.11"
+pyspark = "==3.5.*"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[dependency-groups]
+dev = [
+    "jupyter (>=1.1.1,<2.0.0)",
+    "pytest (>=9.0.2,<10.0.0)",
+    "mkdocs (>=1.6.1,<2.0.0)",
+    "mkdocstrings[python] (>=1.0.4,<2.0.0)",
+    "pyyaml (>=6.0.3,<7.0.0)",
+    "z3-solver (>=4.14.0,<4.15.0)"
+]

pyspark_inspect-0.1.0/pyspark_inspect/__init__.py

@@ -0,0 +1,6 @@
+from .api import inspect_dataframe
+
+
+__all__ = [
+    'inspect_dataframe'
+]

pyspark_inspect-0.1.0/pyspark_inspect/api.py

@@ -0,0 +1,9 @@
+"""High-level public functions and classes."""
+
+from pyspark.sql import DataFrame
+from .plan import Plan
+from .parser.catalyst import get_dataframe_plan, load_catalyst_plan, parse_plan
+
+
+def inspect_dataframe(df: DataFrame) -> Plan:
+    return parse_plan(load_catalyst_plan(get_dataframe_plan(df)))

pyspark_inspect-0.1.0/pyspark_inspect/code_generator.py

@@ -0,0 +1,52 @@
+"""Translate plans into pyspark DataFrames and Columns."""
+
+import typing as tp
+from functools import singledispatch
+from pyspark.sql import SparkSession, DataFrame, Column, functions as F
+from . import plan as P
+from . import expression as E
+
+
+@singledispatch
+def get_spark_expression(expr: tp.Any) -> Column:
+    raise NotImplementedError(f'Cannot translate expression of type {type(expr)}')
+
+
+@get_spark_expression.register
+def _(expr: E.Coalesce) -> Column:
+    return F.coalesce(*map(get_spark_expression, expr.children))
+
+
+@get_spark_expression.register
+def _(expr: E.Alias) -> Column:
+    return get_spark_expression(expr.arg).alias(expr.name)
+
+@get_spark_expression.register
+def _(expr: E.AttributeReference) -> Column:
+    return F.col(f'`{".".join(expr.qualifier)}`.`{expr.name}`')
+
+
+@singledispatch
+def get_plan_dataframe(plan: P.Plan, spark: SparkSession) -> DataFrame:
+    """convert a plan into a Dataframe."""
+    raise NotImplementedError(f'Cannot translate plan of type {type(plan)}')
+
+
+@get_plan_dataframe.register
+def _(plan: P.Project, spark: SparkSession) -> DataFrame:
+    df = get_plan_dataframe(plan.child, spark)
+    return df.select(*[
+        get_spark_expression(expr)
+        for expr in plan.columns
+    ])
+
+
+@get_plan_dataframe.register
+def _(plan: P.Table, spark: SparkSession) -> DataFrame:
+    return spark.read.table(plan.qualified_name)
+
+
+@get_plan_dataframe.register
+def _(plan: P.Alias, spark: SparkSession) -> DataFrame:
+    df = get_plan_dataframe(plan.child, spark)
+    return df.alias('.'.join([*plan.qualifier, plan.alias]))

pyspark_inspect-0.1.0/pyspark_inspect/expression.py

@@ -0,0 +1,201 @@
+"""Representation of Column expressions."""
+
+
+import typing as tp
+import operator as op
+from dataclasses import dataclass
+from pyspark.sql import Column, functions as F
+
+
+@dataclass(frozen=True)
+class Expression:
+    _func: tp.ClassVar[tp.Optional[tp.Callable[..., Column]]] = None
+
+    @property
+    def children(self) -> list['Expression']:
+        raise NotImplementedError()
+
+    def to_column(self) -> Column:
+        raise NotImplementedError()
+
+    def _get_func(self) -> tp.Callable[..., Column]:
+        if self._func is None:
+            raise NotImplementedError()
+        else:
+            return self._func
+
+
+@dataclass(frozen=True)
+class Leaf(Expression):
+    @property
+    def children(self):
+        return []
+
+
+@dataclass(frozen=True)
+class UnaryOperator(Expression):
+    arg: Expression
+
+    @property
+    def children(self):
+        return [self.arg]
+
+    def to_column(self) -> Column:
+        return self._get_func()(self.arg.to_column())
+
+
+@dataclass(frozen=True)
+class BinaryOperator(Expression):
+    left: Expression
+    right: Expression
+
+    @property
+    def children(self):
+        return [self.left, self.right]
+
+    def to_column(self) -> Column:
+        return self._get_func()(self.left.to_column(), self.right.to_column())
+
+
+@dataclass(frozen=True)
+class Variadic(Expression):
+    args: tuple[Expression, ...]
+
+    @property
+    def children(self):
+        return list(self.args)
+
+    def to_column(self):
+        return self._get_func()(*[arg.to_column() for arg in self.args])
+
+
+@dataclass(frozen=True)
+class UnknownExpression(Variadic):
+    class_name: str
+
+    def to_column(self):
+        raise NotImplementedError(f'Cannot translate unknown expression of type {self.class_name}')
+
+
+@dataclass(frozen=True)
+class Literal(Leaf):
+    value: tp.Hashable
+    data_type: str | None = None
+
+    def to_column(self) -> Column:
+        return F.lit(self.value)
+
+
+@dataclass(frozen=True)
+class AttributeReference(Leaf):
+    name: str
+    qualifier: tuple[str, ...]
+    attribute_id: str
+
+    @property
+    def qualified_name(self) -> str:
+        name = f'`{self.name}`'
+        if self.qualifier:
+            qualifier = '.'.join(self.qualifier)
+            return f'`{qualifier}`.{name}'
+        else:
+            return name
+
+    def to_column(self) -> Column:
+        return F.col(self.qualified_name)
+
+
+@dataclass(frozen=True)
+class Alias(UnaryOperator):
+    name: str
+    qualifier: tuple[str, ...]
+
+    @property
+    def qualified_name(self) -> str:
+        name = f'`{self.name}`'
+        if self.qualifier:
+            qualifier = '.'.join(self.qualifier)
+            return f'`{qualifier}`.{name}'
+        else:
+            return name
+
+    def to_column(self) -> Column:
+        return self.arg.to_column().alias(self.qualified_name)
+
+
+@dataclass(frozen=True)
+class Equals(BinaryOperator):
+    _func = op.eq
+
+
+@dataclass(frozen=True)
+class EqNullSafe(BinaryOperator):
+    _func = Column.eqNullSafe
+
+
+@dataclass(frozen=True)
+class GreaterThanOrEqual(BinaryOperator):
+    _func = op.ge
+
+
+@dataclass(frozen=True)
+class GreaterThan(BinaryOperator):
+    _func = op.gt
+
+
+@dataclass(frozen=True)
+class LessThanOrEqual(BinaryOperator):
+    _func = op.le
+
+
+@dataclass(frozen=True)
+class LessThan(BinaryOperator):
+    _func = op.lt
+
+
+@dataclass(frozen=True)
+class And(BinaryOperator):
+    _func = op.and_
+
+
+@dataclass(frozen=True)
+class Or(BinaryOperator):
+    _func = op.or_
+
+
+@dataclass(frozen=True)
+class Not(UnaryOperator):
+    _func = op.neg
+
+
+@dataclass(frozen=True)
+class IsNull(UnaryOperator):
+    _func = F.isnull
+
+
+@dataclass(frozen=True)
+class IsNotNull(UnaryOperator):
+    _func = F.isnotnull
+
+
+@dataclass(frozen=True)
+class Coalesce(Variadic):
+    _func = F.coalesce
+
+
+@dataclass(frozen=True)
+class Concat(Variadic):
+    _func = F.concat
+
+
+@dataclass(frozen=True)
+class When(Variadic):
+    def to_column(self):
+        args = [a.to_column() for a in self.args]
+        when, then = args[:2]
+        c = F.when(when, then)
+        for when, then in zip(args[2::2], args[3::2]):
+            c = c.when(when, then)
+        if len(args) % 2 == 1:
+            c = c.otherwise(args[-1])
+        return c

pyspark_inspect-0.1.0/pyspark_inspect/parser/__init__.py

@@ -0,0 +1 @@
+"""Low-level parsing of Spark plans into Python data structures."""

pyspark_inspect-0.1.0/pyspark_inspect/parser/catalyst.py

@@ -0,0 +1,297 @@
+"""Parse Catalyst JSON plan into Python data structures."""
+
+import typing as tp
+import json
+from dataclasses import dataclass
+from pyspark.errors import PySparkAttributeError
+from pyspark.sql import DataFrame, Column, types as T
+from pyspark_inspect import plan as P, expression as E
+
+
+@dataclass
+class CatalystPlan:
+    children: list['CatalystPlan']
+    data: dict[str, tp.Any]
+
+    @property
+    def class_name(self) -> str:
+        return self.data['class']
+
+
+def add_node(stack, plan: CatalystPlan, num_children: int):
+    if num_children > 0 or not stack:
+        stack.append((plan, num_children))
+    else:
+        top, num_children = stack[-1]
+        top.children.append(plan)
+        if len(top.children) == num_children and len(stack) > 0:
+            stack.pop()
+            add_node(stack, top, 0)
+
+
+def get_dataframe_plan(df: DataFrame) -> str:
+    """Get Catalyst json representation of a DataFrame query plan."""
+    try:
+        jdf = df._jdf
+    except PySparkAttributeError:
+        raise NotImplementedError('Cannot retrieve query plan from JVM. Are you using Spark Connect?')
+    return jdf.logicalPlan().toJSON()
+
+
+def get_column_expression(c: Column) -> str:
+    """Get Catalyst json representation of a column expression."""
+    try:
+        jc = c._jc
+    except PySparkAttributeError:
+        raise NotImplementedError('Cannot retrieve query plan from JVM. Are you using Spark Connect?')
+    return jc.expr().toJSON()
+
+
+def load_catalyst_plan(plan: str | list) -> CatalystPlan:
+    """Load Catalyst plan string or parsed JSON into raw parse tree."""
+    if isinstance(plan, str):
+        plan = tp.cast(list, json.loads(plan))
+    stack: list[tuple[CatalystPlan, int]] = []
+    for item in plan:
+        num_children = item['num-children']
+        node = CatalystPlan(
+            children=[],
+            data=item
+        )
+        add_node(stack, node, num_children)
+    assert len(stack) == 1
+    return stack[0][0]
+
+
+def _parse_logical_relation(plan: CatalystPlan, children: list[P.Plan]) -> P.Table:
+    tid = plan.data['catalogTable']['identifier']
+    qualified_name = '.'.join([tid['catalog'], tid['database'], tid['table']])
+    columns = [parse_expression(load_catalyst_plan(c)) for c in plan.data['output']]
+    return P.Table(qualified_name=qualified_name, columns=tuple(columns))
+
+
+def _parse_subquery_alias(plan: CatalystPlan, children: list[P.Plan]) -> P.Alias:
+    return P.Alias(
+        child=children[0],
+        alias=plan.data['identifier']['name'],
+        qualifier=parse_array_string(plan.data['identifier'].get('qualifier', ''))
+    )
+
+def _parse_project(plan: CatalystPlan, children: list[P.Plan]) -> P.Project:
+    return P.Project(
+        child=children[0],
+        columns=tuple([parse_expression(load_catalyst_plan(p)) for p in plan.data['projectList']])
+    )
+
+
+def _parse_window(plan: CatalystPlan, children: list[P.Plan]) -> P.Project:
+    child = tp.cast(P.Project, children[0])
+    columns = tuple([
+        *child.columns,
+        *[parse_expression(load_catalyst_plan(p)) for p in plan.data['windowExpressions']],
+    ])
+    return P.Project(child=child.child, columns=columns)
+
+
+def _parse_join(plan: CatalystPlan, children: list[P.Plan]) -> P.Join:
+    how = JOIN_TYPE[plan.data['joinType']['object']]
+    on = parse_expression(load_catalyst_plan(plan.data['condition']))
+    return P.Join(left=children[0], right=children[1], on=on, how=how)
+
+
+def _parse_one_row_relation(plan: CatalystPlan, children: list[P.Plan]) -> P.OneRowRelation:
+    return P.OneRowRelation()
+
+
+def _parse_cte_relation_def(plan: CatalystPlan, children: list[P.Plan]) -> P.Plan:
+    return P.CTEDef(
+        cte_id=plan.data['id'],
+        child=children[0]
+    )
+
+
+def _parse_cte_relation_ref(plan: CatalystPlan, children: list[P.Plan]) -> P.Plan:
+    return P.CTERef(cte_id=plan.data['cteId'])
+
+
+def _parse_with_cte(plan: CatalystPlan, children: list[P.Plan]) -> P.Plan:
+    return P.WithCTE(
+        ctes=tuple(children[:-1]),
+        main=children[-1]
+    )
+
+
+def _parse_union(plan: CatalystPlan, children: list[P.Plan]) -> P.Plan:
+    return P.Union(
+        left=children[0],
+        right=children[1],
+        by_name=plan.data['byName'],
+        allow_missing_columns=plan.data['allowMissingCol']
+    )
+
+
+def _parse_except(plan: CatalystPlan, children: list[P.Plan]) -> P.Plan:
+    return P.Except(
+        left=children[0],
+        right=children[1],
+        is_all=plan.data['isAll'],
+    )
+
+
+def _parse_intersect(plan: CatalystPlan, children: list[P.Plan]) -> P.Plan:
+    return P.Intersect(
+        left=children[0],
+        right=children[1],
+        is_all=plan.data['isAll'],
+    )
+
+def _parse_local_limit(plan: CatalystPlan, children: list[P.Plan]) -> P.Plan:
+    limit = int(plan.data['limitExpr'][0]['value'])
+    return P.LocalLimit(children[0], limit=limit)
+
+
+def _parse_global_limit(plan: CatalystPlan, children: list[P.Plan]) -> P.Plan:
+    limit = int(plan.data['limitExpr'][0]['value'])
+    return P.GlobalLimit(children[0], limit=limit)
+
+
+def _parse_aggregate(plan: CatalystPlan, children: list[P.Plan]) -> P.Plan:
+    return P.Aggregate(
+        child=children[0],
+        grouping_expressions=tuple([parse_expression(load_catalyst_plan(p)) for p in plan.data['groupingExpressions']]),
+        columns=tuple([parse_expression(load_catalyst_plan(p)) for p in plan.data['aggregateExpressions']])
+    )
+
+
+def _parse_distinct(plan: CatalystPlan, children: list[P.Plan]) -> P.Plan:
+    return P.Distinct(children[0])
+
+
+def _skip_unary(plan: CatalystPlan, children: list[P.Plan]) -> P.Plan:
+    """Skip a unary plan node, e.g. repartition"""
+    return children[0]
+
+
+LOGICAL_PLAN = 'org.apache.spark.sql.catalyst.plans.logical'
+
+JOIN_TYPE = {
+    'org.apache.spark.sql.catalyst.plans.Inner$': 'inner',
+    'org.apache.spark.sql.catalyst.plans.LeftOuter$': 'left',
+    'org.apache.spark.sql.catalyst.plans.RightOuter$': 'right',
+    'org.apache.spark.sql.catalyst.plans.FullOuter$': 'outer',
+    'org.apache.spark.sql.catalyst.plans.LeftSemi$': 'left-semi',
+    'org.apache.spark.sql.catalyst.plans.LeftAnti$': 'left-anti',
+}
+
+PLAN_PARSER: dict[str, tp.Callable[[CatalystPlan, list[P.Plan]], P.Plan]] = {
+    'org.apache.spark.sql.execution.datasources.LogicalRelation': _parse_logical_relation,
+    f'{LOGICAL_PLAN}.Aggregate': _parse_aggregate,
+    f'{LOGICAL_PLAN}.CTERelationDef': _parse_cte_relation_def,
+    f'{LOGICAL_PLAN}.CTERelationRef': _parse_cte_relation_ref,
+    f'{LOGICAL_PLAN}.Distinct': _parse_distinct,
+    f'{LOGICAL_PLAN}.Except': _parse_except,
+    f'{LOGICAL_PLAN}.GlobalLimit': _parse_global_limit,
+    f'{LOGICAL_PLAN}.Intersect': _parse_intersect,
+    f'{LOGICAL_PLAN}.Join': _parse_join,
+    f'{LOGICAL_PLAN}.LocalLimit': _parse_local_limit,
+    f'{LOGICAL_PLAN}.SubqueryAlias': _parse_subquery_alias,
+    f'{LOGICAL_PLAN}.OneRowRelation': _parse_one_row_relation,
+    f'{LOGICAL_PLAN}.Project': _parse_project,
+    f'{LOGICAL_PLAN}.Union': _parse_union,
+    f'{LOGICAL_PLAN}.Window': _parse_window,
+    f'{LOGICAL_PLAN}.WithCTE': _parse_with_cte,
+    # skipped
+    f'{LOGICAL_PLAN}.Repartition': _skip_unary,
+    f'{LOGICAL_PLAN}.RepartitionByExpression': _skip_unary,
+    f'{LOGICAL_PLAN}.Sort': _skip_unary,
+}
+
+
+def parse_plan(plan: CatalystPlan):
+    children = [parse_plan(c) for c in plan.children]
+    parser = PLAN_PARSER.get(plan.class_name)
+    if not parser:
+        raise ValueError(f'Unsupported plan: {plan.class_name}')
+    return parser(plan, children)
+
+
+def _parse_expression_alias(plan: CatalystPlan, children: list[E.Expression]) -> E.Expression:
+    return E.Alias(
+        arg=children[0],
+        name=plan.data['name'],
+        qualifier=parse_array_string(plan.data['qualifier'])
+    )
+
+
+def _parse_attribute_reference(plan: CatalystPlan, children: list[E.Expression]) -> E.Expression:
+    # TODO consider to parse data type: T._parse_datatype_json_value(plan.data['dataType'])
+    return E.AttributeReference(
+        name=plan.data['name'],
+        qualifier=parse_array_string(plan.data['qualifier']),
+        attribute_id=plan.data['exprId']['id']
+    )
+
+
+def _parse_literal(plan: CatalystPlan, children: list[E.Expression]) -> E.Expression:
+    # TODO: Should cast values
+    raw_value = plan.data['value']
+    raw_data_type = plan.data['dataType']
+    return E.Literal(value=raw_value, data_type=raw_data_type)
+
+
+def _parse_variadic(cls: type) -> tp.Callable:
+    """Generic parser vor variadic expressions without additional arguments."""
+    def parser(plan: CatalystPlan, children: list[E.Expression]):
+        return cls(tuple(children))
+    return parser
+
+
+def _parse_binary(cls: type) -> tp.Callable:
+    """Generic parser for binary expressions without additional arguments."""
+    def parser(plan: CatalystPlan, children: list[E.Expression]):
+        return cls(children[0], children[1])
+    return parser
+
+
+def _parse_unary(cls: type) -> tp.Callable:
+    """Generic parser for unary expressions without additional arguments."""
+    def parser(plan: CatalystPlan, children: list[E.Expression]):
+        return cls(children[0])
+    return parser
+
+
+EXPRESSION_PARSER: dict[str, tp.Callable[[CatalystPlan, list[E.Expression]], E.Expression]] = {
+    'org.apache.spark.sql.catalyst.expressions.Alias': _parse_expression_alias,
+    'org.apache.spark.sql.catalyst.expressions.Coalesce': _parse_variadic(E.Coalesce),
+    'org.apache.spark.sql.catalyst.expressions.AttributeReference': _parse_attribute_reference,
+    'org.apache.spark.sql.catalyst.expressions.And': _parse_binary(E.And),
+    'org.apache.spark.sql.catalyst.expressions.Or': _parse_binary(E.Or),
+    'org.apache.spark.sql.catalyst.expressions.EqualTo': _parse_binary(E.Equals),
+    'org.apache.spark.sql.catalyst.expressions.EqualNullSafe': _parse_binary(E.EqNullSafe),
+    'org.apache.spark.sql.catalyst.expressions.IsNull': _parse_unary(E.IsNull),
+    'org.apache.spark.sql.catalyst.expressions.IsNotNull': _parse_unary(E.IsNotNull),
+    'org.apache.spark.sql.catalyst.expressions.Literal': _parse_literal,
+}
+
+
+def _parse_unknown_expression(plan: CatalystPlan, children: list[E.Expression]) -> E.Expression:
+    """Fallback parser to capture unknown expressions with their children."""
+    return E.UnknownExpression(args=tuple(children), class_name=plan.class_name)
+
+
+def parse_expression(plan: CatalystPlan) -> E.Expression:
+    """Parse a column expression."""
+    children = [parse_expression(c) for c in plan.children]
+    parser = EXPRESSION_PARSER.get(plan.class_name, _parse_unknown_expression)
+    return parser(plan, children)
+
+
+def parse_array_string(s: str) -> tuple[str, ...]:
+    """split comma separated array.
+
+    Note that not all values can be parsed reliably since values are not quoted.
+    """
+    stripped = s[1:-1]
+    if not stripped:
+        return ()
+    return tuple(stripped.split(', '))

pyspark_inspect-0.1.0/pyspark_inspect/plan.py

@@ -0,0 +1,137 @@
+"""Classes representing steps in a query plan."""
+
+import typing as tp
+from dataclasses import dataclass
+from .expression import Expression
+from pyspark.sql import types as T
+
+
+@dataclass(frozen=True)
+class Plan:
+    @property
+    def children(self) -> list['Plan']:
+        return []
+
+
+@dataclass(frozen=True)
+class UnaryPlan(Plan):
+    child: Plan
+
+    @property
+    def children(self) -> list[Plan]:
+        return [self.child]
+
+
+@dataclass(frozen=True)
+class BinaryPlan(Plan):
+    left: Plan
+    right: Plan
+
+    @property
+    def children(self) -> list[Plan]:
+        return [self.left, self.right]
+
+
+@dataclass(frozen=True)
+class Alias(UnaryPlan):
+    alias: str
+    qualifier: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class Project(UnaryPlan):
+    """A projection is a select without row generators like explode."""
+    columns: tuple[Expression, ...]
+
+
+@dataclass(frozen=True)
+class Filter(UnaryPlan):
+    condition: Expression
+
+
+@dataclass(frozen=True)
+class Join(BinaryPlan):
+    """A left/right/inner/outer/cross join."""
+    on: Expression
+    how: tp.Literal['left', 'right', 'inner', 'outer', 'left-anti', 'left-semi']
+
+
+@dataclass(frozen=True)
+class Table(Plan):
+    """A catalog table."""
+    qualified_name: str
+    columns: tuple[Expression, ...]
+
+
+@dataclass(frozen=True)
+class RDD(Plan):
+    """Output of SparkSesssion.createDataFrame.
+
+    Note that we cannot reconstruct data, would need to feed from
+    user.
+    """
+    schema: T.StructType
+
+
+@dataclass(frozen=True)
+class OneRowRelation(Plan):
+    """Dummy input for SQL select without from clause."""
+    pass
+
+
+@dataclass(frozen=True)
+class WithCTE(Plan):
+    """Statement with with-block."""
+    ctes: tuple[Plan, ...]
+    main: Plan
+
+    @property
+    def children(self):
+        return [*self.ctes, self.main]
+
+
+@dataclass(frozen=True)
+class CTEDef(UnaryPlan):
+    cte_id: int
+
+
+@dataclass(frozen=True)
+class CTERef(Plan):
+    cte_id: int
+
+
+@dataclass(frozen=True)
+class Union(BinaryPlan):
+    by_name: bool
+    allow_missing_columns: bool
+
+
+@dataclass(frozen=True)
+class Except(BinaryPlan):
+    is_all: bool
+
+
+@dataclass(frozen=True)
+class Intersect(BinaryPlan):
+    is_all: bool
+
+
+@dataclass(frozen=True)
+class Aggregate(UnaryPlan):
+    grouping_expressions: tuple[Expression, ...]
+    columns: tuple[Expression, ...]
+
+
+@dataclass(frozen=True)
+class Distinct(UnaryPlan):
+    pass
+
+
+@dataclass(frozen=True)
+class GlobalLimit(UnaryPlan):
+    limit: int
+
+
+@dataclass(frozen=True)
+class LocalLimit(UnaryPlan):
+    limit: int