Run Spark workloads

BullSequana AI runs Apache Spark from JupyterLab dev environments and Kubeflow Pipelines. The default JupyterLab image, jupyterlab:python3.12, supports Spark. In a Spark session you can use DataFrames, Spark SQL, Iceberg catalog tables, and pandas UDFs. The platform creates the Spark resources for each session. You stop the session when your work is done.

Before you start

• A JupyterLab dev environment created from a Spark-capable image: jupyterlab:python3.12 (the default) or Databricks Runtime jupyterlab:17.3-lts. Use Databricks Runtime 17.3 LTS when your code needs its package environment. See Choose a notebook image.
• For catalog tables: a warehouse name and object storage credentials from your administrator.
• For pipelines: a BSQAI API key. See Run Spark in Kubeflow Pipelines.

Start a Spark session in a notebook

Spark-capable images include the bsqai_spark helper. Call get_spark_session() to start a session:

from bsqai_spark import SparkSessionConfig, get_spark_session

session = get_spark_session(
    SparkSessionConfig(
        mode="job",
        app_name="orders-summary",
        worker_replicas=2,
        worker_cores=1,
        worker_memory="2g",
    )
)

The call returns when the session is ready. By default, the helper waits up to 5 minutes for the session to start.

Work with DataFrames and SQL

Reuse the same session across notebook cells for related work. Each new call to get_spark_session() starts another set of Spark resources.

from pyspark.sql import functions as F

events = session.createDataFrame(
    [(1, "open"), (2, "close"), (3, "open")],
    ["id", "event"],
)

events.groupBy("event").count().orderBy("event").show()

events.createOrReplaceTempView("events")
session.sql("SELECT event, COUNT(*) AS n FROM events GROUP BY event").show()

When you finish, stop the session:

session.stop()

Read and write catalog tables

Set warehouse to connect the session to the Iceberg catalog. The catalog is registered as lakekeeper. It is not the default Spark catalog, so use fully qualified table names: lakekeeper.<namespace>.<table>.

Provide the warehouse name and storage credentials through configured secrets or environment variables, following your organization's approved secret path. Do not put credentials in notebook source. The helper reads these environment variables:

This example reads the warehouse name from a WAREHOUSE environment variable, writes a table, and reads it back:

import os

from bsqai_spark import SparkSessionConfig, get_spark_session

session = get_spark_session(
    SparkSessionConfig(
        mode="job",
        app_name="catalog-example",
        warehouse=os.environ["WAREHOUSE"],
    )
)

try:
    session.sql("CREATE NAMESPACE IF NOT EXISTS lakekeeper.analytics")

    table = "lakekeeper.analytics.labels"
    labels = session.createDataFrame([(1, "alpha"), (2, "beta")], ["id", "label"])
    labels.writeTo(table).using("iceberg").createOrReplace()

    session.table(table).show()
finally:
    session.stop()

The try/finally block stops the session even if a step fails.

Use pandas UDFs

In jupyterlab:python3.12, Spark sessions support pandas UDFs and mapInPandas:

import pandas as pd
from pyspark.sql.functions import pandas_udf

readings = session.createDataFrame(
    [(1, 50.0), (2, 86.0), (3, 104.0)],
    ["sensor", "fahrenheit"],
)

@pandas_udf("double")
def to_celsius(fahrenheit: pd.Series) -> pd.Series:
    return (fahrenheit - 32) * 5 / 9

readings.select("sensor", to_celsius("fahrenheit").alias("celsius")).show()

def flag_hot(batches):
    for batch in batches:
        batch["hot"] = batch["fahrenheit"] > 80
        yield batch

readings.mapInPandas(
    flag_hot,
    schema="sensor long, fahrenheit double, hot boolean",
).show()

Session lifecycle and time limit

• get_spark_session() asks the platform to create the Spark resources for your session. If startup fails, the helper releases the resources it created.
• The session belongs to the notebook kernel or pipeline step that started it. Reuse it for related work instead of starting a new session per cell.
• Call session.stop() when you finish. This releases the platform resources. Calling session.sparkContext.stop() stops Spark locally but does not release the platform resources.

