rhei_oltp_rusqlite/

engine.rs

//! [`RusqliteEngine`]: the async SQLite OLTP engine.
//!
//! This module provides the central OLTP component of `rhei-oltp-rusqlite`.
//! It wraps `rusqlite` via [`rhei_tokio_rusqlite`] to expose a fully async
//! interface without any `unsafe` code.
//!
//! ## Connection topology
//!
//! | Connection | Count | Purpose |
//! |------------|-------|---------|
//! | Write | 1 | All `INSERT` / `UPDATE` / `DELETE` / DDL |
//! | Read pool | N (default 4) | Concurrent `SELECT` via round-robin |
//! | CDC | 1 (separate) | `_rhei_cdc_log` polling by [`RusqliteCdcProducer`][crate::RusqliteCdcProducer] |
//!
//! The write connection is opened with WAL mode, `synchronous=NORMAL`, and
//! `busy_timeout=5000`.  Read pool connections only set `busy_timeout`.
//!
//! ## Arrow output
//!
//! [`rhei_core::OltpEngine::query`] returns `Vec<RecordBatch>`.  Because
//! SQLite is dynamically typed, the Arrow schema is inferred from actual
//! values rather than declared column types.  A warm-up window of
//! `SCHEMA_WARMUP_ROWS` rows is buffered before the schema is committed, so
//! sparse columns (NULL for many rows, then non-null) are typed correctly.

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

use arrow::array::{ArrayRef, BinaryBuilder, Float64Builder, Int64Builder, StringBuilder};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use tracing::debug;

use crate::error::RusqliteOltpError;

/// Maximum number of rows per Arrow RecordBatch produced by [`query_to_arrow`].
const ARROW_BATCH_ROWS: usize = 8192;

/// Number of rows buffered for type-inference warm-up before the Arrow schema
/// is committed.
///
/// ### Why this exists
///
/// SQLite columns have no enforced declared type — every cell is dynamically
/// typed. `query_to_arrow` must infer the Arrow `DataType` for each column by
/// inspecting actual values. If we commit to a schema after only the first
/// chunk (8 192 rows), *sparse* columns that happen to be NULL in those rows
/// are permanently mis-typed as `Utf8`, even if they contain `Int64` values
/// further into the result set.
///
/// ### The fix: warm-up window
///
/// We accumulate up to `SCHEMA_WARMUP_ROWS` rows before choosing a schema.
/// Each column's type is resolved as soon as its first non-null value appears.
/// Any column still fully NULL after the warm-up window defaults to `Utf8`
/// (safe because there is no later non-null evidence to contradict it).
/// Once the schema is locked we flush the buffered rows in `ARROW_BATCH_ROWS`
/// chunks and continue streaming the remainder of the cursor.
///
/// ### Trade-off
///
/// For a typical row of ~100 bytes the warm-up buffer peaks at ~3 MB — tiny
/// relative to a 10 M-row dataset. Queries that return fewer rows than
/// `SCHEMA_WARMUP_ROWS` incur no extra overhead (schema is locked at EOF).
const SCHEMA_WARMUP_ROWS: usize = 32_768;

/// Rusqlite-backed OLTP engine.
///
/// Uses a single dedicated write connection for DML/DDL and a round-robin pool
/// of read connections for queries. WAL mode is enabled on startup so readers
/// don't block the writer.
pub struct RusqliteEngine {
    /// The path used to open the database (stored for creating new connections).
    db_path: String,
    write_conn: rhei_tokio_rusqlite::Connection,
    read_pool: Vec<rhei_tokio_rusqlite::Connection>,
    read_idx: AtomicUsize,
}

