readstat/
rs_query.rs

1//! SQL query execution via Apache DataFusion.
2
3use std::sync::{Arc, Mutex};
4
5use arrow_array::RecordBatch;
6use arrow_schema::SchemaRef;
7use datafusion::catalog::streaming::StreamingTable;
8use datafusion::datasource::MemTable;
9use datafusion::physical_plan::SendableRecordBatchStream;
10use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
11use datafusion::physical_plan::streaming::PartitionStream;
12use datafusion::prelude::*;
13use futures::StreamExt;
14
15use crate::{ReadStatError, ReadStatWriter, WriteConfig};
16
17/// Error-aware Arrow batch receiver used by streaming SQL queries.
18pub type RecordBatchReceiver = crossbeam::channel::Receiver<Result<RecordBatch, ReadStatError>>;
19/// Sending half of a streaming SQL input channel.
20pub type RecordBatchSender = crossbeam::channel::Sender<Result<RecordBatch, ReadStatError>>;
21/// Async receiving half of a streaming SQL input channel.
22pub type AsyncRecordBatchReceiver = tokio::sync::mpsc::Receiver<Result<RecordBatch, ReadStatError>>;
23/// Async sending half of a streaming SQL input channel.
24pub type AsyncRecordBatchSender = tokio::sync::mpsc::Sender<Result<RecordBatch, ReadStatError>>;
25
26/// Creates a bounded input channel for streaming SQL queries.
27///
28/// A bounded channel applies backpressure to producers; capacity zero creates
29/// a rendezvous channel.
30#[must_use]
31pub fn record_batch_channel(capacity: usize) -> (RecordBatchSender, RecordBatchReceiver) {
32    crossbeam::channel::bounded(capacity)
33}
34
35/// Creates a bounded, executor-friendly input channel for async SQL queries.
36///
37/// # Errors
38///
39/// Returns an error when `capacity` is zero.
40pub fn async_record_batch_channel(
41    capacity: usize,
42) -> Result<(AsyncRecordBatchSender, AsyncRecordBatchReceiver), ReadStatError> {
43    if capacity == 0 {
44        return Err(ReadStatError::Other(
45            "async record batch channel capacity must be greater than zero".into(),
46        ));
47    }
48    Ok(tokio::sync::mpsc::channel(capacity))
49}
50
51fn runtime() -> Result<tokio::runtime::Runtime, ReadStatError> {
52    if tokio::runtime::Handle::try_current().is_ok() {
53        return Err(ReadStatError::SyncSqlInAsyncRuntime);
54    }
55    Ok(tokio::runtime::Runtime::new()?)
56}
57
58/// Synchronously executes SQL against in-memory Arrow batches.
59pub fn execute_sql(
60    batches: Vec<RecordBatch>,
61    schema: SchemaRef,
62    table_name: &str,
63    sql: &str,
64) -> Result<Vec<RecordBatch>, ReadStatError> {
65    runtime()?.block_on(execute_sql_async(batches, schema, table_name, sql))
66}
67
68/// Executes SQL asynchronously against in-memory Arrow batches.
69pub async fn execute_sql_async(
70    batches: Vec<RecordBatch>,
71    schema: SchemaRef,
72    table_name: &str,
73    sql: &str,
74) -> Result<Vec<RecordBatch>, ReadStatError> {
75    let ctx = SessionContext::new();
76    ctx.register_table(
77        table_name,
78        Arc::new(MemTable::try_new(schema, vec![batches])?),
79    )?;
80    collect_with_empty_batch(ctx.sql(sql).await?).await
81}
82
83async fn collect_with_empty_batch(df: DataFrame) -> Result<Vec<RecordBatch>, ReadStatError> {
84    let schema = Arc::new(df.schema().as_arrow().clone());
85    let results = df.collect().await?;
86    Ok(if results.is_empty() {
87        vec![RecordBatch::new_empty(schema)]
88    } else {
89        results
90    })
91}
92
93/// A channel-backed partition is single-execution because receiving consumes
94/// its input. Plans that scan it more than once return an execution error.
95#[derive(Debug)]
96struct ChannelPartitionStream {
97    schema: SchemaRef,
98    receiver: Arc<Mutex<Option<InputReceiver>>>,
99}
100
101impl ChannelPartitionStream {
102    fn new(schema: SchemaRef, receiver: InputReceiver) -> Self {
103        Self {
104            schema,
105            receiver: Arc::new(Mutex::new(Some(receiver))),
106        }
107    }
108}
109
110#[derive(Debug)]
111enum InputReceiver {
112    Blocking(RecordBatchReceiver),
113    Async(AsyncRecordBatchReceiver),
114}
115
116impl PartitionStream for ChannelPartitionStream {
117    fn schema(&self) -> &SchemaRef {
118        &self.schema
119    }
120
121    fn execute(&self, _ctx: Arc<datafusion::execution::TaskContext>) -> SendableRecordBatchStream {
122        let receiver = self
123            .receiver
124            .lock()
125            .unwrap_or_else(|e| e.into_inner())
126            .take();
127        let schema = self.schema.clone();
128        let stream = match receiver {
129            Some(receiver) => {
130                let receiver = match receiver {
131                    InputReceiver::Async(receiver) => receiver,
132                    InputReceiver::Blocking(receiver) => {
133                        let (sender, receiver_async) = tokio::sync::mpsc::channel(2);
134                        // Crossbeam receive is blocking. Bridge it from a dedicated
135                        // thread so polling DataFusion never blocks a Tokio worker.
136                        std::thread::spawn(move || {
137                            loop {
138                                if sender.is_closed() {
139                                    break;
140                                }
141                                match receiver
142                                    .recv_timeout(std::time::Duration::from_millis(100))
143                                {
144                                    Ok(result) => {
145                                        if sender.blocking_send(result).is_err() {
146                                            break;
147                                        }
148                                    }
149                                    Err(crossbeam::channel::RecvTimeoutError::Timeout) => {}
150                                    Err(crossbeam::channel::RecvTimeoutError::Disconnected) => break,
151                                }
152                            }
153                        });
154                        receiver_async
155                    }
156                };
157                futures::stream::unfold(receiver, |mut receiver| async move {
158                    receiver.recv().await.map(|batch| {
159                        let batch = batch.map_err(|error| {
160                            datafusion::error::DataFusionError::External(Box::new(error))
161                        });
162                        (batch, receiver)
163                    })
164                })
165                .left_stream()
166            }
167            None => futures::stream::once(async {
168                Err(datafusion::error::DataFusionError::Execution(
169                    "channel-backed StreamingTable can only be executed once; use execute_sql for plans that scan input multiple times".into(),
170                ))
171            })
172            .right_stream(),
173        };
174        Box::pin(RecordBatchStreamAdapter::new(schema, stream))
175    }
176}
177
178fn streaming_context(
179    receiver: InputReceiver,
180    schema: SchemaRef,
181    table_name: &str,
182) -> Result<SessionContext, ReadStatError> {
183    let ctx = SessionContext::new();
184    let partition = ChannelPartitionStream::new(schema.clone(), receiver);
185    ctx.register_table(
186        table_name,
187        Arc::new(StreamingTable::try_new(schema, vec![Arc::new(partition)])?),
188    )?;
189    Ok(ctx)
190}
191
192/// Synchronously executes SQL from a single-use channel of Arrow batches.
193///
194/// Input is consumed incrementally; query results are collected in memory.
195pub fn execute_sql_stream(
196    receiver: RecordBatchReceiver,
197    schema: SchemaRef,
198    table_name: &str,
199    sql: &str,
200) -> Result<Vec<RecordBatch>, ReadStatError> {
201    runtime()?.block_on(execute_sql_from_input_async(
202        InputReceiver::Blocking(receiver),
203        schema,
204        table_name,
205        sql,
206    ))
207}
208
209/// Asynchronously executes SQL from a single-use channel of Arrow batches.
210///
211/// Input is consumed incrementally without blocking the async executor; query
212/// results are collected in memory. Plans that scan the input more than once
213/// are unsupported.
214pub async fn execute_sql_stream_async(
215    receiver: AsyncRecordBatchReceiver,
216    schema: SchemaRef,
217    table_name: &str,
218    sql: &str,
219) -> Result<Vec<RecordBatch>, ReadStatError> {
220    execute_sql_from_input_async(InputReceiver::Async(receiver), schema, table_name, sql).await
221}
222
223async fn execute_sql_from_input_async(
224    receiver: InputReceiver,
225    schema: SchemaRef,
226    table_name: &str,
227    sql: &str,
228) -> Result<Vec<RecordBatch>, ReadStatError> {
229    let ctx = streaming_context(receiver, schema, table_name)?;
230    collect_with_empty_batch(ctx.sql(sql).await?).await
231}
232
233/// Synchronously streams SQL output directly to a configured writer.
234pub fn execute_sql_and_write_stream(
235    receiver: RecordBatchReceiver,
236    schema: SchemaRef,
237    table_name: &str,
238    sql: &str,
239    config: &WriteConfig,
240) -> Result<usize, ReadStatError> {
241    runtime()?.block_on(execute_sql_and_write_from_input_async(
242        InputReceiver::Blocking(receiver),
243        schema,
244        table_name,
245        sql,
246        config,
247    ))
248}
249
250/// Asynchronously writes each SQL output batch as soon as DataFusion produces it.
251/// Plans that scan the channel-backed table more than once are unsupported.
252pub async fn execute_sql_and_write_stream_async(
253    receiver: AsyncRecordBatchReceiver,
254    schema: SchemaRef,
255    table_name: &str,
256    sql: &str,
257    config: &WriteConfig,
258) -> Result<usize, ReadStatError> {
259    execute_sql_and_write_from_input_async(
260        InputReceiver::Async(receiver),
261        schema,
262        table_name,
263        sql,
264        config,
265    )
266    .await
267}
268
269async fn execute_sql_and_write_from_input_async(
270    receiver: InputReceiver,
271    schema: SchemaRef,
272    table_name: &str,
273    sql: &str,
274    config: &WriteConfig,
275) -> Result<usize, ReadStatError> {
276    let ctx = streaming_context(receiver, schema, table_name)?;
277    let df = ctx.sql(sql).await?;
278    let result_schema = Arc::new(df.schema().as_arrow().clone());
279    let mut stream = df.execute_stream().await?;
280    enum Message {
281        Batch(RecordBatch),
282        Finish,
283    }
284    let (sender, mut receiver) = tokio::sync::mpsc::channel(2);
285    let config = config.clone();
286    let writer_task = tokio::task::spawn_blocking(move || {
287        let mut writer = ReadStatWriter::new(config, result_schema)?;
288        while let Some(message) = receiver.blocking_recv() {
289            match message {
290                Message::Batch(batch) => writer.write(&batch)?,
291                Message::Finish => return writer.finish(),
292            }
293        }
294        Err(ReadStatError::Other(
295            "SQL output was cancelled before the writer finished".into(),
296        ))
297    });
298    while let Some(batch) = stream.next().await {
299        sender
300            .send(Message::Batch(batch?))
301            .await
302            .map_err(|_| ReadStatError::Other("SQL writer stopped unexpectedly".into()))?;
303    }
304    sender
305        .send(Message::Finish)
306        .await
307        .map_err(|_| ReadStatError::Other("SQL writer stopped unexpectedly".into()))?;
308    drop(sender);
309    writer_task
310        .await
311        .map_err(|error| ReadStatError::Other(format!("SQL writer task failed: {error}")))?
312}
313
314/// Reads and validates a SQL query file.
315pub fn read_sql_file(path: &std::path::Path) -> Result<String, ReadStatError> {
316    let sql = std::fs::read_to_string(path)?.trim().to_string();
317    if sql.is_empty() {
318        return Err(ReadStatError::EmptySqlFile(path.to_path_buf()));
319    }
320    Ok(sql)
321}
322
323#[cfg(test)]
324mod tests {
325    use super::*;
326    use arrow_array::{Int32Array, RecordBatch};
327    use arrow_schema::{DataType, Field, Schema};
328
329    fn input() -> (SchemaRef, RecordBatch) {
330        let schema = Arc::new(Schema::new(vec![Field::new(
331            "value",
332            DataType::Int32,
333            false,
334        )]));
335        let batch =
336            RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1, 2]))])
337                .unwrap();
338        (schema, batch)
339    }
340
341    #[tokio::test]
342    async fn async_query_and_sync_runtime_guard() {
343        let (schema, batch) = input();
344        let result = execute_sql_async(vec![batch.clone()], schema.clone(), "t", "select * from t")
345            .await
346            .unwrap();
347        assert_eq!(result.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
348        assert!(matches!(
349            execute_sql(vec![batch], schema, "t", "select * from t"),
350            Err(ReadStatError::SyncSqlInAsyncRuntime)
351        ));
352    }
353
354    #[tokio::test]
355    async fn streaming_propagates_channel_errors_and_preserves_empty_schema() {
356        let (schema, batch) = input();
357        let (sender, receiver) = async_record_batch_channel(1).unwrap();
358        sender
359            .send(Err(ReadStatError::Other("source failed".into())))
360            .await
361            .unwrap();
362        drop(sender);
363        let error = execute_sql_stream_async(receiver, schema.clone(), "t", "select * from t")
364            .await
365            .unwrap_err();
366        assert!(error.to_string().contains("source failed"));
367
368        let (sender, receiver) = async_record_batch_channel(1).unwrap();
369        sender.send(Ok(batch)).await.unwrap();
370        drop(sender);
371        let result = execute_sql_stream_async(receiver, schema, "t", "select * from t where 1=0")
372            .await
373            .unwrap();
374        assert_eq!(result.len(), 1);
375        assert_eq!(result[0].num_rows(), 0);
376        assert_eq!(result[0].num_columns(), 1);
377    }
378
379    #[tokio::test(flavor = "current_thread")]
380    async fn streaming_does_not_block_current_thread_runtime() {
381        let (schema, batch) = input();
382        let (sender, receiver) = async_record_batch_channel(1).unwrap();
383        tokio::spawn(async move {
384            tokio::task::yield_now().await;
385            sender.send(Ok(batch)).await.unwrap();
386        });
387        let result = execute_sql_stream_async(receiver, schema, "t", "select * from t")
388            .await
389            .unwrap();
390        assert_eq!(result.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
391    }
392}
readstat/rs_query.rs

readstat/
rs_query.rs
