infino 0.5.6

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Infino Authors

//! `SupertableReader::query_sql` — DataFusion SQL over a pinned supertable snapshot.
//!
//! ## Public API
//!
//! ```ignore
//! let reader = supertable.reader().expect("reader");
//! let batches: Vec<RecordBatch> =
//!     reader.query_sql("SELECT category, COUNT(*) FROM supertable GROUP BY category")?;
//! ```
//!
//! Sync return type: callers don't need a tokio runtime.
//! Internally the reader drives the async DataFusion plan through the same
//! sync→async bridge used by BM25 and vector search.
//!
//! ## Strategy
//!
//! At `query_sql` time we:
//!
//!   1. Use the reader's already-pinned `Arc<ManifestSnapshot>`.
//!   2. Register a [`SupertableProvider`] as `supertable` in a
//!      fresh `SessionContext`.
//!   3. `ctx.sql(sql).await.collect().await`.
//!
//! The provider's `scan` does the real work — see
//! [`crate::supertable::query::provider`]. In short, it applies
//! **two tiers of pruning**: infino's [`scalar_skip`] drops
//! definitely-irrelevant *superfiles* from the pushed-down `WHERE`
//! predicates, then DataFusion's `ParquetSource` prunes *row
//! groups / pages* and pushes projection + limit into the Parquet
//! reader over the surviving superfiles. This replaces the v1
//! `MemTable` path, which eagerly decoded every row group of every
//! superfile regardless of the query.
//!
//! [`scalar_skip`]: crate::supertable::query::skip::scalar_skip
//! [`SupertableProvider`]: crate::supertable::query::provider::SupertableProvider
//!
//! ## Schema
//!
//! The supertable's *user-visible* schema (`options.scalar_schema`)
//! contains id + scalar columns + FTS columns; vector columns are
//! stored in the embedded vector blob and never exposed via SQL
//! (callers reach them through `vector_search`). The parquet body
//! of each superfile was written with this same scalar schema, so
//! round-trip shape matches without projection or rewrite.
//!
//! **String result type.** String columns are always returned as
//! `LargeUtf8`, regardless of how they are stored or scanned. The scan may
//! run a non-FTS string column internally as `Utf8View` (a comparison
//! optimization), but that view is coerced back to `LargeUtf8` at the plan
//! output and never reaches a caller — a `SELECT`, `GROUP BY` key, or
//! `MIN`/`MAX` over a string column always comes back `LargeUtf8`.

use std::{collections::HashSet, sync::Arc, time::Instant};

use arrow::record_batch::RecordBatch;
use arrow_array::{Array, Decimal128Array};
use arrow_schema::SchemaRef;
use datafusion::{
    datasource::DefaultTableSource,
    error::DataFusionError,
    execution::context::SessionContext,
    logical_expr::{Expr, LogicalPlan},
};

use crate::{
    memory::budgeted_session_context,
    runtime_metrics::op_stats,
    supertable::{
        error::QueryError,
        handle::{Supertable, SupertableReader},
        options::SupertableOptions,
        query::{
            covered_agg::CoveredAggregateRewrite,
            exec::{
                fts_exec::register_bm25, hybrid_exec::register_hybrid_search,
                match_exec::register_match, vector_exec::register_vector_search,
            },
            provider::{SupertableProvider, TABLE_NAME, view_string_schema},
        },
        reader_cache::disk::ForegroundQueryGuard,
    },
};

/// Per-table SQL schemas, built once (`build_sql_schemas`) and cached on the
/// handle instead of recomputed per query. Cheap to clone (fields are `Arc`s).
///
/// - `scalar`: id + scalar + FTS columns, no vectors. What the search TVFs bind to.
/// - `scan`: `scalar` with non-FTS strings viewed as `Utf8View`
///   (`view_string_schema`). What the provider plans against.
#[derive(Clone)]
pub(crate) struct SqlSchemas {
    scalar: SchemaRef,
    scan: SchemaRef,
}

impl SqlSchemas {
    /// Plain scalar schema (id + scalar + FTS, no vectors) the TVFs bind to.
    pub(crate) fn scalar(&self) -> &SchemaRef {
        &self.scalar
    }

    /// String-viewed schema the provider plans against.
    pub(crate) fn scan(&self) -> &SchemaRef {
        &self.scan
    }
}

/// Build the [`SqlSchemas`] for `options`. Called once per table; the result is
/// cached on the handle. This is the one place that walks the full column set,
/// so a wide (thousands of columns) table pays it once, not per query.
pub(crate) fn build_sql_schemas(options: &SupertableOptions) -> SqlSchemas {
    let scalar = options.scalar_schema();
    let fts: HashSet<&str> = options
        .fts_columns
        .iter()
        .map(|c| c.column.as_str())
        .collect();
    let scan = view_string_schema(&scalar, &fts);
    SqlSchemas { scalar, scan }
}

/// Maximum distinct scalar SQL statements cached per manifest snapshot.
const SQL_LOGICAL_PLAN_CACHE_ENTRIES: usize = 64;

/// Cache only plans whose table scans all use [`SupertableProvider`].
///
/// Search TVF providers hold a live reader in their logical plan. Caching
/// those plans on `SupertableInner` would create a reference cycle; scalar
/// providers own only the pinned manifest and storage/cache handles.
fn cacheable_scalar_plan(plan: &LogicalPlan) -> bool {
    fn visit(plan: &LogicalPlan, found_scan: &mut bool) -> bool {
        if let LogicalPlan::TableScan(scan) = plan {
            let Some(source) = scan.source.downcast_ref::<DefaultTableSource>() else {
                return false;
            };
            if source
                .table_provider
                .downcast_ref::<SupertableProvider>()
                .is_none()
            {
                return false;
            }
            *found_scan = true;
        }
        plan.inputs()
            .into_iter()
            .all(|input| visit(input, found_scan))
    }

    let mut found_scan = false;
    visit(plan, &mut found_scan) && found_scan
}

/// Classify a SQL execution error: budget exhaustion -> [`QueryError::OverBudget`]
/// (the catalog surfaces it as `InfinoError::OverBudget`), else an execute error.
fn exec_query_error(e: DataFusionError) -> QueryError {
    match e {
        DataFusionError::ResourcesExhausted(msg) => QueryError::OverBudget(msg),
        other => QueryError::Execute(other.to_string()),
    }
}

impl SupertableReader {
    fn cached_sql_logical_plan(&self, sql: &str) -> Option<LogicalPlan> {
        let guard = self
            .sql_logical_plan_cache()
            .lock()
            .expect("sql logical-plan cache mutex poisoned");
        let (manifest, plans) = guard.as_ref()?;
        Arc::ptr_eq(manifest, self.manifest())
            .then(|| plans.get(sql).cloned())
            .flatten()
    }

    fn cache_sql_logical_plan(&self, sql: String, plan: LogicalPlan) {
        let mut guard = self
            .sql_logical_plan_cache()
            .lock()
            .expect("sql logical-plan cache mutex poisoned");
        if guard
            .as_ref()
            .is_none_or(|(manifest, _)| !Arc::ptr_eq(manifest, self.manifest()))
        {
            *guard = Some((Arc::clone(self.manifest()), Default::default()));
        }
        let (_, plans) = guard.as_mut().expect("cache initialized above");
        if plans.len() >= SQL_LOGICAL_PLAN_CACHE_ENTRIES && !plans.contains_key(&sql) {
            plans.clear();
        }
        plans.insert(sql, plan);
    }

