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 ["tinyint", "smallint", "int", "integer", "bigint"].contains(&decl_type) {
123        return Ok(RemoteType::Sqlite(SqliteType::Integer));
124    }
125    if ["real", "float", "double"].contains(&decl_type) {
126        return Ok(RemoteType::Sqlite(SqliteType::Real));
127    }
128    if decl_type.starts_with("real") {
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 decl_type.starts_with("char")
135        || decl_type.starts_with("varchar")
136        || decl_type.starts_with("text")
137    {
138        return Ok(RemoteType::Sqlite(SqliteType::Text));
139    }
140    if ["binary", "varbinary", "tinyblob", "blob"].contains(&decl_type) {
141        return Ok(RemoteType::Sqlite(SqliteType::Blob));
142    }
143    if decl_type.starts_with("binary") || decl_type.starts_with("varbinary") {
144        return Ok(RemoteType::Sqlite(SqliteType::Blob));
145    }
146    Err(DataFusionError::NotImplemented(format!(
147        "Unsupported sqlite decl type: {decl_type}",
148    )))
149}
150
151fn build_remote_schema(columns: &[OwnedColumn], mut rows: Rows) -> DFResult<RemoteSchema> {
152    let mut remote_field_map = HashMap::with_capacity(columns.len());
153    let mut unknown_cols = vec![];
154    for (col_idx, col) in columns.iter().enumerate() {
155        if let Some(decl_type) = &col.decl_type {
156            let remote_type = decl_type_to_remote_type(&decl_type.to_ascii_lowercase())?;
157            remote_field_map.insert(col_idx, RemoteField::new(&col.name, remote_type, true));
158        } else {
159            unknown_cols.push(col_idx);
160        }
161    }
162
163    if !unknown_cols.is_empty() {
164        while let Some(row) = rows.next().map_err(|e| {
165            DataFusionError::Execution(format!("Failed to get next row from sqlite: {e:?}"))
166        })? {
167            let mut to_be_removed = vec![];
168            for col_idx in unknown_cols.iter() {
169                let value_ref = row.get_ref(*col_idx).map_err(|e| {
170                    DataFusionError::Execution(format!(
171                        "Failed to get value ref for column {col_idx}: {e:?}"
172                    ))
173                })?;
174                match value_ref {
175                    ValueRef::Null => {}
176                    ValueRef::Integer(_) => {
177                        remote_field_map.insert(
178                            *col_idx,
179                            RemoteField::new(
180                                columns[*col_idx].name.clone(),
181                                RemoteType::Sqlite(SqliteType::Integer),
182                                true,
183                            ),
184                        );
185                        to_be_removed.push(*col_idx);
186                    }
187                    ValueRef::Real(_) => {
188                        remote_field_map.insert(
189                            *col_idx,
190                            RemoteField::new(
191                                columns[*col_idx].name.clone(),
192                                RemoteType::Sqlite(SqliteType::Real),
193                                true,
194                            ),
195                        );
196                        to_be_removed.push(*col_idx);
197                    }
198                    ValueRef::Text(_) => {
199                        remote_field_map.insert(
200                            *col_idx,
201                            RemoteField::new(
202                                columns[*col_idx].name.clone(),
203                                RemoteType::Sqlite(SqliteType::Text),
204                                true,
205                            ),
206                        );
207                        to_be_removed.push(*col_idx);
208                    }
209                    ValueRef::Blob(_) => {
210                        remote_field_map.insert(
211                            *col_idx,
212                            RemoteField::new(
213                                columns[*col_idx].name.clone(),
214                                RemoteType::Sqlite(SqliteType::Blob),
215                                true,
216                            ),
217                        );
218                        to_be_removed.push(*col_idx);
219                    }
220                }
221            }
222            for col_idx in to_be_removed.iter() {
223                unknown_cols.retain(|&x| x != *col_idx);
224            }
225            if unknown_cols.is_empty() {
226                break;
227            }
228        }
229    }
230
231    if !unknown_cols.is_empty() {
232        return Err(DataFusionError::NotImplemented(format!(
233            "Failed to infer sqlite decl type for columns: {unknown_cols:?}"
234        )));
235    }
236    let remote_fields = remote_field_map
237        .into_iter()
238        .sorted_by_key(|entry| entry.0)
239        .map(|entry| entry.1)
240        .collect::<Vec<_>>();
241    Ok(RemoteSchema::new(remote_fields))
242}
243
244fn rows_to_batch(
245    mut rows: Rows,
246    table_schema: &SchemaRef,
247    columns: Vec<OwnedColumn>,
248    projection: Option<&Vec<usize>>,
249) -> DFResult<RecordBatch> {
250    let projected_schema = project_schema(table_schema, projection)?;
251    let mut array_builders = vec![];
252    for field in table_schema.fields() {
253        let builder = make_builder(field.data_type(), 1000);
254        array_builders.push(builder);
255    }
256
257    while let Some(row) = rows.next().map_err(|e| {
258        DataFusionError::Execution(format!("Failed to get next row from sqlite: {e:?}"))
259    })? {
260        append_rows_to_array_builders(
261            row,
262            table_schema,
263            &columns,
264            projection,
265            array_builders.as_mut_slice(),
266        )?;
267    }
268
269    let projected_columns = array_builders
270        .into_iter()
271        .enumerate()
272        .filter(|(idx, _)| projections_contains(projection, *idx))
273        .map(|(_, mut builder)| builder.finish())
274        .collect::<Vec<ArrayRef>>();
275    Ok(RecordBatch::try_new(projected_schema, projected_columns)?)
276}
277
278macro_rules! handle_primitive_type {
279    ($builder:expr, $field:expr, $col:expr, $builder_ty:ty, $value_ty:ty, $row:expr, $index:expr) => {{
280        let builder = $builder
281            .as_any_mut()
282            .downcast_mut::<$builder_ty>()
283            .unwrap_or_else(|| {
284                panic!(
285                    "Failed to downcast builder to {} for {:?} and {:?}",
286                    stringify!($builder_ty),
287                    $field,
288                    $col
289                )
290            });
291
292        let v: Option<$value_ty> = $row.get($index).map_err(|e| {
293            DataFusionError::Execution(format!(
294                "Failed to get optional {} value for {:?} and {:?}: {e:?}",
295                stringify!($value_ty),
296                $field,
297                $col
298            ))
299        })?;
300
301        match v {
302            Some(v) => builder.append_value(v),
303            None => builder.append_null(),
304        }
305    }};
306}
307
308fn append_rows_to_array_builders(
309    row: &Row,
310    table_schema: &SchemaRef,
311    columns: &[OwnedColumn],
312    projection: Option<&Vec<usize>>,
313    array_builders: &mut [Box<dyn ArrayBuilder>],
314) -> DFResult<()> {
315    for (idx, field) in table_schema.fields.iter().enumerate() {
316        if !projections_contains(projection, idx) {
317            continue;
318        }
319        let builder = &mut array_builders[idx];
320        let col = columns.get(idx);
321        match field.data_type() {
322            DataType::Null => {
323                let builder = builder
324                    .as_any_mut()
325                    .downcast_mut::<NullBuilder>()
326                    .expect("Failed to downcast builder to NullBuilder");
327                builder.append_null();
328            }
329            DataType::Int64 => {
330                handle_primitive_type!(builder, field, col, Int64Builder, i64, row, idx);
331            }
332            DataType::Float64 => {
333                handle_primitive_type!(builder, field, col, Float64Builder, f64, row, idx);
334            }
335            DataType::Utf8 => {
336                handle_primitive_type!(builder, field, col, StringBuilder, String, row, idx);
337            }
338            DataType::Binary => {
339                handle_primitive_type!(builder, field, col, BinaryBuilder, Vec<u8>, row, idx);
340            }
341            _ => {
342                return Err(DataFusionError::NotImplemented(format!(
343                    "Unsupported data type {:?} for col: {:?}",
344                    field.data_type(),
345                    col
346                )));
347            }
348        }
349    }
350    Ok(())
351}
datafusion_remote_table/connection/sqlite.rs

datafusion_remote_table/connection/
sqlite.rs