impl RusqliteEngine {
    /// Create a new engine backed by a local SQLite database file.
    ///
    /// `read_pool_size` controls how many concurrent read connections are
    /// maintained. Values below 1 are clamped to 1.
    pub async fn new_local(path: &str, read_pool_size: usize) -> Result<Self, RusqliteOltpError> {
        let write_conn = rhei_tokio_rusqlite::Connection::open(path).await?;

        // Enable WAL mode so read connections don't block the writer.
        write_conn
            .call(|conn| {
                conn.execute_batch(
                    "PRAGMA journal_mode=WAL;
                     PRAGMA synchronous=NORMAL;
                     PRAGMA busy_timeout=5000;",
                )?;
                Ok(())
            })
            .await?;

        let pool_size = read_pool_size.max(1);
        let mut read_pool = Vec::with_capacity(pool_size);
        for _ in 0..pool_size {
            let read_conn = rhei_tokio_rusqlite::Connection::open(path).await?;
            read_conn
                .call(|conn| {
                    conn.execute_batch("PRAGMA busy_timeout=5000;")?;
                    Ok(())
                })
                .await?;
            read_pool.push(read_conn);
        }

        Ok(Self {
            db_path: path.to_string(),
            write_conn,
            read_pool,
            read_idx: AtomicUsize::new(0),
        })
    }

    /// Create a fresh, independent connection to the same database file.
    ///
    /// Used to give `RusqliteCdcProducer` its own connection that doesn't compete
    /// with application reads/writes.
    pub async fn new_connection(
        &self,
    ) -> Result<rhei_tokio_rusqlite::Connection, RusqliteOltpError> {
        let conn = rhei_tokio_rusqlite::Connection::open(&self.db_path).await?;
        conn.call(|c| {
            c.execute_batch("PRAGMA busy_timeout=5000;")?;
            Ok(())
        })
        .await?;
        Ok(conn)
    }

    /// Returns a clone of the write connection handle. Used by CDC trigger
    /// setup/teardown (DDL).
    pub fn connection(&self) -> rhei_tokio_rusqlite::Connection {
        self.write_conn.clone()
    }

    fn next_read_conn(&self) -> &rhei_tokio_rusqlite::Connection {
        let idx = self.read_idx.fetch_add(1, Ordering::Relaxed) % self.read_pool.len();
        &self.read_pool[idx]
    }
}

/// Convert a serde_json::Value to a rusqlite Value suitable for binding.
fn json_to_rusqlite(val: &serde_json::Value) -> rusqlite::types::Value {
    match val {
        serde_json::Value::Null => rusqlite::types::Value::Null,
        serde_json::Value::Bool(b) => rusqlite::types::Value::Integer(if *b { 1 } else { 0 }),
        serde_json::Value::Number(n) => {
            if let Some(i) = n.as_i64() {
                rusqlite::types::Value::Integer(i)
            } else if let Some(f) = n.as_f64() {
                rusqlite::types::Value::Real(f)
            } else {
                rusqlite::types::Value::Text(n.to_string())
            }
        }
        serde_json::Value::String(s) => rusqlite::types::Value::Text(s.clone()),
        other => rusqlite::types::Value::Text(other.to_string()),
    }
}