    /// Run a SQL query against this reader's pinned snapshot.
    ///
    /// The snapshot is captured at `query_sql` entry — concurrent
    /// commits don't affect the in-flight query. Returns the
    /// concatenated `Vec<RecordBatch>` from
    /// `DataFrame::collect`.
    ///
    /// The SQL must reference the table as `supertable`. The
    /// available columns are id + scalar + FTS columns; vector
    /// columns are not exposed (use `vector_search` instead).
    ///
    /// Sync API. The first call allocates a tokio Runtime
    /// (single worker thread) cached on the `SupertableInner`;
    /// subsequent calls reuse it.
    ///
    /// Not metered: this entry runs on the cached, collector-detached
    /// [`SessionContext`] (see [`Self::sql_session_context`]), so a
    /// surrounding `with_op_stats` scope reports zero SQL work for it.
    /// The metered SQL surface is the catalog `Connection::query_sql`,
    /// which builds a fresh per-query provider that carries the scope's
    /// collector.
    // Single-table SQL — off the public surface; catalog-level SQL is the
    // public entry point. Reachable from tests/benches via `test-helpers`.
    #[cfg(any(test, feature = "test-helpers"))]
    #[cfg_attr(
        feature = "detailed-tracing",
        tracing::instrument(skip_all, fields(sql = sql))
    )]
    pub fn query_sql(&self, sql: &str) -> Result<Vec<RecordBatch>, QueryError> {
        let _foreground = ForegroundQueryGuard::enter();
        // Read-consistency was applied when `Supertable::reader()` created
        // this pinned reader. SQL therefore observes the same snapshot as
        // `bm25_search` and `vector_search` on this handle.

        // Build (or reuse the cached) SessionContext for the pinned
        // snapshot — the pushdown-aware SupertableProvider plus the
        // search TVFs. See [`SupertableReader::sql_session_context`].
        let ctx = self.sql_session_context()?;
        let tombstone_prefetch = self.tombstone_cache.as_ref().and_then(|cache| {
            let entries = self.manifest().complete_flat_superfiles()?;
            let ids: Vec<_> = entries.iter().map(|entry| entry.superfile_id).collect();
            Some((Arc::clone(cache), ids))
        });
        let cached_plan = self.cached_sql_logical_plan(sql);
        let cache_reader = self.clone();

        let sql = sql.to_owned();
        let drive = async move {
            // The scan runs strings as `Utf8View`; `expand_views_at_output`
            // (set in `budgeted_session_context`) coerces them back to
            // `LargeUtf8` at the plan output, so the result carries no view.
            // Exact manifest statistics can eliminate an unfiltered aggregate
            // before `TableProvider::scan` runs, but only after every
            // superfile's delete view is known. The ordinary scan performs the
            // same prefetch; doing it before planning lets repeated aggregate
            // queries avoid constructing a Parquet plan altogether.
            if let Some((cache, ids)) = tombstone_prefetch {
                cache.prefetch(&ids, Instant::now()).await;
            }
            let df = match cached_plan {
                Some(plan) => ctx
                    .execute_logical_plan(plan)
                    .await
                    .map_err(|e| QueryError::Plan(e.to_string()))?,
                None => {
                    let df = ctx
                        .sql(&sql)
                        .await
                        .map_err(|e| QueryError::Plan(e.to_string()))?;
                    let plan = df.logical_plan().clone();
                    if cacheable_scalar_plan(&plan) {
                        cache_reader.cache_sql_logical_plan(sql.clone(), plan);
                    }
                    df
                }
            };
            df.collect().await.map_err(exec_query_error)
        };

        // Drive through the shared sync→async bridge: ambient
        // runtime → block_in_place on the ambient handle; otherwise
        // the lazily-built owned query_runtime. See
        // [`SupertableReader::block_on`].
        self.block_on(drive)
    }

    /// Build (or reuse the cached) [`SessionContext`] for the
    /// current pinned manifest snapshot: the pushdown-aware
    /// [`SupertableProvider`] registered as `supertable`, plus the
    /// vector / BM25 / hybrid search TVFs.
    ///
    /// The cache keys on the manifest `Arc` — commits publish a new
    /// `Arc`, so any committed state since the last call forces a
    /// rebuild. A hit skips the ~1.5 ms `SessionContext::new()` +
    /// `register_*` setup. Shared by [`query_sql`](Self::query_sql)
    /// (SQL string) and [`scan_ids_matching`](Self::scan_ids_matching)
    /// (programmatic `Expr`), so mutation id-capture gets the same
    /// superfile-skip + row-group/page pruning + lazy tombstone
    /// filtering the read path uses.
    ///
    /// Freshness policy is applied when the reader is created by
    /// [`Supertable::reader`](crate::supertable::handle::Supertable::reader).
    #[cfg_attr(feature = "detailed-tracing", tracing::instrument(skip_all))]
    fn sql_session_context(&self) -> Result<SessionContext, QueryError> {
        // This reader already pins the snapshot; clone is a handful of
        // Arc refcount bumps. Detach any per-query work collector: this
        // context is CACHED across queries, and a collector riding into it
        // would bill later queries into this scope. The whole build below
        // runs under `op_stats::suppressed` for the same reason (provider
        // and TVF constructions capture the thread-local).
        let mut detached = self.clone();
        detached.op_stats = None;
        let reader = Arc::new(detached);
        let manifest = Arc::clone(reader.manifest());

        let mut guard = self
            .sql_session_cache()
            .lock()
            .expect("sql_session_cache mutex poisoned");
        if let Some((cached, ctx)) = &*guard
            && Arc::ptr_eq(cached, &manifest)
        {
            return Ok(ctx.clone());
        }

        let store = Arc::clone(&self.options().store);
        let disk_cache = self.options().disk_cache.as_ref().map(Arc::clone);
        // Cached per-table schemas: the provider scans the string-viewed `scan`
        // schema; the TVFs bind to the plain `scalar` schema.
        let schemas = self.sql_schemas();
        let provider = op_stats::suppressed(|| {
            SupertableProvider::new(
                schemas.scan().clone(),
                Arc::clone(&manifest),
                store,
                disk_cache,
                reader.tombstone_cache.clone(),
            )
        });

        // Gate SQL heap on the connection budget (shared across contexts, so
        // this reader's SQL counts against the same ceiling as the rest).
        let ctx = budgeted_session_context(&self.options().connection_memory_budget)
            .map_err(|e| QueryError::Plan(e.to_string()))?;

        // Covered/residual aggregate rewrite: filter-aligned range
        // aggregates answer covered segments from manifest statistics
        // and scan only the boundary segments. Appended after the
        // built-in rules so it sees pushed-down, normalized plans.
        ctx.add_optimizer_rule(Arc::new(CoveredAggregateRewrite));
        ctx.register_table(TABLE_NAME, Arc::new(provider))
            .map_err(|e| QueryError::Plan(e.to_string()))?;

        // Search TVFs (vector kNN, BM25 FTS, hybrid RRF) bound to
        // the pinned snapshot. They lower to custom `ExecutionPlan`
        // nodes that call the async kernels inside `execute()`.
        register_vector_search(&ctx, Arc::clone(&reader), schemas.scalar().clone());
        register_bm25(&ctx, Arc::clone(&reader), schemas.scalar().clone());
        // Unranked token / exact match TVFs (siblings of bm25_search).
        register_match(&ctx, Arc::clone(&reader), schemas.scalar().clone());
        register_hybrid_search(&ctx, Arc::clone(&reader), schemas.scalar().clone());

        *guard = Some((Arc::clone(&manifest), ctx.clone()));

        Ok(ctx)
    }

    /// Resolve a predicate to the matching `_id` values. Used by
    /// the writer's `delete()` / `update()` entry points to
    /// capture the target-id set at call time (step 0a in the
    /// update / delete pipeline).
    ///
    /// Runs through the same pushdown-aware [`SupertableProvider`]
    /// as `query_sql` (via [`sql_session_context`](Self::sql_session_context)):
    /// `expr` is applied as a `DataFrame::filter` and the result
    /// projected to just `_id`. Superfile skip, row-group / page
    /// pruning, and lazy tombstone filtering all apply, so a
    /// large-table delete/update predicate never materializes every
    /// superfile into memory.
    ///
    /// Note: the resolution is against the **current** manifest
    /// snapshot, exactly like a contemporaneous `query_sql` would
    /// see. Rows that newly match `expr` between this call and
    /// the eventual `commit()` are NOT in the returned set —
    /// captured-at-call semantics match SQL `UPDATE WHERE` /
    /// `DELETE WHERE`.
    pub(crate) fn scan_ids_matching(&self, expr: Expr) -> Result<Vec<i128>, QueryError> {
        let _foreground = ForegroundQueryGuard::enter();
        // Resolve against this reader's pinned snapshot. Callers that need
        // current-state semantics create a fresh reader immediately before
        // invoking this helper.
        let ctx = self.sql_session_context()?;
        let id_column = self.options().id_column.clone();

        let drive = async move {
            let df = ctx
                .table(TABLE_NAME)
                .await
                .map_err(|e| QueryError::Plan(e.to_string()))?
                .filter(expr)
                .map_err(|e| QueryError::Plan(e.to_string()))?
                .select_columns(&[id_column.as_str()])
                .map_err(|e| QueryError::Plan(e.to_string()))?;
            let batches = df.collect().await.map_err(exec_query_error)?;
            extract_id_column(&batches)
        };

        self.block_on(drive)
    }
}

