Python Tasks¶

The Context Model lets you run Python Tasks: Python code that reads input tables, transforms or generates data, and writes results to one or more output tables. pycelonis.ml_workbench.context_model provides @task, TaskConfiguration, and dataset/Spark helpers for interactive ML Workbench notebooks and scheduled executions.

The framework takes care of:

• Creating and cleaning up a SparkSession connected to the Storium Spark Proxy
• Providing Input wrappers in the configured dataset order
• Mapping returned DataFrames to output datasets in their configured order
• Previewing results in dry_run mode instead of writing them

Prerequisites¶

• pyspark needs to be installed, e.g. via pip install pycelonis[mlwb].
• You need a Context Model that contains the tables you want to read, or a Data Pool containing sample input tables.
• Interactive testing is intended to run in ML Workbench, where the Spark connection is automatically set up using your credentials.

Tutorial¶

Writing a Python Task takes three steps:

1. Configure TaskConfiguration for interactive notebook testing.
2. Decorate a synchronous function with @task (or @task()) and return its output.
3. Call the function without arguments.

1. Configure the task¶

TaskConfiguration.set_defaults(...) sets the schema, ordered input/output table lists, and dry-run behavior. schema is required; either table list may be empty. dry_run defaults to True.

For testing against a Context Model, use its managed schema ID. You can find this ID in your Context Model under Settings -> Data Schema.

When scheduled, the Context Model supplies schema, input_tables, output_tables, and dry_run through execution parameters. These values are authoritative, and set_defaults becomes a no-op. Leave this cell in the published notebook; re-running it only replaces interactive defaults.

Scheduled executions receive input and output datasets in the order configured in the Context Model UI. New datasets are appended to their list. Reordering datasets changes their positions; removing a dataset shifts the positions of those after it. Keep the interactive table lists and the code's input indices and return order aligned with the UI. After changing dataset order, check the results with dry_run=True before enabling writes.

Order, duplicate entries, full schema qualification, and execution-slot suffixes are preserved. Path segments containing non-ASCII word characters or punctuation are escaped for Spark; do not add backticks to configured paths.

In [ ]:

from pycelonis.ml_workbench.context_model import TaskConfiguration

# Interactive defaults; the platform supplies authoritative execution parameters.
TaskConfiguration.set_defaults(
    schema="<schema_id>",  # Context Model branch schema ID or Data Pool ID
    input_tables=["input_table", "other_schema.imported_table"],
    output_tables=["output_table"],
    dry_run=True,
)

from pycelonis.ml_workbench.context_model import TaskConfiguration # Interactive defaults; the platform supplies authoritative execution parameters. TaskConfiguration.set_defaults( schema="<schema_id>", # Context Model branch schema ID or Data Pool ID input_tables=["input_table", "other_schema.imported_table"], output_tables=["output_table"], dry_run=True, )

2. Define the task¶

Declare an inputs parameter to receive a list of Input wrappers in configured order. inputs[0] accesses the first table in TaskConfiguration.input_tables, and inputs[1] accesses the second. Input wrappers provide .dataframe(), .count(), .schema(), .columns(), and .show().

Return a Spark or pandas DataFrame for one output, or a tuple of these for multiple outputs. The return count must match output_tables. Spark frames must belong to the execution's session; pandas frames are converted through that session using spark.createDataFrame(...). The complete result is checked and all pandas frames are converted before any automatic preview or write.

In [ ]:

from pycelonis.ml_workbench.context_model import task


@task
def process(inputs):
    df = inputs[0].dataframe()
    print(f"Read {inputs[0].count()} rows from '{inputs[0].table_name}'")
    print(f"Read {inputs[1].count()} rows from '{inputs[1].table_name}'")

    # Apply your transformation logic here.
    return df

from pycelonis.ml_workbench.context_model import task @task def process(inputs): df = inputs[0].dataframe() print(f"Read {inputs[0].count()} rows from '{inputs[0].table_name}'") print(f"Read {inputs[1].count()} rows from '{inputs[1].table_name}'") # Apply your transformation logic here. return df

3. Run the task¶

Calling the decorated function creates a SparkSession, supplies the inputs, runs your code, and sends the returned DataFrames to the configured output tables. The call returns None after output handling. The task closes the sessions it creates when it finishes, including if an error occurs.

