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}