use crate::extended_sql_functions::{json_table_batches_from_text, JsonTableMode};
use crate::{arrow_conv, MongrelQueryError, Result};
use arrow::array::{ArrayRef, Int64Array, StringArray};
use arrow::datatypes::{DataType, Field, Schema as ArrowSchema, SchemaRef};
use arrow::record_batch::{RecordBatch, RecordBatchOptions};
use datafusion::catalog::{Session, TableProvider};
use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
use datafusion::logical_expr::{
BinaryExpr, Expr, Operator, TableProviderFilterPushDown, TableType,
};
use datafusion::physical_plan::ExecutionPlan;
use datafusion_datasource::memory::MemorySourceConfig;
use mongreldb_core::catalog::TableState;
use mongreldb_core::database::{TABLES_DIR, VTAB_DIR};
use mongreldb_core::schema::{
ColumnDef as CoreColumnDef, ColumnFlags, Schema as CoreSchema, TypeId,
};
use mongreldb_core::{Database, ExternalTableEntry, ModuleArg, ModuleCapabilities, Value};
use serde::{Deserialize, Serialize};
use std::any::Any;
use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::Arc;
pub trait ExternalTableModule: Send + Sync {
fn name(&self) -> &str;
fn descriptor(&self) -> ExternalModuleDescriptor;
fn indexes(&self, _entry: &ExternalTableEntry) -> Result<Vec<ExternalModuleIndex>> {
Ok(Vec::new())
}
fn connect(
&self,
ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
) -> Result<Arc<dyn ExternalTable>>;
fn read_rows(
&self,
_ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
) -> Result<Vec<HashMap<u16, Value>>> {
Err(MongrelQueryError::Schema(format!(
"external table {:?} using module {:?} is not row-writable",
entry.name,
self.name()
)))
}
fn prepare_rows(
&self,
_ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
_rows: Vec<HashMap<u16, Value>>,
) -> Result<Vec<u8>> {
Err(MongrelQueryError::Schema(format!(
"external table {:?} using module {:?} is not row-writable",
entry.name,
self.name()
)))
}
fn rows_from_state(&self, state: &[u8]) -> Result<Vec<HashMap<u16, Value>>> {
if state.is_empty() {
Ok(Vec::new())
} else {
decode_state_rows(state)
}
}
fn write(
&self,
ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
op: ExternalWriteOp,
txn: &mut ExternalTxn,
) -> Result<ExternalWriteResult> {
let (rows, changes) = match op {
ExternalWriteOp::Insert { rows: inserted } => {
let changes = inserted.len() as u64;
let mut rows = txn.read_rows()?;
rows.extend(inserted);
(rows, changes)
}
ExternalWriteOp::ReplaceRows { rows, changes } => (rows, changes),
};
txn.replace_state(self.prepare_rows(ctx, entry, rows)?);
Ok(ExternalWriteResult { changes })
}
fn destroy(&self, _ctx: &ModuleConnectCtx<'_>, _entry: &ExternalTableEntry) -> Result<()> {
Ok(())
}
}
pub trait ExternalTable: std::fmt::Debug + Send + Sync {
fn schema(&self) -> SchemaRef;
fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan>;
fn scan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalScan>;
}
#[derive(Clone)]
pub struct ExternalModuleDescriptor {
pub schema: CoreSchema,
pub hidden_columns: Vec<String>,
pub capabilities: ModuleCapabilities,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExternalModuleIndex {
pub name: String,
pub column_ids: Vec<u16>,
pub unique: bool,
pub partial: bool,
}
impl ExternalModuleIndex {
pub fn new(name: impl Into<String>, column_ids: Vec<u16>) -> Self {
Self {
name: name.into(),
column_ids,
unique: false,
partial: false,
}
}
}
pub struct ModuleConnectCtx<'a> {
pub db: &'a Arc<Database>,
}
impl ModuleConnectCtx<'_> {
pub fn raw_state(&self, entry: &ExternalTableEntry) -> Result<Vec<u8>> {
external_table_state_bytes(self.db, entry)
}
pub fn read_state(&self, entry: &ExternalTableEntry, key: &[u8]) -> Result<Option<Vec<u8>>> {
ExternalTxn::read_state_from_bytes(&self.raw_state(entry)?, key)
}
}
pub struct ExternalPlanRequest<'a> {
pub projection: Option<Vec<usize>>,
pub filters: Vec<ExternalFilter>,
pub raw_filters: Vec<&'a Expr>,
pub order_by: Vec<ExternalOrder>,
pub limit: Option<usize>,
pub offset: Option<usize>,
}
pub struct ExternalScan {
pub schema: SchemaRef,
pub batches: Vec<RecordBatch>,
}
#[derive(Clone)]
pub struct ExternalPlan {
pub filter_pushdown: Vec<TableProviderFilterPushDown>,
pub accepted_filters: Vec<AcceptedFilter>,
pub residual_filters_required: bool,
pub estimated_rows: Option<u64>,
pub estimated_cost: f64,
pub order_satisfied: bool,
pub opaque: Arc<dyn Any + Send + Sync>,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ExternalOrder {
pub column_index: usize,
pub descending: bool,
pub nulls_first: Option<bool>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum ExternalFilter {
And(Vec<ExternalFilter>),
Compare {
column_index: usize,
op: ExternalFilterOp,
value: ScalarValue,
},
Unsupported {
expr: String,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ExternalFilterOp {
Eq,
NotEq,
Gt,
GtEq,
Lt,
LtEq,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AcceptedFilter {
pub filter_index: usize,
pub pushdown: TableProviderFilterPushDown,
}
impl ExternalPlan {
pub fn new(
filter_pushdown: Vec<TableProviderFilterPushDown>,
estimated_rows: Option<u64>,
estimated_cost: f64,
order_satisfied: bool,
) -> Self {
let accepted_filters = filter_pushdown
.iter()
.enumerate()
.filter_map(|(filter_index, pushdown)| match pushdown {
TableProviderFilterPushDown::Exact | TableProviderFilterPushDown::Inexact => {
Some(AcceptedFilter {
filter_index,
pushdown: pushdown.clone(),
})
}
TableProviderFilterPushDown::Unsupported => None,
})
.collect();
let residual_filters_required = filter_pushdown.iter().any(|pushdown| {
matches!(
pushdown,
TableProviderFilterPushDown::Inexact | TableProviderFilterPushDown::Unsupported
)
});
Self {
filter_pushdown,
accepted_filters,
residual_filters_required,
estimated_rows,
estimated_cost,
order_satisfied,
opaque: Arc::new(()),
}
}
pub fn with_opaque(mut self, opaque: Arc<dyn Any + Send + Sync>) -> Self {
self.opaque = opaque;
self
}
}
fn external_filter_from_expr(expr: &Expr, schema: &ArrowSchema) -> ExternalFilter {
match expr {
Expr::BinaryExpr(BinaryExpr { left, op, right }) if *op == Operator::And => {
ExternalFilter::And(vec![
external_filter_from_expr(left, schema),
external_filter_from_expr(right, schema),
])
}
Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
let Some(op) = external_filter_op(*op) else {
return ExternalFilter::Unsupported {
expr: expr.to_string(),
};
};
match (left.as_ref(), right.as_ref()) {
(Expr::Column(column), Expr::Literal(value, _)) => schema
.index_of(&column.name)
.map(|column_index| ExternalFilter::Compare {
column_index,
op,
value: value.clone(),
})
.unwrap_or_else(|_| ExternalFilter::Unsupported {
expr: expr.to_string(),
}),
(Expr::Literal(value, _), Expr::Column(column)) => schema
.index_of(&column.name)
.map(|column_index| ExternalFilter::Compare {
column_index,
op: op.flipped(),
value: value.clone(),
})
.unwrap_or_else(|_| ExternalFilter::Unsupported {
expr: expr.to_string(),
}),
_ => ExternalFilter::Unsupported {
expr: expr.to_string(),
},
}
}
_ => ExternalFilter::Unsupported {
expr: expr.to_string(),
},
}
}
fn external_filter_op(op: Operator) -> Option<ExternalFilterOp> {
match op {
Operator::Eq => Some(ExternalFilterOp::Eq),
Operator::NotEq => Some(ExternalFilterOp::NotEq),
Operator::Gt => Some(ExternalFilterOp::Gt),
Operator::GtEq => Some(ExternalFilterOp::GtEq),
Operator::Lt => Some(ExternalFilterOp::Lt),
Operator::LtEq => Some(ExternalFilterOp::LtEq),
_ => None,
}
}
impl ExternalFilterOp {
fn flipped(self) -> Self {
match self {
ExternalFilterOp::Eq => ExternalFilterOp::Eq,
ExternalFilterOp::NotEq => ExternalFilterOp::NotEq,
ExternalFilterOp::Gt => ExternalFilterOp::Lt,
ExternalFilterOp::GtEq => ExternalFilterOp::LtEq,
ExternalFilterOp::Lt => ExternalFilterOp::Gt,
ExternalFilterOp::LtEq => ExternalFilterOp::GtEq,
}
}
}
#[derive(Clone)]
pub enum ExternalWriteOp {
Insert {
rows: Vec<HashMap<u16, Value>>,
},
ReplaceRows {
rows: Vec<HashMap<u16, Value>>,
changes: u64,
},
}
#[derive(Debug, Clone, PartialEq)]
pub enum ExternalBaseWrite {
Put {
table: String,
cells: Vec<(u16, Value)>,
},
Delete {
table: String,
row_id: u64,
},
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ExternalWriteResult {
pub changes: u64,
}
impl ExternalWriteResult {
pub fn new(changes: u64) -> Self {
Self { changes }
}
}
pub struct ExternalTxn {
state: Vec<u8>,
base_writes: Vec<ExternalBaseWrite>,
}
impl ExternalTxn {
pub fn new(state: Vec<u8>) -> Self {
Self {
state,
base_writes: Vec::new(),
}
}
pub fn read_state_from_bytes(state: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>> {
Ok(decode_kv_state(state)?.remove(key))
}
pub fn read_state(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
Self::read_state_from_bytes(&self.state, key)
}
pub fn put_state(&mut self, key: &[u8], value: &[u8]) -> Result<()> {
let mut state = decode_kv_state(&self.state)?;
state.insert(key.to_vec(), value.to_vec());
self.state = encode_kv_state(&state)?;
Ok(())
}
pub fn delete_state(&mut self, key: &[u8]) -> Result<()> {
let mut state = decode_kv_state(&self.state)?;
state.remove(key);
self.state = encode_kv_state(&state)?;
Ok(())
}
pub fn replace_state(&mut self, state: Vec<u8>) {
self.state = state;
}
pub fn read_rows(&self) -> Result<Vec<HashMap<u16, Value>>> {
if self.state.is_empty() {
Ok(Vec::new())
} else {
decode_state_rows(&self.state)
}
}
pub fn replace_rows(
&mut self,
schema: &CoreSchema,
rows: Vec<HashMap<u16, Value>>,
) -> Result<()> {
validate_external_rows(schema, &rows)?;
self.state = encode_state_rows(&rows)?;
Ok(())
}
pub fn emit_base_write(&mut self, op: ExternalBaseWrite) -> Result<()> {
self.base_writes.push(op);
Ok(())
}
fn into_parts(self) -> (Vec<u8>, Vec<ExternalBaseWrite>) {
(self.state, self.base_writes)
}
}
pub struct ExternalModuleRegistry {
modules: parking_lot::RwLock<HashMap<String, Arc<dyn ExternalTableModule>>>,
}
impl Default for ExternalModuleRegistry {
fn default() -> Self {
let registry = Self {
modules: parking_lot::RwLock::new(HashMap::new()),
};
for module in builtin_modules() {
registry.register(module).expect("valid built-in module");
}
registry
}
}
impl ExternalModuleRegistry {
pub fn register(&self, module: Arc<dyn ExternalTableModule>) -> Result<()> {
let name = normalize_module_name(module.name())?;
self.modules.write().insert(name, module);
Ok(())
}
pub fn names(&self) -> Vec<String> {
let mut names = self.modules.read().keys().cloned().collect::<Vec<_>>();
names.sort();
names
}
pub(crate) fn contains(&self, name: &str) -> bool {
normalize_module_name(name)
.ok()
.is_some_and(|name| self.modules.read().contains_key(&name))
}
pub(crate) fn descriptor(&self, name: &str) -> Result<ExternalModuleDescriptor> {
Ok(self.module(name)?.descriptor())
}
pub(crate) fn external_table_provider(
&self,
db: &Arc<Database>,
entry: &ExternalTableEntry,
) -> Result<Arc<dyn TableProvider>> {
let module = self.module(&entry.module)?;
let table = module.connect(&ModuleConnectCtx { db }, entry)?;
Ok(Arc::new(ExternalTableProvider::new(
table,
entry.capabilities,
)))
}
pub(crate) fn external_table_indexes(
&self,
entry: &ExternalTableEntry,
) -> Result<Vec<ExternalModuleIndex>> {
let module = self.module(&entry.module)?;
module.indexes(entry)
}
pub(crate) fn external_table_rows(
&self,
db: &Arc<Database>,
entry: &ExternalTableEntry,
) -> Result<Vec<HashMap<u16, Value>>> {
let module = self.module(&entry.module)?;
module.read_rows(&ModuleConnectCtx { db }, entry)
}
pub(crate) fn external_table_rows_from_state(
&self,
entry: &ExternalTableEntry,
state: &[u8],
) -> Result<Vec<HashMap<u16, Value>>> {
let module = self.module(&entry.module)?;
module.rows_from_state(state)
}
pub(crate) fn external_table_write(
&self,
db: &Arc<Database>,
entry: &ExternalTableEntry,
base_state: Vec<u8>,
op: ExternalWriteOp,
) -> Result<(Vec<u8>, ExternalWriteResult, Vec<ExternalBaseWrite>)> {
let module = self.module(&entry.module)?;
let mut txn = ExternalTxn::new(base_state);
let result = module.write(&ModuleConnectCtx { db }, entry, op, &mut txn)?;
let (state, base_writes) = txn.into_parts();
Ok((state, result, base_writes))
}
pub(crate) fn destroy_external_table(
&self,
db: &Arc<Database>,
entry: &ExternalTableEntry,
) -> Result<()> {
let module = self.module(&entry.module)?;
module.destroy(&ModuleConnectCtx { db }, entry)
}
fn module(&self, name: &str) -> Result<Arc<dyn ExternalTableModule>> {
let name = normalize_module_name(name)?;
self.modules.read().get(&name).cloned().ok_or_else(|| {
MongrelQueryError::Schema(format!("external table module {name:?} is not registered"))
})
}
}
fn normalize_module_name(name: &str) -> Result<String> {
let name = name.trim().to_ascii_lowercase();
if name.is_empty() || !name.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'_') {
return Err(MongrelQueryError::Schema(format!(
"invalid external table module name {name:?}"
)));
}
Ok(name)
}
fn builtin_modules() -> Vec<Arc<dyn ExternalTableModule>> {
vec![
Arc::new(JsonModule {
name: "json_each",
mode: JsonTableMode::Each,
}),
Arc::new(JsonModule {
name: "json_tree",
mode: JsonTableMode::Tree,
}),
Arc::new(JsonModule {
name: "jsonb_each",
mode: JsonTableMode::Each,
}),
Arc::new(JsonModule {
name: "jsonb_tree",
mode: JsonTableMode::Tree,
}),
Arc::new(SchemaTablesModule),
Arc::new(DbStatModule),
Arc::new(FtsDocsModule),
Arc::new(KvStoreModule),
Arc::new(RTreeRectsModule),
Arc::new(SeriesModule),
]
}
struct ExternalTableProvider {
table: Arc<dyn ExternalTable>,
cache_plans: bool,
plan_cache: parking_lot::Mutex<HashMap<ExternalPlanCacheKey, ExternalPlan>>,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct ExternalPlanCacheKey {
projection: Option<Vec<usize>>,
filters: Vec<String>,
order_by: Vec<ExternalOrder>,
limit: Option<usize>,
offset: Option<usize>,
}
impl ExternalPlanCacheKey {
fn from_request(request: &ExternalPlanRequest<'_>) -> Self {
Self {
projection: request.projection.clone(),
filters: request
.filters
.iter()
.map(|filter| format!("{filter:?}"))
.collect(),
order_by: request.order_by.clone(),
limit: request.limit,
offset: request.offset,
}
}
}
impl ExternalTableProvider {
fn new(table: Arc<dyn ExternalTable>, capabilities: ModuleCapabilities) -> Self {
Self {
table,
cache_plans: capabilities.read_only && capabilities.deterministic,
plan_cache: parking_lot::Mutex::new(HashMap::new()),
}
}
fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
if !self.cache_plans {
return self.table.plan(request);
}
let key = ExternalPlanCacheKey::from_request(request);
if let Some(plan) = self.plan_cache.lock().get(&key).cloned() {
return Ok(plan);
}
let plan = self.table.plan(request)?;
self.plan_cache.lock().insert(key, plan.clone());
Ok(plan)
}
}
impl std::fmt::Debug for ExternalTableProvider {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ExternalTableProvider")
.field("table", &self.table)
.finish()
}
}
#[async_trait::async_trait]
impl TableProvider for ExternalTableProvider {
fn schema(&self) -> SchemaRef {
self.table.schema()
}
fn table_type(&self) -> TableType {
TableType::Base
}
fn supports_filters_pushdown(
&self,
filters: &[&Expr],
) -> DFResult<Vec<TableProviderFilterPushDown>> {
let schema = self.table.schema();
let request = ExternalPlanRequest {
projection: None,
filters: filters
.iter()
.map(|filter| external_filter_from_expr(filter, &schema))
.collect(),
raw_filters: filters.to_vec(),
order_by: Vec::new(),
limit: None,
offset: None,
};
let ExternalPlan {
filter_pushdown,
accepted_filters: _accepted_filters,
residual_filters_required: _residual_filters_required,
estimated_rows: _estimated_rows,
estimated_cost: _estimated_cost,
order_satisfied: _order_satisfied,
opaque: _opaque,
} = self.plan(&request)?;
Ok(filter_pushdown)
}
async fn scan(
&self,
_state: &dyn Session,
projection: Option<&Vec<usize>>,
filters: &[Expr],
limit: Option<usize>,
) -> DFResult<Arc<dyn ExecutionPlan>> {
mongreldb_core::trace::QueryTrace::record(|t| {
t.scan_mode = mongreldb_core::trace::ScanMode::ExternalModule;
});
let schema = self.table.schema();
let request = ExternalPlanRequest {
projection: projection.cloned(),
filters: filters
.iter()
.map(|filter| external_filter_from_expr(filter, &schema))
.collect(),
raw_filters: filters.iter().collect(),
order_by: Vec::new(),
limit,
offset: None,
};
let scan = self.table.scan(&request)?;
let partitions = vec![scan.batches];
let exec = MemorySourceConfig::try_new_exec(&partitions, scan.schema, None)?;
Ok(exec)
}
}
struct JsonModule {
name: &'static str,
mode: JsonTableMode,
}
impl ExternalTableModule for JsonModule {
fn name(&self) -> &str {
self.name
}
fn descriptor(&self) -> ExternalModuleDescriptor {
json_descriptor()
}
fn connect(
&self,
_ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
) -> Result<Arc<dyn ExternalTable>> {
let (json, root) = json_args(entry, self.name)?;
let (schema, batches) =
json_table_batches_from_text(self.name, &json, root.as_deref(), self.mode)
.map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
Ok(Arc::new(JsonExternalTable {
module: self.name,
schema,
batches,
}))
}
}
#[derive(Clone)]
struct JsonExternalTable {
module: &'static str,
schema: SchemaRef,
batches: Vec<RecordBatch>,
}
impl std::fmt::Debug for JsonExternalTable {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("JsonExternalTable")
.field("module", &self.module)
.field("schema", &self.schema)
.finish_non_exhaustive()
}
}
impl ExternalTable for JsonExternalTable {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
Ok(ExternalPlan::new(
request
.filters
.iter()
.map(|_| TableProviderFilterPushDown::Unsupported)
.collect(),
None,
1.0,
false,
))
}
fn scan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalScan> {
project_scan(
self.schema.clone(),
self.batches.clone(),
request.projection.as_deref(),
request.limit,
)
}
}
fn json_descriptor() -> ExternalModuleDescriptor {
ExternalModuleDescriptor {
schema: CoreSchema {
schema_id: 0,
columns: vec![
json_column(1, "key", TypeId::Bytes, true),
json_column(2, "value", TypeId::Bytes, true),
json_column(3, "type", TypeId::Bytes, true),
json_column(4, "atom", TypeId::Bytes, true),
json_column(5, "id", TypeId::Int64, false),
json_column(6, "parent", TypeId::Int64, true),
json_column(7, "fullkey", TypeId::Bytes, true),
json_column(8, "path", TypeId::Bytes, true),
json_column(9, "json", TypeId::Bytes, true),
json_column(10, "root", TypeId::Bytes, true),
],
indexes: Vec::new(),
colocation: Vec::new(),
constraints: Default::default(),
},
hidden_columns: vec!["json".to_string(), "root".to_string()],
capabilities: ModuleCapabilities {
read_only: true,
deterministic: true,
trigger_safe: true,
..ModuleCapabilities::default()
},
}
}
fn json_column(id: u16, name: &str, ty: TypeId, nullable: bool) -> CoreColumnDef {
CoreColumnDef {
id,
name: name.to_string(),
ty,
flags: if nullable {
ColumnFlags::empty().with(ColumnFlags::NULLABLE)
} else {
ColumnFlags::empty()
},
}
}
fn json_args(entry: &ExternalTableEntry, module: &str) -> Result<(String, Option<String>)> {
match entry.args.as_slice() {
[json] => Ok((module_arg_string(json).to_string(), None)),
[json, root] => Ok((
module_arg_string(json).to_string(),
Some(module_arg_string(root).to_string()),
)),
_ => Err(MongrelQueryError::Schema(format!(
"{module} external table requires one JSON argument and an optional root path"
))),
}
}
fn module_arg_string(arg: &ModuleArg) -> &str {
match arg {
ModuleArg::Ident(value) | ModuleArg::String(value) | ModuleArg::Number(value) => value,
}
}
struct SchemaTablesModule;
impl ExternalTableModule for SchemaTablesModule {
fn name(&self) -> &str {
"schema_tables"
}
fn descriptor(&self) -> ExternalModuleDescriptor {
ExternalModuleDescriptor {
schema: CoreSchema {
schema_id: 0,
columns: vec![
catalog_column(1, "schema_name", TypeId::Bytes, false),
catalog_column(2, "name", TypeId::Bytes, false),
catalog_column(3, "type", TypeId::Bytes, false),
catalog_column(4, "ncol", TypeId::Int64, false),
catalog_column(5, "module", TypeId::Bytes, true),
catalog_column(6, "created_epoch", TypeId::Int64, false),
],
indexes: Vec::new(),
colocation: Vec::new(),
constraints: Default::default(),
},
hidden_columns: Vec::new(),
capabilities: catalog_capabilities(),
}
}
fn connect(
&self,
ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
) -> Result<Arc<dyn ExternalTable>> {
ensure_no_args(entry, self.name())?;
Ok(Arc::new(SchemaTablesExternalTable {
db: Arc::clone(ctx.db),
schema: schema_tables_arrow_schema(),
}))
}
}
#[derive(Clone)]
struct SchemaTablesExternalTable {
db: Arc<Database>,
schema: SchemaRef,
}
impl std::fmt::Debug for SchemaTablesExternalTable {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SchemaTablesExternalTable")
.field("schema", &self.schema)
.finish_non_exhaustive()
}
}
impl ExternalTable for SchemaTablesExternalTable {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
unsupported_plan(request)
}
fn scan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalScan> {
let batches = schema_tables_batches(&self.db, self.schema.clone())?;
project_scan(
self.schema.clone(),
batches,
request.projection.as_deref(),
request.limit,
)
}
}
struct DbStatModule;
impl ExternalTableModule for DbStatModule {
fn name(&self) -> &str {
"dbstat"
}
fn descriptor(&self) -> ExternalModuleDescriptor {
ExternalModuleDescriptor {
schema: CoreSchema {
schema_id: 0,
columns: vec![
catalog_column(1, "name", TypeId::Bytes, false),
catalog_column(2, "type", TypeId::Bytes, false),
catalog_column(3, "rows", TypeId::Int64, true),
catalog_column(4, "runs", TypeId::Int64, true),
catalog_column(5, "memtable_rows", TypeId::Int64, true),
catalog_column(6, "columns", TypeId::Int64, false),
catalog_column(7, "storage_bytes", TypeId::Int64, false),
],
indexes: Vec::new(),
colocation: Vec::new(),
constraints: Default::default(),
},
hidden_columns: Vec::new(),
capabilities: catalog_capabilities(),
}
}
fn connect(
&self,
ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
) -> Result<Arc<dyn ExternalTable>> {
ensure_no_args(entry, self.name())?;
Ok(Arc::new(DbStatExternalTable {
db: Arc::clone(ctx.db),
schema: dbstat_arrow_schema(),
}))
}
}
#[derive(Clone)]
struct DbStatExternalTable {
db: Arc<Database>,
schema: SchemaRef,
}
impl std::fmt::Debug for DbStatExternalTable {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DbStatExternalTable")
.field("schema", &self.schema)
.finish_non_exhaustive()
}
}
impl ExternalTable for DbStatExternalTable {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
unsupported_plan(request)
}
fn scan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalScan> {
let batches = dbstat_batches(&self.db, self.schema.clone())?;
project_scan(
self.schema.clone(),
batches,
request.projection.as_deref(),
request.limit,
)
}
}
fn catalog_column(id: u16, name: &str, ty: TypeId, nullable: bool) -> CoreColumnDef {
CoreColumnDef {
id,
name: name.to_string(),
ty,
flags: if nullable {
ColumnFlags::empty().with(ColumnFlags::NULLABLE)
} else {
ColumnFlags::empty()
},
}
}
fn catalog_capabilities() -> ModuleCapabilities {
ModuleCapabilities {
read_only: true,
trigger_safe: true,
..ModuleCapabilities::default()
}
}
fn ensure_no_args(entry: &ExternalTableEntry, module: &str) -> Result<()> {
if entry.args.is_empty() {
Ok(())
} else {
Err(MongrelQueryError::Schema(format!(
"{module} external table does not accept arguments"
)))
}
}
fn schema_tables_arrow_schema() -> SchemaRef {
Arc::new(ArrowSchema::new(vec![
Field::new("schema_name", DataType::Utf8, false),
Field::new("name", DataType::Utf8, false),
Field::new("type", DataType::Utf8, false),
Field::new("ncol", DataType::Int64, false),
Field::new("module", DataType::Utf8, true),
Field::new("created_epoch", DataType::Int64, false),
]))
}
fn schema_tables_batches(db: &Arc<Database>, schema: SchemaRef) -> DFResult<Vec<RecordBatch>> {
struct Row {
name: String,
ty: String,
ncol: i64,
module: Option<String>,
created_epoch: i64,
}
let catalog = db.catalog_snapshot();
let mut rows = Vec::new();
for table in catalog
.tables
.into_iter()
.filter(|table| matches!(table.state, TableState::Live))
{
rows.push(Row {
name: table.name,
ty: "table".to_string(),
ncol: saturating_i64(table.schema.columns.len() as u64),
module: None,
created_epoch: saturating_i64(table.created_epoch),
});
}
for table in catalog.external_tables {
rows.push(Row {
name: table.name,
ty: "external".to_string(),
ncol: saturating_i64(table.declared_schema.columns.len() as u64),
module: Some(table.module),
created_epoch: saturating_i64(table.created_epoch),
});
}
rows.sort_by(|a, b| a.name.cmp(&b.name).then_with(|| a.ty.cmp(&b.ty)));
let schema_names = vec!["main".to_string(); rows.len()];
let names = rows.iter().map(|row| row.name.clone()).collect::<Vec<_>>();
let types = rows.iter().map(|row| row.ty.clone()).collect::<Vec<_>>();
let ncols = rows.iter().map(|row| row.ncol).collect::<Vec<_>>();
let modules = rows
.iter()
.map(|row| row.module.clone())
.collect::<Vec<_>>();
let created_epochs = rows.iter().map(|row| row.created_epoch).collect::<Vec<_>>();
Ok(vec![RecordBatch::try_new(
schema,
vec![
Arc::new(StringArray::from(schema_names)) as ArrayRef,
Arc::new(StringArray::from(names)),
Arc::new(StringArray::from(types)),
Arc::new(Int64Array::from(ncols)),
Arc::new(StringArray::from(modules)),
Arc::new(Int64Array::from(created_epochs)),
],
)?])
}
fn dbstat_arrow_schema() -> SchemaRef {
Arc::new(ArrowSchema::new(vec![
Field::new("name", DataType::Utf8, false),
Field::new("type", DataType::Utf8, false),
Field::new("rows", DataType::Int64, true),
Field::new("runs", DataType::Int64, true),
Field::new("memtable_rows", DataType::Int64, true),
Field::new("columns", DataType::Int64, false),
Field::new("storage_bytes", DataType::Int64, false),
]))
}
fn dbstat_batches(db: &Arc<Database>, schema: SchemaRef) -> DFResult<Vec<RecordBatch>> {
struct Row {
name: String,
ty: String,
rows: Option<i64>,
runs: Option<i64>,
memtable_rows: Option<i64>,
columns: i64,
storage_bytes: i64,
}
let catalog = db.catalog_snapshot();
let mut rows = Vec::new();
for table in catalog
.tables
.into_iter()
.filter(|table| matches!(table.state, TableState::Live))
{
let handle = db
.table(&table.name)
.map_err(|e| DataFusionError::Execution(e.to_string()))?;
let table_guard = handle.lock();
let table_dir = db.root().join(TABLES_DIR).join(table.table_id.to_string());
rows.push(Row {
name: table.name,
ty: "table".to_string(),
rows: Some(saturating_i64(table_guard.count())),
runs: Some(saturating_i64(table_guard.run_count() as u64)),
memtable_rows: Some(saturating_i64(table_guard.memtable_len() as u64)),
columns: saturating_i64(table.schema.columns.len() as u64),
storage_bytes: saturating_i64(dir_size(&table_dir)?),
});
}
for table in catalog.external_tables {
let state_dir = db.root().join(VTAB_DIR).join(&table.name);
rows.push(Row {
name: table.name,
ty: "external".to_string(),
rows: None,
runs: None,
memtable_rows: None,
columns: saturating_i64(table.declared_schema.columns.len() as u64),
storage_bytes: saturating_i64(dir_size(&state_dir)?),
});
}
rows.sort_by(|a, b| a.name.cmp(&b.name).then_with(|| a.ty.cmp(&b.ty)));
let names = rows.iter().map(|row| row.name.clone()).collect::<Vec<_>>();
let types = rows.iter().map(|row| row.ty.clone()).collect::<Vec<_>>();
let row_counts = rows.iter().map(|row| row.rows).collect::<Vec<_>>();
let runs = rows.iter().map(|row| row.runs).collect::<Vec<_>>();
let memtable_rows = rows.iter().map(|row| row.memtable_rows).collect::<Vec<_>>();
let columns = rows.iter().map(|row| row.columns).collect::<Vec<_>>();
let storage_bytes = rows.iter().map(|row| row.storage_bytes).collect::<Vec<_>>();
Ok(vec![RecordBatch::try_new(
schema,
vec![
Arc::new(StringArray::from(names)) as ArrayRef,
Arc::new(StringArray::from(types)),
Arc::new(Int64Array::from(row_counts)),
Arc::new(Int64Array::from(runs)),
Arc::new(Int64Array::from(memtable_rows)),
Arc::new(Int64Array::from(columns)),
Arc::new(Int64Array::from(storage_bytes)),
],
)?])
}
fn dir_size(path: &Path) -> DFResult<u64> {
if !path.exists() {
return Ok(0);
}
let metadata = std::fs::symlink_metadata(path)
.map_err(|e| DataFusionError::Execution(format!("stat {:?}: {e}", path)))?;
if metadata.is_file() {
return Ok(metadata.len());
}
if !metadata.is_dir() {
return Ok(0);
}
let mut total = 0_u64;
for entry in std::fs::read_dir(path).map_err(|e| DataFusionError::Execution(format!("{e}")))? {
let entry = entry.map_err(|e| DataFusionError::Execution(format!("{e}")))?;
total = total.saturating_add(dir_size(&entry.path())?);
}
Ok(total)
}
fn saturating_i64(value: u64) -> i64 {
i64::try_from(value).unwrap_or(i64::MAX)
}
fn unsupported_plan(request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
Ok(ExternalPlan::new(
request
.filters
.iter()
.map(|_| TableProviderFilterPushDown::Unsupported)
.collect(),
None,
1.0,
false,
))
}
fn project_scan(
full_schema: SchemaRef,
batches: Vec<RecordBatch>,
projection: Option<&[usize]>,
limit: Option<usize>,
) -> DFResult<ExternalScan> {
let Some(projection) = projection else {
return Ok(ExternalScan {
schema: full_schema,
batches: limit_batches(batches, limit),
});
};
let schema = Arc::new(ArrowSchema::new(
projection
.iter()
.map(|idx| full_schema.field(*idx).clone())
.collect::<Vec<_>>(),
));
let projected = batches
.into_iter()
.map(|batch| {
let columns = projection
.iter()
.map(|idx| batch.column(*idx).clone())
.collect::<Vec<_>>();
if columns.is_empty() {
RecordBatch::try_new_with_options(
schema.clone(),
columns,
&RecordBatchOptions::new().with_row_count(Some(batch.num_rows())),
)
} else {
RecordBatch::try_new(schema.clone(), columns)
}
})
.collect::<std::result::Result<Vec<_>, _>>()?;
Ok(ExternalScan {
schema,
batches: limit_batches(projected, limit),
})
}
fn limit_batches(batches: Vec<RecordBatch>, limit: Option<usize>) -> Vec<RecordBatch> {
let Some(mut remaining) = limit else {
return batches;
};
let mut limited = Vec::new();
for batch in batches {
if remaining == 0 {
break;
}
let take = remaining.min(batch.num_rows());
limited.push(batch.slice(0, take));
remaining -= take;
}
limited
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn project_scan_applies_limit_after_projection() {
let schema = Arc::new(ArrowSchema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("value", DataType::Int64, false),
]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef,
Arc::new(Int64Array::from(vec![10, 11, 12, 13])) as ArrayRef,
],
)
.unwrap();
let scan = project_scan(schema, vec![batch], Some(&[1]), Some(2)).unwrap();
assert_eq!(scan.schema.fields().len(), 1);
assert_eq!(scan.schema.field(0).name(), "value");
assert_eq!(scan.batches.len(), 1);
assert_eq!(scan.batches[0].num_columns(), 1);
assert_eq!(scan.batches[0].num_rows(), 2);
let values = scan.batches[0]
.column(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
assert_eq!(values.value(0), 10);
assert_eq!(values.value(1), 11);
}
#[test]
fn external_plan_derives_accepted_and_residual_filter_metadata() {
let plan = ExternalPlan::new(
vec![
TableProviderFilterPushDown::Exact,
TableProviderFilterPushDown::Inexact,
TableProviderFilterPushDown::Unsupported,
],
Some(42),
3.5,
true,
);
assert_eq!(
plan.accepted_filters,
vec![
AcceptedFilter {
filter_index: 0,
pushdown: TableProviderFilterPushDown::Exact,
},
AcceptedFilter {
filter_index: 1,
pushdown: TableProviderFilterPushDown::Inexact,
},
]
);
assert!(plan.residual_filters_required);
assert_eq!(plan.estimated_rows, Some(42));
assert_eq!(plan.estimated_cost, 3.5);
assert!(plan.order_satisfied);
}
}
struct KvStoreModule;
impl ExternalTableModule for KvStoreModule {
fn name(&self) -> &str {
"kv_store"
}
fn descriptor(&self) -> ExternalModuleDescriptor {
ExternalModuleDescriptor {
schema: kv_store_schema(),
hidden_columns: Vec::new(),
capabilities: ModuleCapabilities {
writable: true,
deterministic: true,
trigger_safe: true,
transaction_safe: true,
..ModuleCapabilities::default()
},
}
}
fn connect(
&self,
ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
) -> Result<Arc<dyn ExternalTable>> {
ensure_no_args(entry, self.name())?;
let rows = read_state_rows(ctx.db, entry)?;
let schema = arrow_conv::arrow_schema(&entry.declared_schema)?;
let batches = core_rows_to_batches(&entry.declared_schema, rows)?;
Ok(Arc::new(KvStoreExternalTable { schema, batches }))
}
fn read_rows(
&self,
ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
) -> Result<Vec<HashMap<u16, Value>>> {
ensure_no_args(entry, self.name())?;
read_state_rows(ctx.db, entry)
}
fn prepare_rows(
&self,
_ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
rows: Vec<HashMap<u16, Value>>,
) -> Result<Vec<u8>> {
ensure_no_args(entry, self.name())?;
validate_external_rows(&entry.declared_schema, &rows)?;
encode_state_rows(&rows)
}
}
#[derive(Clone)]
struct KvStoreExternalTable {
schema: SchemaRef,
batches: Vec<RecordBatch>,
}
impl std::fmt::Debug for KvStoreExternalTable {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("KvStoreExternalTable")
.field("schema", &self.schema)
.finish_non_exhaustive()
}
}
impl ExternalTable for KvStoreExternalTable {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
unsupported_plan(request)
}
fn scan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalScan> {
project_scan(
self.schema.clone(),
self.batches.clone(),
request.projection.as_deref(),
request.limit,
)
}
}
fn kv_store_schema() -> CoreSchema {
CoreSchema {
schema_id: 0,
columns: vec![
CoreColumnDef {
id: 1,
name: "key".to_string(),
ty: TypeId::Bytes,
flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
},
CoreColumnDef {
id: 2,
name: "value".to_string(),
ty: TypeId::Bytes,
flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
},
],
indexes: Vec::new(),
colocation: Vec::new(),
constraints: Default::default(),
}
}
#[derive(Serialize, Deserialize)]
struct ExternalRowState {
cells: Vec<(u16, Value)>,
}
const KV_STATE_MAGIC: &[u8] = b"mongreldb.external.kv.v1\n";
#[derive(Serialize, Deserialize)]
struct ExternalKvState {
entries: Vec<(Vec<u8>, Vec<u8>)>,
}
fn state_file(db: &Arc<Database>, entry: &ExternalTableEntry) -> PathBuf {
db.root()
.join(VTAB_DIR)
.join(&entry.name)
.join("state.json")
}
pub(crate) fn external_table_state_bytes(
db: &Arc<Database>,
entry: &ExternalTableEntry,
) -> Result<Vec<u8>> {
let path = state_file(db, entry);
if !path.exists() {
return Ok(Vec::new());
}
fs::read(&path)
.map_err(|e| MongrelQueryError::Schema(format!("read external state {:?}: {e}", path)))
}
fn read_state_rows(
db: &Arc<Database>,
entry: &ExternalTableEntry,
) -> Result<Vec<HashMap<u16, Value>>> {
let bytes = external_table_state_bytes(db, entry)?;
if bytes.is_empty() {
return Ok(Vec::new());
}
decode_state_rows(&bytes)
}
fn encode_state_rows(rows: &[HashMap<u16, Value>]) -> Result<Vec<u8>> {
let state = rows
.iter()
.map(|row| {
let mut cells = row
.iter()
.map(|(id, value)| (*id, value.clone()))
.collect::<Vec<_>>();
cells.sort_by_key(|(id, _)| *id);
ExternalRowState { cells }
})
.collect::<Vec<_>>();
serde_json::to_vec(&state)
.map_err(|e| MongrelQueryError::Schema(format!("encode external state: {e}")))
}
fn decode_state_rows(state: &[u8]) -> Result<Vec<HashMap<u16, Value>>> {
let rows: Vec<ExternalRowState> = serde_json::from_slice(state)
.map_err(|e| MongrelQueryError::Schema(format!("decode external state: {e}")))?;
Ok(rows
.into_iter()
.map(|row| row.cells.into_iter().collect())
.collect())
}
fn decode_kv_state(state: &[u8]) -> Result<BTreeMap<Vec<u8>, Vec<u8>>> {
if state.is_empty() {
return Ok(BTreeMap::new());
}
let json = state.strip_prefix(KV_STATE_MAGIC).ok_or_else(|| {
MongrelQueryError::Schema(
"external transaction state is not in key/value module format".into(),
)
})?;
let decoded: ExternalKvState = serde_json::from_slice(json)
.map_err(|e| MongrelQueryError::Schema(format!("decode external kv state: {e}")))?;
Ok(decoded.entries.into_iter().collect())
}
fn encode_kv_state(state: &BTreeMap<Vec<u8>, Vec<u8>>) -> Result<Vec<u8>> {
let encoded = ExternalKvState {
entries: state
.iter()
.map(|(key, value)| (key.clone(), value.clone()))
.collect(),
};
let mut out = KV_STATE_MAGIC.to_vec();
out.extend(
serde_json::to_vec(&encoded)
.map_err(|e| MongrelQueryError::Schema(format!("encode external kv state: {e}")))?,
);
Ok(out)
}
fn validate_external_rows(schema: &CoreSchema, rows: &[HashMap<u16, Value>]) -> Result<()> {
let mut known = HashSet::new();
let mut pk_cols = Vec::new();
for column in &schema.columns {
known.insert(column.id);
if column.flags.contains(ColumnFlags::PRIMARY_KEY) {
pk_cols.push(column.id);
}
}
let mut keys = HashSet::new();
for row in rows {
for column_id in row.keys() {
if !known.contains(column_id) {
return Err(MongrelQueryError::Schema(format!(
"external row contains unknown column id {column_id}"
)));
}
}
if !pk_cols.is_empty() {
let mut key = Vec::new();
for column_id in &pk_cols {
let value = row.get(column_id).ok_or_else(|| {
MongrelQueryError::Schema(format!(
"external row is missing primary key column {column_id}"
))
})?;
key.extend_from_slice(&column_id.to_be_bytes());
key.extend_from_slice(&value.encode_key());
key.push(0);
}
if !keys.insert(key) {
return Err(MongrelQueryError::Schema(
"external table primary key conflict".into(),
));
}
}
}
Ok(())
}
fn core_rows_to_batches(
schema: &CoreSchema,
rows: Vec<HashMap<u16, Value>>,
) -> Result<Vec<RecordBatch>> {
let arrow_schema = arrow_conv::arrow_schema(schema)?;
let arrays = schema
.columns
.iter()
.map(|column| {
let values = rows
.iter()
.map(|row| row.get(&column.id).cloned().unwrap_or(Value::Null))
.collect::<Vec<_>>();
arrow_conv::build_array(column.ty, &values)
})
.collect::<Result<Vec<_>>>()?;
Ok(vec![RecordBatch::try_new(arrow_schema, arrays)
.map_err(|e| MongrelQueryError::Arrow(e.to_string()))?])
}
struct FtsDocsModule;
impl ExternalTableModule for FtsDocsModule {
fn name(&self) -> &str {
"fts_docs"
}
fn descriptor(&self) -> ExternalModuleDescriptor {
ExternalModuleDescriptor {
schema: fts_docs_schema(),
hidden_columns: vec!["query".to_string()],
capabilities: ModuleCapabilities {
writable: true,
deterministic: true,
trigger_safe: true,
..ModuleCapabilities::default()
},
}
}
fn indexes(&self, entry: &ExternalTableEntry) -> Result<Vec<ExternalModuleIndex>> {
Ok(vec![ExternalModuleIndex::new(
format!("{}_fts_inverted", entry.name),
vec![2],
)])
}
fn connect(
&self,
ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
) -> Result<Arc<dyn ExternalTable>> {
let options = fts_options(entry)?;
let rows = read_state_rows(ctx.db, entry)?;
let index = FtsInvertedIndex::build(&rows, &options);
Ok(Arc::new(FtsDocsExternalTable {
schema: arrow_conv::arrow_schema(&entry.declared_schema)?,
core_schema: entry.declared_schema.clone(),
options,
index,
rows,
}))
}
fn read_rows(
&self,
ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
) -> Result<Vec<HashMap<u16, Value>>> {
let _ = fts_options(entry)?;
read_state_rows(ctx.db, entry)
}
fn prepare_rows(
&self,
_ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
rows: Vec<HashMap<u16, Value>>,
) -> Result<Vec<u8>> {
let _ = fts_options(entry)?;
let rows = rows
.into_iter()
.map(|mut row| {
row.remove(&3);
row.remove(&4);
row.remove(&5);
row.remove(&6);
row
})
.collect::<Vec<_>>();
validate_external_rows(&entry.declared_schema, &rows)?;
encode_state_rows(&rows)
}
}
#[derive(Clone)]
struct FtsDocsExternalTable {
schema: SchemaRef,
core_schema: CoreSchema,
options: FtsOptions,
index: FtsInvertedIndex,
rows: Vec<HashMap<u16, Value>>,
}
impl std::fmt::Debug for FtsDocsExternalTable {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FtsDocsExternalTable")
.field("schema", &self.schema)
.finish_non_exhaustive()
}
}
impl ExternalTable for FtsDocsExternalTable {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
Ok(ExternalPlan::new(
request
.raw_filters
.iter()
.map(|filter| {
if fts_query_from_expr(filter, &self.options).is_some() {
TableProviderFilterPushDown::Exact
} else {
TableProviderFilterPushDown::Unsupported
}
})
.collect(),
None,
1.0,
false,
))
}
fn scan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalScan> {
let query = request
.raw_filters
.iter()
.filter_map(|filter| fts_query_from_expr(filter, &self.options))
.reduce(FtsQuery::and);
let rows = if let Some(query) = query.as_ref() {
self.index
.candidates(query)
.into_iter()
.filter_map(|idx| {
let row = self.rows.get(idx)?;
let score = fts_match_score(row, query, &self.options)?;
Some(fts_enrich_row(
row.clone(),
Some(query),
&self.options,
Some(score),
))
})
.collect::<Vec<_>>()
} else {
self.rows
.iter()
.cloned()
.map(|row| fts_enrich_row(row, None, &self.options, None))
.collect::<Vec<_>>()
};
project_scan(
self.schema.clone(),
core_rows_to_batches(&self.core_schema, rows)
.map_err(|e| DataFusionError::Execution(e.to_string()))?,
request.projection.as_deref(),
request.limit,
)
}
}
fn fts_docs_schema() -> CoreSchema {
CoreSchema {
schema_id: 0,
columns: vec![
CoreColumnDef {
id: 1,
name: "doc_id".to_string(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
},
CoreColumnDef {
id: 2,
name: "text".to_string(),
ty: TypeId::Bytes,
flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
},
CoreColumnDef {
id: 3,
name: "query".to_string(),
ty: TypeId::Bytes,
flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
},
CoreColumnDef {
id: 4,
name: "rank".to_string(),
ty: TypeId::Float64,
flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
},
CoreColumnDef {
id: 5,
name: "snippet".to_string(),
ty: TypeId::Bytes,
flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
},
CoreColumnDef {
id: 6,
name: "highlight".to_string(),
ty: TypeId::Bytes,
flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
},
],
indexes: Vec::new(),
colocation: Vec::new(),
constraints: Default::default(),
}
}
#[derive(Clone)]
struct FtsOptions {
prefix_queries: bool,
case_sensitive: bool,
min_token_len: usize,
stopwords: HashSet<String>,
}
impl Default for FtsOptions {
fn default() -> Self {
Self {
prefix_queries: false,
case_sensitive: false,
min_token_len: 1,
stopwords: HashSet::new(),
}
}
}
impl FtsOptions {
fn normalize(&self, value: &str) -> String {
if self.case_sensitive {
value.to_string()
} else {
value.to_ascii_lowercase()
}
}
fn keep_term(&self, term: &str) -> bool {
term.len() >= self.min_token_len && !self.stopwords.contains(term)
}
}
#[derive(Clone)]
struct FtsQuery {
groups: Vec<Vec<FtsClause>>,
prohibited: Vec<FtsClause>,
}
impl FtsQuery {
fn and(self, other: Self) -> Self {
let left_groups = if self.groups.is_empty() {
vec![Vec::new()]
} else {
self.groups
};
let right_groups = if other.groups.is_empty() {
vec![Vec::new()]
} else {
other.groups
};
let mut groups = Vec::new();
for left in &left_groups {
for right in &right_groups {
let mut combined = left.clone();
combined.extend(right.clone());
groups.push(combined);
}
}
let mut prohibited = self.prohibited;
prohibited.extend(other.prohibited);
Self { groups, prohibited }
}
fn positive_clauses(&self) -> impl Iterator<Item = &FtsClause> {
self.groups.iter().flatten()
}
}
#[derive(Clone)]
struct FtsInvertedIndex {
terms: HashMap<String, Vec<usize>>,
all_rows: Vec<usize>,
}
impl FtsInvertedIndex {
fn build(rows: &[HashMap<u16, Value>], options: &FtsOptions) -> Self {
let mut term_sets: HashMap<String, BTreeSet<usize>> = HashMap::new();
for (idx, row) in rows.iter().enumerate() {
let Some(Value::Bytes(text)) = row.get(&2) else {
continue;
};
let text = String::from_utf8_lossy(text);
let mut row_terms = HashSet::new();
for span in token_spans_with_options(&text, options) {
if row_terms.insert(span.term.clone()) {
term_sets.entry(span.term).or_default().insert(idx);
}
}
}
let terms = term_sets
.into_iter()
.map(|(term, rows)| (term, rows.into_iter().collect()))
.collect();
Self {
terms,
all_rows: (0..rows.len()).collect(),
}
}
fn candidates(&self, query: &FtsQuery) -> Vec<usize> {
let mut rows = if query.groups.is_empty() {
self.all_rows.iter().copied().collect::<BTreeSet<_>>()
} else {
query
.groups
.iter()
.filter_map(|group| self.group_candidates(group))
.flatten()
.collect::<BTreeSet<_>>()
};
for clause in &query.prohibited {
for idx in self.clause_candidates(clause) {
rows.remove(&idx);
}
}
rows.into_iter().collect()
}
fn group_candidates(&self, group: &[FtsClause]) -> Option<Vec<usize>> {
let mut iter = group.iter().map(|clause| self.clause_candidates(clause));
let first = iter.next()?;
let mut current = first.into_iter().collect::<BTreeSet<_>>();
for rows in iter {
let next = rows.into_iter().collect::<BTreeSet<_>>();
current = current.intersection(&next).copied().collect();
if current.is_empty() {
break;
}
}
Some(current.into_iter().collect())
}
fn clause_candidates(&self, clause: &FtsClause) -> Vec<usize> {
match clause {
FtsClause::Term { term, prefix } => {
if *prefix {
self.terms
.iter()
.filter(|(indexed, _)| indexed.starts_with(term))
.flat_map(|(_, rows)| rows.iter().copied())
.collect::<BTreeSet<_>>()
.into_iter()
.collect()
} else {
self.terms.get(term).cloned().unwrap_or_default()
}
}
FtsClause::Phrase(terms) => {
let clauses = terms
.iter()
.map(|term| FtsClause::Term {
term: term.clone(),
prefix: false,
})
.collect::<Vec<_>>();
self.group_candidates(&clauses).unwrap_or_default()
}
}
}
}
#[derive(Clone, PartialEq, Eq)]
enum FtsClause {
Term { term: String, prefix: bool },
Phrase(Vec<String>),
}
enum FtsQueryToken {
Clause(FtsClause),
Or,
Not,
}
fn fts_options(entry: &ExternalTableEntry) -> Result<FtsOptions> {
let mut options = FtsOptions::default();
for arg in &entry.args {
let raw = module_arg_string(arg).trim();
if raw.is_empty() {
continue;
}
let (key, value) = raw
.split_once('=')
.map_or((raw, "true"), |(key, value)| (key.trim(), value.trim()));
match key.to_ascii_lowercase().as_str() {
"tokenizer" => match value.to_ascii_lowercase().as_str() {
"simple" | "ascii" | "unicode61" => {}
other => {
return Err(MongrelQueryError::Schema(format!(
"fts_docs tokenizer {other:?} is not supported"
)))
}
},
"prefix" | "prefix_queries" => options.prefix_queries = parse_bool_option(value)?,
"case_sensitive" => options.case_sensitive = parse_bool_option(value)?,
"min_token_len" => {
options.min_token_len = value.parse::<usize>().map_err(|e| {
MongrelQueryError::Schema(format!(
"fts_docs min_token_len {value:?} must be an integer: {e}"
))
})?;
if options.min_token_len == 0 {
return Err(MongrelQueryError::Schema(
"fts_docs min_token_len must be at least 1".into(),
));
}
}
"stopwords" => {
options.stopwords = value
.split('|')
.map(str::trim)
.filter(|word| !word.is_empty())
.map(|word| options.normalize(word))
.collect();
}
other => {
return Err(MongrelQueryError::Schema(format!(
"fts_docs option {other:?} is not supported"
)))
}
}
}
Ok(options)
}
fn parse_bool_option(value: &str) -> Result<bool> {
match value.to_ascii_lowercase().as_str() {
"1" | "true" | "yes" | "on" => Ok(true),
"0" | "false" | "no" | "off" => Ok(false),
_ => Err(MongrelQueryError::Schema(format!(
"expected boolean option value, got {value:?}"
))),
}
}
fn fts_query_from_expr(expr: &Expr, options: &FtsOptions) -> Option<FtsQuery> {
match expr {
Expr::BinaryExpr(BinaryExpr { left, op, right }) if *op == Operator::And => {
Some(fts_query_from_expr(left, options)?.and(fts_query_from_expr(right, options)?))
}
Expr::BinaryExpr(BinaryExpr { left, op, right }) if *op == Operator::Eq => {
let literal = match (left.as_ref(), right.as_ref()) {
(Expr::Column(column), Expr::Literal(literal, _)) if column.name == "query" => {
literal_string(literal)?
}
(Expr::Literal(literal, _), Expr::Column(column)) if column.name == "query" => {
literal_string(literal)?
}
_ => return None,
};
parse_fts_query(&literal, options)
}
_ => None,
}
}
fn parse_fts_query(input: &str, options: &FtsOptions) -> Option<FtsQuery> {
let tokens = fts_query_tokens(input, options);
if tokens.is_empty() {
return None;
}
let mut groups: Vec<Vec<FtsClause>> = vec![Vec::new()];
let mut prohibited = Vec::new();
let mut negate = false;
for token in tokens {
match token {
FtsQueryToken::Or => {
if groups.last().is_some_and(|group| !group.is_empty()) {
groups.push(Vec::new());
}
negate = false;
}
FtsQueryToken::Not => negate = true,
FtsQueryToken::Clause(clause) => {
if negate {
prohibited.push(clause);
negate = false;
} else if let Some(group) = groups.last_mut() {
group.push(clause);
}
}
}
}
groups.retain(|group| !group.is_empty());
if groups.is_empty() && prohibited.is_empty() {
None
} else {
Some(FtsQuery { groups, prohibited })
}
}
fn fts_query_tokens(input: &str, options: &FtsOptions) -> Vec<FtsQueryToken> {
let mut tokens = Vec::new();
let mut chars = input.char_indices().peekable();
while let Some((idx, ch)) = chars.next() {
if ch.is_whitespace() {
continue;
}
if ch == '"' {
let start = idx + ch.len_utf8();
let mut end = input.len();
for (next_idx, next_ch) in chars.by_ref() {
if next_ch == '"' {
end = next_idx;
break;
}
}
let terms = token_spans_with_options(&input[start..end], options)
.into_iter()
.map(|span| span.term)
.collect::<Vec<_>>();
if !terms.is_empty() {
tokens.push(FtsQueryToken::Clause(FtsClause::Phrase(terms)));
}
continue;
}
let start = idx;
let mut end = idx + ch.len_utf8();
while let Some((next_idx, next_ch)) = chars.peek().copied() {
if next_ch.is_whitespace() {
break;
}
chars.next();
end = next_idx + next_ch.len_utf8();
}
let raw = &input[start..end];
if raw.eq_ignore_ascii_case("OR") {
tokens.push(FtsQueryToken::Or);
} else if raw.eq_ignore_ascii_case("NOT") {
tokens.push(FtsQueryToken::Not);
} else if let Some(rest) = raw.strip_prefix('-') {
tokens.push(FtsQueryToken::Not);
tokens.extend(
word_clauses(rest, options)
.into_iter()
.map(FtsQueryToken::Clause),
);
} else {
tokens.extend(
word_clauses(raw, options)
.into_iter()
.map(FtsQueryToken::Clause),
);
}
}
tokens
}
fn word_clauses(raw: &str, options: &FtsOptions) -> Vec<FtsClause> {
let raw = raw.trim_matches(|ch: char| !ch.is_ascii_alphanumeric() && ch != '*');
if raw.is_empty() {
return Vec::new();
}
let prefix = options.prefix_queries && raw.ends_with('*');
let raw = raw.trim_end_matches('*');
token_spans_with_options(raw, options)
.into_iter()
.map(|span| FtsClause::Term {
term: span.term,
prefix,
})
.collect()
}
fn fts_match_score(
row: &HashMap<u16, Value>,
query: &FtsQuery,
options: &FtsOptions,
) -> Option<f64> {
let Some(Value::Bytes(text)) = row.get(&2) else {
return None;
};
let text = String::from_utf8_lossy(text);
let tokens = token_spans_with_options(&text, options);
if query
.prohibited
.iter()
.any(|clause| clause_matches(&tokens, clause))
{
return None;
}
let positive_score = if query.groups.is_empty() {
Some(0.0)
} else {
query
.groups
.iter()
.filter(|group| group.iter().all(|clause| clause_matches(&tokens, clause)))
.map(|group| {
group
.iter()
.map(|clause| clause_score(&tokens, clause))
.sum()
})
.max_by(|left: &f64, right: &f64| left.total_cmp(right))
}?;
Some(positive_score)
}
fn clause_matches(tokens: &[TokenSpan], clause: &FtsClause) -> bool {
clause_score(tokens, clause) > 0.0
}
fn clause_score(tokens: &[TokenSpan], clause: &FtsClause) -> f64 {
match clause {
FtsClause::Term { term, prefix } => tokens
.iter()
.filter(|token| token_matches_term(&token.term, term, *prefix))
.count() as f64,
FtsClause::Phrase(terms) => phrase_occurrences(tokens, terms) as f64 * terms.len() as f64,
}
}
fn token_matches_term(token: &str, term: &str, prefix: bool) -> bool {
if prefix {
token.starts_with(term)
} else {
token == term
}
}
fn phrase_occurrences(tokens: &[TokenSpan], terms: &[String]) -> usize {
if terms.is_empty() || tokens.len() < terms.len() {
return 0;
}
tokens
.windows(terms.len())
.filter(|window| {
window
.iter()
.zip(terms)
.all(|(token, term)| token.term == *term)
})
.count()
}
fn token_matches_positive_clause(token: &str, query: &FtsQuery) -> bool {
query.positive_clauses().any(|clause| match clause {
FtsClause::Term { term, prefix } => token_matches_term(token, term, *prefix),
FtsClause::Phrase(terms) => terms.iter().any(|term| term == token),
})
}
fn fts_enrich_row(
mut row: HashMap<u16, Value>,
query: Option<&FtsQuery>,
options: &FtsOptions,
score: Option<f64>,
) -> HashMap<u16, Value> {
row.remove(&3);
row.remove(&4);
row.remove(&5);
row.remove(&6);
let Some(score) = score else {
return row;
};
let text = match row.get(&2) {
Some(Value::Bytes(text)) => String::from_utf8_lossy(text).into_owned(),
_ => String::new(),
};
row.insert(4, Value::Float64(score));
if let Some(query) = query {
row.insert(
5,
Value::Bytes(fts_snippet(&text, query, options).into_bytes()),
);
row.insert(
6,
Value::Bytes(fts_highlight(&text, query, options).into_bytes()),
);
}
row
}
#[derive(Debug, Clone)]
struct TokenSpan {
term: String,
start: usize,
end: usize,
}
fn token_spans_with_options(input: &str, options: &FtsOptions) -> Vec<TokenSpan> {
let mut spans = Vec::new();
let mut start = None;
for (idx, ch) in input.char_indices() {
if ch.is_ascii_alphanumeric() {
start.get_or_insert(idx);
} else if let Some(lo) = start.take() {
push_token_span(input, &mut spans, lo, idx, options);
}
}
if let Some(lo) = start {
push_token_span(input, &mut spans, lo, input.len(), options);
}
spans
}
fn push_token_span(
input: &str,
spans: &mut Vec<TokenSpan>,
start: usize,
end: usize,
options: &FtsOptions,
) {
if start < end {
let term = options.normalize(&input[start..end]);
if options.keep_term(&term) {
spans.push(TokenSpan { term, start, end });
}
}
}
fn fts_snippet(text: &str, query: &FtsQuery, options: &FtsOptions) -> String {
let spans = token_spans_with_options(text, options);
if spans.is_empty() {
return String::new();
}
let first_match = spans
.iter()
.position(|span| token_matches_positive_clause(&span.term, query))
.unwrap_or(0);
let lo = first_match.saturating_sub(3);
let hi = (first_match + 5).min(spans.len());
let mut out = Vec::new();
if lo > 0 {
out.push("...".to_string());
}
for span in &spans[lo..hi] {
let token = &text[span.start..span.end];
if token_matches_positive_clause(&span.term, query) {
out.push(format!("[{token}]"));
} else {
out.push(token.to_string());
}
}
if hi < spans.len() {
out.push("...".to_string());
}
out.join(" ")
}
fn fts_highlight(text: &str, query: &FtsQuery, options: &FtsOptions) -> String {
let spans = token_spans_with_options(text, options);
if spans.is_empty() {
return text.to_string();
}
let mut out = String::new();
let mut cursor = 0;
for span in spans {
out.push_str(&text[cursor..span.start]);
let token = &text[span.start..span.end];
if token_matches_positive_clause(&span.term, query) {
out.push_str("<b>");
out.push_str(token);
out.push_str("</b>");
} else {
out.push_str(token);
}
cursor = span.end;
}
out.push_str(&text[cursor..]);
out
}
struct RTreeRectsModule;
impl ExternalTableModule for RTreeRectsModule {
fn name(&self) -> &str {
"rtree_rects"
}
fn descriptor(&self) -> ExternalModuleDescriptor {
ExternalModuleDescriptor {
schema: rtree_rects_schema(),
hidden_columns: vec![
"query_min_x".to_string(),
"query_max_x".to_string(),
"query_min_y".to_string(),
"query_max_y".to_string(),
],
capabilities: ModuleCapabilities {
writable: true,
deterministic: true,
trigger_safe: true,
..ModuleCapabilities::default()
},
}
}
fn indexes(&self, entry: &ExternalTableEntry) -> Result<Vec<ExternalModuleIndex>> {
Ok(vec![ExternalModuleIndex::new(
format!("{}_rtree_spatial", entry.name),
vec![2, 3, 4, 5],
)])
}
fn connect(
&self,
ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
) -> Result<Arc<dyn ExternalTable>> {
ensure_no_args(entry, self.name())?;
let rows = read_state_rows(ctx.db, entry)?;
let index = RTreeSpatialIndex::build(&rows);
Ok(Arc::new(RTreeRectsExternalTable {
schema: arrow_conv::arrow_schema(&entry.declared_schema)?,
core_schema: entry.declared_schema.clone(),
index,
rows,
}))
}
fn read_rows(
&self,
ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
) -> Result<Vec<HashMap<u16, Value>>> {
ensure_no_args(entry, self.name())?;
read_state_rows(ctx.db, entry)
}
fn prepare_rows(
&self,
_ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
rows: Vec<HashMap<u16, Value>>,
) -> Result<Vec<u8>> {
ensure_no_args(entry, self.name())?;
let rows = rows
.into_iter()
.map(|mut row| {
for column_id in 6..=9 {
row.remove(&column_id);
}
row
})
.collect::<Vec<_>>();
validate_external_rows(&entry.declared_schema, &rows)?;
encode_state_rows(&rows)
}
}
#[derive(Clone)]
struct RTreeRectsExternalTable {
schema: SchemaRef,
core_schema: CoreSchema,
index: RTreeSpatialIndex,
rows: Vec<HashMap<u16, Value>>,
}
impl std::fmt::Debug for RTreeRectsExternalTable {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RTreeRectsExternalTable")
.field("schema", &self.schema)
.finish_non_exhaustive()
}
}
impl ExternalTable for RTreeRectsExternalTable {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
Ok(ExternalPlan::new(
request
.raw_filters
.iter()
.map(|filter| {
if rtree_query_bound(filter).is_some()
|| rtree_query_rect_function(filter).is_some()
{
TableProviderFilterPushDown::Exact
} else {
TableProviderFilterPushDown::Unsupported
}
})
.collect(),
None,
1.0,
false,
))
}
fn scan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalScan> {
let mut hidden_bounds = QueryRect::default();
let mut has_hidden_bounds = false;
for (column, value) in request
.raw_filters
.iter()
.filter_map(|filter| rtree_query_bound(filter))
{
hidden_bounds.set(column, value);
has_hidden_bounds = true;
}
let mut query_rects = request
.raw_filters
.iter()
.filter_map(|filter| rtree_query_rect_function(filter))
.collect::<Vec<_>>();
if has_hidden_bounds || query_rects.is_empty() {
query_rects.push(hidden_bounds);
}
let candidate_rows = query_rects
.iter()
.map(|bounds| {
self.index
.candidates(*bounds)
.into_iter()
.collect::<BTreeSet<_>>()
})
.reduce(|left, right| left.intersection(&right).copied().collect())
.unwrap_or_else(|| self.index.all_rows.iter().copied().collect());
let rows = self
.rows
.iter()
.enumerate()
.filter(|(idx, row)| {
candidate_rows.contains(idx)
&& query_rects
.iter()
.all(|bounds| rtree_row_matches(row, *bounds))
})
.map(|(_, row)| row)
.cloned()
.collect::<Vec<_>>();
project_scan(
self.schema.clone(),
core_rows_to_batches(&self.core_schema, rows)
.map_err(|e| DataFusionError::Execution(e.to_string()))?,
request.projection.as_deref(),
request.limit,
)
}
}
#[derive(Clone, Copy)]
struct QueryRect {
min_x: f64,
max_x: f64,
min_y: f64,
max_y: f64,
}
impl Default for QueryRect {
fn default() -> Self {
Self {
min_x: f64::NEG_INFINITY,
max_x: f64::INFINITY,
min_y: f64::NEG_INFINITY,
max_y: f64::INFINITY,
}
}
}
impl QueryRect {
fn set(&mut self, column: &str, value: f64) {
match column {
"query_min_x" => self.min_x = value,
"query_max_x" => self.max_x = value,
"query_min_y" => self.min_y = value,
"query_max_y" => self.max_y = value,
_ => {}
}
}
}
#[derive(Clone)]
struct RTreeSpatialIndex {
all_rows: Vec<usize>,
min_x: Vec<(f64, usize)>,
max_x: Vec<(f64, usize)>,
min_y: Vec<(f64, usize)>,
max_y: Vec<(f64, usize)>,
}
impl RTreeSpatialIndex {
fn build(rows: &[HashMap<u16, Value>]) -> Self {
let mut all_rows = Vec::new();
let mut min_x = Vec::new();
let mut max_x = Vec::new();
let mut min_y = Vec::new();
let mut max_y = Vec::new();
for (idx, row) in rows.iter().enumerate() {
let Some(x0) = row_f64(row, 2) else {
continue;
};
let Some(x1) = row_f64(row, 3) else {
continue;
};
let Some(y0) = row_f64(row, 4) else {
continue;
};
let Some(y1) = row_f64(row, 5) else {
continue;
};
if !(x0.is_finite() && x1.is_finite() && y0.is_finite() && y1.is_finite()) {
continue;
}
all_rows.push(idx);
min_x.push((x0, idx));
max_x.push((x1, idx));
min_y.push((y0, idx));
max_y.push((y1, idx));
}
for values in [&mut min_x, &mut max_x, &mut min_y, &mut max_y] {
values.sort_by(|left, right| left.0.total_cmp(&right.0));
}
Self {
all_rows,
min_x,
max_x,
min_y,
max_y,
}
}
fn candidates(&self, query: QueryRect) -> Vec<usize> {
let mut rows = self.lte(&self.min_x, query.max_x);
for next in [
self.gte(&self.max_x, query.min_x),
self.lte(&self.min_y, query.max_y),
self.gte(&self.max_y, query.min_y),
] {
rows = rows.intersection(&next).copied().collect();
if rows.is_empty() {
break;
}
}
rows.into_iter().collect()
}
fn lte(&self, values: &[(f64, usize)], bound: f64) -> BTreeSet<usize> {
if bound == f64::INFINITY {
return self.all_rows.iter().copied().collect();
}
if bound.is_nan() {
return BTreeSet::new();
}
let end = values.partition_point(|(value, _)| *value <= bound);
values[..end].iter().map(|(_, idx)| *idx).collect()
}
fn gte(&self, values: &[(f64, usize)], bound: f64) -> BTreeSet<usize> {
if bound == f64::NEG_INFINITY {
return self.all_rows.iter().copied().collect();
}
if bound.is_nan() {
return BTreeSet::new();
}
let start = values.partition_point(|(value, _)| *value < bound);
values[start..].iter().map(|(_, idx)| *idx).collect()
}
}
fn rtree_query_rect_function(expr: &Expr) -> Option<QueryRect> {
let Expr::ScalarFunction(sf) = expr else {
return None;
};
if !sf.func.name().eq_ignore_ascii_case("rtree_intersects") || sf.args.len() != 8 {
return None;
}
for (arg, expected) in sf
.args
.iter()
.take(4)
.zip(["min_x", "max_x", "min_y", "max_y"])
{
let Expr::Column(column) = arg else {
return None;
};
if column.name != expected {
return None;
}
}
Some(QueryRect {
min_x: literal_f64_from_expr(&sf.args[4])?,
max_x: literal_f64_from_expr(&sf.args[5])?,
min_y: literal_f64_from_expr(&sf.args[6])?,
max_y: literal_f64_from_expr(&sf.args[7])?,
})
}
fn literal_f64_from_expr(expr: &Expr) -> Option<f64> {
match expr {
Expr::Literal(literal, _) => literal_f64(literal),
_ => None,
}
}
fn rtree_rects_schema() -> CoreSchema {
CoreSchema {
schema_id: 0,
columns: vec![
CoreColumnDef {
id: 1,
name: "id".to_string(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
},
rect_column(2, "min_x"),
rect_column(3, "max_x"),
rect_column(4, "min_y"),
rect_column(5, "max_y"),
rect_column(6, "query_min_x"),
rect_column(7, "query_max_x"),
rect_column(8, "query_min_y"),
rect_column(9, "query_max_y"),
],
indexes: Vec::new(),
colocation: Vec::new(),
constraints: Default::default(),
}
}
fn rect_column(id: u16, name: &str) -> CoreColumnDef {
CoreColumnDef {
id,
name: name.to_string(),
ty: TypeId::Float64,
flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
}
}
fn rtree_query_bound(expr: &Expr) -> Option<(&str, f64)> {
match expr {
Expr::BinaryExpr(BinaryExpr { left, op, right }) if *op == Operator::Eq => {
match (left.as_ref(), right.as_ref()) {
(Expr::Column(column), Expr::Literal(literal, _)) => {
hidden_rect_column(&column.name).zip(literal_f64(literal))
}
(Expr::Literal(literal, _), Expr::Column(column)) => {
hidden_rect_column(&column.name).zip(literal_f64(literal))
}
_ => None,
}
}
_ => None,
}
}
fn hidden_rect_column(name: &str) -> Option<&str> {
match name {
"query_min_x" | "query_max_x" | "query_min_y" | "query_max_y" => Some(name),
_ => None,
}
}
fn rtree_row_matches(row: &HashMap<u16, Value>, query: QueryRect) -> bool {
let Some(min_x) = row_f64(row, 2) else {
return false;
};
let Some(max_x) = row_f64(row, 3) else {
return false;
};
let Some(min_y) = row_f64(row, 4) else {
return false;
};
let Some(max_y) = row_f64(row, 5) else {
return false;
};
max_x >= query.min_x && min_x <= query.max_x && max_y >= query.min_y && min_y <= query.max_y
}
fn row_f64(row: &HashMap<u16, Value>, column_id: u16) -> Option<f64> {
match row.get(&column_id) {
Some(Value::Float64(value)) => Some(*value),
Some(Value::Int64(value)) => Some(*value as f64),
_ => None,
}
}
struct SeriesModule;
impl ExternalTableModule for SeriesModule {
fn name(&self) -> &str {
"series"
}
fn descriptor(&self) -> ExternalModuleDescriptor {
ExternalModuleDescriptor {
schema: CoreSchema {
schema_id: 0,
columns: vec![
CoreColumnDef {
id: 1,
name: "value".to_string(),
ty: TypeId::Int64,
flags: ColumnFlags::empty(),
},
CoreColumnDef {
id: 2,
name: "start".to_string(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
},
CoreColumnDef {
id: 3,
name: "stop".to_string(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
},
CoreColumnDef {
id: 4,
name: "step".to_string(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::NULLABLE),
},
],
indexes: Vec::new(),
colocation: Vec::new(),
constraints: Default::default(),
},
hidden_columns: vec!["start".to_string(), "stop".to_string(), "step".to_string()],
capabilities: ModuleCapabilities {
read_only: true,
deterministic: true,
trigger_safe: true,
..ModuleCapabilities::default()
},
}
}
fn connect(
&self,
_ctx: &ModuleConnectCtx<'_>,
entry: &ExternalTableEntry,
) -> Result<Arc<dyn ExternalTable>> {
let (start, stop, step) = series_args(entry)?;
let table = SeriesExternalTable::new(start, stop, step)?;
Ok(Arc::new(table))
}
}
struct SeriesExternalTable {
start: i64,
stop: i64,
step: i64,
schema: SchemaRef,
}
impl SeriesExternalTable {
fn new(start: i64, stop: i64, step: i64) -> Result<Self> {
if step == 0 {
return Err(MongrelQueryError::Schema(
"series step must not be 0".into(),
));
}
Ok(Self {
start,
stop,
step,
schema: Arc::new(ArrowSchema::new(vec![Field::new(
"value",
DataType::Int64,
false,
)])),
})
}
fn batches(&self, request: &ExternalPlanRequest<'_>) -> DFResult<Vec<RecordBatch>> {
let mut values = Vec::new();
let mut current = self.start;
while if self.step > 0 {
current <= self.stop
} else {
current >= self.stop
} {
if values.len() >= 1_000_000 {
return Err(DataFusionError::Plan(
"series output is capped at 1,000,000 rows".into(),
));
}
if request
.filters
.iter()
.all(|filter| series_filter_matches(filter, current).unwrap_or(true))
{
values.push(current);
if request.limit.is_some_and(|limit| values.len() >= limit) {
break;
}
}
current = current.saturating_add(self.step);
if (self.step > 0 && current == i64::MAX) || (self.step < 0 && current == i64::MIN) {
break;
}
}
let batch = RecordBatch::try_new(
self.schema.clone(),
vec![Arc::new(Int64Array::from(values)) as ArrayRef],
)?;
Ok(vec![batch])
}
}
impl std::fmt::Debug for SeriesExternalTable {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SeriesExternalTable")
.field("start", &self.start)
.field("stop", &self.stop)
.field("step", &self.step)
.finish_non_exhaustive()
}
}
impl ExternalTable for SeriesExternalTable {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
fn plan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalPlan> {
Ok(ExternalPlan::new(
request
.filters
.iter()
.map(|filter| {
if series_filter_supported(filter) {
TableProviderFilterPushDown::Exact
} else {
TableProviderFilterPushDown::Unsupported
}
})
.collect(),
None,
1.0,
self.step > 0,
))
}
fn scan(&self, request: &ExternalPlanRequest<'_>) -> DFResult<ExternalScan> {
project_scan(
self.schema.clone(),
self.batches(request)?,
request.projection.as_deref(),
request.limit,
)
}
}
fn series_filter_supported(filter: &ExternalFilter) -> bool {
match filter {
ExternalFilter::And(filters) => filters.iter().all(series_filter_supported),
ExternalFilter::Compare {
column_index,
value,
..
} => *column_index == 0 && literal_i64(value).is_some(),
ExternalFilter::Unsupported { .. } => false,
}
}
fn series_filter_matches(filter: &ExternalFilter, value: i64) -> Option<bool> {
match filter {
ExternalFilter::And(filters) => filters.iter().try_fold(true, |matches, filter| {
Some(matches && series_filter_matches(filter, value)?)
}),
ExternalFilter::Compare {
column_index,
op,
value: literal,
} if *column_index == 0 => {
let literal = literal_i64(literal)?;
Some(match op {
ExternalFilterOp::Eq => value == literal,
ExternalFilterOp::NotEq => value != literal,
ExternalFilterOp::Gt => value > literal,
ExternalFilterOp::GtEq => value >= literal,
ExternalFilterOp::Lt => value < literal,
ExternalFilterOp::LtEq => value <= literal,
})
}
_ => None,
}
}
fn literal_i64(value: &ScalarValue) -> Option<i64> {
match value {
ScalarValue::Int8(Some(v)) => Some(*v as i64),
ScalarValue::Int16(Some(v)) => Some(*v as i64),
ScalarValue::Int32(Some(v)) => Some(*v as i64),
ScalarValue::Int64(Some(v)) => Some(*v),
ScalarValue::UInt8(Some(v)) => Some(*v as i64),
ScalarValue::UInt16(Some(v)) => Some(*v as i64),
ScalarValue::UInt32(Some(v)) => Some(*v as i64),
ScalarValue::UInt64(Some(v)) => i64::try_from(*v).ok(),
_ => None,
}
}
fn literal_f64(value: &ScalarValue) -> Option<f64> {
match value {
ScalarValue::Float32(Some(v)) => Some(*v as f64),
ScalarValue::Float64(Some(v)) => Some(*v),
_ => literal_i64(value).map(|value| value as f64),
}
}
fn literal_string(value: &ScalarValue) -> Option<String> {
match value {
ScalarValue::Utf8(Some(v))
| ScalarValue::Utf8View(Some(v))
| ScalarValue::LargeUtf8(Some(v)) => Some(v.clone()),
ScalarValue::Binary(Some(v))
| ScalarValue::BinaryView(Some(v))
| ScalarValue::LargeBinary(Some(v)) => String::from_utf8(v.clone()).ok(),
_ => None,
}
}
fn series_args(entry: &ExternalTableEntry) -> Result<(i64, i64, i64)> {
let values = entry
.args
.iter()
.map(|arg| match arg {
ModuleArg::Ident(value) | ModuleArg::String(value) | ModuleArg::Number(value) => {
value.parse::<i64>().map_err(|e| {
MongrelQueryError::Schema(format!(
"series module argument {value:?} must be an integer: {e}"
))
})
}
})
.collect::<Result<Vec<_>>>()?;
match values.as_slice() {
[] => Ok((0, -1, 1)),
[stop] => Ok((0, *stop, 1)),
[start, stop] => Ok((*start, *stop, 1)),
[start, stop, step] => Ok((*start, *stop, *step)),
_ => Err(MongrelQueryError::Schema(
"series external table accepts at most three arguments".into(),
)),
}
}