skardi 0.4.0

//! Physical execution plan for PostgreSQL full-text search.

use arrow::array::RecordBatch;
use arrow::datatypes::{DataType, SchemaRef};
use datafusion::error::{DataFusionError, Result as DFResult};
use datafusion::execution::{SendableRecordBatchStream, TaskContext};
use datafusion::physical_expr::EquivalenceProperties;
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::{
    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
    execution_plan::{Boundedness, EmissionType},
};
use futures::stream;
use sqlx::PgPool;
use std::any::Any;
use std::fmt;
use std::sync::Arc;

use super::utils::rows_to_batch;

/// Physical execution plan that runs a PostgreSQL full-text search query using
/// `websearch_to_tsquery` and returns matching rows with a `_score` column
/// derived from `ts_rank`.
///
/// Generates SQL of the form:
/// ```sql
/// SELECT <cols>,
///        ts_rank(to_tsvector('<language>', "<text_col>"),
///                websearch_to_tsquery('<language>', $1)) AS _score
/// FROM "<schema>"."<table>"
/// WHERE to_tsvector('<language>', "<text_col>") @@ websearch_to_tsquery('<language>', $1)
///   [AND <filter>]
/// ORDER BY _score DESC
/// LIMIT <limit>
/// ```
#[derive(Debug, Clone)]
pub struct PgFtsExec {
    pool: Arc<PgPool>,
    /// Fully-qualified table identifier (e.g. `"public"."articles"`).
    qualified_table: String,
    /// Name of the text column to search.
    text_col: String,
    /// The user's search query string.
    query: String,
    /// Maximum number of results to return.
    limit: usize,
    /// PostgreSQL text search configuration (e.g. `'english'`, `'simple'`).
    language: String,
    /// Optional SQL WHERE predicate (no "WHERE" keyword).
    /// Safety: generated by `expr_to_pg_sql` which uses DataFusion's `Unparser`
    /// to convert parsed AST expressions — not user-supplied raw strings.
    filter: Option<String>,
    /// Optional scan limit from an outer SQL LIMIT clause.
    scan_limit: Option<usize>,
    /// Output schema: table columns + `_score Float64`.
    schema: SchemaRef,
    /// Cached DataFusion plan metadata.
    plan_properties: PlanProperties,
}

impl PgFtsExec {
    pub fn new(
        pool: Arc<PgPool>,
        qualified_table: String,
        text_col: String,
        query: String,
        limit: usize,
        language: String,
        filter: Option<String>,
        schema: SchemaRef,
    ) -> Self {
        let plan_properties = PlanProperties::new(
            EquivalenceProperties::new(schema.clone()),
            Partitioning::UnknownPartitioning(1),
            EmissionType::Final,
            Boundedness::Bounded,
        );
        Self {
            pool,
            qualified_table,
            text_col,
            query,
            limit,
            language,
            filter,
            scan_limit: None,
            schema,
            plan_properties,
        }
    }

    /// Set the scan limit (from an outer SQL LIMIT clause).
    pub fn with_scan_limit(mut self, limit: usize) -> Self {
        self.scan_limit = Some(limit);
        self
    }

    /// Build the SELECT column list from the output schema (excludes `_score`).
    ///
    /// Decimal128 columns are cast to `float8` so that sqlx can decode them
    /// without requiring the `rust-decimal`/`bigdecimal` feature.
    fn select_columns(&self) -> String {
        self.schema
            .fields()
            .iter()
            .filter(|f| f.name() != "_score")
            .map(|f| {
                let name = format!("\"{}\"", f.name().replace('"', "\"\""));
                match f.data_type() {
                    DataType::Decimal128(_, _) => format!("{name}::float8 AS {name}"),
                    _ => name,
                }
            })
            .collect::<Vec<_>>()
            .join(", ")
    }