/// Build one Arrow [`RecordBatch`] from a row-major chunk of rusqlite values.
fn build_batch(
    chunk: &[Vec<rusqlite::types::Value>],
    schema: &Arc<Schema>,
) -> Result<RecordBatch, rhei_tokio_rusqlite::Error> {
    let col_count = schema.fields().len();
    let mut columns: Vec<ArrayRef> = Vec::with_capacity(col_count);

    for col_idx in 0..col_count {
        let dt = schema.field(col_idx).data_type();
        let array: ArrayRef = match dt {
            DataType::Int64 => {
                let mut b = Int64Builder::with_capacity(chunk.len());
                for row in chunk {
                    match &row[col_idx] {
                        rusqlite::types::Value::Integer(i) => b.append_value(*i),
                        _ => b.append_null(),
                    }
                }
                Arc::new(b.finish())
            }
            DataType::Float64 => {
                let mut b = Float64Builder::with_capacity(chunk.len());
                for row in chunk {
                    match &row[col_idx] {
                        rusqlite::types::Value::Real(f) => b.append_value(*f),
                        rusqlite::types::Value::Integer(i) => b.append_value(*i as f64),
                        _ => b.append_null(),
                    }
                }
                Arc::new(b.finish())
            }
            DataType::Binary => {
                let mut b = BinaryBuilder::with_capacity(chunk.len(), chunk.len() * 16);
                for row in chunk {
                    match &row[col_idx] {
                        rusqlite::types::Value::Blob(bytes) => b.append_value(bytes.as_slice()),
                        _ => b.append_null(),
                    }
                }
                Arc::new(b.finish())
            }
            _ => {
                // Utf8 and all-null fallback
                let mut b = StringBuilder::with_capacity(chunk.len(), chunk.len() * 8);
                for row in chunk {
                    match &row[col_idx] {
                        rusqlite::types::Value::Text(s) => b.append_value(s.as_str()),
                        rusqlite::types::Value::Integer(i) => b.append_value(i.to_string()),
                        rusqlite::types::Value::Real(f) => b.append_value(f.to_string()),
                        rusqlite::types::Value::Null => b.append_null(),
                        rusqlite::types::Value::Blob(_) => b.append_null(),
                    }
                }
                Arc::new(b.finish())
            }
        };
        columns.push(array);
    }

    RecordBatch::try_new(Arc::clone(schema), columns)
        .map_err(|e| rhei_tokio_rusqlite::Error::Other(format!("Arrow error: {e}")))
}

