Skip to main content

mongreldb_query/
external_modules.rs

1use crate::extended_sql_functions::{json_table_batches_from_text, JsonTableMode};
2use crate::{arrow_conv, MongrelQueryError, Result};
3use arrow::array::{ArrayRef, Int64Array, StringArray};
4use arrow::datatypes::{DataType, Field, Schema as ArrowSchema, SchemaRef};
5use arrow::record_batch::{RecordBatch, RecordBatchOptions};
6use datafusion::catalog::{Session, TableProvider};
7use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
8use datafusion::execution::TaskContext;
9use datafusion::logical_expr::{
10    BinaryExpr, Expr, Operator, TableProviderFilterPushDown, TableType,
11};
12use datafusion::physical_expr::EquivalenceProperties;
13use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
14use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
15use datafusion::physical_plan::{
16    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
17    SendableRecordBatchStream,
18};
19use mongreldb_core::catalog::TableState;
20use mongreldb_core::database::{TABLES_DIR, VTAB_DIR};
21use mongreldb_core::schema::{
22    ColumnDef as CoreColumnDef, ColumnFlags, Schema as CoreSchema, TypeId,
23};
24use mongreldb_core::{Database, ExternalTableEntry, ModuleArg, ModuleCapabilities, Value};
25use serde::{Deserialize, Serialize};
26use std::any::Any;
27use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
28use std::fs;
29use std::path::{Path, PathBuf};
30use std::sync::Arc;
31
32use crate::query_registry::SqlTaskContext;
33
34pub trait ExternalTableModule: Send + Sync {
35    fn name(&self) -> &str;
36    fn descriptor(&self) -> ExternalModuleDescriptor;
37    fn indexes(&self, _entry: &ExternalTableEntry) -> Result<Vec<ExternalModuleIndex>> {
38        Ok(Vec::new())
39    }
40    fn connect(
41        &self,
42        ctx: &ModuleConnectCtx<'_>,
43        entry: &ExternalTableEntry,
44    ) -> Result<Arc<dyn ExternalTable>>;
45    fn read_rows(
46        &self,
47        _ctx: &ModuleConnectCtx<'_>,
48        entry: &ExternalTableEntry,
49    ) -> Result<Vec<HashMap<u16, Value>>> {
50        Err(MongrelQueryError::Schema(format!(
51            "external table {:?} using module {:?} is not row-writable",
52            entry.name,
53            self.name()
54        )))
55    }
56    fn prepare_rows(
57        &self,
58        _ctx: &ModuleConnectCtx<'_>,
59        entry: &ExternalTableEntry,
60        _rows: Vec<HashMap<u16, Value>>,
61    ) -> Result<Vec<u8>> {
62        Err(MongrelQueryError::Schema(format!(
63            "external table {:?} using module {:?} is not row-writable",
64            entry.name,
65            self.name()
66        )))
67    }
68    fn rows_from_state(&self, state: &[u8]) -> Result<Vec<HashMap<u16, Value>>> {
69        if state.is_empty() {
70            Ok(Vec::new())
71        } else {
72            decode_state_rows(state)
73        }
74    }
75    fn write(
76        &self,
77        ctx: &ModuleConnectCtx<'_>,
78        entry: &ExternalTableEntry,
79        op: ExternalWriteOp,
80        txn: &mut ExternalTxn,
81    ) -> Result<ExternalWriteResult> {
82        let (rows, changes) = match op {
83            ExternalWriteOp::Insert { rows: inserted } => {
84                let changes = inserted.len() as u64;
85                let mut rows = txn.read_rows()?;
86                rows.extend(inserted);
87                (rows, changes)
88            }
89            ExternalWriteOp::ReplaceRows { rows, changes } => (rows, changes),
90        };
91        txn.replace_state(self.prepare_rows(ctx, entry, rows)?);
92        Ok(ExternalWriteResult { changes })
93    }
94    fn destroy(&self, _ctx: &ModuleConnectCtx<'_>, _entry: &ExternalTableEntry) -> Result<()> {
95        Ok(())
96    }
97}
98
99pub trait ExternalTable: std::fmt::Debug + Send + Sync {
100    fn schema(&self) -> SchemaRef;
101    fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan>;
102    fn scan_with_control(
103        &self,
104        request: &ExternalPlanRequest<'_>,
105        control: &mongreldb_core::ExecutionControl,
106    ) -> DFResult<ExternalScan>;
107}
108
109#[derive(Clone)]
110pub struct ExternalModuleDescriptor {
111    pub schema: CoreSchema,
112    pub hidden_columns: Vec<String>,
113    pub capabilities: ModuleCapabilities,
114}
115
116#[derive(Debug, Clone, PartialEq, Eq)]
117pub struct ExternalModuleIndex {
118    pub name: String,
119    pub column_ids: Vec<u16>,
120    pub unique: bool,
121    pub partial: bool,
122}
123
124impl ExternalModuleIndex {
125    pub fn new(name: impl Into<String>, column_ids: Vec<u16>) -> Self {
126        Self {
127            name: name.into(),
128            column_ids,
129            unique: false,
130            partial: false,
131        }
132    }
133}
134
135pub struct ModuleConnectCtx<'a> {
136    pub db: &'a Arc<Database>,
137}
138
139impl ModuleConnectCtx<'_> {
140    pub fn raw_state(&self, entry: &ExternalTableEntry) -> Result<Vec<u8>> {
141        external_table_state_bytes(self.db, entry)
142    }
143
144    pub fn read_state(&self, entry: &ExternalTableEntry, key: &[u8]) -> Result<Option<Vec<u8>>> {
145        ExternalTxn::read_state_from_bytes(&self.raw_state(entry)?, key)
146    }
147}
148
149pub struct ExternalPlanRequest<'a> {
150    pub projection: Option<Vec<usize>>,
151    pub filters: Vec<ExternalFilter>,
152    pub raw_filters: Vec<&'a Expr>,
153    pub order_by: Vec<ExternalOrder>,
154    pub limit: Option<usize>,
155    pub offset: Option<usize>,
156}
157
158pub struct ExternalScan {
159    pub schema: SchemaRef,
160    pub batches: Vec<RecordBatch>,
161}
162
163#[derive(Clone)]
164pub struct ExternalPlan {
165    pub filter_pushdown: Vec<TableProviderFilterPushDown>,
166    pub accepted_filters: Vec<AcceptedFilter>,
167    pub residual_filters_required: bool,
168    pub estimated_rows: Option<u64>,
169    pub estimated_cost: f64,
170    pub order_satisfied: bool,
171    pub opaque: Arc<dyn Any + Send + Sync>,
172}
173
174#[derive(Debug, Clone, PartialEq, Eq, Hash)]
175pub struct ExternalOrder {
176    pub column_index: usize,
177    pub descending: bool,
178    pub nulls_first: Option<bool>,
179}
180
181#[derive(Debug, Clone, PartialEq)]
182pub enum ExternalFilter {
183    And(Vec<ExternalFilter>),
184    Compare {
185        column_index: usize,
186        op: ExternalFilterOp,
187        value: ScalarValue,
188    },
189    Unsupported {
190        expr: String,
191    },
192}
193
194#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
195pub enum ExternalFilterOp {
196    Eq,
197    NotEq,
198    Gt,
199    GtEq,
200    Lt,
201    LtEq,
202}
203
204#[derive(Debug, Clone, PartialEq, Eq)]
205pub struct AcceptedFilter {
206    pub filter_index: usize,
207    pub pushdown: TableProviderFilterPushDown,
208}
209
210impl ExternalPlan {
211    pub fn new(
212        filter_pushdown: Vec<TableProviderFilterPushDown>,
213        estimated_rows: Option<u64>,
214        estimated_cost: f64,
215        order_satisfied: bool,
216    ) -> Self {
217        let accepted_filters = filter_pushdown
218            .iter()
219            .enumerate()
220            .filter_map(|(filter_index, pushdown)| match pushdown {
221                TableProviderFilterPushDown::Exact | TableProviderFilterPushDown::Inexact => {
222                    Some(AcceptedFilter {
223                        filter_index,
224                        pushdown: pushdown.clone(),
225                    })
226                }
227                TableProviderFilterPushDown::Unsupported => None,
228            })
229            .collect();
230        let residual_filters_required = filter_pushdown.iter().any(|pushdown| {
231            matches!(
232                pushdown,
233                TableProviderFilterPushDown::Inexact | TableProviderFilterPushDown::Unsupported
234            )
235        });
236        Self {
237            filter_pushdown,
238            accepted_filters,
239            residual_filters_required,
240            estimated_rows,
241            estimated_cost,
242            order_satisfied,
243            opaque: Arc::new(()),
244        }
245    }
246
247    pub fn with_opaque(mut self, opaque: Arc<dyn Any + Send + Sync>) -> Self {
248        self.opaque = opaque;
249        self
250    }
251}
252
253fn external_filter_from_expr(expr: &Expr, schema: &ArrowSchema) -> ExternalFilter {
254    match expr {
255        Expr::BinaryExpr(BinaryExpr { left, op, right }) if *op == Operator::And => {
256            ExternalFilter::And(vec![
257                external_filter_from_expr(left, schema),
258                external_filter_from_expr(right, schema),
259            ])
260        }
261        Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
262            let Some(op) = external_filter_op(*op) else {
263                return ExternalFilter::Unsupported {
264                    expr: expr.to_string(),
265                };
266            };
267            match (left.as_ref(), right.as_ref()) {
268                (Expr::Column(column), Expr::Literal(value, _)) => schema
269                    .index_of(&column.name)
270                    .map(|column_index| ExternalFilter::Compare {
271                        column_index,
272                        op,
273                        value: value.clone(),
274                    })
275                    .unwrap_or_else(|_| ExternalFilter::Unsupported {
276                        expr: expr.to_string(),
277                    }),
278                (Expr::Literal(value, _), Expr::Column(column)) => schema
279                    .index_of(&column.name)
280                    .map(|column_index| ExternalFilter::Compare {
281                        column_index,
282                        op: op.flipped(),
283                        value: value.clone(),
284                    })
285                    .unwrap_or_else(|_| ExternalFilter::Unsupported {
286                        expr: expr.to_string(),
287                    }),
288                _ => ExternalFilter::Unsupported {
289                    expr: expr.to_string(),
290                },
291            }
292        }
293        _ => ExternalFilter::Unsupported {
294            expr: expr.to_string(),
295        },
296    }
297}
298
299fn external_filter_op(op: Operator) -> Option<ExternalFilterOp> {
300    match op {
301        Operator::Eq => Some(ExternalFilterOp::Eq),
302        Operator::NotEq => Some(ExternalFilterOp::NotEq),
303        Operator::Gt => Some(ExternalFilterOp::Gt),
304        Operator::GtEq => Some(ExternalFilterOp::GtEq),
305        Operator::Lt => Some(ExternalFilterOp::Lt),
306        Operator::LtEq => Some(ExternalFilterOp::LtEq),
307        _ => None,
308    }
309}
310
311impl ExternalFilterOp {
312    fn flipped(self) -> Self {
313        match self {
314            ExternalFilterOp::Eq => ExternalFilterOp::Eq,
315            ExternalFilterOp::NotEq => ExternalFilterOp::NotEq,
316            ExternalFilterOp::Gt => ExternalFilterOp::Lt,
317            ExternalFilterOp::GtEq => ExternalFilterOp::LtEq,
318            ExternalFilterOp::Lt => ExternalFilterOp::Gt,
319            ExternalFilterOp::LtEq => ExternalFilterOp::GtEq,
320        }
321    }
322}
323
324#[derive(Clone)]
325pub enum ExternalWriteOp {
326    Insert {
327        rows: Vec<HashMap<u16, Value>>,
328    },
329    ReplaceRows {
330        rows: Vec<HashMap<u16, Value>>,
331        changes: u64,
332    },
333}
334
335#[derive(Debug, Clone, PartialEq)]
336pub enum ExternalBaseWrite {
337    Put {
338        table: String,
339        cells: Vec<(u16, Value)>,
340    },
341    Delete {
342        table: String,
343        row_id: u64,
344    },
345}
346
347#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
348pub struct ExternalWriteResult {
349    pub changes: u64,
350}
351
352impl ExternalWriteResult {
353    pub fn new(changes: u64) -> Self {
354        Self { changes }
355    }
356}
357
358pub struct ExternalTxn {
359    state: Vec<u8>,
360    base_writes: Vec<ExternalBaseWrite>,
361}
362
363impl ExternalTxn {
364    pub fn new(state: Vec<u8>) -> Self {
365        Self {
366            state,
367            base_writes: Vec::new(),
368        }
369    }
370
371    pub fn read_state_from_bytes(state: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>> {
372        Ok(decode_kv_state(state)?.remove(key))
373    }
374
375    pub fn read_state(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
376        Self::read_state_from_bytes(&self.state, key)
377    }
378
379    pub fn put_state(&mut self, key: &[u8], value: &[u8]) -> Result<()> {
380        let mut state = decode_kv_state(&self.state)?;
381        state.insert(key.to_vec(), value.to_vec());
382        self.state = encode_kv_state(&state)?;
383        Ok(())
384    }
385
386    pub fn delete_state(&mut self, key: &[u8]) -> Result<()> {
387        let mut state = decode_kv_state(&self.state)?;
388        state.remove(key);
389        self.state = encode_kv_state(&state)?;
390        Ok(())
391    }
392
393    pub fn replace_state(&mut self, state: Vec<u8>) {
394        self.state = state;
395    }
396
397    pub fn read_rows(&self) -> Result<Vec<HashMap<u16, Value>>> {
398        if self.state.is_empty() {
399            Ok(Vec::new())
400        } else {
401            decode_state_rows(&self.state)
402        }
403    }
404
405    pub fn replace_rows(
406        &mut self,
407        schema: &CoreSchema,
408        rows: Vec<HashMap<u16, Value>>,
409    ) -> Result<()> {
410        validate_external_rows(schema, &rows)?;
411        self.state = encode_state_rows(&rows)?;
412        Ok(())
413    }
414
415    pub fn emit_base_write(&mut self, op: ExternalBaseWrite) -> Result<()> {
416        self.base_writes.push(op);
417        Ok(())
418    }
419
420    fn into_parts(self) -> (Vec<u8>, Vec<ExternalBaseWrite>) {
421        (self.state, self.base_writes)
422    }
423}
424
425pub struct ExternalModuleRegistry {
426    modules: parking_lot::RwLock<HashMap<String, Arc<dyn ExternalTableModule>>>,
427}
428
429impl Default for ExternalModuleRegistry {
430    fn default() -> Self {
431        let registry = Self {
432            modules: parking_lot::RwLock::new(HashMap::new()),
433        };
434        for module in builtin_modules() {
435            registry.register(module).expect("valid built-in module");
436        }
437        registry
438    }
439}
440
441impl ExternalModuleRegistry {
442    pub fn register(&self, module: Arc<dyn ExternalTableModule>) -> Result<()> {
443        let name = normalize_module_name(module.name())?;
444        self.modules.write().insert(name, module);
445        Ok(())
446    }
447
448    pub fn names(&self) -> Vec<String> {
449        let mut names = self.modules.read().keys().cloned().collect::<Vec<_>>();
450        names.sort();
451        names
452    }
453
454    pub(crate) fn contains(&self, name: &str) -> bool {
455        normalize_module_name(name)
456            .ok()
457            .is_some_and(|name| self.modules.read().contains_key(&name))
458    }
459
460    pub(crate) fn descriptor(&self, name: &str) -> Result<ExternalModuleDescriptor> {
461        Ok(self.module(name)?.descriptor())
462    }
463
464    pub(crate) fn external_table_provider(
465        &self,
466        db: &Arc<Database>,
467        entry: &ExternalTableEntry,
468    ) -> Result<Arc<dyn TableProvider>> {
469        let module = self.module(&entry.module)?;
470        let table = module.connect(&ModuleConnectCtx { db }, entry)?;
471        Ok(Arc::new(ExternalTableProvider::new(
472            table,
473            entry.capabilities,
474        )))
475    }
476
477    pub(crate) fn external_table_indexes(
478        &self,
479        entry: &ExternalTableEntry,
480    ) -> Result<Vec<ExternalModuleIndex>> {
481        let module = self.module(&entry.module)?;
482        module.indexes(entry)
483    }
484
485    pub(crate) fn external_table_rows(
486        &self,
487        db: &Arc<Database>,
488        entry: &ExternalTableEntry,
489    ) -> Result<Vec<HashMap<u16, Value>>> {
490        let module = self.module(&entry.module)?;
491        module.read_rows(&ModuleConnectCtx { db }, entry)
492    }
493
494    pub(crate) fn external_table_rows_from_state(
495        &self,
496        entry: &ExternalTableEntry,
497        state: &[u8],
498    ) -> Result<Vec<HashMap<u16, Value>>> {
499        let module = self.module(&entry.module)?;
500        module.rows_from_state(state)
501    }
502
503    pub(crate) fn external_table_write(
504        &self,
505        db: &Arc<Database>,
506        entry: &ExternalTableEntry,
507        base_state: Vec<u8>,
508        op: ExternalWriteOp,
509    ) -> Result<(Vec<u8>, ExternalWriteResult, Vec<ExternalBaseWrite>)> {
510        let module = self.module(&entry.module)?;
511        let mut txn = ExternalTxn::new(base_state);
512        let result = module.write(&ModuleConnectCtx { db }, entry, op, &mut txn)?;
513        let (state, base_writes) = txn.into_parts();
514        Ok((state, result, base_writes))
515    }
516
517    pub(crate) fn destroy_external_table(
518        &self,
519        db: &Arc<Database>,
520        entry: &ExternalTableEntry,
521    ) -> Result<()> {
522        let module = self.module(&entry.module)?;
523        module.destroy(&ModuleConnectCtx { db }, entry)
524    }
525
526    fn module(&self, name: &str) -> Result<Arc<dyn ExternalTableModule>> {
527        let name = normalize_module_name(name)?;
528        self.modules.read().get(&name).cloned().ok_or_else(|| {
529            MongrelQueryError::Schema(format!("external table module {name:?} is not registered"))
530        })
531    }
532}
533
534fn normalize_module_name(name: &str) -> Result<String> {
535    let name = name.trim().to_ascii_lowercase();
536    if name.is_empty() || !name.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'_') {
537        return Err(MongrelQueryError::Schema(format!(
538            "invalid external table module name {name:?}"
539        )));
540    }
541    Ok(name)
542}
543
544fn builtin_modules() -> Vec<Arc<dyn ExternalTableModule>> {
545    vec![
546        Arc::new(JsonModule {
547            name: "json_each",
548            mode: JsonTableMode::Each,
549        }),
550        Arc::new(JsonModule {
551            name: "json_tree",
552            mode: JsonTableMode::Tree,
553        }),
554        Arc::new(JsonModule {
555            name: "jsonb_each",
556            mode: JsonTableMode::Each,
557        }),
558        Arc::new(JsonModule {
559            name: "jsonb_tree",
560            mode: JsonTableMode::Tree,
561        }),
562        Arc::new(SchemaTablesModule),
563        Arc::new(DbStatModule),
564        Arc::new(FtsDocsModule),
565        Arc::new(KvStoreModule),
566        Arc::new(RTreeRectsModule),
567        Arc::new(SeriesModule),
568    ]
569}
570
571struct ExternalTableProvider {
572    table: Arc<dyn ExternalTable>,
573    cache_plans: bool,
574    plan_cache: parking_lot::Mutex<HashMap<ExternalPlanCacheKey, ExternalPlan>>,
575}
576
577struct ExternalTableExec {
578    props: Arc<PlanProperties>,
579    schema: SchemaRef,
580    table: Arc<dyn ExternalTable>,
581    projection: Option<Vec<usize>>,
582    filters: Vec<Expr>,
583    limit: Option<usize>,
584}
585
586impl std::fmt::Debug for ExternalTableExec {
587    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
588        f.debug_struct("ExternalTableExec").finish_non_exhaustive()
589    }
590}
591
592impl DisplayAs for ExternalTableExec {
593    fn fmt_as(
594        &self,
595        _format: DisplayFormatType,
596        f: &mut std::fmt::Formatter<'_>,
597    ) -> std::fmt::Result {
598        write!(f, "ExternalTableExec")
599    }
600}
601
602impl ExecutionPlan for ExternalTableExec {
603    fn name(&self) -> &str {
604        "ExternalTableExec"
605    }
606
607    fn properties(&self) -> &Arc<PlanProperties> {
608        &self.props
609    }
610
611    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
612        Vec::new()
613    }
614
615    fn with_new_children(
616        self: Arc<Self>,
617        children: Vec<Arc<dyn ExecutionPlan>>,
618    ) -> DFResult<Arc<dyn ExecutionPlan>> {
619        if children.is_empty() {
620            Ok(self)
621        } else {
622            Err(DataFusionError::Internal(
623                "ExternalTableExec is a leaf node and has no children".into(),
624            ))
625        }
626    }
627
628    fn execute(
629        &self,
630        partition: usize,
631        task_context: Arc<TaskContext>,
632    ) -> DFResult<SendableRecordBatchStream> {
633        use futures::TryStreamExt;
634
635        if partition != 0 {
636            return Err(DataFusionError::Internal(format!(
637                "ExternalTableExec is single-partition; invalid partition {partition}"
638            )));
639        }
640        let query = task_context
641            .session_config()
642            .get_extension::<SqlTaskContext>()
643            .map(|context| context.query().clone());
644        let control = query
645            .as_ref()
646            .map(|query| query.control().child_with_deadline(None))
647            .unwrap_or_else(|| mongreldb_core::ExecutionControl::new(None));
648        control
649            .checkpoint()
650            .map_err(|error| external_control_error(query.as_ref(), error))?;
651        let table = Arc::clone(&self.table);
652        let projection = self.projection.clone();
653        let filters = self.filters.clone();
654        let limit = self.limit;
655        let worker_control = control.clone();
656        let future = async move {
657            let mut task = tokio::task::spawn_blocking(move || {
658                let schema = table.schema();
659                let request = ExternalPlanRequest {
660                    projection,
661                    filters: filters
662                        .iter()
663                        .map(|filter| external_filter_from_expr(filter, &schema))
664                        .collect(),
665                    raw_filters: filters.iter().collect(),
666                    order_by: Vec::new(),
667                    limit,
668                    offset: None,
669                };
670                table.scan_with_control(&request, &worker_control)
671            });
672            let scan = tokio::select! {
673                result = &mut task => result.map_err(|error| {
674                    DataFusionError::Execution(format!("external table worker failed: {error}"))
675                })??,
676                _ = control.cancelled() => {
677                    if tokio::time::timeout(std::time::Duration::from_millis(100), &mut task)
678                        .await
679                        .is_err()
680                    {
681                        eprintln!("external table worker exceeded cancellation grace");
682                    }
683                    return Err(external_control_error_from_control(query.as_ref(), &control));
684                }
685            };
686            control
687                .checkpoint()
688                .map_err(|error| external_control_error(query.as_ref(), error))?;
689            Ok(futures::stream::iter(scan.batches.into_iter().map(Ok)))
690        };
691        let stream = futures::stream::once(future).try_flatten();
692        Ok(Box::pin(RecordBatchStreamAdapter::new(
693            Arc::clone(&self.schema),
694            stream,
695        )))
696    }
697}
698
699fn external_control_error(
700    query: Option<&crate::RegisteredSqlQuery>,
701    error: mongreldb_core::MongrelError,
702) -> DataFusionError {
703    if let Some(error) = query.and_then(|query| query.checkpoint().err()) {
704        DataFusionError::External(Box::new(error))
705    } else {
706        DataFusionError::Execution(error.to_string())
707    }
708}
709
710fn external_control_error_from_control(
711    query: Option<&crate::RegisteredSqlQuery>,
712    control: &mongreldb_core::ExecutionControl,
713) -> DataFusionError {
714    external_control_error(
715        query,
716        control
717            .checkpoint()
718            .expect_err("cancelled external control must fail its checkpoint"),
719    )
720}
721
722#[derive(Debug, Clone, PartialEq, Eq, Hash)]
723struct ExternalPlanCacheKey {
724    projection: Option<Vec<usize>>,
725    filters: Vec<String>,
726    order_by: Vec<ExternalOrder>,
727    limit: Option<usize>,
728    offset: Option<usize>,
729}
730
731impl ExternalPlanCacheKey {
732    fn from_request(request: &ExternalPlanRequest<'_>) -> Self {
733        Self {
734            projection: request.projection.clone(),
735            filters: request
736                .filters
737                .iter()
738                .map(|filter| format!("{filter:?}"))
739                .collect(),
740            order_by: request.order_by.clone(),
741            limit: request.limit,
742            offset: request.offset,
743        }
744    }
745}
746
747impl ExternalTableProvider {
748    fn new(table: Arc<dyn ExternalTable>, capabilities: ModuleCapabilities) -> Self {
749        Self {
750            table,
751            cache_plans: capabilities.read_only && capabilities.deterministic,
752            plan_cache: parking_lot::Mutex::new(HashMap::new()),
753        }
754    }
755
756    fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
757        if !self.cache_plans {
758            return self.table.plan(request);
759        }
760        let key = ExternalPlanCacheKey::from_request(request);
761        if let Some(plan) = self.plan_cache.lock().get(&key).cloned() {
762            return Ok(plan);
763        }
764        let plan = self.table.plan(request)?;
765        self.plan_cache.lock().insert(key, plan.clone());
766        Ok(plan)
767    }
768}
769
770impl std::fmt::Debug for ExternalTableProvider {
771    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
772        f.debug_struct("ExternalTableProvider")
773            .field("table", &self.table)
774            .finish()
775    }
776}
777
778#[async_trait::async_trait]
779impl TableProvider for ExternalTableProvider {
780    fn schema(&self) -> SchemaRef {
781        self.table.schema()
782    }
783
784    fn table_type(&self) -> TableType {
785        TableType::Base
786    }
787
788    fn supports_filters_pushdown(
789        &self,
790        filters: &[&Expr],
791    ) -> DFResult<Vec<TableProviderFilterPushDown>> {
792        let schema = self.table.schema();
793        let request = ExternalPlanRequest {
794            projection: None,
795            filters: filters
796                .iter()
797                .map(|filter| external_filter_from_expr(filter, &schema))
798                .collect(),
799            raw_filters: filters.to_vec(),
800            order_by: Vec::new(),
801            limit: None,
802            offset: None,
803        };
804        let ExternalPlan {
805            filter_pushdown,
806            accepted_filters: _accepted_filters,
807            residual_filters_required: _residual_filters_required,
808            estimated_rows: _estimated_rows,
809            estimated_cost: _estimated_cost,
810            order_satisfied: _order_satisfied,
811            opaque: _opaque,
812        } = self.plan(&request)?;
813        Ok(filter_pushdown)
814    }
815
816    async fn scan(
817        &self,
818        _state: &dyn Session,
819        projection: Option<&Vec<usize>>,
820        filters: &[Expr],
821        limit: Option<usize>,
822    ) -> DFResult<Arc<dyn ExecutionPlan>> {
823        mongreldb_core::trace::QueryTrace::record(|t| {
824            t.scan_mode = mongreldb_core::trace::ScanMode::ExternalModule;
825        });
826        let schema = self.table.schema();
827        let request = ExternalPlanRequest {
828            projection: projection.cloned(),
829            filters: filters
830                .iter()
831                .map(|filter| external_filter_from_expr(filter, &schema))
832                .collect(),
833            raw_filters: filters.iter().collect(),
834            order_by: Vec::new(),
835            limit,
836            offset: None,
837        };
838        let _plan = self.plan(&request)?;
839        let output_schema = match projection {
840            Some(projection) => Arc::new(schema.project(projection)?),
841            None => schema,
842        };
843        let props = Arc::new(PlanProperties::new(
844            EquivalenceProperties::new(Arc::clone(&output_schema)),
845            Partitioning::UnknownPartitioning(1),
846            EmissionType::Incremental,
847            Boundedness::Bounded,
848        ));
849        Ok(Arc::new(ExternalTableExec {
850            props,
851            schema: output_schema,
852            table: Arc::clone(&self.table),
853            projection: projection.cloned(),
854            filters: filters.to_vec(),
855            limit,
856        }))
857    }
858}
859
860struct JsonModule {
861    name: &'static str,
862    mode: JsonTableMode,
863}
864
865impl ExternalTableModule for JsonModule {
866    fn name(&self) -> &str {
867        self.name
868    }
869
870    fn descriptor(&self) -> ExternalModuleDescriptor {
871        json_descriptor()
872    }
873
874    fn connect(
875        &self,
876        _ctx: &ModuleConnectCtx<'_>,
877        entry: &ExternalTableEntry,
878    ) -> Result<Arc<dyn ExternalTable>> {
879        let (json, root) = json_args(entry, self.name)?;
880        let (schema, batches) =
881            json_table_batches_from_text(self.name, &json, root.as_deref(), self.mode)
882                .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
883        Ok(Arc::new(JsonExternalTable {
884            module: self.name,
885            schema,
886            batches,
887        }))
888    }
889}
890
891#[derive(Clone)]
892struct JsonExternalTable {
893    module: &'static str,
894    schema: SchemaRef,
895    batches: Vec<RecordBatch>,
896}
897
898impl std::fmt::Debug for JsonExternalTable {
899    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
900        f.debug_struct("JsonExternalTable")
901            .field("module", &self.module)
902            .field("schema", &self.schema)
903            .finish_non_exhaustive()
904    }
905}
906
907impl ExternalTable for JsonExternalTable {
908    fn schema(&self) -> SchemaRef {
909        self.schema.clone()
910    }
911
912    fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
913        Ok(ExternalPlan::new(
914            request
915                .filters
916                .iter()
917                .map(|_| TableProviderFilterPushDown::Unsupported)
918                .collect(),
919            None,
920            1.0,
921            false,
922        ))
923    }
924
925    fn scan_with_control(
926        &self,
927        request: &ExternalPlanRequest<'_>,
928        control: &mongreldb_core::ExecutionControl,
929    ) -> DFResult<ExternalScan> {
930        project_scan_with_control(
931            self.schema.clone(),
932            self.batches.clone(),
933            request.projection.as_deref(),
934            request.limit,
935            Some(control),
936        )
937    }
938}
939
940fn json_descriptor() -> ExternalModuleDescriptor {
941    ExternalModuleDescriptor {
942        schema: CoreSchema {
943            schema_id: 0,
944            columns: vec![
945                json_column(1, "key", TypeId::Bytes, true),
946                json_column(2, "value", TypeId::Bytes, true),
947                json_column(3, "type", TypeId::Bytes, true),
948                json_column(4, "atom", TypeId::Bytes, true),
949                json_column(5, "id", TypeId::Int64, false),
950                json_column(6, "parent", TypeId::Int64, true),
951                json_column(7, "fullkey", TypeId::Bytes, true),
952                json_column(8, "path", TypeId::Bytes, true),
953                json_column(9, "json", TypeId::Bytes, true),
954                json_column(10, "root", TypeId::Bytes, true),
955            ],
956            indexes: Vec::new(),
957            colocation: Vec::new(),
958            constraints: Default::default(),
959            clustered: false,
960        },
961        hidden_columns: vec!["json".to_string(), "root".to_string()],
962        capabilities: ModuleCapabilities {
963            read_only: true,
964            deterministic: true,
965            trigger_safe: true,
966            ..ModuleCapabilities::default()
967        },
968    }
969}
970
971fn json_column(id: u16, name: &str, ty: TypeId, nullable: bool) -> CoreColumnDef {
972    CoreColumnDef {
973        id,
974        name: name.to_string(),
975        ty,
976        flags: if nullable {
977            ColumnFlags::empty().with(ColumnFlags::NULLABLE)
978        } else {
979            ColumnFlags::empty()
980        },
981        default_value: None,
982    }
983}
984
985fn json_args(entry: &ExternalTableEntry, module: &str) -> Result<(String, Option<String>)> {
986    match entry.args.as_slice() {
987        [json] => Ok((module_arg_string(json).to_string(), None)),
988        [json, root] => Ok((
989            module_arg_string(json).to_string(),
990            Some(module_arg_string(root).to_string()),
991        )),
992        _ => Err(MongrelQueryError::Schema(format!(
993            "{module} external table requires one JSON argument and an optional root path"
994        ))),
995    }
996}
997
998fn module_arg_string(arg: &ModuleArg) -> &str {
999    match arg {
1000        ModuleArg::Ident(value) | ModuleArg::String(value) | ModuleArg::Number(value) => value,
1001    }
1002}
1003
1004struct SchemaTablesModule;
1005
1006impl ExternalTableModule for SchemaTablesModule {
1007    fn name(&self) -> &str {
1008        "schema_tables"
1009    }
1010
1011    fn descriptor(&self) -> ExternalModuleDescriptor {
1012        ExternalModuleDescriptor {
1013            schema: CoreSchema {
1014                schema_id: 0,
1015                columns: vec![
1016                    catalog_column(1, "schema_name", TypeId::Bytes, false),
1017                    catalog_column(2, "name", TypeId::Bytes, false),
1018                    catalog_column(3, "type", TypeId::Bytes, false),
1019                    catalog_column(4, "ncol", TypeId::Int64, false),
1020                    catalog_column(5, "module", TypeId::Bytes, true),
1021                    catalog_column(6, "created_epoch", TypeId::Int64, false),
1022                ],
1023                indexes: Vec::new(),
1024                colocation: Vec::new(),
1025                constraints: Default::default(),
1026                clustered: false,
1027            },
1028            hidden_columns: Vec::new(),
1029            capabilities: catalog_capabilities(),
1030        }
1031    }
1032
1033    fn connect(
1034        &self,
1035        ctx: &ModuleConnectCtx<'_>,
1036        entry: &ExternalTableEntry,
1037    ) -> Result<Arc<dyn ExternalTable>> {
1038        ensure_no_args(entry, self.name())?;
1039        Ok(Arc::new(SchemaTablesExternalTable {
1040            db: Arc::clone(ctx.db),
1041            schema: schema_tables_arrow_schema(),
1042        }))
1043    }
1044}
1045
1046#[derive(Clone)]
1047struct SchemaTablesExternalTable {
1048    db: Arc<Database>,
1049    schema: SchemaRef,
1050}
1051
1052impl std::fmt::Debug for SchemaTablesExternalTable {
1053    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1054        f.debug_struct("SchemaTablesExternalTable")
1055            .field("schema", &self.schema)
1056            .finish_non_exhaustive()
1057    }
1058}
1059
1060impl ExternalTable for SchemaTablesExternalTable {
1061    fn schema(&self) -> SchemaRef {
1062        self.schema.clone()
1063    }
1064
1065    fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
1066        unsupported_plan(request)
1067    }
1068
1069    fn scan_with_control(
1070        &self,
1071        request: &ExternalPlanRequest<'_>,
1072        control: &mongreldb_core::ExecutionControl,
1073    ) -> DFResult<ExternalScan> {
1074        external_checkpoint(Some(control), 0)?;
1075        let batches = schema_tables_batches(&self.db, self.schema.clone(), Some(control))?;
1076        project_scan_with_control(
1077            self.schema.clone(),
1078            batches,
1079            request.projection.as_deref(),
1080            request.limit,
1081            Some(control),
1082        )
1083    }
1084}
1085
1086struct DbStatModule;
1087
1088impl ExternalTableModule for DbStatModule {
1089    fn name(&self) -> &str {
1090        "dbstat"
1091    }
1092
1093    fn descriptor(&self) -> ExternalModuleDescriptor {
1094        ExternalModuleDescriptor {
1095            schema: CoreSchema {
1096                schema_id: 0,
1097                columns: vec![
1098                    catalog_column(1, "name", TypeId::Bytes, false),
1099                    catalog_column(2, "type", TypeId::Bytes, false),
1100                    catalog_column(3, "rows", TypeId::Int64, true),
1101                    catalog_column(4, "runs", TypeId::Int64, true),
1102                    catalog_column(5, "memtable_rows", TypeId::Int64, true),
1103                    catalog_column(6, "columns", TypeId::Int64, false),
1104                    catalog_column(7, "storage_bytes", TypeId::Int64, false),
1105                ],
1106                indexes: Vec::new(),
1107                colocation: Vec::new(),
1108                constraints: Default::default(),
1109                clustered: false,
1110            },
1111            hidden_columns: Vec::new(),
1112            capabilities: catalog_capabilities(),
1113        }
1114    }
1115
1116    fn connect(
1117        &self,
1118        ctx: &ModuleConnectCtx<'_>,
1119        entry: &ExternalTableEntry,
1120    ) -> Result<Arc<dyn ExternalTable>> {
1121        ensure_no_args(entry, self.name())?;
1122        Ok(Arc::new(DbStatExternalTable {
1123            db: Arc::clone(ctx.db),
1124            schema: dbstat_arrow_schema(),
1125        }))
1126    }
1127}
1128
1129#[derive(Clone)]
1130struct DbStatExternalTable {
1131    db: Arc<Database>,
1132    schema: SchemaRef,
1133}
1134
1135impl std::fmt::Debug for DbStatExternalTable {
1136    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1137        f.debug_struct("DbStatExternalTable")
1138            .field("schema", &self.schema)
1139            .finish_non_exhaustive()
1140    }
1141}
1142
1143impl ExternalTable for DbStatExternalTable {
1144    fn schema(&self) -> SchemaRef {
1145        self.schema.clone()
1146    }
1147
1148    fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
1149        unsupported_plan(request)
1150    }
1151
1152    fn scan_with_control(
1153        &self,
1154        request: &ExternalPlanRequest<'_>,
1155        control: &mongreldb_core::ExecutionControl,
1156    ) -> DFResult<ExternalScan> {
1157        external_checkpoint(Some(control), 0)?;
1158        let batches = dbstat_batches(&self.db, self.schema.clone(), Some(control))?;
1159        project_scan_with_control(
1160            self.schema.clone(),
1161            batches,
1162            request.projection.as_deref(),
1163            request.limit,
1164            Some(control),
1165        )
1166    }
1167}
1168
1169fn catalog_column(id: u16, name: &str, ty: TypeId, nullable: bool) -> CoreColumnDef {
1170    CoreColumnDef {
1171        id,
1172        name: name.to_string(),
1173        ty,
1174        flags: if nullable {
1175            ColumnFlags::empty().with(ColumnFlags::NULLABLE)
1176        } else {
1177            ColumnFlags::empty()
1178        },
1179        default_value: None,
1180    }
1181}
1182
1183fn catalog_capabilities() -> ModuleCapabilities {
1184    ModuleCapabilities {
1185        read_only: true,
1186        trigger_safe: true,
1187        ..ModuleCapabilities::default()
1188    }
1189}
1190
1191fn ensure_no_args(entry: &ExternalTableEntry, module: &str) -> Result<()> {
1192    if entry.args.is_empty() {
1193        Ok(())
1194    } else {
1195        Err(MongrelQueryError::Schema(format!(
1196            "{module} external table does not accept arguments"
1197        )))
1198    }
1199}
1200
1201fn schema_tables_arrow_schema() -> SchemaRef {
1202    Arc::new(ArrowSchema::new(vec![
1203        Field::new("schema_name", DataType::Utf8, false),
1204        Field::new("name", DataType::Utf8, false),
1205        Field::new("type", DataType::Utf8, false),
1206        Field::new("ncol", DataType::Int64, false),
1207        Field::new("module", DataType::Utf8, true),
1208        Field::new("created_epoch", DataType::Int64, false),
1209    ]))
1210}
1211
1212fn schema_tables_batches(
1213    db: &Arc<Database>,
1214    schema: SchemaRef,
1215    control: Option<&mongreldb_core::ExecutionControl>,
1216) -> DFResult<Vec<RecordBatch>> {
1217    struct Row {
1218        name: String,
1219        ty: String,
1220        ncol: i64,
1221        module: Option<String>,
1222        created_epoch: i64,
1223    }
1224
1225    let catalog = db.catalog_snapshot();
1226    let mut rows = Vec::new();
1227    for (index, table) in catalog
1228        .tables
1229        .into_iter()
1230        .filter(|table| matches!(table.state, TableState::Live))
1231        .enumerate()
1232    {
1233        external_checkpoint(control, index)?;
1234        rows.push(Row {
1235            name: table.name,
1236            ty: "table".to_string(),
1237            ncol: saturating_i64(table.schema.columns.len() as u64),
1238            module: None,
1239            created_epoch: saturating_i64(table.created_epoch),
1240        });
1241    }
1242    for (index, table) in catalog.external_tables.into_iter().enumerate() {
1243        external_checkpoint(control, index)?;
1244        rows.push(Row {
1245            name: table.name,
1246            ty: "external".to_string(),
1247            ncol: saturating_i64(table.declared_schema.columns.len() as u64),
1248            module: Some(table.module),
1249            created_epoch: saturating_i64(table.created_epoch),
1250        });
1251    }
1252    rows.sort_by(|a, b| a.name.cmp(&b.name).then_with(|| a.ty.cmp(&b.ty)));
1253
1254    let schema_names = vec!["main".to_string(); rows.len()];
1255    let names = rows.iter().map(|row| row.name.clone()).collect::<Vec<_>>();
1256    let types = rows.iter().map(|row| row.ty.clone()).collect::<Vec<_>>();
1257    let ncols = rows.iter().map(|row| row.ncol).collect::<Vec<_>>();
1258    let modules = rows
1259        .iter()
1260        .map(|row| row.module.clone())
1261        .collect::<Vec<_>>();
1262    let created_epochs = rows.iter().map(|row| row.created_epoch).collect::<Vec<_>>();
1263
1264    Ok(vec![RecordBatch::try_new(
1265        schema,
1266        vec![
1267            Arc::new(StringArray::from(schema_names)) as ArrayRef,
1268            Arc::new(StringArray::from(names)),
1269            Arc::new(StringArray::from(types)),
1270            Arc::new(Int64Array::from(ncols)),
1271            Arc::new(StringArray::from(modules)),
1272            Arc::new(Int64Array::from(created_epochs)),
1273        ],
1274    )?])
1275}
1276
1277fn dbstat_arrow_schema() -> SchemaRef {
1278    Arc::new(ArrowSchema::new(vec![
1279        Field::new("name", DataType::Utf8, false),
1280        Field::new("type", DataType::Utf8, false),
1281        Field::new("rows", DataType::Int64, true),
1282        Field::new("runs", DataType::Int64, true),
1283        Field::new("memtable_rows", DataType::Int64, true),
1284        Field::new("columns", DataType::Int64, false),
1285        Field::new("storage_bytes", DataType::Int64, false),
1286    ]))
1287}
1288
1289fn dbstat_batches(
1290    db: &Arc<Database>,
1291    schema: SchemaRef,
1292    control: Option<&mongreldb_core::ExecutionControl>,
1293) -> DFResult<Vec<RecordBatch>> {
1294    struct Row {
1295        name: String,
1296        ty: String,
1297        rows: Option<i64>,
1298        runs: Option<i64>,
1299        memtable_rows: Option<i64>,
1300        columns: i64,
1301        storage_bytes: i64,
1302    }
1303
1304    let catalog = db.catalog_snapshot();
1305    let mut rows = Vec::new();
1306    for (index, table) in catalog
1307        .tables
1308        .into_iter()
1309        .filter(|table| matches!(table.state, TableState::Live))
1310        .enumerate()
1311    {
1312        external_checkpoint(control, index)?;
1313        let handle = db
1314            .table(&table.name)
1315            .map_err(|e| DataFusionError::Execution(e.to_string()))?;
1316        let table_guard = handle.lock();
1317        let table_dir = db.root().join(TABLES_DIR).join(table.table_id.to_string());
1318        rows.push(Row {
1319            name: table.name,
1320            ty: "table".to_string(),
1321            rows: Some(saturating_i64(table_guard.count())),
1322            runs: Some(saturating_i64(table_guard.run_count() as u64)),
1323            memtable_rows: Some(saturating_i64(table_guard.memtable_len() as u64)),
1324            columns: saturating_i64(table.schema.columns.len() as u64),
1325            storage_bytes: saturating_i64(dir_size(&table_dir, control)?),
1326        });
1327    }
1328    for (index, table) in catalog.external_tables.into_iter().enumerate() {
1329        external_checkpoint(control, index)?;
1330        let state_dir = db.root().join(VTAB_DIR).join(&table.name);
1331        rows.push(Row {
1332            name: table.name,
1333            ty: "external".to_string(),
1334            rows: None,
1335            runs: None,
1336            memtable_rows: None,
1337            columns: saturating_i64(table.declared_schema.columns.len() as u64),
1338            storage_bytes: saturating_i64(dir_size(&state_dir, control)?),
1339        });
1340    }
1341    rows.sort_by(|a, b| a.name.cmp(&b.name).then_with(|| a.ty.cmp(&b.ty)));
1342
1343    let names = rows.iter().map(|row| row.name.clone()).collect::<Vec<_>>();
1344    let types = rows.iter().map(|row| row.ty.clone()).collect::<Vec<_>>();
1345    let row_counts = rows.iter().map(|row| row.rows).collect::<Vec<_>>();
1346    let runs = rows.iter().map(|row| row.runs).collect::<Vec<_>>();
1347    let memtable_rows = rows.iter().map(|row| row.memtable_rows).collect::<Vec<_>>();
1348    let columns = rows.iter().map(|row| row.columns).collect::<Vec<_>>();
1349    let storage_bytes = rows.iter().map(|row| row.storage_bytes).collect::<Vec<_>>();
1350
1351    Ok(vec![RecordBatch::try_new(
1352        schema,
1353        vec![
1354            Arc::new(StringArray::from(names)) as ArrayRef,
1355            Arc::new(StringArray::from(types)),
1356            Arc::new(Int64Array::from(row_counts)),
1357            Arc::new(Int64Array::from(runs)),
1358            Arc::new(Int64Array::from(memtable_rows)),
1359            Arc::new(Int64Array::from(columns)),
1360            Arc::new(Int64Array::from(storage_bytes)),
1361        ],
1362    )?])
1363}
1364
1365fn dir_size(path: &Path, control: Option<&mongreldb_core::ExecutionControl>) -> DFResult<u64> {
1366    external_checkpoint(control, 0)?;
1367    if !path.exists() {
1368        return Ok(0);
1369    }
1370    let metadata = std::fs::symlink_metadata(path)
1371        .map_err(|e| DataFusionError::Execution(format!("stat {:?}: {e}", path)))?;
1372    if metadata.is_file() {
1373        return Ok(metadata.len());
1374    }
1375    if !metadata.is_dir() {
1376        return Ok(0);
1377    }
1378
1379    let mut total = 0_u64;
1380    for (index, entry) in std::fs::read_dir(path)
1381        .map_err(|e| DataFusionError::Execution(format!("{e}")))?
1382        .enumerate()
1383    {
1384        external_checkpoint(control, index)?;
1385        let entry = entry.map_err(|e| DataFusionError::Execution(format!("{e}")))?;
1386        total = total.saturating_add(dir_size(&entry.path(), control)?);
1387    }
1388    Ok(total)
1389}
1390
1391fn saturating_i64(value: u64) -> i64 {
1392    i64::try_from(value).unwrap_or(i64::MAX)
1393}
1394
1395fn unsupported_plan(request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
1396    Ok(ExternalPlan::new(
1397        request
1398            .filters
1399            .iter()
1400            .map(|_| TableProviderFilterPushDown::Unsupported)
1401            .collect(),
1402        None,
1403        1.0,
1404        false,
1405    ))
1406}
1407
1408#[cfg(test)]
1409fn project_scan(
1410    full_schema: SchemaRef,
1411    batches: Vec<RecordBatch>,
1412    projection: Option<&[usize]>,
1413    limit: Option<usize>,
1414) -> DFResult<ExternalScan> {
1415    project_scan_with_control(full_schema, batches, projection, limit, None)
1416}
1417
1418fn project_scan_with_control(
1419    full_schema: SchemaRef,
1420    batches: Vec<RecordBatch>,
1421    projection: Option<&[usize]>,
1422    limit: Option<usize>,
1423    control: Option<&mongreldb_core::ExecutionControl>,
1424) -> DFResult<ExternalScan> {
1425    external_checkpoint(control, 0)?;
1426    let Some(projection) = projection else {
1427        return Ok(ExternalScan {
1428            schema: full_schema,
1429            batches: limit_batches(batches, limit, control)?,
1430        });
1431    };
1432    let schema = Arc::new(ArrowSchema::new(
1433        projection
1434            .iter()
1435            .map(|idx| full_schema.field(*idx).clone())
1436            .collect::<Vec<_>>(),
1437    ));
1438    let projected = batches
1439        .into_iter()
1440        .enumerate()
1441        .map(|(index, batch)| {
1442            external_checkpoint(control, index)?;
1443            let columns = projection
1444                .iter()
1445                .map(|idx| batch.column(*idx).clone())
1446                .collect::<Vec<_>>();
1447            let batch = if columns.is_empty() {
1448                RecordBatch::try_new_with_options(
1449                    schema.clone(),
1450                    columns,
1451                    &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())),
1452                )
1453            } else {
1454                RecordBatch::try_new(schema.clone(), columns)
1455            };
1456            batch.map_err(DataFusionError::from)
1457        })
1458        .collect::<DFResult<Vec<_>>>()?;
1459    Ok(ExternalScan {
1460        schema,
1461        batches: limit_batches(projected, limit, control)?,
1462    })
1463}
1464
1465fn limit_batches(
1466    batches: Vec<RecordBatch>,
1467    limit: Option<usize>,
1468    control: Option<&mongreldb_core::ExecutionControl>,
1469) -> DFResult<Vec<RecordBatch>> {
1470    let Some(mut remaining) = limit else {
1471        return Ok(batches);
1472    };
1473    let mut limited = Vec::new();
1474    for (index, batch) in batches.into_iter().enumerate() {
1475        external_checkpoint(control, index)?;
1476        if remaining == 0 {
1477            break;
1478        }
1479        let take = remaining.min(batch.num_rows());
1480        limited.push(batch.slice(0, take));
1481        remaining -= take;
1482    }
1483    Ok(limited)
1484}
1485
1486#[inline]
1487fn external_checkpoint(
1488    control: Option<&mongreldb_core::ExecutionControl>,
1489    index: usize,
1490) -> DFResult<()> {
1491    if index % 256 == 0 {
1492        control
1493            .map(mongreldb_core::ExecutionControl::checkpoint)
1494            .transpose()
1495            .map_err(|error| DataFusionError::Execution(error.to_string()))?;
1496    }
1497    Ok(())
1498}
1499
1500#[cfg(test)]
1501mod tests {
1502    use super::*;
1503
1504    #[test]
1505    fn project_scan_applies_limit_after_projection() {
1506        let schema = Arc::new(ArrowSchema::new(vec![
1507            Field::new("id", DataType::Int64, false),
1508            Field::new("value", DataType::Int64, false),
1509        ]));
1510        let batch = RecordBatch::try_new(
1511            schema.clone(),
1512            vec![
1513                Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef,
1514                Arc::new(Int64Array::from(vec![10, 11, 12, 13])) as ArrayRef,
1515            ],
1516        )
1517        .unwrap();
1518
1519        let scan = project_scan(schema, vec![batch], Some(&[1]), Some(2)).unwrap();
1520
1521        assert_eq!(scan.schema.fields().len(), 1);
1522        assert_eq!(scan.schema.field(0).name(), "value");
1523        assert_eq!(scan.batches.len(), 1);
1524        assert_eq!(scan.batches[0].num_columns(), 1);
1525        assert_eq!(scan.batches[0].num_rows(), 2);
1526        let values = scan.batches[0]
1527            .column(0)
1528            .as_any()
1529            .downcast_ref::<Int64Array>()
1530            .unwrap();
1531        assert_eq!(values.value(0), 10);
1532        assert_eq!(values.value(1), 11);
1533    }
1534
1535    #[test]
1536    fn external_plan_derives_accepted_and_residual_filter_metadata() {
1537        let plan = ExternalPlan::new(
1538            vec![
1539                TableProviderFilterPushDown::Exact,
1540                TableProviderFilterPushDown::Inexact,
1541                TableProviderFilterPushDown::Unsupported,
1542            ],
1543            Some(42),
1544            3.5,
1545            true,
1546        );
1547
1548        assert_eq!(
1549            plan.accepted_filters,
1550            vec![
1551                AcceptedFilter {
1552                    filter_index: 0,
1553                    pushdown: TableProviderFilterPushDown::Exact,
1554                },
1555                AcceptedFilter {
1556                    filter_index: 1,
1557                    pushdown: TableProviderFilterPushDown::Inexact,
1558                },
1559            ]
1560        );
1561        assert!(plan.residual_filters_required);
1562        assert_eq!(plan.estimated_rows, Some(42));
1563        assert_eq!(plan.estimated_cost, 3.5);
1564        assert!(plan.order_satisfied);
1565    }
1566}
1567
1568struct KvStoreModule;
1569
1570impl ExternalTableModule for KvStoreModule {
1571    fn name(&self) -> &str {
1572        "kv_store"
1573    }
1574
1575    fn descriptor(&self) -> ExternalModuleDescriptor {
1576        ExternalModuleDescriptor {
1577            schema: kv_store_schema(),
1578            hidden_columns: Vec::new(),
1579            capabilities: ModuleCapabilities {
1580                writable: true,
1581                deterministic: true,
1582                trigger_safe: true,
1583                transaction_safe: true,
1584                ..ModuleCapabilities::default()
1585            },
1586        }
1587    }
1588
1589    fn connect(
1590        &self,
1591        ctx: &ModuleConnectCtx<'_>,
1592        entry: &ExternalTableEntry,
1593    ) -> Result<Arc<dyn ExternalTable>> {
1594        ensure_no_args(entry, self.name())?;
1595        let rows = read_state_rows(ctx.db, entry)?;
1596        let schema = arrow_conv::arrow_schema(&entry.declared_schema)?;
1597        let batches = core_rows_to_batches(&entry.declared_schema, rows)?;
1598        Ok(Arc::new(KvStoreExternalTable { schema, batches }))
1599    }
1600
1601    fn read_rows(
1602        &self,
1603        ctx: &ModuleConnectCtx<'_>,
1604        entry: &ExternalTableEntry,
1605    ) -> Result<Vec<HashMap<u16, Value>>> {
1606        ensure_no_args(entry, self.name())?;
1607        read_state_rows(ctx.db, entry)
1608    }
1609
1610    fn prepare_rows(
1611        &self,
1612        _ctx: &ModuleConnectCtx<'_>,
1613        entry: &ExternalTableEntry,
1614        rows: Vec<HashMap<u16, Value>>,
1615    ) -> Result<Vec<u8>> {
1616        ensure_no_args(entry, self.name())?;
1617        validate_external_rows(&entry.declared_schema, &rows)?;
1618        encode_state_rows(&rows)
1619    }
1620}
1621
1622#[derive(Clone)]
1623struct KvStoreExternalTable {
1624    schema: SchemaRef,
1625    batches: Vec<RecordBatch>,
1626}
1627
1628impl std::fmt::Debug for KvStoreExternalTable {
1629    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1630        f.debug_struct("KvStoreExternalTable")
1631            .field("schema", &self.schema)
1632            .finish_non_exhaustive()
1633    }
1634}
1635
1636impl ExternalTable for KvStoreExternalTable {
1637    fn schema(&self) -> SchemaRef {
1638        self.schema.clone()
1639    }
1640
1641    fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
1642        unsupported_plan(request)
1643    }
1644
1645    fn scan_with_control(
1646        &self,
1647        request: &ExternalPlanRequest<'_>,
1648        control: &mongreldb_core::ExecutionControl,
1649    ) -> DFResult<ExternalScan> {
1650        project_scan_with_control(
1651            self.schema.clone(),
1652            self.batches.clone(),
1653            request.projection.as_deref(),
1654            request.limit,
1655            Some(control),
1656        )
1657    }
1658}
1659
1660fn kv_store_schema() -> CoreSchema {
1661    CoreSchema {
1662        schema_id: 0,
1663        columns: vec![
1664            CoreColumnDef {
1665                id: 1,
1666                name: "key".to_string(),
1667                ty: TypeId::Bytes,
1668                flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
1669                default_value: None,
1670            },
1671            CoreColumnDef {
1672                id: 2,
1673                name: "value".to_string(),
1674                ty: TypeId::Bytes,
1675                flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
1676                default_value: None,
1677            },
1678        ],
1679        indexes: Vec::new(),
1680        colocation: Vec::new(),
1681        constraints: Default::default(),
1682        clustered: false,
1683    }
1684}
1685
1686#[derive(Serialize, Deserialize)]
1687struct ExternalRowState {
1688    cells: Vec<(u16, Value)>,
1689}
1690
1691const KV_STATE_MAGIC: &[u8] = b"mongreldb.external.kv.v1\n";
1692
1693#[derive(Serialize, Deserialize)]
1694struct ExternalKvState {
1695    entries: Vec<(Vec<u8>, Vec<u8>)>,
1696}
1697
1698fn state_file(db: &Arc<Database>, entry: &ExternalTableEntry) -> PathBuf {
1699    db.root()
1700        .join(VTAB_DIR)
1701        .join(&entry.name)
1702        .join("state.json")
1703}
1704
1705pub(crate) fn external_table_state_bytes(
1706    db: &Arc<Database>,
1707    entry: &ExternalTableEntry,
1708) -> Result<Vec<u8>> {
1709    let path = state_file(db, entry);
1710    if !path.exists() {
1711        return Ok(Vec::new());
1712    }
1713    fs::read(&path)
1714        .map_err(|e| MongrelQueryError::Schema(format!("read external state {:?}: {e}", path)))
1715}
1716
1717fn read_state_rows(
1718    db: &Arc<Database>,
1719    entry: &ExternalTableEntry,
1720) -> Result<Vec<HashMap<u16, Value>>> {
1721    let bytes = external_table_state_bytes(db, entry)?;
1722    if bytes.is_empty() {
1723        return Ok(Vec::new());
1724    }
1725    decode_state_rows(&bytes)
1726}
1727
1728fn encode_state_rows(rows: &[HashMap<u16, Value>]) -> Result<Vec<u8>> {
1729    let state = rows
1730        .iter()
1731        .map(|row| {
1732            let mut cells = row
1733                .iter()
1734                .map(|(id, value)| (*id, value.clone()))
1735                .collect::<Vec<_>>();
1736            cells.sort_by_key(|(id, _)| *id);
1737            ExternalRowState { cells }
1738        })
1739        .collect::<Vec<_>>();
1740    serde_json::to_vec(&state)
1741        .map_err(|e| MongrelQueryError::Schema(format!("encode external state: {e}")))
1742}
1743
1744fn decode_state_rows(state: &[u8]) -> Result<Vec<HashMap<u16, Value>>> {
1745    let rows: Vec<ExternalRowState> = serde_json::from_slice(state)
1746        .map_err(|e| MongrelQueryError::Schema(format!("decode external state: {e}")))?;
1747    Ok(rows
1748        .into_iter()
1749        .map(|row| row.cells.into_iter().collect())
1750        .collect())
1751}
1752
1753fn decode_kv_state(state: &[u8]) -> Result<BTreeMap<Vec<u8>, Vec<u8>>> {
1754    if state.is_empty() {
1755        return Ok(BTreeMap::new());
1756    }
1757    let json = state.strip_prefix(KV_STATE_MAGIC).ok_or_else(|| {
1758        MongrelQueryError::Schema(
1759            "external transaction state is not in key/value module format".into(),
1760        )
1761    })?;
1762    let decoded: ExternalKvState = serde_json::from_slice(json)
1763        .map_err(|e| MongrelQueryError::Schema(format!("decode external kv state: {e}")))?;
1764    Ok(decoded.entries.into_iter().collect())
1765}
1766
1767fn encode_kv_state(state: &BTreeMap<Vec<u8>, Vec<u8>>) -> Result<Vec<u8>> {
1768    let encoded = ExternalKvState {
1769        entries: state
1770            .iter()
1771            .map(|(key, value)| (key.clone(), value.clone()))
1772            .collect(),
1773    };
1774    let mut out = KV_STATE_MAGIC.to_vec();
1775    out.extend(
1776        serde_json::to_vec(&encoded)
1777            .map_err(|e| MongrelQueryError::Schema(format!("encode external kv state: {e}")))?,
1778    );
1779    Ok(out)
1780}
1781
1782fn validate_external_rows(schema: &CoreSchema, rows: &[HashMap<u16, Value>]) -> Result<()> {
1783    let mut known = HashSet::new();
1784    let mut pk_cols = Vec::new();
1785    for column in &schema.columns {
1786        known.insert(column.id);
1787        if column.flags.contains(ColumnFlags::PRIMARY_KEY) {
1788            pk_cols.push(column.id);
1789        }
1790    }
1791    let mut keys = HashSet::new();
1792    for row in rows {
1793        for column_id in row.keys() {
1794            if !known.contains(column_id) {
1795                return Err(MongrelQueryError::Schema(format!(
1796                    "external row contains unknown column id {column_id}"
1797                )));
1798            }
1799        }
1800        if !pk_cols.is_empty() {
1801            let mut key = Vec::new();
1802            for column_id in &pk_cols {
1803                let value = row.get(column_id).ok_or_else(|| {
1804                    MongrelQueryError::Schema(format!(
1805                        "external row is missing primary key column {column_id}"
1806                    ))
1807                })?;
1808                key.extend_from_slice(&column_id.to_be_bytes());
1809                key.extend_from_slice(&value.encode_key());
1810                key.push(0);
1811            }
1812            if !keys.insert(key) {
1813                return Err(MongrelQueryError::Schema(
1814                    "external table primary key conflict".into(),
1815                ));
1816            }
1817        }
1818    }
1819    Ok(())
1820}
1821
1822fn core_rows_to_batches(
1823    schema: &CoreSchema,
1824    rows: Vec<HashMap<u16, Value>>,
1825) -> Result<Vec<RecordBatch>> {
1826    let arrow_schema = arrow_conv::arrow_schema(schema)?;
1827    let arrays = schema
1828        .columns
1829        .iter()
1830        .map(|column| {
1831            let values = rows
1832                .iter()
1833                .map(|row| row.get(&column.id).cloned().unwrap_or(Value::Null))
1834                .collect::<Vec<_>>();
1835            arrow_conv::build_array(column.ty.clone(), &values)
1836        })
1837        .collect::<Result<Vec<_>>>()?;
1838    Ok(vec![RecordBatch::try_new(arrow_schema, arrays)
1839        .map_err(|e| MongrelQueryError::Arrow(e.to_string()))?])
1840}
1841
1842struct FtsDocsModule;
1843
1844impl ExternalTableModule for FtsDocsModule {
1845    fn name(&self) -> &str {
1846        "fts_docs"
1847    }
1848
1849    fn descriptor(&self) -> ExternalModuleDescriptor {
1850        ExternalModuleDescriptor {
1851            schema: fts_docs_schema(),
1852            hidden_columns: vec!["query".to_string()],
1853            capabilities: ModuleCapabilities {
1854                writable: true,
1855                deterministic: true,
1856                trigger_safe: true,
1857                ..ModuleCapabilities::default()
1858            },
1859        }
1860    }
1861
1862    fn indexes(&self, entry: &ExternalTableEntry) -> Result<Vec<ExternalModuleIndex>> {
1863        Ok(vec![ExternalModuleIndex::new(
1864            format!("{}_fts_inverted", entry.name),
1865            vec![2],
1866        )])
1867    }
1868
1869    fn connect(
1870        &self,
1871        ctx: &ModuleConnectCtx<'_>,
1872        entry: &ExternalTableEntry,
1873    ) -> Result<Arc<dyn ExternalTable>> {
1874        let options = fts_options(entry)?;
1875        let rows = read_state_rows(ctx.db, entry)?;
1876        let index = FtsInvertedIndex::build(&rows, &options);
1877        Ok(Arc::new(FtsDocsExternalTable {
1878            schema: arrow_conv::arrow_schema(&entry.declared_schema)?,
1879            core_schema: entry.declared_schema.clone(),
1880            options,
1881            index,
1882            rows,
1883        }))
1884    }
1885
1886    fn read_rows(
1887        &self,
1888        ctx: &ModuleConnectCtx<'_>,
1889        entry: &ExternalTableEntry,
1890    ) -> Result<Vec<HashMap<u16, Value>>> {
1891        let _ = fts_options(entry)?;
1892        read_state_rows(ctx.db, entry)
1893    }
1894
1895    fn prepare_rows(
1896        &self,
1897        _ctx: &ModuleConnectCtx<'_>,
1898        entry: &ExternalTableEntry,
1899        rows: Vec<HashMap<u16, Value>>,
1900    ) -> Result<Vec<u8>> {
1901        let _ = fts_options(entry)?;
1902        let rows = rows
1903            .into_iter()
1904            .map(|mut row| {
1905                row.remove(&3);
1906                row.remove(&4);
1907                row.remove(&5);
1908                row.remove(&6);
1909                row
1910            })
1911            .collect::<Vec<_>>();
1912        validate_external_rows(&entry.declared_schema, &rows)?;
1913        encode_state_rows(&rows)
1914    }
1915}
1916
1917#[derive(Clone)]
1918struct FtsDocsExternalTable {
1919    schema: SchemaRef,
1920    core_schema: CoreSchema,
1921    options: FtsOptions,
1922    index: FtsInvertedIndex,
1923    rows: Vec<HashMap<u16, Value>>,
1924}
1925
1926impl std::fmt::Debug for FtsDocsExternalTable {
1927    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1928        f.debug_struct("FtsDocsExternalTable")
1929            .field("schema", &self.schema)
1930            .finish_non_exhaustive()
1931    }
1932}
1933
1934impl ExternalTable for FtsDocsExternalTable {
1935    fn schema(&self) -> SchemaRef {
1936        self.schema.clone()
1937    }
1938
1939    fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
1940        Ok(ExternalPlan::new(
1941            request
1942                .raw_filters
1943                .iter()
1944                .map(|filter| {
1945                    if fts_query_from_expr(filter, &self.options).is_some() {
1946                        TableProviderFilterPushDown::Exact
1947                    } else {
1948                        TableProviderFilterPushDown::Unsupported
1949                    }
1950                })
1951                .collect(),
1952            None,
1953            1.0,
1954            false,
1955        ))
1956    }
1957
1958    fn scan_with_control(
1959        &self,
1960        request: &ExternalPlanRequest<'_>,
1961        control: &mongreldb_core::ExecutionControl,
1962    ) -> DFResult<ExternalScan> {
1963        external_checkpoint(Some(control), 0)?;
1964        let query = request
1965            .raw_filters
1966            .iter()
1967            .filter_map(|filter| fts_query_from_expr(filter, &self.options))
1968            .reduce(FtsQuery::and);
1969        let mut rows = Vec::new();
1970        if let Some(query) = query.as_ref() {
1971            for (index, row_index) in self.index.candidates(query).into_iter().enumerate() {
1972                external_checkpoint(Some(control), index)?;
1973                let Some(row) = self.rows.get(row_index) else {
1974                    continue;
1975                };
1976                let Some(score) = fts_match_score(row, query, &self.options) else {
1977                    continue;
1978                };
1979                rows.push(fts_enrich_row(
1980                    row.clone(),
1981                    Some(query),
1982                    &self.options,
1983                    Some(score),
1984                ));
1985            }
1986        } else {
1987            rows.reserve(self.rows.len());
1988            for (index, row) in self.rows.iter().cloned().enumerate() {
1989                external_checkpoint(Some(control), index)?;
1990                rows.push(fts_enrich_row(row, None, &self.options, None));
1991            }
1992        }
1993        external_checkpoint(Some(control), 0)?;
1994        project_scan_with_control(
1995            self.schema.clone(),
1996            core_rows_to_batches(&self.core_schema, rows)
1997                .map_err(|e| DataFusionError::Execution(e.to_string()))?,
1998            request.projection.as_deref(),
1999            request.limit,
2000            Some(control),
2001        )
2002    }
2003}
2004
2005fn fts_docs_schema() -> CoreSchema {
2006    CoreSchema {
2007        schema_id: 0,
2008        columns: vec![
2009            CoreColumnDef {
2010                id: 1,
2011                name: "doc_id".to_string(),
2012                ty: TypeId::Int64,
2013                flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
2014                default_value: None,
2015            },
2016            CoreColumnDef {
2017                id: 2,
2018                name: "text".to_string(),
2019                ty: TypeId::Bytes,
2020                flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
2021                default_value: None,
2022            },
2023            CoreColumnDef {
2024                id: 3,
2025                name: "query".to_string(),
2026                ty: TypeId::Bytes,
2027                flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
2028                default_value: None,
2029            },
2030            CoreColumnDef {
2031                id: 4,
2032                name: "rank".to_string(),
2033                ty: TypeId::Float64,
2034                flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
2035                default_value: None,
2036            },
2037            CoreColumnDef {
2038                id: 5,
2039                name: "snippet".to_string(),
2040                ty: TypeId::Bytes,
2041                flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
2042                default_value: None,
2043            },
2044            CoreColumnDef {
2045                id: 6,
2046                name: "highlight".to_string(),
2047                ty: TypeId::Bytes,
2048                flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
2049                default_value: None,
2050            },
2051        ],
2052        indexes: Vec::new(),
2053        colocation: Vec::new(),
2054        constraints: Default::default(),
2055        clustered: false,
2056    }
2057}
2058
2059#[derive(Clone)]
2060struct FtsOptions {
2061    prefix_queries: bool,
2062    case_sensitive: bool,
2063    min_token_len: usize,
2064    stopwords: HashSet<String>,
2065}
2066
2067impl Default for FtsOptions {
2068    fn default() -> Self {
2069        Self {
2070            prefix_queries: false,
2071            case_sensitive: false,
2072            min_token_len: 1,
2073            stopwords: HashSet::new(),
2074        }
2075    }
2076}
2077
2078impl FtsOptions {
2079    fn normalize(&self, value: &str) -> String {
2080        if self.case_sensitive {
2081            value.to_string()
2082        } else {
2083            value.to_ascii_lowercase()
2084        }
2085    }
2086
2087    fn keep_term(&self, term: &str) -> bool {
2088        term.len() >= self.min_token_len && !self.stopwords.contains(term)
2089    }
2090}
2091
2092#[derive(Clone)]
2093struct FtsQuery {
2094    groups: Vec<Vec<FtsClause>>,
2095    prohibited: Vec<FtsClause>,
2096}
2097
2098impl FtsQuery {
2099    fn and(self, other: Self) -> Self {
2100        let left_groups = if self.groups.is_empty() {
2101            vec![Vec::new()]
2102        } else {
2103            self.groups
2104        };
2105        let right_groups = if other.groups.is_empty() {
2106            vec![Vec::new()]
2107        } else {
2108            other.groups
2109        };
2110        let mut groups = Vec::new();
2111        for left in &left_groups {
2112            for right in &right_groups {
2113                let mut combined = left.clone();
2114                combined.extend(right.clone());
2115                groups.push(combined);
2116            }
2117        }
2118        let mut prohibited = self.prohibited;
2119        prohibited.extend(other.prohibited);
2120        Self { groups, prohibited }
2121    }
2122
2123    fn positive_clauses(&self) -> impl Iterator<Item = &FtsClause> {
2124        self.groups.iter().flatten()
2125    }
2126}
2127
2128#[derive(Clone)]
2129struct FtsInvertedIndex {
2130    terms: HashMap<String, Vec<usize>>,
2131    all_rows: Vec<usize>,
2132}
2133
2134impl FtsInvertedIndex {
2135    fn build(rows: &[HashMap<u16, Value>], options: &FtsOptions) -> Self {
2136        let mut term_sets: HashMap<String, BTreeSet<usize>> = HashMap::new();
2137        for (idx, row) in rows.iter().enumerate() {
2138            let Some(Value::Bytes(text)) = row.get(&2) else {
2139                continue;
2140            };
2141            let text = String::from_utf8_lossy(text);
2142            let mut row_terms = HashSet::new();
2143            for span in token_spans_with_options(&text, options) {
2144                if row_terms.insert(span.term.clone()) {
2145                    term_sets.entry(span.term).or_default().insert(idx);
2146                }
2147            }
2148        }
2149        let terms = term_sets
2150            .into_iter()
2151            .map(|(term, rows)| (term, rows.into_iter().collect()))
2152            .collect();
2153        Self {
2154            terms,
2155            all_rows: (0..rows.len()).collect(),
2156        }
2157    }
2158
2159    fn candidates(&self, query: &FtsQuery) -> Vec<usize> {
2160        let mut rows = if query.groups.is_empty() {
2161            self.all_rows.iter().copied().collect::<BTreeSet<_>>()
2162        } else {
2163            query
2164                .groups
2165                .iter()
2166                .filter_map(|group| self.group_candidates(group))
2167                .flatten()
2168                .collect::<BTreeSet<_>>()
2169        };
2170        for clause in &query.prohibited {
2171            for idx in self.clause_candidates(clause) {
2172                rows.remove(&idx);
2173            }
2174        }
2175        rows.into_iter().collect()
2176    }
2177
2178    fn group_candidates(&self, group: &[FtsClause]) -> Option<Vec<usize>> {
2179        let mut iter = group.iter().map(|clause| self.clause_candidates(clause));
2180        let first = iter.next()?;
2181        let mut current = first.into_iter().collect::<BTreeSet<_>>();
2182        for rows in iter {
2183            let next = rows.into_iter().collect::<BTreeSet<_>>();
2184            current = current.intersection(&next).copied().collect();
2185            if current.is_empty() {
2186                break;
2187            }
2188        }
2189        Some(current.into_iter().collect())
2190    }
2191
2192    fn clause_candidates(&self, clause: &FtsClause) -> Vec<usize> {
2193        match clause {
2194            FtsClause::Term { term, prefix } => {
2195                if *prefix {
2196                    self.terms
2197                        .iter()
2198                        .filter(|(indexed, _)| indexed.starts_with(term))
2199                        .flat_map(|(_, rows)| rows.iter().copied())
2200                        .collect::<BTreeSet<_>>()
2201                        .into_iter()
2202                        .collect()
2203                } else {
2204                    self.terms.get(term).cloned().unwrap_or_default()
2205                }
2206            }
2207            FtsClause::Phrase(terms) => {
2208                let clauses = terms
2209                    .iter()
2210                    .map(|term| FtsClause::Term {
2211                        term: term.clone(),
2212                        prefix: false,
2213                    })
2214                    .collect::<Vec<_>>();
2215                self.group_candidates(&clauses).unwrap_or_default()
2216            }
2217        }
2218    }
2219}
2220
2221#[derive(Clone, PartialEq, Eq)]
2222enum FtsClause {
2223    Term { term: String, prefix: bool },
2224    Phrase(Vec<String>),
2225}
2226
2227enum FtsQueryToken {
2228    Clause(FtsClause),
2229    Or,
2230    Not,
2231}
2232
2233fn fts_options(entry: &ExternalTableEntry) -> Result<FtsOptions> {
2234    let mut options = FtsOptions::default();
2235    for arg in &entry.args {
2236        let raw = module_arg_string(arg).trim();
2237        if raw.is_empty() {
2238            continue;
2239        }
2240        let (key, value) = raw
2241            .split_once('=')
2242            .map_or((raw, "true"), |(key, value)| (key.trim(), value.trim()));
2243        match key.to_ascii_lowercase().as_str() {
2244            "tokenizer" => match value.to_ascii_lowercase().as_str() {
2245                "simple" | "ascii" | "unicode61" => {}
2246                other => {
2247                    return Err(MongrelQueryError::Schema(format!(
2248                        "fts_docs tokenizer {other:?} is not supported"
2249                    )))
2250                }
2251            },
2252            "prefix" | "prefix_queries" => options.prefix_queries = parse_bool_option(value)?,
2253            "case_sensitive" => options.case_sensitive = parse_bool_option(value)?,
2254            "min_token_len" => {
2255                options.min_token_len = value.parse::<usize>().map_err(|e| {
2256                    MongrelQueryError::Schema(format!(
2257                        "fts_docs min_token_len {value:?} must be an integer: {e}"
2258                    ))
2259                })?;
2260                if options.min_token_len == 0 {
2261                    return Err(MongrelQueryError::Schema(
2262                        "fts_docs min_token_len must be at least 1".into(),
2263                    ));
2264                }
2265            }
2266            "stopwords" => {
2267                options.stopwords = value
2268                    .split('|')
2269                    .map(str::trim)
2270                    .filter(|word| !word.is_empty())
2271                    .map(|word| options.normalize(word))
2272                    .collect();
2273            }
2274            other => {
2275                return Err(MongrelQueryError::Schema(format!(
2276                    "fts_docs option {other:?} is not supported"
2277                )))
2278            }
2279        }
2280    }
2281    Ok(options)
2282}
2283
2284fn parse_bool_option(value: &str) -> Result<bool> {
2285    match value.to_ascii_lowercase().as_str() {
2286        "1" | "true" | "yes" | "on" => Ok(true),
2287        "0" | "false" | "no" | "off" => Ok(false),
2288        _ => Err(MongrelQueryError::Schema(format!(
2289            "expected boolean option value, got {value:?}"
2290        ))),
2291    }
2292}
2293
2294fn fts_query_from_expr(expr: &Expr, options: &FtsOptions) -> Option<FtsQuery> {
2295    match expr {
2296        Expr::BinaryExpr(BinaryExpr { left, op, right }) if *op == Operator::And => {
2297            Some(fts_query_from_expr(left, options)?.and(fts_query_from_expr(right, options)?))
2298        }
2299        Expr::BinaryExpr(BinaryExpr { left, op, right }) if *op == Operator::Eq => {
2300            let literal = match (left.as_ref(), right.as_ref()) {
2301                (Expr::Column(column), Expr::Literal(literal, _)) if column.name == "query" => {
2302                    literal_string(literal)?
2303                }
2304                (Expr::Literal(literal, _), Expr::Column(column)) if column.name == "query" => {
2305                    literal_string(literal)?
2306                }
2307                _ => return None,
2308            };
2309            parse_fts_query(&literal, options)
2310        }
2311        _ => None,
2312    }
2313}
2314
2315fn parse_fts_query(input: &str, options: &FtsOptions) -> Option<FtsQuery> {
2316    let tokens = fts_query_tokens(input, options);
2317    if tokens.is_empty() {
2318        return None;
2319    }
2320    let mut groups: Vec<Vec<FtsClause>> = vec![Vec::new()];
2321    let mut prohibited = Vec::new();
2322    let mut negate = false;
2323    for token in tokens {
2324        match token {
2325            FtsQueryToken::Or => {
2326                if groups.last().is_some_and(|group| !group.is_empty()) {
2327                    groups.push(Vec::new());
2328                }
2329                negate = false;
2330            }
2331            FtsQueryToken::Not => negate = true,
2332            FtsQueryToken::Clause(clause) => {
2333                if negate {
2334                    prohibited.push(clause);
2335                    negate = false;
2336                } else if let Some(group) = groups.last_mut() {
2337                    group.push(clause);
2338                }
2339            }
2340        }
2341    }
2342    groups.retain(|group| !group.is_empty());
2343    if groups.is_empty() && prohibited.is_empty() {
2344        None
2345    } else {
2346        Some(FtsQuery { groups, prohibited })
2347    }
2348}
2349
2350fn fts_query_tokens(input: &str, options: &FtsOptions) -> Vec<FtsQueryToken> {
2351    let mut tokens = Vec::new();
2352    let mut chars = input.char_indices().peekable();
2353    while let Some((idx, ch)) = chars.next() {
2354        if ch.is_whitespace() {
2355            continue;
2356        }
2357        if ch == '"' {
2358            let start = idx + ch.len_utf8();
2359            let mut end = input.len();
2360            for (next_idx, next_ch) in chars.by_ref() {
2361                if next_ch == '"' {
2362                    end = next_idx;
2363                    break;
2364                }
2365            }
2366            let terms = token_spans_with_options(&input[start..end], options)
2367                .into_iter()
2368                .map(|span| span.term)
2369                .collect::<Vec<_>>();
2370            if !terms.is_empty() {
2371                tokens.push(FtsQueryToken::Clause(FtsClause::Phrase(terms)));
2372            }
2373            continue;
2374        }
2375        let start = idx;
2376        let mut end = idx + ch.len_utf8();
2377        while let Some((next_idx, next_ch)) = chars.peek().copied() {
2378            if next_ch.is_whitespace() {
2379                break;
2380            }
2381            chars.next();
2382            end = next_idx + next_ch.len_utf8();
2383        }
2384        let raw = &input[start..end];
2385        if raw.eq_ignore_ascii_case("OR") {
2386            tokens.push(FtsQueryToken::Or);
2387        } else if raw.eq_ignore_ascii_case("NOT") {
2388            tokens.push(FtsQueryToken::Not);
2389        } else if let Some(rest) = raw.strip_prefix('-') {
2390            tokens.push(FtsQueryToken::Not);
2391            tokens.extend(
2392                word_clauses(rest, options)
2393                    .into_iter()
2394                    .map(FtsQueryToken::Clause),
2395            );
2396        } else {
2397            tokens.extend(
2398                word_clauses(raw, options)
2399                    .into_iter()
2400                    .map(FtsQueryToken::Clause),
2401            );
2402        }
2403    }
2404    tokens
2405}
2406
2407fn word_clauses(raw: &str, options: &FtsOptions) -> Vec<FtsClause> {
2408    let raw = raw.trim_matches(|ch: char| !ch.is_ascii_alphanumeric() && ch != '*');
2409    if raw.is_empty() {
2410        return Vec::new();
2411    }
2412    let prefix = options.prefix_queries && raw.ends_with('*');
2413    let raw = raw.trim_end_matches('*');
2414    token_spans_with_options(raw, options)
2415        .into_iter()
2416        .map(|span| FtsClause::Term {
2417            term: span.term,
2418            prefix,
2419        })
2420        .collect()
2421}
2422
2423fn fts_match_score(
2424    row: &HashMap<u16, Value>,
2425    query: &FtsQuery,
2426    options: &FtsOptions,
2427) -> Option<f64> {
2428    let Some(Value::Bytes(text)) = row.get(&2) else {
2429        return None;
2430    };
2431    let text = String::from_utf8_lossy(text);
2432    let tokens = token_spans_with_options(&text, options);
2433    if query
2434        .prohibited
2435        .iter()
2436        .any(|clause| clause_matches(&tokens, clause))
2437    {
2438        return None;
2439    }
2440    let positive_score = if query.groups.is_empty() {
2441        Some(0.0)
2442    } else {
2443        query
2444            .groups
2445            .iter()
2446            .filter(|group| group.iter().all(|clause| clause_matches(&tokens, clause)))
2447            .map(|group| {
2448                group
2449                    .iter()
2450                    .map(|clause| clause_score(&tokens, clause))
2451                    .sum()
2452            })
2453            .max_by(|left: &f64, right: &f64| left.total_cmp(right))
2454    }?;
2455    Some(positive_score)
2456}
2457
2458fn clause_matches(tokens: &[TokenSpan], clause: &FtsClause) -> bool {
2459    clause_score(tokens, clause) > 0.0
2460}
2461
2462fn clause_score(tokens: &[TokenSpan], clause: &FtsClause) -> f64 {
2463    match clause {
2464        FtsClause::Term { term, prefix } => tokens
2465            .iter()
2466            .filter(|token| token_matches_term(&token.term, term, *prefix))
2467            .count() as f64,
2468        FtsClause::Phrase(terms) => phrase_occurrences(tokens, terms) as f64 * terms.len() as f64,
2469    }
2470}
2471
2472fn token_matches_term(token: &str, term: &str, prefix: bool) -> bool {
2473    if prefix {
2474        token.starts_with(term)
2475    } else {
2476        token == term
2477    }
2478}
2479
2480fn phrase_occurrences(tokens: &[TokenSpan], terms: &[String]) -> usize {
2481    if terms.is_empty() || tokens.len() < terms.len() {
2482        return 0;
2483    }
2484    tokens
2485        .windows(terms.len())
2486        .filter(|window| {
2487            window
2488                .iter()
2489                .zip(terms)
2490                .all(|(token, term)| token.term == *term)
2491        })
2492        .count()
2493}
2494
2495fn token_matches_positive_clause(token: &str, query: &FtsQuery) -> bool {
2496    query.positive_clauses().any(|clause| match clause {
2497        FtsClause::Term { term, prefix } => token_matches_term(token, term, *prefix),
2498        FtsClause::Phrase(terms) => terms.iter().any(|term| term == token),
2499    })
2500}
2501
2502fn fts_enrich_row(
2503    mut row: HashMap<u16, Value>,
2504    query: Option<&FtsQuery>,
2505    options: &FtsOptions,
2506    score: Option<f64>,
2507) -> HashMap<u16, Value> {
2508    row.remove(&3);
2509    row.remove(&4);
2510    row.remove(&5);
2511    row.remove(&6);
2512    let Some(score) = score else {
2513        return row;
2514    };
2515    let text = match row.get(&2) {
2516        Some(Value::Bytes(text)) => String::from_utf8_lossy(text).into_owned(),
2517        _ => String::new(),
2518    };
2519    row.insert(4, Value::Float64(score));
2520    if let Some(query) = query {
2521        row.insert(
2522            5,
2523            Value::Bytes(fts_snippet(&text, query, options).into_bytes()),
2524        );
2525        row.insert(
2526            6,
2527            Value::Bytes(fts_highlight(&text, query, options).into_bytes()),
2528        );
2529    }
2530    row
2531}
2532
2533#[derive(Debug, Clone)]
2534struct TokenSpan {
2535    term: String,
2536    start: usize,
2537    end: usize,
2538}
2539
2540fn token_spans_with_options(input: &str, options: &FtsOptions) -> Vec<TokenSpan> {
2541    let mut spans = Vec::new();
2542    let mut start = None;
2543    for (idx, ch) in input.char_indices() {
2544        if ch.is_ascii_alphanumeric() {
2545            start.get_or_insert(idx);
2546        } else if let Some(lo) = start.take() {
2547            push_token_span(input, &mut spans, lo, idx, options);
2548        }
2549    }
2550    if let Some(lo) = start {
2551        push_token_span(input, &mut spans, lo, input.len(), options);
2552    }
2553    spans
2554}
2555
2556fn push_token_span(
2557    input: &str,
2558    spans: &mut Vec<TokenSpan>,
2559    start: usize,
2560    end: usize,
2561    options: &FtsOptions,
2562) {
2563    if start < end {
2564        let term = options.normalize(&input[start..end]);
2565        if options.keep_term(&term) {
2566            spans.push(TokenSpan { term, start, end });
2567        }
2568    }
2569}
2570
2571fn fts_snippet(text: &str, query: &FtsQuery, options: &FtsOptions) -> String {
2572    let spans = token_spans_with_options(text, options);
2573    if spans.is_empty() {
2574        return String::new();
2575    }
2576    let first_match = spans
2577        .iter()
2578        .position(|span| token_matches_positive_clause(&span.term, query))
2579        .unwrap_or(0);
2580    let lo = first_match.saturating_sub(3);
2581    let hi = (first_match + 5).min(spans.len());
2582    let mut out = Vec::new();
2583    if lo > 0 {
2584        out.push("...".to_string());
2585    }
2586    for span in &spans[lo..hi] {
2587        let token = &text[span.start..span.end];
2588        if token_matches_positive_clause(&span.term, query) {
2589            out.push(format!("[{token}]"));
2590        } else {
2591            out.push(token.to_string());
2592        }
2593    }
2594    if hi < spans.len() {
2595        out.push("...".to_string());
2596    }
2597    out.join(" ")
2598}
2599
2600fn fts_highlight(text: &str, query: &FtsQuery, options: &FtsOptions) -> String {
2601    let spans = token_spans_with_options(text, options);
2602    if spans.is_empty() {
2603        return text.to_string();
2604    }
2605    let mut out = String::new();
2606    let mut cursor = 0;
2607    for span in spans {
2608        out.push_str(&text[cursor..span.start]);
2609        let token = &text[span.start..span.end];
2610        if token_matches_positive_clause(&span.term, query) {
2611            out.push_str("<b>");
2612            out.push_str(token);
2613            out.push_str("</b>");
2614        } else {
2615            out.push_str(token);
2616        }
2617        cursor = span.end;
2618    }
2619    out.push_str(&text[cursor..]);
2620    out
2621}
2622
2623struct RTreeRectsModule;
2624
2625impl ExternalTableModule for RTreeRectsModule {
2626    fn name(&self) -> &str {
2627        "rtree_rects"
2628    }
2629
2630    fn descriptor(&self) -> ExternalModuleDescriptor {
2631        ExternalModuleDescriptor {
2632            schema: rtree_rects_schema(),
2633            hidden_columns: vec![
2634                "query_min_x".to_string(),
2635                "query_max_x".to_string(),
2636                "query_min_y".to_string(),
2637                "query_max_y".to_string(),
2638            ],
2639            capabilities: ModuleCapabilities {
2640                writable: true,
2641                deterministic: true,
2642                trigger_safe: true,
2643                ..ModuleCapabilities::default()
2644            },
2645        }
2646    }
2647
2648    fn indexes(&self, entry: &ExternalTableEntry) -> Result<Vec<ExternalModuleIndex>> {
2649        Ok(vec![ExternalModuleIndex::new(
2650            format!("{}_rtree_spatial", entry.name),
2651            vec![2, 3, 4, 5],
2652        )])
2653    }
2654
2655    fn connect(
2656        &self,
2657        ctx: &ModuleConnectCtx<'_>,
2658        entry: &ExternalTableEntry,
2659    ) -> Result<Arc<dyn ExternalTable>> {
2660        ensure_no_args(entry, self.name())?;
2661        let rows = read_state_rows(ctx.db, entry)?;
2662        let index = RTreeSpatialIndex::build(&rows);
2663        Ok(Arc::new(RTreeRectsExternalTable {
2664            schema: arrow_conv::arrow_schema(&entry.declared_schema)?,
2665            core_schema: entry.declared_schema.clone(),
2666            index,
2667            rows,
2668        }))
2669    }
2670
2671    fn read_rows(
2672        &self,
2673        ctx: &ModuleConnectCtx<'_>,
2674        entry: &ExternalTableEntry,
2675    ) -> Result<Vec<HashMap<u16, Value>>> {
2676        ensure_no_args(entry, self.name())?;
2677        read_state_rows(ctx.db, entry)
2678    }
2679
2680    fn prepare_rows(
2681        &self,
2682        _ctx: &ModuleConnectCtx<'_>,
2683        entry: &ExternalTableEntry,
2684        rows: Vec<HashMap<u16, Value>>,
2685    ) -> Result<Vec<u8>> {
2686        ensure_no_args(entry, self.name())?;
2687        let rows = rows
2688            .into_iter()
2689            .map(|mut row| {
2690                for column_id in 6..=9 {
2691                    row.remove(&column_id);
2692                }
2693                row
2694            })
2695            .collect::<Vec<_>>();
2696        validate_external_rows(&entry.declared_schema, &rows)?;
2697        encode_state_rows(&rows)
2698    }
2699}
2700
2701#[derive(Clone)]
2702struct RTreeRectsExternalTable {
2703    schema: SchemaRef,
2704    core_schema: CoreSchema,
2705    index: RTreeSpatialIndex,
2706    rows: Vec<HashMap<u16, Value>>,
2707}
2708
2709impl std::fmt::Debug for RTreeRectsExternalTable {
2710    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2711        f.debug_struct("RTreeRectsExternalTable")
2712            .field("schema", &self.schema)
2713            .finish_non_exhaustive()
2714    }
2715}
2716
2717impl ExternalTable for RTreeRectsExternalTable {
2718    fn schema(&self) -> SchemaRef {
2719        self.schema.clone()
2720    }
2721
2722    fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
2723        Ok(ExternalPlan::new(
2724            request
2725                .raw_filters
2726                .iter()
2727                .map(|filter| {
2728                    if rtree_query_bound(filter).is_some()
2729                        || rtree_query_rect_function(filter).is_some()
2730                    {
2731                        TableProviderFilterPushDown::Exact
2732                    } else {
2733                        TableProviderFilterPushDown::Unsupported
2734                    }
2735                })
2736                .collect(),
2737            None,
2738            1.0,
2739            false,
2740        ))
2741    }
2742
2743    fn scan_with_control(
2744        &self,
2745        request: &ExternalPlanRequest<'_>,
2746        control: &mongreldb_core::ExecutionControl,
2747    ) -> DFResult<ExternalScan> {
2748        external_checkpoint(Some(control), 0)?;
2749        let mut hidden_bounds = QueryRect::default();
2750        let mut has_hidden_bounds = false;
2751        for (column, value) in request
2752            .raw_filters
2753            .iter()
2754            .filter_map(|filter| rtree_query_bound(filter))
2755        {
2756            hidden_bounds.set(column, value);
2757            has_hidden_bounds = true;
2758        }
2759        let mut query_rects = request
2760            .raw_filters
2761            .iter()
2762            .filter_map(|filter| rtree_query_rect_function(filter))
2763            .collect::<Vec<_>>();
2764        if has_hidden_bounds || query_rects.is_empty() {
2765            query_rects.push(hidden_bounds);
2766        }
2767        let mut candidate_rows: Option<BTreeSet<usize>> = None;
2768        for bounds in &query_rects {
2769            let mut next = BTreeSet::new();
2770            for (index, row) in self.index.candidates(*bounds).into_iter().enumerate() {
2771                external_checkpoint(Some(control), index)?;
2772                next.insert(row);
2773            }
2774            candidate_rows = Some(match candidate_rows.take() {
2775                Some(current) => {
2776                    let mut intersection = BTreeSet::new();
2777                    for (index, row) in current.into_iter().enumerate() {
2778                        external_checkpoint(Some(control), index)?;
2779                        if next.contains(&row) {
2780                            intersection.insert(row);
2781                        }
2782                    }
2783                    intersection
2784                }
2785                None => next,
2786            });
2787        }
2788        let candidate_rows =
2789            candidate_rows.unwrap_or_else(|| self.index.all_rows.iter().copied().collect());
2790        let mut rows = Vec::new();
2791        for (index, row) in self.rows.iter().enumerate() {
2792            external_checkpoint(Some(control), index)?;
2793            if candidate_rows.contains(&index)
2794                && query_rects
2795                    .iter()
2796                    .all(|bounds| rtree_row_matches(row, *bounds))
2797            {
2798                rows.push(row.clone());
2799            }
2800        }
2801        external_checkpoint(Some(control), 0)?;
2802        project_scan_with_control(
2803            self.schema.clone(),
2804            core_rows_to_batches(&self.core_schema, rows)
2805                .map_err(|e| DataFusionError::Execution(e.to_string()))?,
2806            request.projection.as_deref(),
2807            request.limit,
2808            Some(control),
2809        )
2810    }
2811}
2812
2813#[derive(Clone, Copy)]
2814struct QueryRect {
2815    min_x: f64,
2816    max_x: f64,
2817    min_y: f64,
2818    max_y: f64,
2819}
2820
2821impl Default for QueryRect {
2822    fn default() -> Self {
2823        Self {
2824            min_x: f64::NEG_INFINITY,
2825            max_x: f64::INFINITY,
2826            min_y: f64::NEG_INFINITY,
2827            max_y: f64::INFINITY,
2828        }
2829    }
2830}
2831
2832impl QueryRect {
2833    fn set(&mut self, column: &str, value: f64) {
2834        match column {
2835            "query_min_x" => self.min_x = value,
2836            "query_max_x" => self.max_x = value,
2837            "query_min_y" => self.min_y = value,
2838            "query_max_y" => self.max_y = value,
2839            _ => {}
2840        }
2841    }
2842}
2843
2844#[derive(Clone)]
2845struct RTreeSpatialIndex {
2846    all_rows: Vec<usize>,
2847    min_x: Vec<(f64, usize)>,
2848    max_x: Vec<(f64, usize)>,
2849    min_y: Vec<(f64, usize)>,
2850    max_y: Vec<(f64, usize)>,
2851}
2852
2853impl RTreeSpatialIndex {
2854    fn build(rows: &[HashMap<u16, Value>]) -> Self {
2855        let mut all_rows = Vec::new();
2856        let mut min_x = Vec::new();
2857        let mut max_x = Vec::new();
2858        let mut min_y = Vec::new();
2859        let mut max_y = Vec::new();
2860        for (idx, row) in rows.iter().enumerate() {
2861            let Some(x0) = row_f64(row, 2) else {
2862                continue;
2863            };
2864            let Some(x1) = row_f64(row, 3) else {
2865                continue;
2866            };
2867            let Some(y0) = row_f64(row, 4) else {
2868                continue;
2869            };
2870            let Some(y1) = row_f64(row, 5) else {
2871                continue;
2872            };
2873            if !(x0.is_finite() && x1.is_finite() && y0.is_finite() && y1.is_finite()) {
2874                continue;
2875            }
2876            all_rows.push(idx);
2877            min_x.push((x0, idx));
2878            max_x.push((x1, idx));
2879            min_y.push((y0, idx));
2880            max_y.push((y1, idx));
2881        }
2882        for values in [&mut min_x, &mut max_x, &mut min_y, &mut max_y] {
2883            values.sort_by(|left, right| left.0.total_cmp(&right.0));
2884        }
2885        Self {
2886            all_rows,
2887            min_x,
2888            max_x,
2889            min_y,
2890            max_y,
2891        }
2892    }
2893
2894    fn candidates(&self, query: QueryRect) -> Vec<usize> {
2895        let mut rows = self.lte(&self.min_x, query.max_x);
2896        for next in [
2897            self.gte(&self.max_x, query.min_x),
2898            self.lte(&self.min_y, query.max_y),
2899            self.gte(&self.max_y, query.min_y),
2900        ] {
2901            rows = rows.intersection(&next).copied().collect();
2902            if rows.is_empty() {
2903                break;
2904            }
2905        }
2906        rows.into_iter().collect()
2907    }
2908
2909    fn lte(&self, values: &[(f64, usize)], bound: f64) -> BTreeSet<usize> {
2910        if bound == f64::INFINITY {
2911            return self.all_rows.iter().copied().collect();
2912        }
2913        if bound.is_nan() {
2914            return BTreeSet::new();
2915        }
2916        let end = values.partition_point(|(value, _)| *value <= bound);
2917        values[..end].iter().map(|(_, idx)| *idx).collect()
2918    }
2919
2920    fn gte(&self, values: &[(f64, usize)], bound: f64) -> BTreeSet<usize> {
2921        if bound == f64::NEG_INFINITY {
2922            return self.all_rows.iter().copied().collect();
2923        }
2924        if bound.is_nan() {
2925            return BTreeSet::new();
2926        }
2927        let start = values.partition_point(|(value, _)| *value < bound);
2928        values[start..].iter().map(|(_, idx)| *idx).collect()
2929    }
2930}
2931
2932fn rtree_query_rect_function(expr: &Expr) -> Option<QueryRect> {
2933    let Expr::ScalarFunction(sf) = expr else {
2934        return None;
2935    };
2936    if !sf.func.name().eq_ignore_ascii_case("rtree_intersects") || sf.args.len() != 8 {
2937        return None;
2938    }
2939    for (arg, expected) in sf
2940        .args
2941        .iter()
2942        .take(4)
2943        .zip(["min_x", "max_x", "min_y", "max_y"])
2944    {
2945        let Expr::Column(column) = arg else {
2946            return None;
2947        };
2948        if column.name != expected {
2949            return None;
2950        }
2951    }
2952    Some(QueryRect {
2953        min_x: literal_f64_from_expr(&sf.args[4])?,
2954        max_x: literal_f64_from_expr(&sf.args[5])?,
2955        min_y: literal_f64_from_expr(&sf.args[6])?,
2956        max_y: literal_f64_from_expr(&sf.args[7])?,
2957    })
2958}
2959
2960fn literal_f64_from_expr(expr: &Expr) -> Option<f64> {
2961    match expr {
2962        Expr::Literal(literal, _) => literal_f64(literal),
2963        _ => None,
2964    }
2965}
2966
2967fn rtree_rects_schema() -> CoreSchema {
2968    CoreSchema {
2969        schema_id: 0,
2970        columns: vec![
2971            CoreColumnDef {
2972                id: 1,
2973                name: "id".to_string(),
2974                ty: TypeId::Int64,
2975                flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
2976                default_value: None,
2977            },
2978            rect_column(2, "min_x"),
2979            rect_column(3, "max_x"),
2980            rect_column(4, "min_y"),
2981            rect_column(5, "max_y"),
2982            rect_column(6, "query_min_x"),
2983            rect_column(7, "query_max_x"),
2984            rect_column(8, "query_min_y"),
2985            rect_column(9, "query_max_y"),
2986        ],
2987        indexes: Vec::new(),
2988        colocation: Vec::new(),
2989        constraints: Default::default(),
2990        clustered: false,
2991    }
2992}
2993
2994fn rect_column(id: u16, name: &str) -> CoreColumnDef {
2995    CoreColumnDef {
2996        id,
2997        name: name.to_string(),
2998        ty: TypeId::Float64,
2999        flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
3000        default_value: None,
3001    }
3002}
3003
3004fn rtree_query_bound(expr: &Expr) -> Option<(&str, f64)> {
3005    match expr {
3006        Expr::BinaryExpr(BinaryExpr { left, op, right }) if *op == Operator::Eq => {
3007            match (left.as_ref(), right.as_ref()) {
3008                (Expr::Column(column), Expr::Literal(literal, _)) => {
3009                    hidden_rect_column(&column.name).zip(literal_f64(literal))
3010                }
3011                (Expr::Literal(literal, _), Expr::Column(column)) => {
3012                    hidden_rect_column(&column.name).zip(literal_f64(literal))
3013                }
3014                _ => None,
3015            }
3016        }
3017        _ => None,
3018    }
3019}
3020
3021fn hidden_rect_column(name: &str) -> Option<&str> {
3022    match name {
3023        "query_min_x" | "query_max_x" | "query_min_y" | "query_max_y" => Some(name),
3024        _ => None,
3025    }
3026}
3027
3028fn rtree_row_matches(row: &HashMap<u16, Value>, query: QueryRect) -> bool {
3029    let Some(min_x) = row_f64(row, 2) else {
3030        return false;
3031    };
3032    let Some(max_x) = row_f64(row, 3) else {
3033        return false;
3034    };
3035    let Some(min_y) = row_f64(row, 4) else {
3036        return false;
3037    };
3038    let Some(max_y) = row_f64(row, 5) else {
3039        return false;
3040    };
3041    max_x >= query.min_x && min_x <= query.max_x && max_y >= query.min_y && min_y <= query.max_y
3042}
3043
3044fn row_f64(row: &HashMap<u16, Value>, column_id: u16) -> Option<f64> {
3045    match row.get(&column_id) {
3046        Some(Value::Float64(value)) => Some(*value),
3047        Some(Value::Int64(value)) => Some(*value as f64),
3048        _ => None,
3049    }
3050}
3051
3052struct SeriesModule;
3053
3054impl ExternalTableModule for SeriesModule {
3055    fn name(&self) -> &str {
3056        "series"
3057    }
3058
3059    fn descriptor(&self) -> ExternalModuleDescriptor {
3060        ExternalModuleDescriptor {
3061            schema: CoreSchema {
3062                schema_id: 0,
3063                columns: vec![
3064                    CoreColumnDef {
3065                        id: 1,
3066                        name: "value".to_string(),
3067                        ty: TypeId::Int64,
3068                        flags: ColumnFlags::empty(),
3069                        default_value: None,
3070                    },
3071                    CoreColumnDef {
3072                        id: 2,
3073                        name: "start".to_string(),
3074                        ty: TypeId::Int64,
3075                        flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
3076                        default_value: None,
3077                    },
3078                    CoreColumnDef {
3079                        id: 3,
3080                        name: "stop".to_string(),
3081                        ty: TypeId::Int64,
3082                        flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
3083                        default_value: None,
3084                    },
3085                    CoreColumnDef {
3086                        id: 4,
3087                        name: "step".to_string(),
3088                        ty: TypeId::Int64,
3089                        flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
3090                        default_value: None,
3091                    },
3092                ],
3093                indexes: Vec::new(),
3094                colocation: Vec::new(),
3095                constraints: Default::default(),
3096                clustered: false,
3097            },
3098            hidden_columns: vec!["start".to_string(), "stop".to_string(), "step".to_string()],
3099            capabilities: ModuleCapabilities {
3100                read_only: true,
3101                deterministic: true,
3102                trigger_safe: true,
3103                ..ModuleCapabilities::default()
3104            },
3105        }
3106    }
3107
3108    fn connect(
3109        &self,
3110        _ctx: &ModuleConnectCtx<'_>,
3111        entry: &ExternalTableEntry,
3112    ) -> Result<Arc<dyn ExternalTable>> {
3113        let (start, stop, step) = series_args(entry)?;
3114        let table = SeriesExternalTable::new(start, stop, step)?;
3115        Ok(Arc::new(table))
3116    }
3117}
3118
3119struct SeriesExternalTable {
3120    start: i64,
3121    stop: i64,
3122    step: i64,
3123    schema: SchemaRef,
3124}
3125
3126impl SeriesExternalTable {
3127    fn new(start: i64, stop: i64, step: i64) -> Result<Self> {
3128        if step == 0 {
3129            return Err(MongrelQueryError::Schema(
3130                "series step must not be 0".into(),
3131            ));
3132        }
3133        Ok(Self {
3134            start,
3135            stop,
3136            step,
3137            schema: Arc::new(ArrowSchema::new(vec![Field::new(
3138                "value",
3139                DataType::Int64,
3140                false,
3141            )])),
3142        })
3143    }
3144
3145    fn batches(
3146        &self,
3147        request: &ExternalPlanRequest<'_>,
3148        control: &mongreldb_core::ExecutionControl,
3149    ) -> DFResult<Vec<RecordBatch>> {
3150        let mut values = Vec::new();
3151        let mut current = self.start;
3152        let mut scanned = 0;
3153        while if self.step > 0 {
3154            current <= self.stop
3155        } else {
3156            current >= self.stop
3157        } {
3158            external_checkpoint(Some(control), scanned)?;
3159            scanned += 1;
3160            if values.len() >= 1_000_000 {
3161                return Err(DataFusionError::Plan(
3162                    "series output is capped at 1,000,000 rows".into(),
3163                ));
3164            }
3165            if request
3166                .filters
3167                .iter()
3168                .all(|filter| series_filter_matches(filter, current).unwrap_or(true))
3169            {
3170                values.push(current);
3171                if request.limit.is_some_and(|limit| values.len() >= limit) {
3172                    break;
3173                }
3174            }
3175            current = current.saturating_add(self.step);
3176            if (self.step > 0 && current == i64::MAX) || (self.step < 0 && current == i64::MIN) {
3177                break;
3178            }
3179        }
3180        let batch = RecordBatch::try_new(
3181            self.schema.clone(),
3182            vec![Arc::new(Int64Array::from(values)) as ArrayRef],
3183        )?;
3184        Ok(vec![batch])
3185    }
3186}
3187
3188impl std::fmt::Debug for SeriesExternalTable {
3189    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3190        f.debug_struct("SeriesExternalTable")
3191            .field("start", &self.start)
3192            .field("stop", &self.stop)
3193            .field("step", &self.step)
3194            .finish_non_exhaustive()
3195    }
3196}
3197
3198impl ExternalTable for SeriesExternalTable {
3199    fn schema(&self) -> SchemaRef {
3200        self.schema.clone()
3201    }
3202
3203    fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
3204        Ok(ExternalPlan::new(
3205            request
3206                .filters
3207                .iter()
3208                .map(|filter| {
3209                    if series_filter_supported(filter) {
3210                        TableProviderFilterPushDown::Exact
3211                    } else {
3212                        TableProviderFilterPushDown::Unsupported
3213                    }
3214                })
3215                .collect(),
3216            None,
3217            1.0,
3218            self.step > 0,
3219        ))
3220    }
3221
3222    fn scan_with_control(
3223        &self,
3224        request: &ExternalPlanRequest<'_>,
3225        control: &mongreldb_core::ExecutionControl,
3226    ) -> DFResult<ExternalScan> {
3227        external_checkpoint(Some(control), 0)?;
3228        let batches = self.batches(request, control)?;
3229        external_checkpoint(Some(control), 0)?;
3230        project_scan_with_control(
3231            self.schema.clone(),
3232            batches,
3233            request.projection.as_deref(),
3234            request.limit,
3235            Some(control),
3236        )
3237    }
3238}
3239
3240fn series_filter_supported(filter: &ExternalFilter) -> bool {
3241    match filter {
3242        ExternalFilter::And(filters) => filters.iter().all(series_filter_supported),
3243        ExternalFilter::Compare {
3244            column_index,
3245            value,
3246            ..
3247        } => *column_index == 0 && literal_i64(value).is_some(),
3248        ExternalFilter::Unsupported { .. } => false,
3249    }
3250}
3251
3252fn series_filter_matches(filter: &ExternalFilter, value: i64) -> Option<bool> {
3253    match filter {
3254        ExternalFilter::And(filters) => filters.iter().try_fold(true, |matches, filter| {
3255            Some(matches && series_filter_matches(filter, value)?)
3256        }),
3257        ExternalFilter::Compare {
3258            column_index,
3259            op,
3260            value: literal,
3261        } if *column_index == 0 => {
3262            let literal = literal_i64(literal)?;
3263            Some(match op {
3264                ExternalFilterOp::Eq => value == literal,
3265                ExternalFilterOp::NotEq => value != literal,
3266                ExternalFilterOp::Gt => value > literal,
3267                ExternalFilterOp::GtEq => value >= literal,
3268                ExternalFilterOp::Lt => value < literal,
3269                ExternalFilterOp::LtEq => value <= literal,
3270            })
3271        }
3272        _ => None,
3273    }
3274}
3275
3276fn literal_i64(value: &ScalarValue) -> Option<i64> {
3277    match value {
3278        ScalarValue::Int8(Some(v)) => Some(*v as i64),
3279        ScalarValue::Int16(Some(v)) => Some(*v as i64),
3280        ScalarValue::Int32(Some(v)) => Some(*v as i64),
3281        ScalarValue::Int64(Some(v)) => Some(*v),
3282        ScalarValue::UInt8(Some(v)) => Some(*v as i64),
3283        ScalarValue::UInt16(Some(v)) => Some(*v as i64),
3284        ScalarValue::UInt32(Some(v)) => Some(*v as i64),
3285        ScalarValue::UInt64(Some(v)) => i64::try_from(*v).ok(),
3286        _ => None,
3287    }
3288}
3289
3290fn literal_f64(value: &ScalarValue) -> Option<f64> {
3291    match value {
3292        ScalarValue::Float32(Some(v)) => Some(*v as f64),
3293        ScalarValue::Float64(Some(v)) => Some(*v),
3294        _ => literal_i64(value).map(|value| value as f64),
3295    }
3296}
3297
3298fn literal_string(value: &ScalarValue) -> Option<String> {
3299    match value {
3300        ScalarValue::Utf8(Some(v))
3301        | ScalarValue::Utf8View(Some(v))
3302        | ScalarValue::LargeUtf8(Some(v)) => Some(v.clone()),
3303        ScalarValue::Binary(Some(v))
3304        | ScalarValue::BinaryView(Some(v))
3305        | ScalarValue::LargeBinary(Some(v)) => String::from_utf8(v.clone()).ok(),
3306        _ => None,
3307    }
3308}
3309
3310fn series_args(entry: &ExternalTableEntry) -> Result<(i64, i64, i64)> {
3311    let values = entry
3312        .args
3313        .iter()
3314        .map(|arg| match arg {
3315            ModuleArg::Ident(value) | ModuleArg::String(value) | ModuleArg::Number(value) => {
3316                value.parse::<i64>().map_err(|e| {
3317                    MongrelQueryError::Schema(format!(
3318                        "series module argument {value:?} must be an integer: {e}"
3319                    ))
3320                })
3321            }
3322        })
3323        .collect::<Result<Vec<_>>>()?;
3324    match values.as_slice() {
3325        [] => Ok((0, -1, 1)),
3326        [stop] => Ok((0, *stop, 1)),
3327        [start, stop] => Ok((*start, *stop, 1)),
3328        [start, stop, step] => Ok((*start, *stop, *step)),
3329        _ => Err(MongrelQueryError::Schema(
3330            "series external table accepts at most three arguments".into(),
3331        )),
3332    }
3333}