With dry_run=True, each returned frame is previewed with df.show(). With dry_run=False, the target table's data and schema are replaced using overwrite mode and overwriteSchema=true. Multi-output writes are sequential, not transactional: a later Spark failure does not roll back earlier writes.

In [ ]:

process()

process()

Other supported patterns¶

These examples define alternative tasks; configure the corresponding input/output lists before invoking them.

First-input shorthand. input_ always aliases inputs[0], even when more inputs are configured. Both parameters may appear in the same function; declaring input_ requires at least one input.

@task()
def first_input(input_):
    return input_.dataframe()

Multiple outputs. The first returned DataFrame goes to the first configured output table, the second to the second table, and so on. Matching counts and valid DataFrame types do not detect swapped results: returning them in the wrong order can overwrite the wrong table. Return a tuple, not a list or dictionary:

TaskConfiguration.set_defaults(
    schema="<schema_id>",
    input_tables=["orders", "customers"],
    output_tables=["orders_clean", "customers_clean"],
    dry_run=True,
)

@task
def split_outputs(inputs):
    orders_df = inputs[0].dataframe()
    customers_df = inputs[1].dataframe()
    return orders_df, customers_df  # orders_clean, then customers_clean

Pandas outputs. Configure one output table for this task. Pandas results and direct Output(...).write_dataframe(pandas_df) calls use Spark's schema inference and the same overwrite/schema-replacement behavior. Dry runs also convert to Spark and call show(). Tuples may mix pandas and Spark frames. Empty frames or unsupported values may fail Spark conversion; inference errors propagate without automatic output writes.

Use pandas outputs for small to medium results. With Spark Connect, conversion transfers the pandas data from the notebook to the server over gRPC, which can be slow or encounter configured transport limits for large results. For large results, keep processing in Spark and return a Spark DataFrame.

import pandas as pd

@task
def pandas_task():
    return pd.DataFrame({"customer_id": [1, 2]})

Direct table access. A no-parameter task can use Input and Output by name if explicit table references are preferred over positions. Returning None skips automatic output routing, even if output tables are configured; manual-only tasks may use empty lists.

from pycelonis.ml_workbench.context_model import Input, Output

@task
def manual_task():
    df = Input("input_table").dataframe()
    output = Output("output_table")
    print(f"Writing to {output.resolved_table_name}")
    output.write_dataframe(df)

Inside a task, Output("output_table") selects the corresponding configured output dataset. When a scheduled execution uses an output slot such as output_table@<execution-id>, the wrapper resolves to that destination automatically. Its resolved_table_name contains the destination with Spark escaping, which you can inspect before writing. The table_name attribute and representation retain the path you supplied. You can also supply a full slot path directly; Output escapes path segments containing punctuation or non-ASCII characters.

Quoted Spark identifiers keep their meaning: a dot inside a quoted table name is not a schema separator, and doubled backticks preserve a literal backtick.

An exact configured path takes precedence. A schema-qualified logical name restricts matching to that path; an unqualified name must match a single destination. If multiple destinations match, construction raises ValueError; use a qualified name or the full slot path to disambiguate. Repeated entries for the same destination are not ambiguous. A name that does not match a configured output is escaped and used directly. When the task has configured outputs, this fallback logs a warning at construction so typos are visible; manual access to other tables is still supported. Only canonical UUID suffixes identify execution slots; names ending in @union, @latest, or other suffixes are treated literally.

Output captures its destination when constructed and retains it across nested tasks, session resets, and configuration changes. An Output constructed outside a task is escaped but is not remapped when later used inside one. When an explicit Spark session is supplied, configured-destination resolution applies only if it is the task's current active session. Other explicit sessions receive path escaping without the task's output mapping.

Inside a task, Input and Output use the task's Spark session automatically. When constructed without dry_run, Output captures the task's dry-run setting. Pass dry_run=True explicitly for previews, or dry_run=False to write. Outside a task, Output defaults to previews; using the same wrapper later inside a task does not change its captured setting.

You may configure the same table as an input and output, but Spark can reject overwriting a table that the returned DataFrame reads from. Such errors propagate to the caller.

Advanced: Spark sessions and concurrent execution¶