impl Supertable {
    /// Register this supertable's pushdown-aware provider into `ctx`
    /// under `name`, applying the read-consistency policy first. The
    /// catalog's multi-table [`Connection::query_sql`] calls this once
    /// per referenced table. Returns the pinned reader so the caller can
    /// later wire the same snapshot into search TVFs.
    ///
    /// [`Connection::query_sql`]: crate::Connection::query_sql
    pub(crate) fn register_into(
        &self,
        ctx: &SessionContext,
        name: &str,
    ) -> Result<Arc<SupertableReader>, QueryError> {
        // `reader()` applies the read-consistency freshness check itself (and,
        // under Strong, fails rather than serving a stale snapshot), so no
        // separate `ensure_fresh` call is needed here.
        let reader = Arc::new(self.reader().map_err(QueryError::ManifestLoad)?);
        let manifest = Arc::clone(reader.manifest());
        let store = Arc::clone(&self.options().store);
        let disk_cache = self.options().disk_cache.as_ref().map(Arc::clone);
        // Provider scans the cached string-viewed schema.
        let provider = SupertableProvider::new(
            self.sql_schemas().scan().clone(),
            manifest,
            store,
            disk_cache,
            reader.tombstone_cache.clone(),
        );
        ctx.register_table(name, Arc::new(provider))
            .map_err(|e| QueryError::Plan(e.to_string()))?;
        Ok(reader)
    }
}

