datafusion_remote_table/connection/
sqlite.rs

1use crate::connection::{RemoteDbType, projections_contains};
2use crate::{
3    Connection, ConnectionOptions, DFResult, Pool, RemoteField, RemoteSchema, RemoteSchemaRef,
4    RemoteType, SqliteType,
5};
6use datafusion::arrow::array::{
7    ArrayBuilder, ArrayRef, BinaryBuilder, Float64Builder, Int64Builder, NullBuilder, RecordBatch,
8    StringBuilder, make_builder,
9};
10use datafusion::arrow::datatypes::{DataType, SchemaRef};
11use datafusion::common::{DataFusionError, project_schema};
12use datafusion::execution::SendableRecordBatchStream;
13use datafusion::physical_plan::memory::MemoryStream;
14use datafusion::prelude::Expr;
15use itertools::Itertools;
16use rusqlite::types::ValueRef;
17use rusqlite::{Column, Row, Rows};
18use std::collections::HashMap;
19use std::path::PathBuf;
20use std::sync::Arc;
21
22#[derive(Debug)]
23pub struct SqlitePool {
24    pool: tokio_rusqlite::Connection,
25}
26
27pub async fn connect_sqlite(path: &PathBuf) -> DFResult<SqlitePool> {
28    let pool = tokio_rusqlite::Connection::open(path).await.map_err(|e| {
29        DataFusionError::Execution(format!("Failed to open sqlite connection: {e:?}"))
30    })?;
31    Ok(SqlitePool { pool })
32}
33
34#[async_trait::async_trait]
35impl Pool for SqlitePool {
36    async fn get(&self) -> DFResult<Arc<dyn Connection>> {
37        let conn = self.pool.clone();
38        Ok(Arc::new(SqliteConnection { conn }))
39    }
40}
41
42#[derive(Debug)]
43pub struct SqliteConnection {
44    conn: tokio_rusqlite::Connection,
45}
46
47#[async_trait::async_trait]
48impl Connection for SqliteConnection {
49    async fn infer_schema(&self, sql: &str) -> DFResult<(RemoteSchemaRef, SchemaRef)> {
50        let sql = sql.to_string();
51        self.conn
52            .call(move |conn| {
53                let mut stmt = conn.prepare(&sql)?;
54                let columns: Vec<OwnedColumn> =
55                    stmt.columns().iter().map(sqlite_col_to_owned_col).collect();
56                let rows = stmt.query([])?;
57
58                let remote_schema = Arc::new(
59                    build_remote_schema(columns.as_slice(), rows)
60                        .map_err(|e| tokio_rusqlite::Error::Other(Box::new(e)))?,
61                );
62                let arrow_schema = Arc::new(remote_schema.to_arrow_schema());
63                Ok((remote_schema, arrow_schema))
64            })
65            .await
66            .map_err(|e| DataFusionError::Execution(format!("Failed to infer schema: {e:?}")))
67    }
68
69    async fn query(
70        &self,
71        _conn_options: &ConnectionOptions,
72        sql: &str,
73        table_schema: SchemaRef,
74        projection: Option<&Vec<usize>>,
75        filters: &[Expr],
76        limit: Option<usize>,
77    ) -> DFResult<SendableRecordBatchStream> {
78        let projected_schema = project_schema(&table_schema, projection)?;
79        let sql = RemoteDbType::Sqlite
80            .try_rewrite_query(sql, filters, limit)
81            .unwrap_or_else(|| sql.to_string());
82        let sql_clone = sql.clone();
83        let projection = projection.cloned();
84        let batch = self
85            .conn
86            .call(move |conn| {
87                let mut stmt = conn.prepare(&sql)?;
88                let columns: Vec<OwnedColumn> =
89                    stmt.columns().iter().map(sqlite_col_to_owned_col).collect();
90                let rows = stmt.query([])?;
91
92                let batch = rows_to_batch(rows, &table_schema, columns, projection.as_ref())
93                    .map_err(|e| tokio_rusqlite::Error::Other(e.into()))?;
94                Ok(batch)
95            })
96            .await
97            .map_err(|e| {
98                DataFusionError::Execution(format!(
99                    "Failed to execute query {sql_clone} on sqlite: {e:?}"
100                ))
101            })?;
102
103        let memory_stream = MemoryStream::try_new(vec![batch], projected_schema, None)?;
104        Ok(Box::pin(memory_stream))
105    }
106}
107
108#[derive(Debug)]
109struct OwnedColumn {
110    name: String,
111    decl_type: Option<String>,
112}
113
114fn sqlite_col_to_owned_col(sqlite_col: &Column) -> OwnedColumn {
115    OwnedColumn {
116        name: sqlite_col.name().to_string(),
117        decl_type: sqlite_col.decl_type().map(|x| x.to_string()),
118    }
119}
120
121fn decl_type_to_remote_type(decl_type: &str) -> DFResult<RemoteType> {
122    if "null".eq(decl_type) {
123        return Ok(RemoteType::Sqlite(SqliteType::Null));
124    }
125    if ["tinyint", "smallint", "int", "integer", "bigint"].contains(&decl_type) {
126        return Ok(RemoteType::Sqlite(SqliteType::Integer));
127    }
128    if ["real", "float", "double"].contains(&decl_type) {
129        return Ok(RemoteType::Sqlite(SqliteType::Real));
130    }
131    if ["text", "varchar", "char", "string"].contains(&decl_type) {
132        return Ok(RemoteType::Sqlite(SqliteType::Text));
133    }
134    if ["binary", "varbinary", "tinyblob", "blob"].contains(&decl_type) {
135        return Ok(RemoteType::Sqlite(SqliteType::Blob));
136    }
137    Err(DataFusionError::NotImplemented(format!(
138        "Unsupported sqlite decl type: {decl_type}",
139    )))
140}
141
142fn build_remote_schema(columns: &[OwnedColumn], mut rows: Rows) -> DFResult<RemoteSchema> {
143    let mut remote_field_map = HashMap::with_capacity(columns.len());
144    let mut unknown_cols = vec![];
145    for (col_idx, col) in columns.iter().enumerate() {
146        if let Some(decl_type) = &col.decl_type {
147            let remote_type = decl_type_to_remote_type(&decl_type.to_ascii_lowercase())?;
148            remote_field_map.insert(col_idx, RemoteField::new(&col.name, remote_type, true));
149        } else {
150            unknown_cols.push(col_idx);
151        }
152    }
153
154    if !unknown_cols.is_empty() {
155        while let Some(row) = rows.next().map_err(|e| {
156            DataFusionError::Execution(format!("Failed to get next row from sqlite: {e:?}"))
157        })? {
158            let mut to_be_removed = vec![];
159            for col_idx in unknown_cols.iter() {
160                let value_ref = row.get_ref(*col_idx).map_err(|e| {
161                    DataFusionError::Execution(format!(
162                        "Failed to get value ref for column {col_idx}: {e:?}"
163                    ))
164                })?;
165                match value_ref {
166                    ValueRef::Null => {}
167                    ValueRef::Integer(_) => {
168                        remote_field_map.insert(
169                            *col_idx,
170                            RemoteField::new(
171                                columns[*col_idx].name.clone(),
172                                RemoteType::Sqlite(SqliteType::Integer),
173                                true,
174                            ),
175                        );
176                        to_be_removed.push(*col_idx);
177                    }
178                    ValueRef::Real(_) => {
179                        remote_field_map.insert(
180                            *col_idx,
181                            RemoteField::new(
182                                columns[*col_idx].name.clone(),
183                                RemoteType::Sqlite(SqliteType::Real),
184                                true,
185                            ),
186                        );
187                        to_be_removed.push(*col_idx);
188                    }
189                    ValueRef::Text(_) => {
190                        remote_field_map.insert(
191                            *col_idx,
192                            RemoteField::new(
193                                columns[*col_idx].name.clone(),
194                                RemoteType::Sqlite(SqliteType::Text),
195                                true,
196                            ),
197                        );
198                        to_be_removed.push(*col_idx);
199                    }
200                    ValueRef::Blob(_) => {
201                        remote_field_map.insert(
202                            *col_idx,
203                            RemoteField::new(
204                                columns[*col_idx].name.clone(),
205                                RemoteType::Sqlite(SqliteType::Blob),
206                                true,
207                            ),
208                        );
209                        to_be_removed.push(*col_idx);
210                    }
211                }
212            }
213            for col_idx in to_be_removed.iter() {
214                unknown_cols.retain(|&x| x != *col_idx);
215            }
216            if unknown_cols.is_empty() {
217                break;
218            }
219        }
220    }
221
222    if !unknown_cols.is_empty() {
223        return Err(DataFusionError::NotImplemented(format!(
224            "Failed to infer sqlite decl type for columns: {unknown_cols:?}"
225        )));
226    }
227    let remote_fields = remote_field_map
228        .into_iter()
229        .sorted_by_key(|entry| entry.0)
230        .map(|entry| entry.1)
231        .collect::<Vec<_>>();
232    Ok(RemoteSchema::new(remote_fields))
233}
234
235fn rows_to_batch(
236    mut rows: Rows,
237    table_schema: &SchemaRef,
238    columns: Vec<OwnedColumn>,
239    projection: Option<&Vec<usize>>,
240) -> DFResult<RecordBatch> {
241    let projected_schema = project_schema(table_schema, projection)?;
242    let mut array_builders = vec![];
243    for field in table_schema.fields() {
244        let builder = make_builder(field.data_type(), 1000);
245        array_builders.push(builder);
246    }
247
248    while let Some(row) = rows.next().map_err(|e| {
249        DataFusionError::Execution(format!("Failed to get next row from sqlite: {e:?}"))
250    })? {
251        append_rows_to_array_builders(
252            row,
253            table_schema,
254            &columns,
255            projection,
256            array_builders.as_mut_slice(),
257        )?;
258    }
259
260    let projected_columns = array_builders
261        .into_iter()
262        .enumerate()
263        .filter(|(idx, _)| projections_contains(projection, *idx))
264        .map(|(_, mut builder)| builder.finish())
265        .collect::<Vec<ArrayRef>>();
266    Ok(RecordBatch::try_new(projected_schema, projected_columns)?)
267}
268
269macro_rules! handle_primitive_type {
270    ($builder:expr, $field:expr, $col:expr, $builder_ty:ty, $value_ty:ty, $row:expr, $index:expr) => {{
271        let builder = $builder
272            .as_any_mut()
273            .downcast_mut::<$builder_ty>()
274            .unwrap_or_else(|| {
275                panic!(
276                    "Failed to downcast builder to {} for {:?} and {:?}",
277                    stringify!($builder_ty),
278                    $field,
279                    $col
280                )
281            });
282
283        let v: Option<$value_ty> = $row.get($index).map_err(|e| {
284            DataFusionError::Execution(format!(
285                "Failed to get optional {} value for {:?} and {:?}: {e:?}",
286                stringify!($value_ty),
287                $field,
288                $col
289            ))
290        })?;
291
292        match v {
293            Some(v) => builder.append_value(v),
294            None => builder.append_null(),
295        }
296    }};
297}
298
299fn append_rows_to_array_builders(
300    row: &Row,
301    table_schema: &SchemaRef,
302    columns: &[OwnedColumn],
303    projection: Option<&Vec<usize>>,
304    array_builders: &mut [Box<dyn ArrayBuilder>],
305) -> DFResult<()> {
306    for (idx, field) in table_schema.fields.iter().enumerate() {
307        if !projections_contains(projection, idx) {
308            continue;
309        }
310        let builder = &mut array_builders[idx];
311        let col = columns.get(idx);
312        match field.data_type() {
313            DataType::Null => {
314                let builder = builder
315                    .as_any_mut()
316                    .downcast_mut::<NullBuilder>()
317                    .expect("Failed to downcast builder to NullBuilder");
318                builder.append_null();
319            }
320            DataType::Int64 => {
321                handle_primitive_type!(builder, field, col, Int64Builder, i64, row, idx);
322            }
323            DataType::Float64 => {
324                handle_primitive_type!(builder, field, col, Float64Builder, f64, row, idx);
325            }
326            DataType::Utf8 => {
327                handle_primitive_type!(builder, field, col, StringBuilder, String, row, idx);
328            }
329            DataType::Binary => {
330                handle_primitive_type!(builder, field, col, BinaryBuilder, Vec<u8>, row, idx);
331            }
332            _ => {
333                return Err(DataFusionError::NotImplemented(format!(
334                    "Unsupported data type {:?} for col: {:?}",
335                    field.data_type(),
336                    col
337                )));
338            }
339        }
340    }
341    Ok(())
342}
datafusion_remote_table/connection/sqlite.rs

datafusion_remote_table/connection/
sqlite.rs
