Skip to main content

exoware_sql/
server.rs

1//! Connect-backed server for `sql.v1`.
2//!
3//! [`SqlServer`] builds a DataFusion session over a [`KvSchema`] and exposes:
4//! - [`Service::query`] unary SQL against that session.
5//! - [`Service::subscribe`] streaming: for every atomic ingest batch that
6//!   touches a registered table's primary-key family, decode its rows
7//!   and re-run the subscriber's SQL `WHERE` predicate against just those
8//!   rows. Each matching batch produces one [`SubscribeResponse`] carrying
9//!   only the rows that satisfied the predicate.
10//!
11//! The streaming path builds a small transient [`MemTable`] per batch and
12//! runs `SELECT * FROM <table> WHERE <where_sql>` against it, so any SQL
13//! expression DataFusion accepts (referring to the table's columns) works
14//! as the predicate.
15
16#![allow(refining_impl_trait)]
17
18use std::collections::HashMap;
19use std::future::Future;
20use std::pin::Pin;
21use std::sync::Arc;
22
23use crate::proto::sql::v1::{
24    cell::Kind as ProtoCellKind, Cell as ProtoCell, Column as ProtoColumn, Index as ProtoIndex,
25    IndexLayout as ProtoIndexLayout, ListValue as ProtoListValue, Null as ProtoNull,
26    QueryRequestView, QueryResponse, Row as ProtoRow, Service, ServiceServer, SubscribeRequestView,
27    SubscribeResponse, Table as ProtoTable, TablesRequestView, TablesResponse,
28};
29use bytes::Bytes;
30use connectrpc::{ConnectError, ConnectRpcService, RequestContext as Context};
31use datafusion::arrow::array::{
32    Array, ArrayRef, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, Date64Array,
33    Decimal128Array, Decimal256Array, FixedSizeBinaryArray, Float32Array, Float64Array, Int32Array,
34    Int64Array, LargeBinaryArray, LargeListArray, LargeStringArray, ListArray, StringArray,
35    StringViewArray, TimestampMicrosecondArray, TimestampMillisecondArray,
36    TimestampNanosecondArray, TimestampSecondArray, UInt32Array, UInt64Array,
37};
38use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit};
39use datafusion::arrow::record_batch::RecordBatch;
40use datafusion::common::{DataFusionError, Result as DataFusionResult};
41use datafusion::datasource::MemTable;
42use datafusion::prelude::SessionContext;
43use exoware_sdk::keys::Key;
44use exoware_sdk::kv_codec::{decode_stored_row, Utf8};
45use exoware_sdk::selector::Selector;
46use exoware_sdk::stream_filter::StreamFilter;
47use exoware_sdk::{PrefixedStoreClient, StreamSubscription};
48use futures::future::BoxFuture;
49use futures::stream::Stream;
50use futures::FutureExt;
51
52use crate::builder::{projected_column_indices, ProjectedBatchBuilder};
53use crate::codec::decode_primary_key_selected;
54use crate::filter::ScanAccessPlan;
55use crate::predicate::QueryPredicate;
56use crate::schema::KvSchema;
57use crate::types::{IndexLayout, ResolvedIndexSpec, TableModel};
58
59const MAX_CONNECTRPC_BODY_BYTES: usize = 256 * 1024 * 1024;
60
61type SubscribeStream = Pin<Box<dyn Stream<Item = Result<SubscribeResponse, ConnectError>> + Send>>;
62
63/// One registered table's streaming-decode state.
64#[derive(Clone)]
65struct TableStream {
66    model: Arc<TableModel>,
67    schema: SchemaRef,
68    access_plan: Arc<ScanAccessPlan>,
69    selector: Selector,
70    indexes: Arc<Vec<ResolvedIndexSpec>>,
71}
72
73impl TableStream {
74    fn new(model: Arc<TableModel>, indexes: Vec<ResolvedIndexSpec>) -> Self {
75        let projection: Option<Vec<usize>> = Some((0..model.columns.len()).collect());
76        let access_plan = Arc::new(ScanAccessPlan::new(
77            &model,
78            &projection,
79            &QueryPredicate::default(),
80        ));
81        let selector = Selector {
82            prefix: model.primary_key_prefix.as_bytes().clone(),
83            payload_regex: Utf8::from("(?s-u).*"),
84        };
85        Self {
86            schema: model.schema.clone(),
87            access_plan,
88            model,
89            selector,
90            indexes: Arc::new(indexes),
91        }
92    }
93
94    fn decode_batch(&self, entries: &[(Key, Bytes)]) -> DataFusionResult<RecordBatch> {
95        let mut builder = ProjectedBatchBuilder::from_access_plan(&self.model, &self.access_plan);
96        for (key, value) in entries {
97            if !self.model.primary_key_prefix.matches(key) {
98                continue;
99            }
100            let Some(pk_values) = decode_primary_key_selected(
101                self.model.table_prefix,
102                key,
103                &self.model,
104                &self.access_plan.required_pk_mask,
105            ) else {
106                continue;
107            };
108            let Ok(archived) = decode_stored_row(value) else {
109                continue;
110            };
111            if archived.values.len() != self.model.columns.len() {
112                continue;
113            }
114            let _ = builder.append_archived_row(&pk_values, &archived)?;
115        }
116        builder.finish(&self.schema)
117    }
118}
119
120/// SQL server bound to a single [`KvSchema`].
121///
122/// Construct with [`SqlServer::new`], pass to [`sql_connect_stack`] to mount
123/// on an axum router.
124pub struct SqlServer {
125    ctx: Arc<SessionContext>,
126    streams: HashMap<String, TableStream>,
127    // Registration order, preserved for the `Tables` RPC so clients see
128    // tables in the same order the operator declared them.
129    table_names: Vec<String>,
130    store: PrefixedStoreClient,
131}
132
133impl SqlServer {
134    /// Build a server from a [`KvSchema`]. The schema's tables are registered
135    /// in a new [`SessionContext`] that drives both unary `Query` and the
136    /// per-batch predicate evaluation on `Subscribe`.
137    pub fn new(schema: KvSchema) -> DataFusionResult<Self> {
138        let store = schema.client().clone();
139        let mut streams = HashMap::with_capacity(schema.tables().len());
140        let mut table_names = Vec::with_capacity(schema.tables().len());
141        for (name, config) in schema.tables() {
142            let model =
143                Arc::new(TableModel::from_config(config).map_err(|e| {
144                    DataFusionError::Execution(format!("invalid table config: {e}"))
145                })?);
146            let indexes = model
147                .resolve_index_specs(&config.index_specs)
148                .map_err(|e| DataFusionError::Execution(format!("invalid index specs: {e}")))?;
149            streams.insert(name.clone(), TableStream::new(model, indexes));
150            table_names.push(name.clone());
151        }
152        let ctx = SessionContext::new();
153        schema.register_all(&ctx)?;
154        Ok(Self {
155            ctx: Arc::new(ctx),
156            streams,
157            table_names,
158            store,
159        })
160    }
161
162    /// Borrow the underlying DataFusion session, e.g. to `INSERT` seed rows
163    /// without going through the connect API.
164    pub fn session(&self) -> &SessionContext {
165        &self.ctx
166    }
167
168    #[allow(clippy::result_large_err)]
169    fn stream(&self, table: &str) -> Result<&TableStream, ConnectError> {
170        self.streams
171            .get(table)
172            .ok_or_else(|| ConnectError::not_found(format!("unknown table '{table}'")))
173    }
174
175    fn describe_tables(&self) -> Vec<ProtoTable> {
176        self.table_names
177            .iter()
178            .filter_map(|name| {
179                let stream = self.streams.get(name)?;
180                let columns = stream
181                    .schema
182                    .fields()
183                    .iter()
184                    .map(|field| ProtoColumn {
185                        name: field.name().clone(),
186                        data_type: format!("{}", field.data_type()),
187                        nullable: field.is_nullable(),
188                        ..Default::default()
189                    })
190                    .collect();
191                let primary_key_columns = stream
192                    .model
193                    .primary_key_indices
194                    .iter()
195                    .map(|&idx| idx as u32)
196                    .collect();
197                let indexes = stream
198                    .indexes
199                    .iter()
200                    .map(|spec| {
201                        let key_set: std::collections::HashSet<usize> =
202                            spec.key_columns.iter().copied().collect();
203                        ProtoIndex {
204                            name: spec.name.clone(),
205                            layout: proto_index_layout(spec.layout).into(),
206                            key_columns: spec.key_columns.iter().map(|&idx| idx as u32).collect(),
207                            // Columns stored in the index payload beyond the
208                            // key itself ("covered" columns that let point
209                            // lookups skip the base-row fetch).
210                            cover_columns: spec
211                                .value_column_mask
212                                .iter()
213                                .enumerate()
214                                .filter_map(|(idx, covered)| {
215                                    (*covered && !key_set.contains(&idx)).then_some(idx as u32)
216                                })
217                                .collect(),
218                            ..Default::default()
219                        }
220                    })
221                    .collect();
222                Some(ProtoTable {
223                    name: name.clone(),
224                    columns,
225                    primary_key_columns,
226                    indexes,
227                    ..Default::default()
228                })
229            })
230            .collect()
231    }
232}
233
234fn proto_index_layout(layout: IndexLayout) -> ProtoIndexLayout {
235    match layout {
236        IndexLayout::Lexicographic => ProtoIndexLayout::INDEX_LAYOUT_LEXICOGRAPHIC,
237        IndexLayout::ZOrder => ProtoIndexLayout::INDEX_LAYOUT_Z_ORDER,
238    }
239}
240
241/// Turn a [`SqlServer`] into a mounted Connect service stack ready to hand to
242/// axum's `fallback_service`.
243pub fn sql_connect_stack(server: Arc<SqlServer>) -> ConnectRpcService<ServiceServer<SqlConnect>> {
244    ConnectRpcService::new(ServiceServer::new(SqlConnect::new(server)))
245        .with_limits(
246            connectrpc::Limits::default()
247                .max_request_body_size(MAX_CONNECTRPC_BODY_BYTES)
248                .max_message_size(MAX_CONNECTRPC_BODY_BYTES),
249        )
250        .with_compression(exoware_sdk::connect_compression_registry())
251}
252
253/// Connect handler implementing `sql.v1.Service`.
254#[derive(Clone)]
255pub struct SqlConnect {
256    server: Arc<SqlServer>,
257}
258
259impl SqlConnect {
260    pub fn new(server: Arc<SqlServer>) -> Self {
261        Self { server }
262    }
263}
264
265impl Service for SqlConnect {
266    fn subscribe(
267        &self,
268        _ctx: Context,
269        request: buffa::view::OwnedView<SubscribeRequestView<'static>>,
270    ) -> impl Future<Output = connectrpc::ServiceResult<SubscribeStream>> + Send {
271        let server = self.server.clone();
272        async move {
273            let table_name = request.table.to_string();
274            let where_sql = request.where_sql.trim().to_string();
275            let since = request.since_sequence_number.filter(|seq| *seq != 0);
276            let stream = server.stream(&table_name)?.clone();
277
278            let filter = StreamFilter {
279                selectors: vec![stream.selector.clone()],
280                value_filters: vec![],
281            };
282            let sub = server
283                .store
284                .stream()
285                .subscribe(filter, since)
286                .await
287                .map_err(client_error_to_connect)?;
288
289            let output = Box::pin(BatchPredicateStream::new(
290                sub, stream, table_name, where_sql,
291            ));
292            Ok(connectrpc::Response::stream(output as SubscribeStream))
293        }
294    }
295
296    fn tables(
297        &self,
298        _ctx: Context,
299        _request: buffa::view::OwnedView<TablesRequestView<'static>>,
300    ) -> impl Future<Output = connectrpc::ServiceResult<TablesResponse>> + Send {
301        let server = self.server.clone();
302        async move {
303            connectrpc::Response::ok(TablesResponse {
304                tables: server.describe_tables(),
305                ..Default::default()
306            })
307        }
308    }
309
310    fn query(
311        &self,
312        _ctx: Context,
313        request: buffa::view::OwnedView<QueryRequestView<'static>>,
314    ) -> impl Future<Output = connectrpc::ServiceResult<QueryResponse>> + Send {
315        let server = self.server.clone();
316        async move {
317            let sql = request.sql.to_string();
318            let df = server
319                .ctx
320                .sql(&sql)
321                .await
322                .map_err(datafusion_error_to_connect)?;
323            let schema = df.schema().clone();
324            let batches = df.collect().await.map_err(datafusion_error_to_connect)?;
325            let columns: Vec<String> = schema.fields().iter().map(|f| f.name().clone()).collect();
326            let rows =
327                record_batches_to_proto_rows(&batches).map_err(datafusion_error_to_connect)?;
328            connectrpc::Response::ok(QueryResponse {
329                column: columns,
330                rows,
331                ..Default::default()
332            })
333        }
334    }
335}
336
337struct BatchPredicateStream {
338    sub: StreamSubscription,
339    state: TableStream,
340    table_name: String,
341    where_sql: String,
342    building: Option<BoxFuture<'static, Result<Option<SubscribeResponse>, ConnectError>>>,
343}
344
345impl BatchPredicateStream {
346    fn new(
347        sub: StreamSubscription,
348        state: TableStream,
349        table_name: String,
350        where_sql: String,
351    ) -> Self {
352        Self {
353            sub,
354            state,
355            table_name,
356            where_sql,
357            building: None,
358        }
359    }
360}
361
362impl Stream for BatchPredicateStream {
363    type Item = Result<SubscribeResponse, ConnectError>;
364
365    fn poll_next(
366        self: Pin<&mut Self>,
367        cx: &mut std::task::Context<'_>,
368    ) -> std::task::Poll<Option<Self::Item>> {
369        let this = self.get_mut();
370        loop {
371            if let Some(fut) = this.building.as_mut() {
372                match fut.as_mut().poll(cx) {
373                    std::task::Poll::Pending => return std::task::Poll::Pending,
374                    std::task::Poll::Ready(Ok(Some(resp))) => {
375                        this.building = None;
376                        return std::task::Poll::Ready(Some(Ok(resp)));
377                    }
378                    std::task::Poll::Ready(Ok(None)) => {
379                        this.building = None;
380                    }
381                    std::task::Poll::Ready(Err(err)) => {
382                        this.building = None;
383                        return std::task::Poll::Ready(Some(Err(err)));
384                    }
385                }
386            }
387
388            let frame = {
389                let next_fut = this.sub.next();
390                tokio::pin!(next_fut);
391                match next_fut.as_mut().poll(cx) {
392                    std::task::Poll::Ready(Ok(Some(frame))) => frame,
393                    std::task::Poll::Ready(Ok(None)) => return std::task::Poll::Ready(None),
394                    std::task::Poll::Ready(Err(err)) => {
395                        return std::task::Poll::Ready(Some(Err(client_error_to_connect(err))));
396                    }
397                    std::task::Poll::Pending => return std::task::Poll::Pending,
398                }
399            };
400
401            let sequence_number = frame.sequence_number;
402            let entries: Vec<(Key, Bytes)> = frame
403                .entries
404                .into_iter()
405                .map(|entry| (entry.key, entry.value))
406                .collect();
407            let state = this.state.clone();
408            let table_name = this.table_name.clone();
409            let where_sql = this.where_sql.clone();
410            this.building = Some(
411                async move {
412                    evaluate_batch(state, table_name, where_sql, sequence_number, entries).await
413                }
414                .boxed(),
415            );
416        }
417    }
418}
419
420async fn evaluate_batch(
421    state: TableStream,
422    table_name: String,
423    where_sql: String,
424    sequence_number: u64,
425    entries: Vec<(Key, Bytes)>,
426) -> Result<Option<SubscribeResponse>, ConnectError> {
427    let batch = state
428        .decode_batch(&entries)
429        .map_err(datafusion_error_to_connect)?;
430    if batch.num_rows() == 0 {
431        return Ok(None);
432    }
433
434    let filtered = if where_sql.is_empty() {
435        batch
436    } else {
437        apply_where(state.schema.clone(), batch, &table_name, &where_sql)
438            .await
439            .map_err(datafusion_error_to_connect)?
440    };
441    if filtered.num_rows() == 0 {
442        return Ok(None);
443    }
444
445    let columns: Vec<String> = filtered
446        .schema()
447        .fields()
448        .iter()
449        .map(|f| f.name().clone())
450        .collect();
451    let rows = record_batches_to_proto_rows(std::slice::from_ref(&filtered))
452        .map_err(datafusion_error_to_connect)?;
453    Ok(Some(SubscribeResponse {
454        sequence_number,
455        column: columns,
456        rows,
457        ..Default::default()
458    }))
459}
460
461async fn apply_where(
462    schema: SchemaRef,
463    batch: RecordBatch,
464    table_name: &str,
465    where_sql: &str,
466) -> DataFusionResult<RecordBatch> {
467    let ctx = SessionContext::new();
468    let mem = MemTable::try_new(schema.clone(), vec![vec![batch]])?;
469    ctx.register_table(table_name, Arc::new(mem))?;
470    let sql = format!("SELECT * FROM {table_name} WHERE {where_sql}");
471    let df = ctx.sql(&sql).await?;
472    let batches = df.collect().await?;
473    if batches.is_empty() {
474        return Ok(RecordBatch::new_empty(schema));
475    }
476    datafusion::arrow::compute::concat_batches(&schema, batches.iter())
477        .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))
478}
479
480fn record_batches_to_proto_rows(batches: &[RecordBatch]) -> DataFusionResult<Vec<ProtoRow>> {
481    let mut out = Vec::with_capacity(batches.iter().map(|b| b.num_rows()).sum());
482    for batch in batches {
483        for row_idx in 0..batch.num_rows() {
484            let mut cells = Vec::with_capacity(batch.num_columns());
485            for col_idx in 0..batch.num_columns() {
486                cells.push(arrow_value_to_cell(batch.column(col_idx), row_idx)?);
487            }
488            out.push(ProtoRow {
489                cells,
490                ..Default::default()
491            });
492        }
493    }
494    Ok(out)
495}
496
497fn arrow_value_to_cell(array: &ArrayRef, row: usize) -> DataFusionResult<ProtoCell> {
498    let kind = if array.is_null(row) {
499        ProtoCellKind::NullValue(Box::<ProtoNull>::default())
500    } else {
501        arrow_value_to_kind(array, row)?
502    };
503    Ok(ProtoCell {
504        kind: Some(kind),
505        ..Default::default()
506    })
507}
508
509fn arrow_value_to_kind(array: &ArrayRef, row: usize) -> DataFusionResult<ProtoCellKind> {
510    match array.data_type() {
511        DataType::Int64 => Ok(ProtoCellKind::Int64Value(
512            array
513                .as_any()
514                .downcast_ref::<Int64Array>()
515                .unwrap()
516                .value(row),
517        )),
518        DataType::Int32 => Ok(ProtoCellKind::Int64Value(
519            array
520                .as_any()
521                .downcast_ref::<Int32Array>()
522                .unwrap()
523                .value(row) as i64,
524        )),
525        DataType::UInt64 => Ok(ProtoCellKind::Uint64Value(
526            array
527                .as_any()
528                .downcast_ref::<UInt64Array>()
529                .unwrap()
530                .value(row),
531        )),
532        DataType::UInt32 => Ok(ProtoCellKind::Uint64Value(
533            array
534                .as_any()
535                .downcast_ref::<UInt32Array>()
536                .unwrap()
537                .value(row) as u64,
538        )),
539        DataType::Float64 => Ok(ProtoCellKind::Float64Value(
540            array
541                .as_any()
542                .downcast_ref::<Float64Array>()
543                .unwrap()
544                .value(row),
545        )),
546        DataType::Float32 => Ok(ProtoCellKind::Float64Value(
547            array
548                .as_any()
549                .downcast_ref::<Float32Array>()
550                .unwrap()
551                .value(row) as f64,
552        )),
553        DataType::Boolean => Ok(ProtoCellKind::BooleanValue(
554            array
555                .as_any()
556                .downcast_ref::<BooleanArray>()
557                .unwrap()
558                .value(row),
559        )),
560        DataType::Utf8 => Ok(ProtoCellKind::Utf8Value(
561            array
562                .as_any()
563                .downcast_ref::<StringArray>()
564                .unwrap()
565                .value(row)
566                .to_string(),
567        )),
568        DataType::LargeUtf8 => Ok(ProtoCellKind::Utf8Value(
569            array
570                .as_any()
571                .downcast_ref::<LargeStringArray>()
572                .unwrap()
573                .value(row)
574                .to_string(),
575        )),
576        DataType::Utf8View => Ok(ProtoCellKind::Utf8Value(
577            array
578                .as_any()
579                .downcast_ref::<StringViewArray>()
580                .unwrap()
581                .value(row)
582                .to_string(),
583        )),
584        DataType::FixedSizeBinary(_) => {
585            Ok(ProtoCellKind::FixedSizeBinaryValue(Bytes::copy_from_slice(
586                array
587                    .as_any()
588                    .downcast_ref::<FixedSizeBinaryArray>()
589                    .unwrap()
590                    .value(row),
591            )))
592        }
593        DataType::Binary => Ok(ProtoCellKind::BinaryValue(Bytes::copy_from_slice(
594            array
595                .as_any()
596                .downcast_ref::<BinaryArray>()
597                .unwrap()
598                .value(row),
599        ))),
600        DataType::LargeBinary => Ok(ProtoCellKind::BinaryValue(Bytes::copy_from_slice(
601            array
602                .as_any()
603                .downcast_ref::<LargeBinaryArray>()
604                .unwrap()
605                .value(row),
606        ))),
607        DataType::BinaryView => Ok(ProtoCellKind::BinaryValue(Bytes::copy_from_slice(
608            array
609                .as_any()
610                .downcast_ref::<BinaryViewArray>()
611                .unwrap()
612                .value(row),
613        ))),
614        DataType::Date32 => Ok(ProtoCellKind::Date32Value(
615            array
616                .as_any()
617                .downcast_ref::<Date32Array>()
618                .unwrap()
619                .value(row),
620        )),
621        DataType::Date64 => Ok(ProtoCellKind::Date64Value(
622            array
623                .as_any()
624                .downcast_ref::<Date64Array>()
625                .unwrap()
626                .value(row),
627        )),
628        DataType::Timestamp(unit, _) => {
629            let v = match unit {
630                TimeUnit::Second => array
631                    .as_any()
632                    .downcast_ref::<TimestampSecondArray>()
633                    .unwrap()
634                    .value(row),
635                TimeUnit::Millisecond => array
636                    .as_any()
637                    .downcast_ref::<TimestampMillisecondArray>()
638                    .unwrap()
639                    .value(row),
640                TimeUnit::Microsecond => array
641                    .as_any()
642                    .downcast_ref::<TimestampMicrosecondArray>()
643                    .unwrap()
644                    .value(row),
645                TimeUnit::Nanosecond => array
646                    .as_any()
647                    .downcast_ref::<TimestampNanosecondArray>()
648                    .unwrap()
649                    .value(row),
650            };
651            Ok(ProtoCellKind::TimestampValue(v))
652        }
653        DataType::Decimal128(_, _) => {
654            let v = array
655                .as_any()
656                .downcast_ref::<Decimal128Array>()
657                .unwrap()
658                .value(row);
659            Ok(ProtoCellKind::Decimal128Value(Bytes::copy_from_slice(
660                &v.to_be_bytes(),
661            )))
662        }
663        DataType::Decimal256(_, _) => {
664            let v = array
665                .as_any()
666                .downcast_ref::<Decimal256Array>()
667                .unwrap()
668                .value(row);
669            Ok(ProtoCellKind::Decimal256Value(Bytes::copy_from_slice(
670                &v.to_be_bytes(),
671            )))
672        }
673        DataType::List(_) => {
674            let list = array.as_any().downcast_ref::<ListArray>().unwrap();
675            Ok(ProtoCellKind::ListValue(Box::new(list_array_to_proto(
676                &list.value(row),
677            )?)))
678        }
679        DataType::LargeList(_) => {
680            let list = array.as_any().downcast_ref::<LargeListArray>().unwrap();
681            Ok(ProtoCellKind::ListValue(Box::new(list_array_to_proto(
682                &list.value(row),
683            )?)))
684        }
685        other => Err(DataFusionError::NotImplemented(format!(
686            "cell conversion for arrow type {other:?}"
687        ))),
688    }
689}
690
691fn list_array_to_proto(elements: &ArrayRef) -> DataFusionResult<ProtoListValue> {
692    let mut cells = Vec::with_capacity(elements.len());
693    for idx in 0..elements.len() {
694        cells.push(arrow_value_to_cell(elements, idx)?);
695    }
696    Ok(ProtoListValue {
697        elements: cells,
698        ..Default::default()
699    })
700}
701
702fn datafusion_error_to_connect(err: DataFusionError) -> ConnectError {
703    match err {
704        DataFusionError::Plan(msg)
705        | DataFusionError::SQL(_, Some(msg))
706        | DataFusionError::Configuration(msg)
707        | DataFusionError::NotImplemented(msg) => ConnectError::invalid_argument(msg),
708        DataFusionError::SchemaError(err, _) => ConnectError::invalid_argument(err.to_string()),
709        other => ConnectError::internal(other.to_string()),
710    }
711}
712
713fn client_error_to_connect(err: exoware_sdk::ClientError) -> ConnectError {
714    if let Some(rpc) = err.rpc_error() {
715        ConnectError::new(rpc.code, rpc.message.clone().unwrap_or_default())
716    } else {
717        ConnectError::internal(err.to_string())
718    }
719}
720
721// Silence unused import when no tests reference them.
722#[allow(dead_code)]
723fn _assert_projected_column_indices_visible() {
724    let _ = projected_column_indices;
725}
726
727#[cfg(test)]
728mod tests {
729    use super::*;
730
731    /// Result cells for Binary columns can arrive as any of the three arrow
732    /// binary encodings; every one becomes a `binary_value` cell.
733    #[test]
734    fn binary_arrays_convert_to_binary_cells() {
735        let body: &[u8] = &[0x00, 0xFF, 0x42];
736        let arrays: Vec<ArrayRef> = vec![
737            Arc::new(BinaryArray::from_iter_values([body])),
738            Arc::new(LargeBinaryArray::from_iter_values([body])),
739            Arc::new(BinaryViewArray::from_iter_values([body])),
740        ];
741        for array in arrays {
742            let body_type = array.data_type().clone();
743            let cell = arrow_value_to_cell(&array, 0).expect("cell conversion");
744            match cell.kind {
745                Some(ProtoCellKind::BinaryValue(bytes)) => {
746                    assert_eq!(bytes.as_ref(), body, "wrong bytes for {body_type:?}")
747                }
748                other => panic!("expected binary_value for {body_type:?}, got {other:?}"),
749            }
750        }
751    }
752}