/// Drain `_id`-only batches into a `Vec<i128>`. The supertable's
/// `_id` is a Decimal128(38, 0) column; we read the raw 128-bit
/// integer value directly.
fn extract_id_column(batches: &[RecordBatch]) -> Result<Vec<i128>, QueryError> {
    let mut out: Vec<i128> = Vec::new();
    for batch in batches {
        if batch.num_columns() != 1 {
            return Err(QueryError::Plan(format!(
                "scan_ids_matching: expected 1-column batch, got {}",
                batch.num_columns()
            )));
        }
        let col = batch.column(0);
        let arr = col
            .as_any()
            .downcast_ref::<Decimal128Array>()
            .ok_or_else(|| {
                QueryError::Plan("scan_ids_matching: _id column not Decimal128".into())
            })?;
        for i in 0..arr.len() {
            if arr.is_null(i) {
                continue;
            }
            out.push(arr.value(i));
        }
    }
    Ok(out)
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use arrow_array::{
        Array, Decimal128Array, FixedSizeListArray, Float32Array, Int64Array, LargeStringArray,
        RecordBatch, StringArray, StringViewArray,
    };
    use arrow_schema::{DataType, Field, Schema};

    use crate::{
        memory::ConnectionMemoryBudget,
        storage::{LocalFsStorageProvider, StorageProvider},
        superfile::{
            builder::{FtsConfig, VectorConfig},
            vector::{distance::Metric, rerank_codec::RerankCodec},
        },
        supertable::{
            Supertable, SupertableOptions, error::QueryError, query::sql::build_sql_schemas,
        },
        test_helpers::default_tokenizer as tok,
    };

    /// One more than the manifest's exact-value cardinality cap.
    const HIGH_CARDINALITY_ROWS: usize = 257;

    /// Schema with id + scalar + FTS column. No vector; query_sql
    /// is scalar-only by design.
    fn schema_id_cat_title() -> Arc<Schema> {
        Arc::new(Schema::new(vec![
            Field::new("category", DataType::LargeUtf8, false),
            Field::new("title", DataType::LargeUtf8, false),
        ]))
    }

    fn options_id_cat_title() -> SupertableOptions {
        // Single-threaded writer pool so each commit produces
        // exactly one superfile — keeps assertions on per-superfile
        // counts deterministic.
        let pool = Arc::new(
            rayon::ThreadPoolBuilder::new()
                .num_threads(1)
                .build()
                .expect("rayon pool"),
        );
        SupertableOptions::new(
            schema_id_cat_title(),
            vec![FtsConfig {
                column: "title".into(),
                positions: false,
            }],
            vec![],
            Some(tok()),
        )
        .expect("valid options")
        .with_writer_pool(pool)
    }

    // Ingest `batch` on a measured supertable, then return a second handle over
    // the same durable storage under a 0-byte gate. Ingest is gated by the
    // budget too, so a query-gating test can't reuse one tiny-budget handle for
    // both; this does the setup on a measured handle and hands back the gated
    // reader. The returned `TempDir` guard must be held: dropping it deletes the
    // store the reader is still reading through.
    fn zero_gate_reader_after_ingest(batch: &RecordBatch) -> (tempfile::TempDir, Supertable) {
        let dir = tempfile::tempdir().expect("tempdir");
        let storage: Arc<dyn StorageProvider> =
            Arc::new(LocalFsStorageProvider::new(dir.path()).expect("localfs"));

        let ingest = Supertable::create(options_id_cat_title().with_storage(Arc::clone(&storage)))
            .expect("create");
        let mut w = ingest.writer().expect("writer");
        w.append(batch).expect("append");
        w.commit().expect("commit");

        let mut qopts = options_id_cat_title().with_storage(storage);
        qopts.connection_memory_budget = ConnectionMemoryBudget::with_limit(1);
        (dir, Supertable::open(qopts).expect("open"))
    }

    /// Build a small categorical batch — start id sequence at
    /// `start`, plant `cats[i] / titles[i]` per row.
    fn build_cat_batch(_start: u64, cats: &[&str], titles: &[&str]) -> RecordBatch {
        assert_eq!(cats.len(), titles.len());
        let cat_arr = LargeStringArray::from(cats.to_vec());
        let title_arr = LargeStringArray::from(titles.to_vec());
        RecordBatch::try_new(
            schema_id_cat_title(),
            vec![Arc::new(cat_arr), Arc::new(title_arr)],
        )
        .expect("build batch")
    }

    /// A single-superfile table seeded with one committed batch of
    /// `cats`/`titles`. Collapses the create + append + commit boilerplate the
    /// string-view tests share.
    fn seeded(cats: &[&str], titles: &[&str]) -> Supertable {
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let mut w = st.writer().expect("writer");
        w.append(&build_cat_batch(0, cats, titles)).expect("append");
        w.commit().expect("commit");
        st
    }

    fn rating_table(ratings: &[i64]) -> Supertable {
        let schema = Arc::new(Schema::new(vec![Field::new(
            "rating",
            DataType::Int64,
            false,
        )]));
        let pool = Arc::new(
            rayon::ThreadPoolBuilder::new()
                .num_threads(1)
                .build()
                .expect("rayon pool"),
        );
        let options = SupertableOptions::new(schema.clone(), vec![], vec![], None)
            .expect("rating options")
            .with_writer_pool(pool);
        let table = Supertable::create(options).expect("create rating table");
        let mut writer = table.writer().expect("rating writer");
        let batch =
            RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(ratings.to_vec()))])
                .expect("rating batch");
        writer.append(&batch).expect("append ratings");
        writer.commit().expect("commit ratings");
        drop(writer);
        table
    }

    /// Convenience: run a query and pull a single `Int64` aggregate
    /// value from cell (0,0).
    fn run_count(st: &Supertable, sql: &str) -> i64 {
        let batches = st
            .reader()
            .expect("reader")
            .query_sql(sql)
            .expect("query_sql ok");
        assert!(!batches.is_empty(), "expected at least one result batch");
        let n = batches[0]
            .column(0)
            .as_any()
            .downcast_ref::<Int64Array>()
            .expect("count column is Int64");
        n.value(0)
    }

    /// `extract_id_column` collects non-null Decimal128 `_id`s from single-column
    /// batches and rejects a batch that isn't exactly one column.
    #[test]
    fn extract_id_column_reads_decimal128_and_rejects_multi_column() {
        use arrow_array::ArrayRef;
        let ids: ArrayRef = Arc::new(
            Decimal128Array::from(vec![Some(1i128), Some(2), None, Some(3)])
                .with_precision_and_scale(38, 0)
                .expect("decimal"),
        );
        let schema = Arc::new(Schema::new(vec![Field::new(
            "_id",
            DataType::Decimal128(38, 0),
            true,
        )]));
        let batch = RecordBatch::try_new(schema, vec![ids]).expect("batch");
        assert_eq!(
            super::extract_id_column(&[batch]).expect("ids"),
            vec![1i128, 2, 3]
        );

        let two = RecordBatch::try_new(
            Arc::new(Schema::new(vec![
                Field::new("a", DataType::Int64, false),
                Field::new("b", DataType::Int64, false),
            ])),
            vec![
                Arc::new(Int64Array::from(vec![1])) as ArrayRef,
                Arc::new(Int64Array::from(vec![2])) as ArrayRef,
            ],
        )
        .expect("two-col batch");
        assert!(super::extract_id_column(&[two]).is_err());
    }

    #[test]
    fn query_sql_count_star_returns_zero_on_empty_supertable() {
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let n = run_count(&st, "SELECT COUNT(*) FROM supertable");
        assert_eq!(n, 0);
    }

    #[test]
    fn query_sql_count_star_returns_total_doc_count() {
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let mut w = st.writer().expect("writer");
        w.append(&build_cat_batch(
            0,
            &["rust", "rust", "python"],
            &["a", "b", "c"],
        ))
        .expect("append");
        w.commit().expect("commit");

        let n = run_count(&st, "SELECT COUNT(*) FROM supertable");
        assert_eq!(n, 3);
    }

    #[test]
    fn query_sql_caches_scalar_plan_but_not_search_tvf_plan() {
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let mut writer = st.writer().expect("writer");
        writer
            .append(&build_cat_batch(0, &["rust"], &["searchable"]))
            .expect("append");
        writer.commit().expect("commit");
        let reader = st.reader().expect("reader");
        let scalar_sql = "SELECT COUNT(*) FROM supertable";

        reader.query_sql(scalar_sql).expect("first scalar query");
        reader.query_sql(scalar_sql).expect("cached scalar query");
        {
            let guard = reader
                .sql_logical_plan_cache()
                .lock()
                .expect("plan cache lock");
            let (_, plans) = guard.as_ref().expect("scalar plan cached");
            assert_eq!(plans.len(), 1);
            assert!(plans.contains_key(scalar_sql));
        }

        reader
            .query_sql("SELECT _id FROM bm25_search('title', 'searchable', 10)")
            .expect("search TVF query");
        let guard = reader
            .sql_logical_plan_cache()
            .lock()
            .expect("plan cache lock");
        let (_, plans) = guard.as_ref().expect("scalar plan remains cached");
        assert_eq!(
            plans.len(),
            1,
            "search TVF plans hold readers and must not enter the inner cache"
        );
    }

    /// Regression test for the cold-reopen consumer leak. Running
    /// `query_sql` builds and caches a `SessionContext` on the
    /// `SupertableInner`, and that context registers the search TVFs.
    /// When the TVFs held a strong `Arc<SupertableReader>` (which holds
    /// the `Arc<SupertableInner>`), the chain
    /// `inner -> cached SessionContext -> TVF -> reader -> inner` formed a
    /// reference cycle that pinned the whole consumer — every fresh
    /// consumer reopen (the cold query path) leaked one, OOMing at scale.
    /// With the TVFs holding a `WeakReader`, dropping the last external
    /// handle releases the inner; a `Weak` to it must fail to upgrade.
    #[test]
    fn query_sql_session_cache_does_not_leak_consumer() {
        let weak = {
            let st = Supertable::create(options_id_cat_title()).expect("create");
            let mut w = st.writer().expect("writer");
            w.append(&build_cat_batch(0, &["rust"], &["a"]))
                .expect("append");
            w.commit().expect("commit");

            // Populate the cached SessionContext (registers the TVFs).
            assert_eq!(run_count(&st, "SELECT COUNT(*) FROM supertable"), 1);

            let weak = Arc::downgrade(st.inner());
            drop(w);
            drop(st);
            weak
        };

        assert!(
            weak.upgrade().is_none(),
            "SQL session cache leaked the consumer — the \
             inner -> SessionContext -> TVF -> reader -> inner cycle was not broken",
        );
    }

    #[test]
    fn query_sql_filter_predicate_applied_above_mem_table() {
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let mut w = st.writer().expect("writer");
        w.append(&build_cat_batch(
            0,
            &["rust", "rust", "python", "rust", "go"],
            &["a", "b", "c", "d", "e"],
        ))
        .expect("append");
        w.commit().expect("commit");

        let n = run_count(
            &st,
            "SELECT COUNT(*) FROM supertable WHERE category = 'rust'",
        );
        assert_eq!(n, 3);
    }

    #[test]
    fn query_sql_range_count_uses_exact_value_frequencies() {
        let table = rating_table(&[0, 5, 9, 10, 10, 20, 99]);
        assert_eq!(
            run_count(&table, "SELECT COUNT(*) FROM supertable WHERE rating < 10"),
            3
        );
        assert_eq!(
            run_count(
                &table,
                "SELECT COUNT(*) FROM supertable WHERE rating BETWEEN 10 AND 20"
            ),
            3
        );
        assert_eq!(
            run_count(&table, "SELECT COUNT(*) FROM supertable WHERE rating > 100"),
            0
        );
    }

    #[test]
    fn query_sql_range_count_falls_back_above_value_count_cap() {
        let ratings: Vec<i64> = (0..HIGH_CARDINALITY_ROWS)
            .map(|value| value as i64)
            .collect();
        let table = rating_table(&ratings);
        assert_eq!(
            run_count(&table, "SELECT COUNT(*) FROM supertable WHERE rating < 10"),
            10
        );
    }

    #[test]
    fn query_sql_group_by_over_budget_is_refused() {
        // The reader path (second production ctx site) is gated too: a 0-byte
        // gate refuses an aggregate that cannot fold from exact manifest
        // value counts and surfaces as QueryError::OverBudget. High
        // cardinality keeps the count-fold fast path out of play: a
        // low-cardinality batch would fold COUNT DISTINCT from the
        // manifest's exact value counts and never hit the gate.
        let categories: Vec<String> = (0..HIGH_CARDINALITY_ROWS)
            .map(|value| format!("category-{value}"))
            .collect();
        let category_refs: Vec<&str> = categories.iter().map(String::as_str).collect();
        let titles = vec!["title"; HIGH_CARDINALITY_ROWS];
        let (_dir, st) =
            zero_gate_reader_after_ingest(&build_cat_batch(0, &category_refs, &titles));

        let err = st
            .reader()
            .expect("reader")
            .query_sql("SELECT category, COUNT(*) FROM supertable GROUP BY category")
            .expect_err("0-byte gate refuses the aggregate");

        assert!(matches!(err, QueryError::OverBudget(_)), "got {err:?}");
    }

    #[test]
    fn query_sql_streaming_scan_is_not_refused_under_a_zero_gate() {
        // A projection streams (no buffering), so it runs even at a 0-byte gate:
        // the budget bounds sort/aggregate/join, not scans.
        let (_dir, st) =
            zero_gate_reader_after_ingest(&build_cat_batch(0, &["rust", "python"], &["a", "b"]));

        let rows: usize = st
            .reader()
            .expect("reader")
            .query_sql("SELECT title FROM supertable")
            .expect("a streaming scan is not gated")
            .iter()
            .map(|b| b.num_rows())
            .sum();

        assert_eq!(rows, 2);
    }

    #[test]
    fn query_sql_group_by_returns_correct_per_category_counts() {
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let mut w = st.writer().expect("writer");
        w.append(&build_cat_batch(
            0,
            &["rust", "rust", "python", "rust", "python", "go"],
            &["a", "b", "c", "d", "e", "f"],
        ))
        .expect("append");
        w.commit().expect("commit");

        let batches = st
            .reader()
            .expect("reader")
            .query_sql(
                "SELECT category, COUNT(*) AS n FROM supertable \
                 GROUP BY category ORDER BY category",
            )
            .expect("group-by query");
        assert_eq!(batches.len(), 1);

        let cat_col = batches[0].column(0);
        let counts = batches[0]
            .column(1)
            .as_any()
            .downcast_ref::<Int64Array>()
            .expect("count is Int64");
        // DataFusion may materialize the GROUP BY key as Utf8,
        // LargeUtf8, or StringView depending on hash-aggregate
        // type promotion; accept all three.
        let extract = |i: usize| -> String {
            if let Some(a) = cat_col.as_any().downcast_ref::<LargeStringArray>() {
                a.value(i).to_string()
            } else if let Some(a) = cat_col.as_any().downcast_ref::<StringArray>() {
                a.value(i).to_string()
            } else if let Some(a) = cat_col.as_any().downcast_ref::<StringViewArray>() {
                a.value(i).to_string()
            } else {
                panic!("unexpected category column type: {:?}", cat_col.data_type())
            }
        };
        let mut got: Vec<(String, i64)> = (0..cat_col.len())
            .map(|i| (extract(i), counts.value(i)))
            .collect();
        got.sort();
        assert_eq!(
            got,
            vec![
                ("go".to_string(), 1),
                ("python".to_string(), 2),
                ("rust".to_string(), 3),
            ]
        );
    }

    // ---- Utf8View scan ----------------------------------------------------

    /// The scan runs strings as `Utf8View`, but `expand_views_at_output` coerces
    /// them to `LargeUtf8` at the plan output, so no view leaks to a caller: a
    /// GROUP BY key on a `LargeUtf8` column comes back `LargeUtf8`, not a view.
    #[test]
    fn query_sql_string_group_by_key_is_large_utf8_not_view() {
        let st = seeded(&["rust", "go", "rust"], &["a", "b", "c"]);

        let batches = st
            .reader()
            .expect("reader")
            .query_sql("SELECT category FROM supertable GROUP BY category")
            .expect("group-by");
        let col = batches[0].column(0);
        assert_eq!(
            col.data_type(),
            &DataType::LargeUtf8,
            "public result must be LargeUtf8, not Utf8View"
        );
        assert!(
            col.as_any().downcast_ref::<LargeStringArray>().is_some(),
            "category should downcast to LargeStringArray"
        );
        assert!(
            col.as_any().downcast_ref::<StringViewArray>().is_none(),
            "Utf8View must not leak to the caller"
        );
    }

    /// A projected + `ORDER BY` string column returns `LargeUtf8` and the
    /// values are correctly sorted (the view compare ran during the sort).
    #[test]
    fn query_sql_ordered_string_projection_is_large_utf8_and_sorted() {
        let st = seeded(&["rust", "go", "python"], &["a", "b", "c"]);
        let batches = st
            .reader()
            .expect("reader")
            .query_sql("SELECT category FROM supertable ORDER BY category")
            .expect("order-by");
        let col = batches[0]
            .column(0)
            .as_any()
            .downcast_ref::<LargeStringArray>()
            .expect("category is LargeUtf8");
        let got: Vec<&str> = (0..col.len()).map(|i| col.value(i)).collect();
        assert_eq!(got, vec!["go", "python", "rust"]);
    }

    /// Grouped `MIN(string)` aggregates on the view and returns `LargeUtf8`,
    /// with correct per-group minima.
    #[test]
    fn query_sql_grouped_min_string_is_large_utf8() {
        let st = seeded(&["rust", "rust", "go", "go"], &["b", "a", "d", "c"]);
        let batches = st
            .reader()
            .expect("reader")
            .query_sql(
                "SELECT category, MIN(title) AS m FROM supertable \
                 GROUP BY category ORDER BY category",
            )
            .expect("grouped min");
        let cat = batches[0]
            .column(0)
            .as_any()
            .downcast_ref::<LargeStringArray>()
            .expect("category is LargeUtf8");
        let m = batches[0]
            .column(1)
            .as_any()
            .downcast_ref::<LargeStringArray>()
            .expect("MIN(title) is LargeUtf8");
        let got: Vec<(&str, &str)> = (0..cat.len()).map(|i| (cat.value(i), m.value(i))).collect();
        assert_eq!(got, vec![("go", "c"), ("rust", "a")]);
    }

    /// Ungrouped `MIN(string)` over a viewed column. On its own this trips a
    /// DataFusion `ProjectionPushdown` schema mismatch (`Utf8View` vs
    /// `LargeUtf8`); `expand_views_at_output` (set in `budgeted_session_context`)
    /// coerces the view at the plan output and sidesteps it. Returns `LargeUtf8`.
    #[test]
    fn query_sql_ungrouped_min_string() {
        let st = seeded(&["rust", "go", "python"], &["a", "b", "c"]);
        let batches = st
            .reader()
            .expect("reader")
            .query_sql("SELECT MIN(category) AS m FROM supertable")
            .expect("ungrouped min");
        let col = batches[0]
            .column(0)
            .as_any()
            .downcast_ref::<LargeStringArray>()
            .expect("MIN(string) is LargeUtf8");
        assert_eq!(col.value(0), "go");
    }

    /// Unit: `build_sql_schemas` views the scan schema (non-FTS strings ->
    /// `Utf8View`, FTS kept) and keeps the plain `scalar`. The walk done once
    /// per table.
    #[test]
    fn build_sql_schemas_views_scan_and_keeps_scalar() {
        let s = build_sql_schemas(&options_id_cat_title());
        // scan: `category` (non-FTS string) viewed; `title` (FTS) kept.
        assert_eq!(
            s.scan()
                .field_with_name("category")
                .expect("category")
                .data_type(),
            &DataType::Utf8View,
        );
        assert_eq!(
            s.scan()
                .field_with_name("title")
                .expect("title")
                .data_type(),
            &DataType::LargeUtf8,
            "FTS column stays LargeUtf8 in the scan schema",
        );
        // scalar: no viewing.
        assert_eq!(
            s.scalar()
                .field_with_name("category")
                .expect("category")
                .data_type(),
            &DataType::LargeUtf8,
        );
    }

    /// The per-table schemas are built once and memoized on the handle, not
    /// rebuilt per query (the whole point of the cache for wide tables).
    #[test]
    fn sql_schemas_is_memoized_across_calls() {
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let a = st.sql_schemas();
        let b = st.sql_schemas();
        assert!(
            Arc::ptr_eq(&a, &b),
            "sql_schemas must be cached (same Arc), not recomputed per call",
        );
    }

    /// A NULL string value survives the view scan + output coercion.
    #[test]
    fn query_sql_null_string_survives() {
        let schema = Arc::new(Schema::new(vec![
            Field::new("category", DataType::LargeUtf8, true), // nullable, so we can plant a NULL
            Field::new("title", DataType::LargeUtf8, false),
        ]));
        let pool = Arc::new(
            rayon::ThreadPoolBuilder::new()
                .num_threads(1)
                .build()
                .expect("rayon pool"),
        );
        let opts = SupertableOptions::new(
            Arc::clone(&schema),
            vec![FtsConfig {
                column: "title".into(),
                positions: false,
            }],
            vec![],
            Some(tok()),
        )
        .expect("valid options")
        .with_writer_pool(pool);

        let st = Supertable::create(opts).expect("create");
        let mut w = st.writer().expect("writer");
        let batch = RecordBatch::try_new(
            schema,
            vec![
                Arc::new(LargeStringArray::from(vec![Some("rust"), None, Some("go")])),
                Arc::new(LargeStringArray::from(vec!["a", "b", "c"])),
            ],
        )
        .expect("batch");
        w.append(&batch).expect("append");
        w.commit().expect("commit");

        let batches = st
            .reader()
            .expect("reader")
            .query_sql("SELECT category FROM supertable")
            .expect("select");
        let col = batches[0]
            .column(0)
            .as_any()
            .downcast_ref::<LargeStringArray>()
            .expect("category is LargeUtf8");
        assert_eq!(col.null_count(), 1, "the NULL survives the view + coercion");
    }

    /// A column the user declared `Utf8View` comes back `LargeUtf8`: the view
    /// is an internal scan type, and `expand_views_at_output` coerces every
    /// view to `LargeUtf8` at the plan output, so SQL results never expose one.
    #[test]
    fn query_sql_declared_utf8view_column_returns_large_utf8() {
        let schema = Arc::new(Schema::new(vec![
            Field::new("category", DataType::Utf8View, false), // user declares a view
            Field::new("title", DataType::LargeUtf8, false),   // FTS column must be LargeUtf8
        ]));
        let pool = Arc::new(
            rayon::ThreadPoolBuilder::new()
                .num_threads(1)
                .build()
                .expect("rayon pool"),
        );
        let opts = SupertableOptions::new(
            Arc::clone(&schema),
            vec![FtsConfig {
                column: "title".into(),
                positions: false,
            }],
            vec![],
            Some(tok()),
        )
        .expect("valid options")
        .with_writer_pool(pool);

        let st = Supertable::create(opts).expect("create");
        let mut w = st.writer().expect("writer");
        let batch = RecordBatch::try_new(
            schema,
            vec![
                Arc::new(StringViewArray::from(vec!["rust", "go", "rust"])),
                Arc::new(LargeStringArray::from(vec!["a", "b", "c"])),
            ],
        )
        .expect("batch");
        w.append(&batch).expect("append");
        w.commit().expect("commit");

        let batches = st
            .reader()
            .expect("reader")
            .query_sql("SELECT category FROM supertable GROUP BY category")
            .expect("group-by");
        assert_eq!(
            batches[0].column(0).data_type(),
            &DataType::LargeUtf8,
            "views are internal; SQL results expose LargeUtf8, not Utf8View"
        );
    }

    /// Alias on a viewed string column: the aliased output name has no declared
    /// type, so it defaults to `LargeUtf8`; values stay correct.
    #[test]
    fn query_sql_aliased_string_column_is_large_utf8() {
        let st = seeded(&["rust", "go", "rust"], &["a", "b", "c"]);

        let batches = st
            .reader()
            .expect("reader")
            .query_sql("SELECT category AS c FROM supertable GROUP BY c ORDER BY c")
            .expect("alias");
        let col = batches[0]
            .column(0)
            .as_any()
            .downcast_ref::<LargeStringArray>()
            .expect("aliased column is LargeUtf8");
        let got: Vec<&str> = (0..col.len()).map(|i| col.value(i)).collect();
        assert_eq!(got, vec!["go", "rust"]);
    }

    /// String column projected through a CTE: the name survives, so it is
    /// returned as `LargeUtf8`.
    #[test]
    fn query_sql_cte_string_column_is_declared_type() {
        let st = seeded(&["rust", "go", "rust"], &["a", "b", "c"]);

        let batches = st
            .reader()
            .expect("reader")
            .query_sql(
                "WITH t AS (SELECT category FROM supertable) \
                 SELECT category FROM t GROUP BY category ORDER BY category",
            )
            .expect("cte");
        let col = batches[0]
            .column(0)
            .as_any()
            .downcast_ref::<LargeStringArray>()
            .expect("CTE column is LargeUtf8");
        let got: Vec<&str> = (0..col.len()).map(|i| col.value(i)).collect();
        assert_eq!(got, vec!["go", "rust"]);
    }

    /// String column projected through a FROM-subquery.
    #[test]
    fn query_sql_subquery_string_column_is_declared_type() {
        let st = seeded(&["rust", "go", "rust"], &["a", "b", "c"]);

        let batches = st
            .reader()
            .expect("reader")
            .query_sql(
                "SELECT category FROM (SELECT category FROM supertable) sub \
                 GROUP BY category ORDER BY category",
            )
            .expect("subquery");
        let col = batches[0]
            .column(0)
            .as_any()
            .downcast_ref::<LargeStringArray>()
            .expect("subquery column is LargeUtf8");
        let got: Vec<&str> = (0..col.len()).map(|i| col.value(i)).collect();
        assert_eq!(got, vec!["go", "rust"]);
    }

    /// No data loss through the view + coercion for values that stress
    /// `Utf8View`'s layout: strings past the 12-byte inline limit (stored
    /// out-of-line), values sharing a 4-byte prefix (the view compares the
    /// prefix first, so it must fall through to the full bytes and keep them
    /// distinct), an empty string, and multi-byte unicode. GROUP BY exercises
    /// both the comparison (distinct groups) and the coercion (exact values).
    #[test]
    fn query_sql_string_values_survive_view_and_coercion() {
        let vals = [
            "",                        // empty
            "short",                   // inline (<= 12 bytes)
            "sixteen_byte_val",        // 16 bytes, out-of-line
            "prefabricated_alpha",     // shares "pref" 4-byte prefix ...
            "prefabricated_omega",     // ... differs later, must stay distinct
            "café_ünïcode_日本語_str", // multi-byte unicode, out-of-line
            "sixteen_byte_val",        // duplicate: must fold to one group
        ];
        let titles: Vec<&str> = (0..vals.len()).map(|_| "t").collect();
        let st = seeded(&vals, &titles);
        let batches = st
            .reader()
            .expect("reader")
            .query_sql("SELECT category FROM supertable GROUP BY category ORDER BY category")
            .expect("group-by over layout-stressing values");
        let col = batches[0]
            .column(0)
            .as_any()
            .downcast_ref::<LargeStringArray>()
            .expect("category is LargeUtf8");
        let mut got: Vec<&str> = (0..col.len()).map(|i| col.value(i)).collect();
        got.sort_unstable();

        // Every distinct value survives byte-for-byte; the duplicate folds to
        // one; the prefix-sharing pair stays as two.
        let mut want: Vec<&str> = vec![
            "",
            "café_ünïcode_日本語_str",
            "prefabricated_alpha",
            "prefabricated_omega",
            "short",
            "sixteen_byte_val",
        ];
        want.sort_unstable();
        assert_eq!(got, want);
    }

    /// `SELECT DISTINCT` on a viewed string column: dedup compares on the view,
    /// result comes back `LargeUtf8`.
    #[test]
    fn query_sql_distinct_string_is_declared_type() {
        let st = seeded(&["rust", "go", "rust"], &["a", "b", "c"]);

        let batches = st
            .reader()
            .expect("reader")
            .query_sql("SELECT DISTINCT category FROM supertable ORDER BY category")
            .expect("distinct");
        let col = batches[0]
            .column(0)
            .as_any()
            .downcast_ref::<LargeStringArray>()
            .expect("distinct column is LargeUtf8");
        let got: Vec<&str> = (0..col.len()).map(|i| col.value(i)).collect();
        assert_eq!(got, vec!["go", "rust"]);
    }

    /// Self-join whose join key is a viewed string column: the equality runs on
    /// `Utf8View`, and the projected key comes back `LargeUtf8`.
    #[test]
    fn query_sql_self_join_on_string_key() {
        let st = seeded(&["rust", "go", "rust"], &["a", "b", "c"]);

        let batches = st
            .reader()
            .expect("reader")
            .query_sql(
                "SELECT a.category AS cat FROM supertable a \
                 JOIN supertable b ON a.category = b.category \
                 GROUP BY a.category ORDER BY a.category",
            )
            .expect("self-join on string key");
        let col = batches[0]
            .column(0)
            .as_any()
            .downcast_ref::<LargeStringArray>()
            .expect("join key projects as LargeUtf8");
        let got: Vec<&str> = (0..col.len()).map(|i| col.value(i)).collect();
        assert_eq!(got, vec!["go", "rust"]);
    }

    #[test]
    fn query_sql_group_by_with_nulls_falls_back_and_keeps_null_group() {
        let schema = Arc::new(Schema::new(vec![Field::new(
            "category",
            DataType::LargeUtf8,
            true,
        )]));
        let options = SupertableOptions::new(schema.clone(), vec![], vec![], None)
            .expect("nullable category options");
        let table = Supertable::create(options).expect("create");
        let mut writer = table.writer().expect("writer");
        let categories = LargeStringArray::from(vec![Some("rust"), None, Some("rust")]);
        let batch = RecordBatch::try_new(schema, vec![Arc::new(categories)]).expect("batch");
        writer.append(&batch).expect("append");
        writer.commit().expect("commit");

        let batches = table
            .reader()
            .expect("reader")
            .query_sql(
                "SELECT category, COUNT(*) AS n FROM supertable \
                 GROUP BY category ORDER BY category NULLS FIRST",
            )
            .expect("group by nullable category");
        let categories = batches[0]
            .column(0)
            .as_any()
            .downcast_ref::<LargeStringArray>()
            .expect("large utf8 category");
        let counts = batches[0]
            .column(1)
            .as_any()
            .downcast_ref::<Int64Array>()
            .expect("int64 count");
        assert!(categories.is_null(0));
        assert_eq!(counts.value(0), 1);
        assert_eq!(categories.value(1), "rust");
        assert_eq!(counts.value(1), 2);
    }

    #[test]
    fn query_sql_scans_across_multiple_superfiles() {
        // Three commits → three superfiles. SQL must aggregate across
        // all of them.
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let mut w = st.writer().expect("writer");
        w.append(&build_cat_batch(0, &["rust", "rust"], &["a", "b"]))
            .expect("a1");
        w.commit().expect("c1");
        w.append(&build_cat_batch(10, &["python"], &["c"]))
            .expect("a2");
        w.commit().expect("c2");
        w.append(&build_cat_batch(20, &["rust", "go"], &["d", "e"]))
            .expect("a3");
        w.commit().expect("c3");

        assert_eq!(st.reader().expect("reader").n_superfiles(), 3);

        let n_total = run_count(&st, "SELECT COUNT(*) FROM supertable");
        assert_eq!(n_total, 5);

        let n_rust = run_count(
            &st,
            "SELECT COUNT(*) FROM supertable WHERE category = 'rust'",
        );
        assert_eq!(n_rust, 3);
    }

    #[test]
    fn query_sql_equality_on_fts_column_across_superfiles_is_correct() {
        // Equality on the FTS-indexed `title` column drives the new
        // term-bloom prune leaf (plus the scalar min/max leaf). The two
        // superfiles whose bloom lacks "bravo" may be pruned, but the
        // result must still be exactly the one matching row — proving
        // the bloom prune never drops a match.
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let mut w = st.writer().expect("writer");
        w.append(&build_cat_batch(0, &["x"], &["alpha"]))
            .expect("a1");
        w.commit().expect("c1");
        w.append(&build_cat_batch(10, &["y"], &["bravo"]))
            .expect("a2");
        w.commit().expect("c2");
        w.append(&build_cat_batch(20, &["z"], &["charlie"]))
            .expect("a3");
        w.commit().expect("c3");
        assert_eq!(st.reader().expect("reader").n_superfiles(), 3);

        assert_eq!(
            run_count(&st, "SELECT COUNT(*) FROM supertable WHERE title = 'bravo'"),
            1
        );
        assert_eq!(
            run_count(
                &st,
                "SELECT COUNT(*) FROM supertable WHERE title = 'nonexistent'"
            ),
            0
        );
    }

    #[test]
    fn query_sql_multiword_equality_on_fts_column_is_correct() {
        // Multi-word literal: the equality lowers to a `TermPresence`
        // leaf over {rust, async, runtime} (AND). The second superfile's
        // bloom lacks those tokens and is pruned, yet results are exact
        // — DataFusion's FilterExec re-applies the full string equality.
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let mut w = st.writer().expect("writer");
        w.append(&build_cat_batch(0, &["lang"], &["rust async runtime"]))
            .expect("a1");
        w.commit().expect("c1");
        w.append(&build_cat_batch(10, &["lang"], &["python data science"]))
            .expect("a2");
        w.commit().expect("c2");
        assert_eq!(st.reader().expect("reader").n_superfiles(), 2);

        assert_eq!(
            run_count(
                &st,
                "SELECT COUNT(*) FROM supertable WHERE title = 'rust async runtime'"
            ),
            1
        );
        // Tokens present in superfile 1, but no row equals this exact
        // string — the prune is an optimization, correctness holds.
        assert_eq!(
            run_count(
                &st,
                "SELECT COUNT(*) FROM supertable WHERE title = 'rust async'"
            ),
            0
        );
    }

    #[test]
    fn query_sql_fts_equality_superset_is_narrowed_to_exact_match() {
        // Index-driven row selection: the candidate plan resolves
        // `WHERE title = 'rust async'` to the term-AND posting set, which
        // within one superfile is a *superset* — both rows below contain
        // {rust, async}. The FilterExec above the scan must narrow that
        // candidate superset to the single exact-equality row, proving
        // the row-level prune never over-returns.
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let mut w = st.writer().expect("writer");
        w.append(&build_cat_batch(
            0,
            &["x", "y"],
            &["rust async", "rust async runtime"],
        ))
        .expect("append");
        w.commit().expect("commit");

        assert_eq!(
            run_count(
                &st,
                "SELECT COUNT(*) FROM supertable WHERE title = 'rust async'",
            ),
            1,
        );
        let batches = st
            .reader()
            .expect("reader")
            .query_sql("SELECT title FROM supertable WHERE title = 'rust async'")
            .expect("query");
        let total: usize = batches.iter().map(|b| b.num_rows()).sum();
        assert_eq!(total, 1);
    }

    #[test]
    fn query_sql_fts_or_and_in_are_exact() {
        // OR of two FTS equalities, AND with a non-FTS conjunct, and IN —
        // all index-bounded except where a branch is un-boundable, and
        // all verified exact by FilterExec.
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let mut w = st.writer().expect("writer");
        w.append(&build_cat_batch(
            0,
            &["rust", "python", "rust", "go"],
            &["alpha", "beta", "gamma", "delta"],
        ))
        .expect("append");
        w.commit().expect("commit");

        // OR of two FTS equalities → union, exact.
        assert_eq!(
            run_count(
                &st,
                "SELECT COUNT(*) FROM supertable WHERE title = 'alpha' OR title = 'beta'",
            ),
            2,
        );
        // AND with a non-FTS conjunct: FTS branch bounds candidates, the
        // category check is verified in pass 2.
        assert_eq!(
            run_count(
                &st,
                "SELECT COUNT(*) FROM supertable \
                 WHERE title = 'alpha' AND category = 'rust'",
            ),
            1,
        );
        assert_eq!(
            run_count(
                &st,
                "SELECT COUNT(*) FROM supertable \
                 WHERE title = 'alpha' AND category = 'python'",
            ),
            0,
        );
        // IN on the FTS column → OR of equalities.
        assert_eq!(
            run_count(
                &st,
                "SELECT COUNT(*) FROM supertable WHERE title IN ('alpha', 'delta', 'zzz')",
            ),
            2,
        );
    }

    #[test]
    fn query_sql_not_predicates_are_exact() {
        // NOT / != aren't index-prefiltered (Unbounded → scan), but must
        // still be exact; and `= AND !=` prefilters on the `=` branch
        // while FilterExec applies the negation.
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let mut w = st.writer().expect("writer");
        w.append(&build_cat_batch(
            0,
            &["rust", "python", "rust", "go"],
            &["alpha", "beta", "alpha", "delta"],
        ))
        .expect("append");
        w.commit().expect("commit");

        // Standalone NOT (scan fallback): 4 rows, 2 are 'alpha' → 2 left.
        assert_eq!(
            run_count(
                &st,
                "SELECT COUNT(*) FROM supertable WHERE NOT (title = 'alpha')",
            ),
            2,
        );
        // `!=` (NotEq) likewise.
        assert_eq!(
            run_count(
                &st,
                "SELECT COUNT(*) FROM supertable WHERE title != 'alpha'"
            ),
            2,
        );
        // `= AND !=`: candidates from the `title='alpha'` branch (2 rows),
        // then FilterExec drops category='rust' → 1 remains.
        assert_eq!(
            run_count(
                &st,
                "SELECT COUNT(*) FROM supertable \
                 WHERE title = 'alpha' AND category != 'rust'",
            ),
            0,
        );
        assert_eq!(
            run_count(
                &st,
                "SELECT COUNT(*) FROM supertable \
                 WHERE title = 'alpha' AND category != 'python'",
            ),
            2,
        );
    }

    #[test]
    fn query_sql_or_with_non_fts_branch_matches_full_scan() {
        // `title = 'alpha' OR category = 'go'` is un-boundable (the
        // category branch could match any row), so the planner falls back
        // to a full scan + FilterExec — and must still be exact.
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let mut w = st.writer().expect("writer");
        w.append(&build_cat_batch(
            0,
            &["rust", "python", "go", "go"],
            &["alpha", "beta", "gamma", "delta"],
        ))
        .expect("append");
        w.commit().expect("commit");

        // alpha (1 row) ∪ category=go (2 rows), disjoint → 3.
        assert_eq!(
            run_count(
                &st,
                "SELECT COUNT(*) FROM supertable WHERE title = 'alpha' OR category = 'go'",
            ),
            3,
        );
    }

    #[test]
    fn query_sql_select_orders_ids_across_superfiles() {
        // Verifies row identity round-trips through MemTable +
        // DataFusion: rows planted across two superfiles come back
        // in monotonic _id order under ORDER BY. The _id values
        // are auto-injected by the supertable (timestamp +
        // worker + counter), so we don't assert specific
        // values — only strict-increasing order.
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let mut w = st.writer().expect("writer");
        w.append(&build_cat_batch(100, &["a", "b"], &["t1", "t2"]))
            .expect("a1");
        w.commit().expect("c1");
        w.append(&build_cat_batch(200, &["c"], &["t3"]))
            .expect("a2");
        w.commit().expect("c2");

        let batches = st
            .reader()
            .expect("reader")
            .query_sql("SELECT _id FROM supertable ORDER BY _id")
            .expect("query");
        let ids: Vec<i128> = batches
            .iter()
            .flat_map(|b| {
                let a = b
                    .column(0)
                    .as_any()
                    .downcast_ref::<Decimal128Array>()
                    .expect("_id is Decimal128");
                (0..a.len()).map(|i| a.value(i)).collect::<Vec<_>>()
            })
            .collect();
        assert_eq!(ids.len(), 3);
        for w in ids.windows(2) {
            assert!(w[0] < w[1], "expected strictly increasing _id");
        }
    }

    #[test]
    fn query_sql_select_star_exposes_only_user_columns_plus_id() {
        // The supertable is a thin SQL skin over scalar columns —
        // `inf.*` KV metadata stays invisible. The injected `_id`
        // column is part of the visible schema.
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let mut w = st.writer().expect("writer");
        w.append(&build_cat_batch(0, &["x"], &["t"])).expect("a");
        w.commit().expect("c");

        let batches = st
            .reader()
            .expect("reader")
            .query_sql("SELECT * FROM supertable LIMIT 1")
            .expect("query");
        let schema = batches[0].schema();
        let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
        assert_eq!(names, vec!["_id", "category", "title"]);
    }

    #[test]
    fn query_sql_runtime_is_cached_across_calls() {
        // Two queries on the same supertable must share one
        // Runtime — the OnceLock guarantees this; we assert by
        // checking that both calls succeed without spawning a
        // fresh Runtime per call (observed indirectly via the
        // `.await` over `block_on` not double-allocating; if the
        // cache regressed, tests would still pass but would leak
        // a Runtime per call. The functional check below is
        // adequate for correctness; benchmarks would catch leak).
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let mut w = st.writer().expect("writer");
        w.append(&build_cat_batch(0, &["x"], &["t"])).expect("a");
        w.commit().expect("c");
        for _ in 0..3 {
            let n = run_count(&st, "SELECT COUNT(*) FROM supertable");
            assert_eq!(n, 1);
        }
    }

    #[test]
    fn query_sql_invalid_sql_returns_plan_error() {
        let st = Supertable::create(options_id_cat_title()).expect("create");
        let err = st
            .reader()
            .expect("reader")
            .query_sql("SELECT NOT_A_REAL_FN(*) FROM supertable")
            .expect_err("expected a plan error");
        assert!(
            matches!(err, QueryError::Plan(_)),
            "expected Plan variant; got {err:?}"
        );
    }

    // ---- vector schema integration ----------------------------------

    /// Build a schema that includes a vector column. The supertable
    /// strips it at commit time; SQL surface only sees the scalar
    /// columns. `query_sql` SELECTing the vector column must error
    /// (DataFusion's planner rejects unknown column).
    fn schema_with_vector(dim: usize) -> Arc<Schema> {
        Arc::new(Schema::new(vec![
            Field::new("title", DataType::LargeUtf8, false),
            Field::new(
                "emb",
                DataType::FixedSizeList(
                    Arc::new(Field::new("item", DataType::Float32, true)),
                    dim as i32,
                ),
                false,
            ),
        ]))
    }

    fn options_with_vector(dim: usize) -> SupertableOptions {
        let pool = Arc::new(
            rayon::ThreadPoolBuilder::new()
                .num_threads(1)
                .build()
                .expect("rayon pool"),
        );
        SupertableOptions::new(
            schema_with_vector(dim),
            vec![FtsConfig {
                column: "title".into(),
                positions: false,
            }],
            vec![VectorConfig {
                column: "emb".into(),
                dim,
                rot_seed: 0,
                metric: Metric::Cosine,
                rerank_codec: RerankCodec::Fp32,
                provided_centroids: None,
            }],
            Some(tok()),
        )
        .expect("valid options")
        .with_writer_pool(pool)
    }

    fn build_vector_batch(_start: u64, n: usize, dim: usize) -> RecordBatch {
        let titles = LargeStringArray::from((0..n).map(|i| format!("doc {i}")).collect::<Vec<_>>());
        let mut flat = Vec::<f32>::with_capacity(n * dim);
        for i in 0..n {
            for d in 0..dim {
                flat.push(((i + d) as f32) / 100.0);
            }
        }
        let item_field = Arc::new(Field::new("item", DataType::Float32, true));
        let values = Float32Array::from(flat);
        let emb = FixedSizeListArray::try_new(
            item_field,
            dim as i32,
            Arc::new(values) as Arc<dyn Array>,
            None,
        )
        .expect("FixedSizeList build");
        RecordBatch::try_new(
            schema_with_vector(dim),
            vec![Arc::new(titles), Arc::new(emb)],
        )
        .expect("build batch")
    }

    #[test]
    fn query_sql_hides_vector_columns_from_sql_surface() {
        let st = Supertable::create(options_with_vector(16)).expect("create");
        let mut w = st.writer().expect("writer");
        // n=8 ≥ n_cent=4 so kmeans has data to cluster.
        w.append(&build_vector_batch(0, 8, 16)).expect("append");
        w.commit().expect("commit");

        let batches = st
            .reader()
            .expect("reader")
            .query_sql("SELECT * FROM supertable LIMIT 1")
            .expect("query");
        let schema = batches[0].schema();
        let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
        // `emb` was stripped by `vector_split` at commit time and
        // lives in the embedded vector blob — not visible to SQL.
        // The supertable-injected `_id` is visible.
        assert_eq!(names, vec!["_id", "title"]);
    }

    #[test]
    fn query_sql_referencing_vector_column_returns_plan_error() {
        let st = Supertable::create(options_with_vector(16)).expect("create");
        let mut w = st.writer().expect("writer");
        w.append(&build_vector_batch(0, 8, 16)).expect("append");
        w.commit().expect("commit");

        let err = st
            .reader()
            .expect("reader")
            .query_sql("SELECT emb FROM supertable")
            .expect_err("vector column should not be in the SQL schema");
        assert!(
            matches!(err, QueryError::Plan(_)),
            "expected Plan variant; got {err:?}"
        );
    }
}