/// Execute a query inside a `conn.call()` closure and convert results to Arrow
/// [`RecordBatch`]es.
///
/// Rows are accumulated into a row buffer and flushed into Arrow
/// [`RecordBatch`]es every [`ARROW_BATCH_ROWS`] rows so that large result sets
/// never materialise as a single in-memory allocation.
///
/// ### Schema inference
///
/// SQLite is dynamically typed, so the Arrow schema cannot be derived from the
/// column declarations alone — it must be inferred from the actual values.
/// To handle *sparse* columns (NULL for many consecutive rows, non-null later),
/// we use a **warm-up window**: the first [`SCHEMA_WARMUP_ROWS`] rows are
/// buffered. Each column's Arrow type is resolved as soon as its first non-null
/// value appears. Any column still fully NULL when the window is exhausted or
/// the cursor is drained defaults to `Utf8`. The schema is then locked and the
/// buffered rows are flushed in `ARROW_BATCH_ROWS`-sized chunks; subsequent
/// rows stream through normally. See [`SCHEMA_WARMUP_ROWS`] for the rationale.
fn query_to_arrow(
    conn: &mut rusqlite::Connection,
    sql: &str,
    params: &[rusqlite::types::Value],
) -> Result<Vec<RecordBatch>, rhei_tokio_rusqlite::Error> {
    let mut stmt = conn.prepare(sql)?;

    let col_count = stmt.column_count();
    if col_count == 0 {
        // Statement produces no columns (e.g., a write statement called via query).
        return Ok(vec![]);
    }

    let col_names: Vec<String> = stmt
        .column_names()
        .into_iter()
        .map(|s| s.to_string())
        .collect();

    let params_refs: Vec<&dyn rusqlite::types::ToSql> = params
        .iter()
        .map(|v| v as &dyn rusqlite::types::ToSql)
        .collect();

    let mut rows = stmt.query(params_refs.as_slice())?;

    // --- Phase 1: warm-up buffer -------------------------------------------------
    // Accumulate up to SCHEMA_WARMUP_ROWS rows while tracking per-column type hints.
    // `None` means "still fully NULL for this column".
    let mut hints: Vec<Option<DataType>> = vec![None; col_count];
    let mut warmup: Vec<Vec<rusqlite::types::Value>> = Vec::with_capacity(SCHEMA_WARMUP_ROWS);

    while warmup.len() < SCHEMA_WARMUP_ROWS {
        match rows.next()? {
            None => break,
            Some(row) => {
                let mut vals = Vec::with_capacity(col_count);
                for (col_idx, hint) in hints.iter_mut().enumerate() {
                    let val: rusqlite::types::Value = row.get(col_idx)?;
                    // Lock in the type as soon as we see a non-null value.
                    if hint.is_none() {
                        *hint = match &val {
                            rusqlite::types::Value::Integer(_) => Some(DataType::Int64),
                            rusqlite::types::Value::Real(_) => Some(DataType::Float64),
                            rusqlite::types::Value::Text(_) => Some(DataType::Utf8),
                            rusqlite::types::Value::Blob(_) => Some(DataType::Binary),
                            rusqlite::types::Value::Null => None,
                        };
                    }
                    vals.push(val);
                }
                warmup.push(vals);
            }
        }
    }

    // Commit the schema: any column still None (all-NULL in warmup) defaults to Utf8.
    let schema = Arc::new(Schema::new(
        hints
            .into_iter()
            .zip(col_names.iter())
            .map(|(hint, name)| {
                let dt = hint.unwrap_or(DataType::Utf8);
                Field::new(name, dt, true)
            })
            .collect::<Vec<_>>(),
    ));

    // --- Phase 2: flush warmup buffer + stream remainder ------------------------
    let mut batches: Vec<RecordBatch> = Vec::new();
    let mut chunk: Vec<Vec<rusqlite::types::Value>> = Vec::with_capacity(ARROW_BATCH_ROWS);

    // Helper: flush `chunk` into a batch and clear it.
    let flush_chunk = |chunk: &mut Vec<Vec<rusqlite::types::Value>>,
                       schema: &Arc<Schema>,
                       batches: &mut Vec<RecordBatch>|
     -> Result<(), rhei_tokio_rusqlite::Error> {
        if chunk.is_empty() {
            return Ok(());
        }
        batches.push(build_batch(chunk, schema)?);
        chunk.clear();
        Ok(())
    };

    // Flush the warm-up rows first.
    for row_vals in warmup {
        chunk.push(row_vals);
        if chunk.len() >= ARROW_BATCH_ROWS {
            flush_chunk(&mut chunk, &schema, &mut batches)?;
        }
    }

    // Stream the remainder of the cursor.
    while let Some(row) = rows.next()? {
        let mut vals = Vec::with_capacity(col_count);
        for i in 0..col_count {
            let val: rusqlite::types::Value = row.get(i)?;
            vals.push(val);
        }
        chunk.push(vals);

        if chunk.len() >= ARROW_BATCH_ROWS {
            flush_chunk(&mut chunk, &schema, &mut batches)?;
        }
    }

    // Flush any remaining rows.
    flush_chunk(&mut chunk, &schema, &mut batches)?;

    if batches.is_empty() {
        // Query returned zero rows: emit one empty batch so callers always get a schema.
        batches.push(RecordBatch::new_empty(Arc::clone(&schema)));
    }

    Ok(batches)
}

impl rhei_core::OltpEngine for RusqliteEngine {
    type Error = RusqliteOltpError;

    /// Execute a read-only SQL query and return the results as Arrow
    /// [`RecordBatch`]es.
    ///
    /// Dispatches to the next connection in the round-robin read pool.
    /// The Arrow schema is inferred from actual cell values (see the
    /// `SCHEMA_WARMUP_ROWS` constant for the sparse-column strategy).
    ///
    /// # Errors
    ///
    /// Returns [`RusqliteOltpError`] if statement preparation, execution, or
    /// Arrow conversion fails.
    async fn query(
        &self,
        sql: &str,
        params: &[serde_json::Value],
    ) -> Result<Vec<RecordBatch>, Self::Error> {
        debug!(sql, params_count = params.len(), "OLTP rusqlite query");
        let rusqlite_params: Vec<rusqlite::types::Value> =
            params.iter().map(json_to_rusqlite).collect();
        let sql_owned = sql.to_string();

        let conn = self.next_read_conn();
        let batches = conn
            .call(move |c| query_to_arrow(c, &sql_owned, &rusqlite_params))
            .await?;
        Ok(batches)
    }