    /// Build the full-text search SELECT query.
    ///
    /// Uses `$1` as the bind parameter for the search query string.
    ///
    /// NOTE: `to_tsvector(lang, col)` appears in both the WHERE and ts_rank() expressions.
    /// PostgreSQL does not CSE these, so the conversion runs twice per matched row.
    /// We intentionally keep the raw expression in the WHERE clause (rather than a
    /// subquery alias) so that PostgreSQL can use a GIN index on
    /// `to_tsvector('lang', col)` when one exists. With the index the double
    /// computation is negligible (ts_rank only runs on matched rows); without one
    /// the real fix is to create the index, not to restructure the query.
    fn build_query(&self) -> String {
        let cols = self.select_columns();
        let text_col = format!("\"{}\"", self.text_col.replace('"', "\"\""));
        let lang = &self.language;
        let tsvec = format!("to_tsvector('{lang}', {text_col})");
        let tsquery = format!("websearch_to_tsquery('{lang}', $1)");
        // ts_rank returns float4; cast to float8 to match the Arrow Float64 _score column.
        let rank_expr = format!("ts_rank({tsvec}, {tsquery})::float8");
        let match_expr = format!("{tsvec} @@ {tsquery}");

        let mut where_parts = vec![match_expr];
        if let Some(ref f) = self.filter {
            where_parts.push(f.clone());
        }
        let where_clause = format!(" WHERE {}", where_parts.join(" AND "));

        let effective_limit = self
            .scan_limit
            .map(|sl| sl.min(self.limit))
            .unwrap_or(self.limit);

        let score_expr = format!("{rank_expr} AS _score");
        let select_list = if cols.is_empty() {
            score_expr
        } else {
            format!("{cols}, {score_expr}")
        };

        format!(
            "SELECT {select_list} \
             FROM {table}{where_clause} \
             ORDER BY _score DESC \
             LIMIT {limit}",
            table = self.qualified_table,
            limit = effective_limit,
        )
    }

    /// Execute the query and return all rows as a single `RecordBatch`.
    async fn run(&self) -> DFResult<RecordBatch> {
        let sql = self.build_query();
        tracing::debug!("pg_fts SQL: {}", sql);

        let rows = sqlx::query(&sql)
            .bind(&self.query)
            .fetch_all(self.pool.as_ref())
            .await
            .map_err(|e| DataFusionError::External(Box::new(e)))?;

        rows_to_batch(&rows, &self.schema)
    }
}

impl DisplayAs for PgFtsExec {
    fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "PgFtsExec: table={}, text_col={}, limit={}",
            self.qualified_table, self.text_col, self.limit
        )
    }
}

impl ExecutionPlan for PgFtsExec {
    fn name(&self) -> &str {
        "PgFtsExec"
    }
    fn as_any(&self) -> &dyn Any {
        self
    }
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
    fn properties(&self) -> &PlanProperties {
        &self.plan_properties
    }

    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
        vec![]
    }

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> DFResult<Arc<dyn ExecutionPlan>> {
        if children.is_empty() {
            Ok(self)
        } else {
            Err(DataFusionError::Internal(
                "PgFtsExec expects 0 children".to_string(),
            ))
        }
    }

    fn execute(
        &self,
        _partition: usize,
        _context: Arc<TaskContext>,
    ) -> DFResult<SendableRecordBatchStream> {
        let exec = self.clone();
        let schema = self.schema.clone();
        let fut = async move { exec.run().await };
        Ok(Box::pin(RecordBatchStreamAdapter::new(
            schema,
            stream::once(fut),
        )))
    }
}