Every session has a time limit. The default is 1 hour. The limit counts from when the session started, regardless of activity. When it is reached, the platform removes the session's resources, even if work is still running. The time limit also releases resources when you cannot call session.stop(), for example after a kernel crash.

Set the limit with the idle_timeout option. Use a number followed by s, m, h, or d, for example 30m, 4h, or 2d. The maximum is 365d. Choose a limit that covers your expected work:

session = get_spark_session(
    SparkSessionConfig(mode="job", app_name="nightly-aggregation", idle_timeout="4h")
)

Your administrator sets per-user limits on the number of Spark sessions and on total worker cores. Sessions in notebooks and pipeline steps count toward the same limits.

Run Spark in Kubeflow Pipelines

A pipeline step can start a Spark session with the same helper. Each step that uses Spark creates its own session and stops it before it ends. Steps pass results to each other through standard pipeline outputs and artifacts. Pipeline steps can also persist results as catalog tables. The example below passes results as pipeline outputs only. For general pipeline authoring, see Pipelines.

Set up an API key

A pipeline step authenticates with a BSQAI API key.

1. In the Portal, open Settings → API Keys and create a key. The full key, starting with sk-bsq-v1-, is shown only once. Copy it.
2. Store the key as a Kubernetes Secret in your Kubeflow profile. Name the Secret spark-api-key, and store the key under the data key key. If you do not have permission to create Secrets, ask your workspace administrator.

Do not put the API key in pipeline source.

Define the pipeline

Each Spark step needs to:

• Use a Spark-capable notebook image as the step image. The image picker shows a short image name. Ask your administrator for the full image path.
• Receive the API key as the SPARK_API_KEY environment variable.
• Turn off caching. A cached Spark step reports success without running Spark.
• Write every output it declares.

Compiling this pipeline requires the kfp and kfp-kubernetes packages.

from kfp import dsl, kubernetes

SPARK_IMAGE = "<registry>/<path>/jupyterlab:python3.12"  # Ask your administrator

PRODUCE = r"""
import pathlib, sys
from bsqai_spark import SparkSessionConfig, get_spark_session

rows, out = int(sys.argv[1]), sys.argv[2]
session = get_spark_session(SparkSessionConfig(mode="job", app_name="produce", worker_replicas=1))
try:
    total = session.range(0, rows).selectExpr("id * 2 AS v").groupBy().sum("v").first()[0]
    pathlib.Path(out).write_text(str(total))
finally:
    session.stop()
"""

CONSUME = r"""
import pathlib, sys
from bsqai_spark import SparkSessionConfig, get_spark_session

total, out = int(sys.argv[1]), sys.argv[2]
session = get_spark_session(SparkSessionConfig(mode="job", app_name="consume", worker_replicas=1))
try:
    df = session.createDataFrame([(i, total + i) for i in range(5)], ["id", "v"])
    result = df.groupBy().sum("v").first()[0]
    pathlib.Path(out).write_text(str(result))
finally:
    session.stop()
"""


@dsl.container_component
def produce(rows: int, total: dsl.OutputPath(int)):
    return dsl.ContainerSpec(image=SPARK_IMAGE, command=["python3", "-c", PRODUCE], args=[rows, total])


@dsl.container_component
def consume(total: int, result: dsl.OutputPath(int)):
    return dsl.ContainerSpec(image=SPARK_IMAGE, command=["python3", "-c", CONSUME], args=[total, result])


def spark_step(task: dsl.PipelineTask) -> dsl.PipelineTask:
    kubernetes.use_secret_as_env(
        task,
        secret_name="spark-api-key",
        secret_key_to_env={"key": "SPARK_API_KEY"},
    )
    task.set_caching_options(False)
    return task


@dsl.pipeline(name="spark-two-steps")
def spark_two_steps(rows: int = 20):
    first = spark_step(produce(rows=rows))
    spark_step(consume(total=first.outputs["total"]))

To see a step's output, open the run in the Kubeflow Pipelines UI and select the step. Its logs appear in the side panel.

Steps that run in parallel each start their own session and count toward your Spark limits.

Troubleshooting