    /// Execute a single write statement (`INSERT`, `UPDATE`, `DELETE`, or DDL)
    /// and return the number of rows affected.
    ///
    /// Always dispatches to the single dedicated write connection.
    ///
    /// # Errors
    ///
    /// Returns [`RusqliteOltpError`] if the statement fails (e.g., constraint
    /// violation, syntax error, or the database is locked beyond
    /// `busy_timeout`).
    async fn execute(&self, sql: &str, params: &[serde_json::Value]) -> Result<u64, Self::Error> {
        debug!(sql, params_count = params.len(), "OLTP rusqlite execute");
        let rusqlite_params: Vec<rusqlite::types::Value> =
            params.iter().map(json_to_rusqlite).collect();
        let sql_owned = sql.to_string();

        let rows_affected = self
            .write_conn
            .call(move |c| {
                let params_refs: Vec<&dyn rusqlite::types::ToSql> = rusqlite_params
                    .iter()
                    .map(|v| v as &dyn rusqlite::types::ToSql)
                    .collect();
                let changed = c.execute(&sql_owned, params_refs.as_slice())?;
                Ok(changed as u64)
            })
            .await?;
        Ok(rows_affected)
    }

    /// Execute a batch of write statements inside a single `BEGIN`…`COMMIT`
    /// transaction.
    ///
    /// All statements succeed or the entire transaction is rolled back.
    /// Dispatches to the dedicated write connection.
    ///
    /// # Errors
    ///
    /// Returns [`RusqliteOltpError`] if any statement fails or the transaction
    /// cannot be committed.
    async fn execute_batch(
        &self,
        statements: &[(String, Vec<serde_json::Value>)],
    ) -> Result<(), Self::Error> {
        debug!(count = statements.len(), "OLTP rusqlite execute_batch");
        // Clone the statements so we can move them into the closure
        let stmts: Vec<(String, Vec<rusqlite::types::Value>)> = statements
            .iter()
            .map(|(sql, params)| {
                let rusqlite_params: Vec<rusqlite::types::Value> =
                    params.iter().map(json_to_rusqlite).collect();
                (sql.clone(), rusqlite_params)
            })
            .collect();

        self.write_conn
            .call(move |c| {
                let tx = c.transaction()?;
                for (sql, params) in &stmts {
                    let params_refs: Vec<&dyn rusqlite::types::ToSql> = params
                        .iter()
                        .map(|v| v as &dyn rusqlite::types::ToSql)
                        .collect();
                    tx.execute(sql, params_refs.as_slice())?;
                }
                tx.commit()?;
                Ok(())
            })
            .await?;
        Ok(())
    }

