Crate spark_connect

Expand description

§spark-connect

spark-connect

An idiomatic, SQL-first Rust client for Apache Spark Connect.

This crate is an async layer over the official apache-spark-connect client, whose API is synchronous. It brings the whole official API into async code and adds a sqlx-style interface for binding parameters safely.

§✨ Features

⚙️ Spark-compatible connection builder (sc://host:port format);
🪶 Async execution on tokio, for the whole official API;
🧩 Parameterized queries, bound through the ToLiteral trait;
🧾 Results as Arrow RecordBatches, typed Rows, or a row stream;

§Getting Started

use spark_connect::SparkSession;

// 1️⃣ Connect to a Spark Connect endpoint
let session = SparkSession::builder()
    .remote("sc://localhost:15002")
    .get_or_create()
    .await?;

// 2️⃣ Execute a simple SQL query and receive a Vec<RecordBatch>
let batches = session
    .query("SELECT ? AS rule, ? AS text")
    .bind(42)
    .bind("world")
    .execute()
    .await?;

It’s that simple!

§🧩 Parameterized Queries

The SparkSession::query method uses the ToLiteral trait to bind each parameter as a Spark literal, so values never reach the server as SQL text:

// This
let batches = session
    .query("SELECT ? AS id, ? AS text")
    .bind(42)
    .bind("world")
    .execute()
    .await?;

// is the same as this
use spark_connect::{prelude::*, Expression};

let batches = session
    .run(|spark| {
        spark
            .sql_with_args(
                "SELECT ? AS id, ? AS text",
                vec![
                    Expression::Literal(42.to_literal()),
                    Expression::Literal("world".to_literal()),
                ],
                Default::default(),
            )?
            .collect_record_batches()
    })
    .await?;

bind accepts Rust primitives, String/&str, Vec<u8> and, with the chrono feature, NaiveDate and NaiveDateTime. Pass a LiteralExpression directly for decimals, arrays, maps, structs and typed nulls.

Named parameters and the rest of the official SQL API are reached through sql_with_args, as above.

§🧰 DataFrames

The official client’s DataFrame API is used as it is: transformations only build a plan, and each action has an _async counterpart in the prelude, run off the async executor:

use spark_connect::prelude::*;

let df = spark.sql("SELECT * FROM people")?.filter(col("age").gt(lit(17)));

let adults = df.count_async().await?;
let rows = df.collect_async().await?;

This crate re-exports the official crate’s items - DataFrame, Column, functions, col, lit, … - alongside its own. Anything else in the official API is reachable through SparkSession::run.

§⚖️ Trade-offs

The official client blocks on a runtime of its own, so every call that reaches the server runs on tokio’s blocking thread pool. In practice:

Dropping a future does not cancel the query. A timed-out or abandoned action runs to completion on the server and its result is discarded. Use SparkSession::interrupt_all or SparkSession::interrupt_tag to stop work on the server.
Each in-flight action holds a blocking thread. Tokio’s pool allows 512 by default, far above typical Spark concurrency.
The official synchronous actions remain callable - df.collect() next to df.collect_async() - and panic when called from an async task.
Large results can be streamed row by row with to_local_iterator_async.

§🧠 Concepts

SparkSession - the main entry point for executing SQL queries and managing a session.
SqlQueryBuilder - helper for binding parameters and executing queries.
ext - async counterparts of the official types’ actions.

§⚙️ Requirements

A running Spark Connect server (Spark 4.0+);
Network access to the configured sc:// endpoint;
tokio runtime;
protoc at build time.

§Installing protoc

apache-spark-connect-proto compiles the Connect .proto files while it builds, using the Protocol Buffers compiler:

apt-get install -y protobuf-compiler  # Debian, Ubuntu
brew install protobuf                 # macOS
winget install protobuf               # Windows

If it is not on the PATH, point PROTOC at it:

PROTOC=/path/to/protoc cargo build

In GitHub Actions:

- uses: arduino/setup-protoc@v3
  with:
    repo-token: ${{ secrets.GITHUB_TOKEN }}

§🔒 Example Connection Strings

sc://localhost:15002
sc://spark-cluster:15002/;user_id=francisco
sc://10.0.0.5:15002/;session_id=abc123;user_agent=my-app

§📘 Learn More

apache-spark-connect API reference - the DataFrame, column and function APIs this crate builds on;
Apache Spark Connect documentation;
Spark Connect client connection string.

This project is not affiliated with, endorsed by, or sponsored by the Apache Software Foundation. “Apache”, “Apache Spark”, and “Spark Connect” are trademarks of the Apache Software Foundation.

Re-exports§

pub use apache_spark_connect;

Modules§

catalog: Catalog API mirroring pyspark.sql.connect.catalog.Catalog.
column: Column API mirroring PySpark’s pyspark.sql.connect.column.Column.
conf: RuntimeConf implementation mirroring pyspark.sql.conf.RuntimeConfig.
dataframe: DataFrame implementation mirroring pyspark.sql.DataFrame.
datasource: Data source registration support.
expression: Expression tree mirroring PySpark’s pyspark.sql.connect.expressions.
ext: Async counterparts of the official types’ blocking actions.
functions: SQL functions mirroring pyspark.sql.connect.functions.builtin.
group: GroupedData implementation for aggregations.
merge: MergeIntoWriter mirroring pyspark.sql.connect.merge.MergeIntoWriter.
ml: ML module mirroring pyspark.ml.connect.base and related ML classes.
observation: Observation implementation for collecting metrics.
pipelines: Spark Declarative Pipelines (SDP) command execution.
plan: Logical plan nodes mirroring pyspark.sql.connect.plan.
prelude: The session, the ToLiteral trait, every ext trait, and the official col and lit.
profiler: Client-side profiler collector for UDF and plan profiling results.
query: Provides support for building and executing parameterized SQL queries through SparkSession::query.
readwriter: DataFrameReader and DataFrameWriter implementation mirroring pyspark.sql.connect.readwriter.
resource: Resource profile for specifying executor and task resource requirements.
row: Row type mirroring pyspark.sql.Row.
session: SparkSession implementation mirroring pyspark.sql.SparkSession.
storage: Storage-level presets, mirroring pyspark.StorageLevel.
streaming: Structured Streaming support mirroring pyspark.sql.connect.streaming.
table_arg: Table argument support mirroring pyspark.sql.connect.table_arg.TableArg.
tvf: Table-valued functions (TVF) implementation mirroring pyspark.sql.connect.tvf.TableValuedFunction.
types: DataType type system mirroring PySpark’s pyspark.sql.types.
udf: User-defined function (UDF) support.
window: Window functions and specifications mirroring PySpark’s pyspark.sql.window.

Structs§

Alias: pyspark.sql.connect.expressions.Alias / ColumnAlias
CaseWhen: pyspark.sql.connect.expressions.CaseWhen
Cast: pyspark.sql.connect.expressions.CastExpression / Cast
CoGroupedData: A pair of cogrouped GroupedData, mirroring pyspark.sql.PandasCogroupedOps.
Column: pyspark.sql.connect.column.Column
ColumnReference: pyspark.sql.connect.expressions.ColumnReference Represents a reference to a column by name (unresolved attribute).
CommonInlineUserDefinedDataSourceExpression: Represents a CommonInlineUserDefinedDataSource command. Wraps a Python data source with its name.
DataFrame: A Spark DataFrame, lazily evaluated.
DataFrameReader: DataFrameReader for reading data from various sources.
DataFrameWriter: DataFrameWriter for writing DataFrames to various destinations.
DataFrameWriterV2: DataFrameWriterV2 for the v2 write API (DataFrame.writeTo).
DataStreamReader: DataStreamReader for reading streaming data from various sources.
DataStreamWriter: DataStreamWriter for writing streaming data to various sinks.
ExecutorResourceRequests: Builder for executor resource requests.
GroupedData: Grouped data for performing aggregations.
LocalRowIterator: Iterator over rows from a DataFrame, yielded lazily as the server streams results.
MergeIntoWriter: Fluent builder for a MERGE INTO command. Create it via crate::dataframe::DataFrame::merge_into.
Observation: An Observation for collecting metrics from a DataFrame.
ProfilerCollector: Client-side collector for accumulated profiler results.
PythonDataSourcePayload: Represents a Python data source with its serialized command and metadata.
ResourceProfile: A Spark ResourceProfile, optionally with a server-assigned id.
ResourceProfileBuilder: Builder for ResourceProfile.
SortOrder: pyspark.sql.connect.expressions.SortOrder
SparkError: Re-exported so fallible APIs (e.g. the #[spark_wasm_udf]-generated UDF constructors) can name the client’s Result/error types. A structured error carrying a PySpark error class and its message parameters.
SparkSession: Represents a logical connection to a Spark Connect backend.
SparkSessionBuilder: Builder for creating SparkSession instances.
StorageLevel: Re-exported so persist(...) / storage_level() have a public storage-level type. StorageLevel for persisting Datasets/Tables.
StreamingQuery: A handle to an active streaming query.
StreamingQueryManager: Manager for active streaming queries.
StructField: A field in a StructType, mirroring pyspark.sql.types.StructField.
TaskResourceRequests: Builder for task resource requests.
UnresolvedFunction: pyspark.sql.connect.expressions.UnresolvedFunction Represents a function call with a name and arguments.
Window: pyspark.sql.window.Window
WindowSpec: pyspark.sql.window.WindowSpec

Enums§

DataType: The base DataType representation, mirroring pyspark.sql.types.DataType.
Expression: The base Expression type, mirroring pyspark.sql.connect.expressions.Expression.
FrameBound: Frame bound value used in window frame definitions.
FrameType: Frame type for window specifications.
LiteralExpression: pyspark.sql.connect.expressions.LiteralExpression
ReadType: ReadType variant for Read relation.
SaveMode: SaveMode for write operations.
TableSaveMethod: TableSaveMethod for WriteOperation.
Trigger: Trigger for streaming queries.

Traits§

StorageLevelExt: Named StorageLevel presets (MEMORY_AND_DISK, DISK_ONLY, …) mirroring pyspark.StorageLevel. Named storage-level presets on StorageLevel, matching the constants on pyspark.StorageLevel (MEMORY_AND_DISK, DISK_ONLY, OFF_HEAP, …).
ToLiteral: A trait that allows automatic conversion of Rust primitives and complex types into Spark data types.

Functions§

col: Mirrors pyspark.sql.functions.col.
lit: Mirrors pyspark.sql.functions.lit for an integer value.
lit_boolean: Create a literal from a boolean.
lit_double: Create a literal from a double.
lit_string: Create a literal from a string.

Type Aliases§

Result: Re-exported so fallible APIs (e.g. the #[spark_wasm_udf]-generated UDF constructors) can name the client’s Result/error types.