// ─── Tests ────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use arrow::datatypes::{DataType, Field, Schema};

    /// Build a `PgFtsExec` with a lazy (never-connecting) pool for query-building tests.
    fn make_exec(
        cols: Vec<(&str, DataType)>,
        text_col: &str,
        query: &str,
        filter: Option<&str>,
        limit: usize,
    ) -> PgFtsExec {
        make_exec_with_lang(cols, text_col, query, filter, limit, "english")
    }

    fn make_exec_with_lang(
        cols: Vec<(&str, DataType)>,
        text_col: &str,
        query: &str,
        filter: Option<&str>,
        limit: usize,
        language: &str,
    ) -> PgFtsExec {
        let pool =
            Arc::new(sqlx::PgPool::connect_lazy("postgresql://localhost/test").expect("lazy pool"));
        let mut fields: Vec<Field> = cols
            .into_iter()
            .map(|(name, dt)| Field::new(name, dt, true))
            .collect();
        fields.push(Field::new("_score", DataType::Float64, true));
        let schema = Arc::new(Schema::new(fields));
        PgFtsExec::new(
            pool,
            "\"public\".\"articles\"".to_string(),
            text_col.to_string(),
            query.to_string(),
            limit,
            language.to_string(),
            filter.map(str::to_string),
            schema,
        )
    }

    #[tokio::test]
    async fn test_build_query_uses_parameterized_search() {
        let exec = make_exec(
            vec![("id", DataType::Int64), ("content", DataType::Utf8)],
            "content",
            "test query",
            None,
            10,
        );
        let sql = exec.build_query();
        assert!(
            sql.contains("$1"),
            "query should use $1 bind parameter; sql={sql}"
        );
        assert!(
            sql.contains("websearch_to_tsquery('english', $1)"),
            "should use websearch_to_tsquery; sql={sql}"
        );
    }

    #[tokio::test]
    async fn test_build_query_uses_ts_rank() {
        let exec = make_exec(
            vec![("id", DataType::Int64)],
            "body",
            "search terms",
            None,
            10,
        );
        let sql = exec.build_query();
        assert!(
            sql.contains("ts_rank("),
            "should use ts_rank for scoring; sql={sql}"
        );
        assert!(
            sql.contains("AS _score"),
            "score column should be aliased; sql={sql}"
        );
    }

    #[tokio::test]
    async fn test_build_query_orders_by_score_desc() {
        let exec = make_exec(vec![("id", DataType::Int64)], "body", "test", None, 10);
        let sql = exec.build_query();
        assert!(
            sql.contains("ORDER BY _score DESC"),
            "should order by _score descending; sql={sql}"
        );
    }

    #[tokio::test]
    async fn test_build_query_with_filter() {
        let exec = make_exec(
            vec![("id", DataType::Int64)],
            "body",
            "test",
            Some("category = 'news'"),
            10,
        );
        let sql = exec.build_query();
        assert!(
            sql.contains("category = 'news'"),
            "filter should appear in WHERE clause; sql={sql}"
        );
        assert!(
            sql.contains("@@") && sql.contains("AND"),
            "filter should be ANDed with FTS match; sql={sql}"
        );
    }

    #[tokio::test]
    async fn test_build_query_limit() {
        for limit in [1, 5, 100] {
            let exec = make_exec(vec![("id", DataType::Int64)], "body", "q", None, limit);
            let sql = exec.build_query();
            assert!(sql.contains(&format!("LIMIT {limit}")));
        }
    }

    #[tokio::test]
    async fn test_build_query_quotes_text_column() {
        let exec = make_exec(vec![("id", DataType::Int64)], "full text", "test", None, 10);
        let sql = exec.build_query();
        assert!(
            sql.contains("\"full text\""),
            "text column should be quoted; sql={sql}"
        );
    }

    #[tokio::test]
    async fn test_select_columns_excludes_score() {
        let exec = make_exec(
            vec![("id", DataType::Int64), ("title", DataType::Utf8)],
            "body",
            "q",
            None,
            10,
        );
        let cols = exec.select_columns();
        assert!(cols.contains("\"id\""));
        assert!(cols.contains("\"title\""));
        assert!(!cols.contains("_score"));
    }

    #[tokio::test]
    async fn test_select_columns_casts_decimal_to_float8() {
        let exec = make_exec(
            vec![("price", DataType::Decimal128(10, 2))],
            "body",
            "q",
            None,
            10,
        );
        let cols = exec.select_columns();
        assert!(
            cols.contains("::float8"),
            "Decimal128 columns must be cast to float8"
        );
    }

    #[tokio::test]
    async fn test_scan_limit_takes_minimum() {
        let exec =
            make_exec(vec![("id", DataType::Int64)], "body", "q", None, 100).with_scan_limit(5);
        let sql = exec.build_query();
        assert!(
            sql.contains("LIMIT 5"),
            "scan_limit < limit should win; sql={sql}"
        );
    }

    #[tokio::test]
    async fn test_scan_limit_does_not_exceed_function_limit() {
        let exec =
            make_exec(vec![("id", DataType::Int64)], "body", "q", None, 10).with_scan_limit(50);
        let sql = exec.build_query();
        assert!(
            sql.contains("LIMIT 10"),
            "function limit < scan_limit should win; sql={sql}"
        );
    }

    #[tokio::test]
    async fn test_build_query_uses_custom_language() {
        let exec = make_exec_with_lang(
            vec![("id", DataType::Int64)],
            "body",
            "test",
            None,
            10,
            "simple",
        );
        let sql = exec.build_query();
        assert!(
            sql.contains("to_tsvector('simple',"),
            "should use custom language; sql={sql}"
        );
        assert!(
            sql.contains("websearch_to_tsquery('simple',"),
            "should use custom language; sql={sql}"
        );
        assert!(
            !sql.contains("english"),
            "should not contain default language; sql={sql}"
        );
    }
}