spark_connect/
query.rs

1//! Provides support for building and executing parameterized SQL queries through [`SparkSession::query`].
2//!
3//! # Overview
4//!
5//! This module defines the internal [`SqlQueryBuilder`] type used by [`SparkSession::query`] to
6//! support a fluent, type-safe API for parameterized SQL queries.  
7//!
8//! Users are not expected to instantiate [`SqlQueryBuilder`] directly; instead, call
9//! [`SparkSession::query`], and then chain `.bind()` calls to attach
10//! parameters before executing the query.
11//!
12//! # Example
13//!
14//! ```
15//! use spark_connect::SparkSessionBuilder;
16//! use arrow::array::RecordBatch;
17//!
18//! # tokio_test::block_on(async {
19//! let session = SparkSessionBuilder::new("sc://localhost:15002").build().await.unwrap();
20//!
21//! // Build and execute a parameterized query fluently
22//! let results: Vec<RecordBatch> = session
23//!     .query("SELECT ? AS id, ? AS name")
24//!     .bind(42)
25//!     .bind("Alice")
26//!     .execute()
27//!     .await
28//!     .unwrap();
29//!
30//! assert!(!results.is_empty());
31//! # });
32//! ```
33//!
34//! # How it works
35//!
36//! - [`SparkSession::query`] creates an internal [`SqlQueryBuilder`] instance tied to the session
37//!   and initializes it with a SQL query string containing `?` placeholders.
38//! - `.bind()` attaches parameter values, converting each Rust type into a Spark [`Literal`] via
39//!   the [`ToLiteral`] trait.
40//! - `.execute()` runs the query asynchronously and collects the resulting Arrow
41//!   [`RecordBatch`]es into memory.
42//!
43//! # See also
44//! - [`ToLiteral`] — converts native Rust types into Spark literals.
45//! - [`SparkSession::sql`] — executes parameterized SQL queries directly.
46//!
47//! # Errors
48//!
49//! Returns a [`SparkError`] if query preparation or execution fails.
50
51use crate::{SparkSession, error::SparkError};
52use crate::spark::expression::Literal;
53use crate::ToLiteral;
54
55use arrow::array::RecordBatch;
56
57
58pub struct SqlQueryBuilder<'a> {
59    session: &'a SparkSession,
60    query: String,
61    params: Vec<Literal>,
62}
63
64impl<'a> SqlQueryBuilder<'a> {
65    pub(crate) fn new(session: &'a SparkSession, query: &str) -> Self {
66        Self {
67            session,
68            query: query.to_string(),
69            params: Vec::new(),
70        }
71    }
72
73    pub fn bind<T: ToLiteral>(mut self, value: T) -> Self {
74        self.params.push(value.to_literal());
75        self
76    }
77
78    pub async fn execute(self) -> Result<Vec<RecordBatch>, SparkError> {
79        let plan = self.session.sql(&self.query, self.params).await?;
80        self.session.collect(plan).await
81    }
82}
spark_connect/query.rs

spark_connect/
query.rs
