use std::any::Any;
use arrow::datatypes::DataType;
use datafusion::common::{Column, ScalarValue, not_impl_err, plan_err};
use datafusion::error::Result as DFResult;
use datafusion::logical_expr::{
ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
};
use datafusion::prelude::SessionContext;
pub const FTS_MATCH: &str = "fts_match";
#[derive(Debug, PartialEq, Eq, Hash)]
struct FtsMatchUdf {
signature: Signature,
}
impl FtsMatchUdf {
fn new() -> Self {
Self {
signature: Signature::user_defined(Volatility::Immutable),
}
}
}
fn indexable(data_type: &DataType) -> bool {
match data_type {
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View | DataType::LargeBinary => true,
DataType::List(field) | DataType::LargeList(field) => {
matches!(field.data_type(), DataType::Utf8 | DataType::LargeUtf8)
}
_ => false,
}
}
impl ScalarUDFImpl for FtsMatchUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
FTS_MATCH
}
fn signature(&self) -> &Signature {
&self.signature
}
fn coerce_types(&self, arg_types: &[DataType]) -> DFResult<Vec<DataType>> {
let [column, query] = arg_types else {
return plan_err!(
"{FTS_MATCH} takes two arguments, a column and a query, not {}",
arg_types.len()
);
};
if !indexable(column) && *column != DataType::Null {
return plan_err!(
"{FTS_MATCH} searches a text, binary or list-of-text column, not {column}"
);
}
let query = match query {
DataType::Null => DataType::Null,
q if q.is_numeric()
|| matches!(q, DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View) =>
{
DataType::Utf8
}
other => return plan_err!("{FTS_MATCH}'s query is text, not {other}"),
};
Ok(vec![column.clone(), query])
}
fn return_type(&self, _arg_types: &[DataType]) -> DFResult<DataType> {
Ok(DataType::Boolean)
}
fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> DFResult<ColumnarValue> {
not_impl_err!(
"{FTS_MATCH} is answered by a Lance provider's full-text indexes; this query \
was not pushed down to one"
)
}
}
pub fn register_fts_match_udf(ctx: &SessionContext) {
ctx.register_udf(ScalarUDF::new_from_impl(FtsMatchUdf::new()));
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct FtsMatch<'a> {
pub column: &'a Column,
pub query: &'a Expr,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FtsQueryText<'a> {
Text(&'a str),
MatchesNothing,
NotALiteral,
}
impl<'a> FtsMatch<'a> {
pub fn query_text(&self) -> FtsQueryText<'a> {
match self.query {
Expr::Literal(
ScalarValue::Utf8(Some(text))
| ScalarValue::LargeUtf8(Some(text))
| ScalarValue::Utf8View(Some(text)),
_,
) if !text.trim().is_empty() => FtsQueryText::Text(text),
Expr::Literal(
ScalarValue::Utf8(_)
| ScalarValue::LargeUtf8(_)
| ScalarValue::Utf8View(_)
| ScalarValue::Null,
_,
) => FtsQueryText::MatchesNothing,
_ => FtsQueryText::NotALiteral,
}
}
}
pub fn as_fts_match(expr: &Expr) -> Option<FtsMatch<'_>> {
let Expr::ScalarFunction(call) = expr else {
return None;
};
if call.func.name() != FTS_MATCH || !call.func.inner().as_any().is::<FtsMatchUdf>() {
return None;
}
match call.args.as_slice() {
[Expr::Column(column), query] => Some(FtsMatch { column, query }),
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
use arrow::array::{Int64Array, RecordBatch, StringArray};
use arrow::datatypes::{Field, Schema};
use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
use datafusion::datasource::MemTable;
use datafusion::execution::FunctionRegistry;
use datafusion::logical_expr::LogicalPlan;
use datafusion::logical_expr::utils::split_conjunction;
use datafusion::prelude::{col, lit};
use std::sync::Arc;
async fn ctx_with_corpus() -> SessionContext {
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("body", DataType::Utf8, true),
]));
let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![
Arc::new(Int64Array::from(vec![1, 2])),
Arc::new(StringArray::from(vec![Some("onboarding notes"), None])),
],
)
.unwrap();
let ctx = SessionContext::new();
ctx.register_table(
"corpus",
Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()),
)
.unwrap();
register_fts_match_udf(&ctx);
ctx
}
fn filter_predicates(plan: &LogicalPlan) -> Vec<Expr> {
let mut out = Vec::new();
plan.apply(|node| {
if let LogicalPlan::Filter(filter) = node {
out.push(filter.predicate.clone());
}
Ok(TreeNodeRecursion::Continue)
})
.unwrap();
out
}
#[tokio::test]
async fn the_predicate_plans_and_survives_optimization_intact() {
let ctx = ctx_with_corpus().await;
let plan = ctx
.sql("SELECT id FROM corpus WHERE fts_match(body, 'onboarding') AND id > 0")
.await
.expect("fts_match plans")
.into_optimized_plan()
.unwrap();
let predicates = filter_predicates(&plan);
let found = predicates.iter().any(|p| {
split_conjunction(p)
.into_iter()
.any(|conjunct| as_fts_match(conjunct).is_some())
});
assert!(
found,
"fts_match must reach the scan's filter: {predicates:?}"
);
}
#[tokio::test]
async fn evaluating_it_locally_is_an_error_not_an_empty_result() {
let ctx = ctx_with_corpus().await;
let err = ctx
.sql("SELECT id FROM corpus WHERE fts_match(body, 'onboarding')")
.await
.unwrap()
.collect()
.await
.expect_err("a MemTable cannot answer fts_match");
let msg = err.to_string();
assert!(
msg.contains("fts_match") && msg.contains("not pushed down"),
"{msg}"
);
}
#[tokio::test]
async fn it_returns_boolean_and_refuses_the_wrong_arity() {
let ctx = ctx_with_corpus().await;
let df = ctx
.sql("SELECT fts_match(body, 'q') AS m FROM corpus")
.await
.unwrap();
assert_eq!(
df.schema().field_with_name(None, "m").unwrap().data_type(),
&DataType::Boolean
);
let one_argument = match ctx.sql("SELECT id FROM corpus WHERE fts_match(body)").await {
Err(e) => e.to_string(),
Ok(df) => df
.into_optimized_plan()
.expect_err("one argument must not plan")
.to_string(),
};
assert!(
one_argument.contains("fts_match") && !one_argument.contains("not pushed down"),
"{one_argument}"
);
}
#[tokio::test]
async fn as_fts_match_takes_a_column_and_nothing_else() {
let ctx = ctx_with_corpus().await;
let udf = ctx.udf(FTS_MATCH).unwrap();
let on_column = udf.call(vec![col("body"), lit("onboarding")]);
let m = as_fts_match(&on_column).expect("a column is recognised");
assert_eq!(m.column.name, "body");
assert_eq!(m.query, &lit("onboarding"));
let on_expression = udf.call(vec![lit("body"), lit("onboarding")]);
assert!(
as_fts_match(&on_expression).is_none(),
"a literal is not an index"
);
assert!(as_fts_match(&col("body").eq(lit("x"))).is_none());
}
#[derive(Debug)]
struct Claiming {
inner: MemTable,
exact: bool,
}
#[async_trait::async_trait]
impl datafusion::catalog::TableProvider for Claiming {
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> arrow::datatypes::SchemaRef {
self.inner.schema()
}
fn table_type(&self) -> datafusion::datasource::TableType {
datafusion::datasource::TableType::Base
}
fn supports_filters_pushdown(
&self,
filters: &[&Expr],
) -> DFResult<Vec<datafusion::logical_expr::TableProviderFilterPushDown>> {
use datafusion::logical_expr::TableProviderFilterPushDown as P;
Ok(filters
.iter()
.map(|f| match (as_fts_match(f), self.exact) {
(Some(_), true) => P::Exact,
(Some(_), false) => P::Inexact,
(None, _) => P::Unsupported,
})
.collect())
}
async fn scan(
&self,
state: &dyn datafusion::catalog::Session,
projection: Option<&Vec<usize>>,
_filters: &[Expr],
limit: Option<usize>,
) -> DFResult<Arc<dyn datafusion::physical_plan::ExecutionPlan>> {
self.inner.scan(state, projection, &[], limit).await
}
}
async fn ctx_with_claiming(exact: bool) -> SessionContext {
let ctx = ctx_with_corpus().await;
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("body", DataType::Utf8, true),
]));
let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![
Arc::new(Int64Array::from(vec![1])),
Arc::new(StringArray::from(vec![Some("onboarding notes")])),
],
)
.unwrap();
let table = Claiming {
inner: MemTable::try_new(schema, vec![vec![batch]]).unwrap(),
exact,
};
ctx.register_table("claiming", Arc::new(table)).unwrap();
ctx
}
fn scan_filters(plan: &LogicalPlan) -> Vec<Expr> {
let mut out = Vec::new();
plan.apply(|node| {
if let LogicalPlan::TableScan(scan) = node {
out.extend(scan.filters.iter().cloned());
}
Ok(TreeNodeRecursion::Continue)
})
.unwrap();
out
}
#[tokio::test]
async fn a_provider_that_claims_it_exact_receives_it_and_the_query_runs() {
let ctx = ctx_with_claiming(true).await;
let df = ctx
.sql("SELECT id FROM claiming WHERE fts_match(body, 'onboarding')")
.await
.unwrap();
let plan = df.clone().into_optimized_plan().unwrap();
assert!(
filter_predicates(&plan).is_empty(),
"no Filter left: {plan}"
);
assert!(
scan_filters(&plan)
.iter()
.any(|f| as_fts_match(f).is_some()),
"the scan received the predicate: {plan}"
);
df.collect().await.expect("invoke was never reached");
}
#[tokio::test]
async fn compound_predicates_and_an_inexact_provider_fail_loudly() {
let exact = ctx_with_claiming(true).await;
for sql in [
"SELECT id FROM claiming WHERE fts_match(body, 'q') OR id > 0",
"SELECT id FROM claiming WHERE NOT fts_match(body, 'q')",
] {
let err = exact
.sql(sql)
.await
.unwrap()
.collect()
.await
.expect_err(sql);
assert!(err.to_string().contains("not pushed down"), "{sql}: {err}");
}
let inexact = ctx_with_claiming(false).await;
let err = inexact
.sql("SELECT id FROM claiming WHERE fts_match(body, 'q')")
.await
.unwrap()
.collect()
.await
.expect_err("Inexact keeps the Filter");
assert!(err.to_string().contains("not pushed down"), "{err}");
}
#[tokio::test]
async fn the_column_keeps_its_type_and_the_query_becomes_text() {
let schema = Arc::new(Schema::new(vec![
Field::new("body", DataType::Utf8, true),
Field::new(
"tags",
DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
true,
),
]));
let ctx = SessionContext::new();
ctx.register_table(
"t",
Arc::new(MemTable::try_new(Arc::clone(&schema), vec![vec![]]).unwrap()),
)
.unwrap();
register_fts_match_udf(&ctx);
for (sql, text) in [
(
"SELECT 1 FROM t WHERE fts_match(body, arrow_cast('q', 'Utf8View'))",
"q",
),
("SELECT 1 FROM t WHERE fts_match(tags, 'q')", "q"),
("SELECT 1 FROM t WHERE fts_match(body, 2024)", "2024"),
] {
let plan = ctx.sql(sql).await.unwrap().into_optimized_plan().unwrap();
let found = filter_predicates(&plan)
.iter()
.flat_map(|p| {
split_conjunction(p)
.into_iter()
.cloned()
.collect::<Vec<_>>()
})
.find_map(|c| {
as_fts_match(&c).map(|m| {
(
m.column.name.clone(),
m.query_text() == FtsQueryText::Text(text),
)
})
});
let (column, query_ok) =
found.unwrap_or_else(|| panic!("{sql}: not recognised: {plan}"));
assert!(query_ok, "{sql}: query is not {text:?}: {plan}");
assert!(column == "body" || column == "tags", "{sql}");
}
let err = match ctx.sql("SELECT 1 FROM t WHERE fts_match(1, 'q')").await {
Err(e) => e.to_string(),
Ok(df) => df
.into_optimized_plan()
.expect_err("a number is not a searchable column")
.to_string(),
};
assert!(err.contains("searches a text"), "{err}");
}
#[tokio::test]
async fn query_text_classifies_the_three_cases() {
let ctx = ctx_with_corpus().await;
let udf = ctx.udf(FTS_MATCH).unwrap();
let text = |query: Expr| {
let call = udf.call(vec![col("body"), query]);
as_fts_match(&call).unwrap().query_text().clone_static()
};
assert_eq!(text(lit("onboarding")), Static::Text("onboarding".into()));
assert_eq!(text(lit(" ")), Static::MatchesNothing);
assert_eq!(text(lit(ScalarValue::Utf8(None))), Static::MatchesNothing);
assert_eq!(text(lit(ScalarValue::Null)), Static::MatchesNothing);
assert_eq!(text(col("id")), Static::NotALiteral);
}
#[derive(Debug, PartialEq)]
enum Static {
Text(String),
MatchesNothing,
NotALiteral,
}
trait CloneStatic {
fn clone_static(&self) -> Static;
}
impl CloneStatic for FtsQueryText<'_> {
fn clone_static(&self) -> Static {
match self {
FtsQueryText::Text(t) => Static::Text(t.to_string()),
FtsQueryText::MatchesNothing => Static::MatchesNothing,
FtsQueryText::NotALiteral => Static::NotALiteral,
}
}
}
#[tokio::test]
async fn an_all_literal_call_survives_folding_and_fails_at_execution() {
let ctx = ctx_with_corpus().await;
let df = ctx
.sql("SELECT id FROM corpus WHERE fts_match('body', 'q')")
.await
.unwrap();
df.clone()
.into_optimized_plan()
.expect("folding failed quietly and kept the call");
let err = df.collect().await.expect_err("and it errors when run");
assert!(err.to_string().contains("not pushed down"), "{err}");
}
#[test]
fn a_different_function_of_the_same_name_is_not_recognised() {
#[derive(Debug, PartialEq, Eq, Hash)]
struct Impostor(Signature);
impl ScalarUDFImpl for Impostor {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
FTS_MATCH
}
fn signature(&self) -> &Signature {
&self.0
}
fn return_type(&self, _: &[DataType]) -> DFResult<DataType> {
Ok(DataType::Boolean)
}
fn invoke_with_args(&self, _: ScalarFunctionArgs) -> DFResult<ColumnarValue> {
Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))))
}
}
let impostor = ScalarUDF::new_from_impl(Impostor(Signature::any(2, Volatility::Immutable)));
let expr = impostor.call(vec![col("body"), lit("q")]);
assert!(as_fts_match(&expr).is_none());
}
}