    /// Return `true` if a table named `table_name` exists in `sqlite_master`.
    ///
    /// Uses a read-pool connection.
    ///
    /// # Errors
    ///
    /// Returns [`RusqliteOltpError`] if the metadata query fails.
    async fn table_exists(&self, table_name: &str) -> Result<bool, Self::Error> {
        let tbl = table_name.to_string();
        let conn = self.next_read_conn();
        let exists = conn
            .call(move |c| {
                let count: i64 = c.query_row(
                    "SELECT count(*) FROM sqlite_master WHERE type='table' AND name=?1",
                    rusqlite::params![tbl],
                    |row| row.get(0),
                )?;
                Ok(count > 0)
            })
            .await?;
        Ok(exists)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow::array::{Array, Int64Array};
    use rhei_core::OltpEngine;

    #[tokio::test]
    async fn test_basic_crud() {
        let dir = tempfile::TempDir::new().unwrap();
        let path = dir.path().join("test.db");
        let engine = RusqliteEngine::new_local(path.to_str().unwrap(), 2)
            .await
            .unwrap();

        // Create table
        engine
            .execute(
                "CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT, age INTEGER)",
                &[],
            )
            .await
            .unwrap();

        // Insert
        engine
            .execute(
                "INSERT INTO users (id, name, age) VALUES (?1, ?2, ?3)",
                &[
                    serde_json::json!(1),
                    serde_json::json!("Alice"),
                    serde_json::json!(30),
                ],
            )
            .await
            .unwrap();

        // Query
        let batches = engine.query("SELECT * FROM users", &[]).await.unwrap();
        assert_eq!(batches.len(), 1);
        assert_eq!(batches[0].num_rows(), 1);
        assert_eq!(batches[0].num_columns(), 3);

        // table_exists
        assert!(engine.table_exists("users").await.unwrap());
        assert!(!engine.table_exists("nonexistent").await.unwrap());
    }

    #[tokio::test]
    async fn test_execute_batch() {
        let dir = tempfile::TempDir::new().unwrap();
        let path = dir.path().join("test_batch.db");
        let engine = RusqliteEngine::new_local(path.to_str().unwrap(), 1)
            .await
            .unwrap();

        engine
            .execute("CREATE TABLE items (id INTEGER PRIMARY KEY, val TEXT)", &[])
            .await
            .unwrap();

        let stmts: Vec<(String, Vec<serde_json::Value>)> = vec![
            (
                "INSERT INTO items (id, val) VALUES (?1, ?2)".to_string(),
                vec![serde_json::json!(1), serde_json::json!("a")],
            ),
            (
                "INSERT INTO items (id, val) VALUES (?1, ?2)".to_string(),
                vec![serde_json::json!(2), serde_json::json!("b")],
            ),
            (
                "INSERT INTO items (id, val) VALUES (?1, ?2)".to_string(),
                vec![serde_json::json!(3), serde_json::json!("c")],
            ),
        ];

        engine.execute_batch(&stmts).await.unwrap();

        let batches = engine
            .query("SELECT COUNT(*) FROM items", &[])
            .await
            .unwrap();
        let count = batches[0]
            .column(0)
            .as_any()
            .downcast_ref::<Int64Array>()
            .unwrap()
            .value(0);
        assert_eq!(count, 3);
    }

    #[tokio::test]
    async fn test_chunked_query() {
        let dir = tempfile::TempDir::new().unwrap();
        let path = dir.path().join("test_chunks.db");
        let engine = RusqliteEngine::new_local(path.to_str().unwrap(), 1)
            .await
            .unwrap();

        engine
            .execute("CREATE TABLE big (id INTEGER PRIMARY KEY, val TEXT)", &[])
            .await
            .unwrap();

        // Insert 20 000 rows via a single execute_batch to keep the test fast.
        let stmts: Vec<(String, Vec<serde_json::Value>)> = (0..20_000u64)
            .map(|i| {
                (
                    "INSERT INTO big (id, val) VALUES (?1, ?2)".to_string(),
                    vec![serde_json::json!(i), serde_json::json!(format!("v{i}"))],
                )
            })
            .collect();
        engine.execute_batch(&stmts).await.unwrap();

        let batches = engine.query("SELECT * FROM big", &[]).await.unwrap();

        // With ARROW_BATCH_ROWS = 8192 and 20 000 rows we expect 3 batches
        // (8192 + 8192 + 3616).
        assert!(
            batches.len() > 1,
            "expected multiple RecordBatches, got {}",
            batches.len()
        );

        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
        assert_eq!(total_rows, 20_000, "total row count mismatch");

        // Verify all batches use the same schema.
        let schema = batches[0].schema();
        for (i, batch) in batches.iter().enumerate() {
            assert_eq!(batch.schema(), schema, "schema mismatch on batch {i}");
        }
    }

    #[tokio::test]
    async fn test_new_connection() {
        let dir = tempfile::TempDir::new().unwrap();
        let path = dir.path().join("test_conn.db");
        let engine = RusqliteEngine::new_local(path.to_str().unwrap(), 1)
            .await
            .unwrap();

        engine
            .execute("CREATE TABLE t (id INTEGER PRIMARY KEY)", &[])
            .await
            .unwrap();

        // new_connection should be able to see the table
        let conn = engine.new_connection().await.unwrap();
        let exists: bool = conn
            .call(|c| {
                let count: i64 = c.query_row(
                    "SELECT count(*) FROM sqlite_master WHERE type='table' AND name='t'",
                    [],
                    |row| row.get(0),
                )?;
                Ok(count > 0)
            })
            .await
            .unwrap();
        assert!(exists);
    }

    /// Regression test for sparse-column type inference.
    ///
    /// Inserts 10 000 rows with `score IS NULL` followed by rows with integer
    /// values for `score`. Before the warm-up fix, the schema was locked after
    /// the first 8 192 rows, all of which had a NULL `score` — permanently
    /// mis-typing the column as `Utf8`. The correct type is `Int64`.
    #[tokio::test]
    async fn test_sparse_column_type_inference() {
        use arrow::array::Int64Array;
        use arrow::datatypes::DataType;

        let dir = tempfile::TempDir::new().unwrap();
        let path = dir.path().join("test_sparse.db");
        let engine = RusqliteEngine::new_local(path.to_str().unwrap(), 1)
            .await
            .unwrap();

        engine
            .execute(
                "CREATE TABLE sparse (id INTEGER PRIMARY KEY, score INTEGER)",
                &[],
            )
            .await
            .unwrap();

        // Insert 10 000 NULL-score rows, then 100 rows with integer scores.
        // With the old code the first ARROW_BATCH_ROWS (8 192) rows were all-NULL
        // for `score`, locking it as Utf8. We need at least SCHEMA_WARMUP_ROWS
        // (32 768) non-null rows to exercise the boundary, but 10 100 rows is
        // sufficient to demonstrate the old bug (NULL for the first two 8 192-row
        // chunks, then non-null) and confirm the fix works.
        const NULL_ROWS: usize = 10_000;
        const INT_ROWS: usize = 100;

        let null_stmts: Vec<(String, Vec<serde_json::Value>)> = (0..NULL_ROWS)
            .map(|i| {
                (
                    "INSERT INTO sparse (id, score) VALUES (?1, NULL)".to_string(),
                    vec![serde_json::json!(i)],
                )
            })
            .collect();
        engine.execute_batch(&null_stmts).await.unwrap();

        let int_stmts: Vec<(String, Vec<serde_json::Value>)> = (NULL_ROWS..NULL_ROWS + INT_ROWS)
            .map(|i| {
                (
                    "INSERT INTO sparse (id, score) VALUES (?1, ?2)".to_string(),
                    vec![serde_json::json!(i), serde_json::json!(i as i64)],
                )
            })
            .collect();
        engine.execute_batch(&int_stmts).await.unwrap();

        let batches = engine.query("SELECT * FROM sparse", &[]).await.unwrap();
        assert!(!batches.is_empty(), "expected at least one batch");

        // All batches must share the same schema.
        let schema = batches[0].schema();
        for (idx, batch) in batches.iter().enumerate() {
            assert_eq!(batch.schema(), schema, "schema mismatch on batch {idx}");
        }

        // `score` column must be Int64, not Utf8.
        let score_field = schema
            .field_with_name("score")
            .expect("score field missing");
        assert_eq!(
            score_field.data_type(),
            &DataType::Int64,
            "sparse column 'score' should be Int64; \
             was it mis-typed as Utf8 due to early schema lock?"
        );

        // Spot-check: the non-null integer values should be readable as Int64.
        let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
        assert_eq!(total_rows, NULL_ROWS + INT_ROWS);

        // Find a batch that contains the integer rows and verify a value.
        let mut found_int = false;
        for batch in &batches {
            let score_col = batch.column_by_name("score").expect("score column missing");
            if let Some(arr) = score_col.as_any().downcast_ref::<Int64Array>() {
                for row in 0..arr.len() {
                    if arr.is_valid(row) {
                        // The integer value must be >= NULL_ROWS (our insert offset).
                        assert!(
                            arr.value(row) >= NULL_ROWS as i64,
                            "unexpected integer value {}",
                            arr.value(row)
                        );
                        found_int = true;
                    }
                }
            }
        }
        assert!(
            found_int,
            "no non-null Int64 values found in 'score' column"
        );
    }
}