Spark session access¶

Inside a task, get_spark_session() returns the task's Spark session without requiring a schema or credentials. Reuse is the default whenever an active session is available, including during notebook cell re-runs with matching configuration. Conflicting explicit connection values raise ValueError. If no session is available, creation requires a schema and connection settings. Pass reuse=False to deliberately create a fresh session.

For interactive use outside a task, create a session first and reset it when finished:

from pycelonis.ml_workbench.context_model import Input, get_spark_session, reset_spark_session

spark = get_spark_session(schema="<schema_id>")
try:
    df = Input("input_table").dataframe()
    df.show()
finally:
    reset_spark_session(stop=True)

Input and Output can also receive an explicit spark argument, which takes precedence over the task's session. Spark DataFrames must belong to the session used by Output; a frame from another session is rejected before previewing or writing.

Cleanup and nesting¶

The default reset_spark_session(stop=False) clears the active reference without stopping the session. Every task invocation starts with a fresh session, including repeated and nested calls; subsequent get_spark_session() calls inside it reuse that session. A task stops all sessions it creates when it finishes, including sessions replaced or reset during execution. If a reset fails to stop a session, task cleanup retries. Sessions supplied explicitly by the caller are not stopped automatically unless the task created them.

Outside a task, matching calls reuse the available session by default. If you request another session with reuse=False, keep references to both sessions and stop them explicitly when finished. Replacing the active reference does not stop the previous session, which may still be needed by existing DataFrames.

Each task has its own session scope: its active session and dry-run defaults are isolated for that execution. A nested task restores the outer task's active session and defaults when it finishes, including on errors. Reuse performs no network health check; call reset_spark_session after manually stopping a session.

Caching and concurrency¶

Each Input caches its latest session/DataFrame pair in the current Python execution context. Repeated access with the same session reuses the cached frame. Child contexts inherit existing entries, but later updates do not replace their parent's or siblings' entries. Cache entries are released when their Input is garbage-collected.

New threads need an explicit session or must invoke their own task. Async work that inherits a task's context must finish before that task exits, because cleanup stops the task's sessions. Functions decorated with @task must be synchronous and cannot be generators.

From notebook to Python Task¶

Once you are happy with the result:

1. Publish this ML Workbench.
2. Create a Transformation resource in the ML Workbench UI.
3. Select that ML Workbench Transformation when creating a Python Task in the Context Model.
4. Configure input and output datasets in the order expected by inputs[i] and the returned DataFrames. Match this order in your interactive TaskConfiguration defaults.

The Context Model supplies execution parameters, including dry_run=False, so interactive defaults are ignored.

Conclusion¶

You have learned to configure TaskConfiguration, access ordered inputs with @task, return one or more output DataFrames, and test safely with dry-run previews before publishing.

Migrating from @transformation¶

This section is for users of the legacy API. pycelonis.ml_workbench.transformation remains supported but deprecated. The @transformation decorator will be removed in a future release; no removal version has been announced. Use @task for new code.

• Replace TransformationConfig with TaskConfiguration. The two configurations are independent: setting one does not configure functions using the other.
• Replace @transformation(...) with @task or @task(), and rewrite name-matched function parameters as ordered inputs or input_.
• Return output DataFrame(s) instead of receiving output parameters. For manual writes, import Input and Output from pycelonis.ml_workbench.context_model.

Legacy helper imports still resolve to the canonical Input/Output classes. Each legacy module emits a deprecation warning on its first import using the shared warning utilities. The decorator and configuration modules also warn now that task and TaskConfiguration are available; importing from the canonical package remains warning-free.

Manual Outputs constructed inside @transformation also resolve logical names against its captured output table configuration, including execution slots. The same construction-time capture and explicit-session rules apply.

Breaking change for manual outputs: This also applies to legacy helper imports. An Output constructed inside @transformation captures TransformationConfig.dry_run when dry_run is omitted, even if an explicit Spark session is supplied. Previously, omission always meant preview-only. Pass dry_run=True explicitly to preserve previews when the transformation is configured to write. The goal is for manual writes to follow the same notebook-preview/scheduled-write setting as injected outputs. Outputs constructed outside a decorated execution still default to previews; using them later inside a task or transformation does not change their captured setting.