paimon-datafusion 0.3.0

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use std::collections::HashMap;
use std::sync::Arc;

mod common;

use common::string_value;
use datafusion::arrow::array::{Array, Int32Array};
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::arrow::util::display::array_value_to_string;
use datafusion::catalog::CatalogProvider;
use datafusion::datasource::TableProvider;
use datafusion::logical_expr::{col, lit, TableProviderFilterPushDown};
use datafusion::physical_plan::{displayable, ExecutionPlan};
use datafusion::prelude::SessionConfig;
use paimon::catalog::Identifier;
use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options};
use paimon_datafusion::{PaimonCatalogProvider, PaimonTableProvider, SQLContext};

fn get_test_warehouse() -> String {
    std::env::var("PAIMON_TEST_WAREHOUSE").unwrap_or_else(|_| "/tmp/paimon-warehouse".to_string())
}

fn create_catalog() -> FileSystemCatalog {
    let warehouse = get_test_warehouse();
    let mut options = Options::new();
    options.set(CatalogOptions::WAREHOUSE, warehouse);
    FileSystemCatalog::new(options).expect("Failed to create catalog")
}

async fn create_context() -> SQLContext {
    let catalog = create_catalog();
    let catalog: Arc<dyn Catalog> = Arc::new(catalog);
    let mut ctx = SQLContext::new();
    ctx.register_catalog("paimon", catalog)
        .await
        .expect("Failed to register catalog");
    ctx
}

async fn create_provider(table_name: &str) -> PaimonTableProvider {
    let catalog = create_catalog();
    let identifier = Identifier::new("default", table_name);
    let table = catalog
        .get_table(&identifier)
        .await
        .expect("Failed to get table");

    PaimonTableProvider::try_new(table).expect("Failed to create table provider")
}

async fn create_provider_with_options(
    table_name: &str,
    extra_options: HashMap<String, String>,
) -> PaimonTableProvider {
    let catalog = create_catalog();
    let identifier = Identifier::new("default", table_name);
    let table = catalog
        .get_table(&identifier)
        .await
        .expect("Failed to get table")
        .copy_with_options(extra_options);

    PaimonTableProvider::try_new(table).expect("Failed to create table provider")
}

async fn read_rows(table_name: &str) -> Vec<(i32, String)> {
    let sql = format!("SELECT id, name FROM paimon.default.{table_name}");
    let batches = collect_query(&sql)
        .await
        .expect("Failed to collect query result");

    assert!(
        !batches.is_empty(),
        "Expected at least one batch from table {table_name}"
    );

    let mut actual_rows = extract_id_name_rows(&batches);
    actual_rows.sort_by_key(|(id, _)| *id);
    actual_rows
}

async fn collect_query(sql: &str) -> datafusion::error::Result<Vec<RecordBatch>> {
    let ctx = create_context().await;
    ctx.sql(sql).await?.collect().await
}

fn collect_rows(batches: &[RecordBatch]) -> Vec<Vec<String>> {
    let mut rows = Vec::new();
    for batch in batches {
        for row_index in 0..batch.num_rows() {
            let mut row = Vec::with_capacity(batch.num_columns());
            for column in batch.columns() {
                if column.is_null(row_index) {
                    row.push("NULL".to_string());
                } else {
                    row.push(
                        array_value_to_string(column.as_ref(), row_index)
                            .expect("query result value should format"),
                    );
                }
            }
            rows.push(row);
        }
    }
    rows
}

async fn assert_sql_rows(sql: &str, expected: &[&[&str]]) {
    let batches = collect_query(sql).await.expect("query should succeed");
    let actual = collect_rows(&batches);
    let expected: Vec<Vec<String>> = expected
        .iter()
        .map(|row| row.iter().map(|value| value.to_string()).collect())
        .collect();
    assert_eq!(actual, expected, "unexpected result for SQL: {sql}");
}

async fn assert_time_travel_schema_evolution_rows(
    tag: &str,
    select_list: &str,
    where_clause: Option<&str>,
    expected: &[&[&str]],
) {
    let sql = match where_clause {
        Some(where_clause) => format!(
            "SELECT {select_list} FROM paimon.default.time_travel_schema_evolution VERSION AS OF '{tag}' WHERE {where_clause} ORDER BY id"
        ),
        None => format!(
            "SELECT {select_list} FROM paimon.default.time_travel_schema_evolution VERSION AS OF '{tag}' ORDER BY id"
        ),
    };
    assert_sql_rows(&sql, expected).await;
}

async fn create_physical_plan(sql: &str) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
    let ctx = create_context().await;
    ctx.sql(sql).await?.create_physical_plan().await
}

fn extract_id_name_rows(
    batches: &[datafusion::arrow::record_batch::RecordBatch],
) -> Vec<(i32, String)> {
    let mut rows = Vec::new();
    for batch in batches {
        let id_array = batch
            .column_by_name("id")
            .and_then(|column| column.as_any().downcast_ref::<Int32Array>())
            .expect("Expected Int32Array for id column");
        let name_array = batch
            .column_by_name("name")
            .expect("Expected string array for name column");

        for row_index in 0..batch.num_rows() {
            rows.push((
                id_array.value(row_index),
                string_value(name_array.as_ref(), row_index).to_string(),
            ));
        }
    }
    rows
}

fn format_physical_plan(plan: &Arc<dyn ExecutionPlan>) -> String {
    displayable(plan.as_ref()).indent(true).to_string()
}

fn paimon_scan_lines(plan_text: &str) -> Vec<&str> {
    plan_text
        .lines()
        .filter(|line| line.contains("PaimonTableScan:"))
        .collect()
}

#[tokio::test]
async fn test_read_log_table_via_datafusion() {
    let actual_rows = read_rows("simple_log_table").await;
    let expected_rows = vec![
        (1, "alice".to_string()),
        (2, "bob".to_string()),
        (3, "carol".to_string()),
    ];

    assert_eq!(
        actual_rows, expected_rows,
        "Rows should match expected values"
    );
}

#[tokio::test]
async fn test_read_primary_key_table_via_datafusion() {
    let actual_rows = read_rows("simple_dv_pk_table").await;
    let expected_rows = vec![
        (1, "alice-v2".to_string()),
        (2, "bob-v2".to_string()),
        (3, "carol-v2".to_string()),
        (4, "dave-v2".to_string()),
        (5, "eve-v2".to_string()),
        (6, "frank-v1".to_string()),
    ];

    assert_eq!(
        actual_rows, expected_rows,
        "Primary key table rows should match expected values"
    );
}

#[tokio::test]
async fn test_projection_via_datafusion() {
    let batches = collect_query("SELECT id FROM paimon.default.simple_log_table")
        .await
        .expect("Subset projection should succeed");

    assert!(
        !batches.is_empty(),
        "Expected at least one batch from projected query"
    );

    let mut actual_ids = Vec::new();
    for batch in &batches {
        let schema = batch.schema();
        let field_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
        assert_eq!(
            field_names,
            vec!["id"],
            "Projected query should only return 'id' column"
        );

        let id_array = batch
            .column_by_name("id")
            .and_then(|col| col.as_any().downcast_ref::<Int32Array>())
            .expect("Expected Int32Array for id column");
        for i in 0..id_array.len() {
            actual_ids.push(id_array.value(i));
        }
    }

    actual_ids.sort();
    assert_eq!(
        actual_ids,
        vec![1, 2, 3],
        "Projected id values should match"
    );
}

#[tokio::test]
async fn test_supports_partition_filters_pushdown() {
    let provider = create_provider("multi_partitioned_log_table").await;
    let partition_filter = col("dt").eq(lit("2024-01-01"));
    let mixed_and_filter = col("dt").eq(lit("2024-01-01")).and(col("id").gt(lit(1)));
    let data_filter = col("id").gt(lit(1));

    let supports = provider
        .supports_filters_pushdown(&[&partition_filter, &mixed_and_filter, &data_filter])
        .expect("supports_filters_pushdown should succeed");

    assert_eq!(
        supports,
        vec![
            TableProviderFilterPushDown::Exact,
            TableProviderFilterPushDown::Inexact,
            TableProviderFilterPushDown::Inexact,
        ]
    );
}

/// Verifies that `PaimonTableProvider::scan()` produces more than one
/// execution partition for a multi-partition table, and that the reported
/// partition count is still capped by `target_partitions`.
#[tokio::test]
async fn test_scan_partition_count_respects_session_config() {
    let provider = create_provider("partitioned_log_table").await;

    // With generous target_partitions, the plan should expose more than one partition.
    let config = SessionConfig::new().with_target_partitions(8);
    let ctx = datafusion::prelude::SessionContext::new_with_config(config);
    let state = ctx.state();
    let plan = provider
        .scan(&state, None, &[], None)
        .await
        .expect("scan() should succeed");

    let partition_count = plan.properties().output_partitioning().partition_count();
    assert!(
        partition_count > 1,
        "partitioned_log_table should produce >1 partitions, got {partition_count}"
    );

    // With target_partitions=1, all splits must be coalesced into a single partition
    let config_single = SessionConfig::new().with_target_partitions(1);
    let ctx_single = datafusion::prelude::SessionContext::new_with_config(config_single);
    let state_single = ctx_single.state();
    let plan_single = provider
        .scan(&state_single, None, &[], None)
        .await
        .expect("scan() should succeed with target_partitions=1");

    assert_eq!(
        plan_single
            .properties()
            .output_partitioning()
            .partition_count(),
        1,
        "target_partitions=1 should coalesce all splits into exactly 1 partition"
    );
}

#[tokio::test]
async fn test_partition_filter_query_via_datafusion() {
    let batches = collect_query(
        "SELECT id, name FROM paimon.default.partitioned_log_table WHERE dt = '2024-01-01'",
    )
    .await
    .expect("Partition filter query should succeed");

    let mut actual_rows = extract_id_name_rows(&batches);
    actual_rows.sort_by_key(|(id, _)| *id);
    assert_eq!(
        actual_rows,
        vec![(1, "alice".to_string()), (2, "bob".to_string())]
    );
}

#[tokio::test]
async fn test_multi_partition_filter_query_via_datafusion() {
    let batches = collect_query(
        "SELECT id, name FROM paimon.default.multi_partitioned_log_table WHERE dt = '2024-01-01' AND hr = 10",
    )
    .await
    .expect("Multi-partition filter query should succeed");

    let mut actual_rows = extract_id_name_rows(&batches);
    actual_rows.sort_by_key(|(id, _)| *id);
    assert_eq!(
        actual_rows,
        vec![(1, "alice".to_string()), (2, "bob".to_string())]
    );
}

#[tokio::test]
async fn test_mixed_and_filter_keeps_residual_datafusion_filter() {
    let batches = collect_query(
        "SELECT id, name FROM paimon.default.partitioned_log_table WHERE dt = '2024-01-01' AND id > 1",
    )
    .await
    .expect("Mixed filter query should succeed");

    let actual_rows = extract_id_name_rows(&batches);

    assert_eq!(actual_rows, vec![(2, "bob".to_string())]);
}

#[tokio::test]
async fn test_partially_translated_filter_keeps_partition_pruning_and_correctness() {
    let sql = "SELECT id, name FROM paimon.default.multi_partitioned_log_table WHERE dt = '2024-01-01' AND hr + 1 > 20 LIMIT 1";
    let plan = create_physical_plan(sql)
        .await
        .expect("Physical plan creation should succeed");
    let plan_text = format_physical_plan(&plan);
    let scan_lines = paimon_scan_lines(&plan_text);

    assert!(
        !scan_lines.is_empty(),
        "plan should contain a PaimonTableScan, plan:\n{plan_text}"
    );
    assert!(
        scan_lines
            .iter()
            .any(|line| line.contains("predicate=dt = '2024-01-01'")),
        "The translated partition predicate should still be pushed into PaimonTableScan, plan:\n{plan_text}"
    );
    assert!(
        scan_lines.iter().all(|line| !line.contains("fetch=")),
        "Partially translated filters should not revive the removed fetch contract, plan:\n{plan_text}"
    );

    let batches = collect_query(sql)
        .await
        .expect("Partially translated filter + LIMIT query should succeed");
    let rows = extract_id_name_rows(&batches);

    assert_eq!(
        rows,
        vec![(3, "carol".to_string())],
        "The residual filter should still be enforced above the scan"
    );
}

#[tokio::test]
async fn test_temporal_filter_pushdown_via_datafusion_scan() {
    let (_tmp, sql_context) = common::setup_sql_context().await;
    sql_context
        .sql(
            "CREATE TABLE paimon.test_db.temporal_filter_pushdown (
                id INT,
                name STRING,
                ts TIMESTAMP(6),
                lzts TIMESTAMP(6) WITH TIME ZONE
            )",
        )
        .await
        .expect("CREATE TABLE should succeed")
        .collect()
        .await
        .expect("CREATE TABLE should collect");
    sql_context
        .sql(
            "INSERT INTO paimon.test_db.temporal_filter_pushdown VALUES
                (1, 'alice', TIMESTAMP '2024-01-01 00:00:00.123456', TIMESTAMP '2024-01-01 00:00:00.123456+00:00'),
                (2, 'bob', TIMESTAMP '2024-01-01 00:00:00.654321', TIMESTAMP '2024-01-01 00:00:00.654321+00:00'),
                (3, 'carol', TIMESTAMP '2024-01-02 00:00:00.123456', TIMESTAMP '2024-01-02 00:00:00.123456+00:00')",
        )
        .await
        .expect("INSERT should succeed")
        .collect()
        .await
        .expect("INSERT should collect");

    let timestamp_sql = "SELECT id, name FROM paimon.test_db.temporal_filter_pushdown \
        WHERE ts = TIMESTAMP '2024-01-01 00:00:00.123456' AND id + 1 > 2";
    let plan = sql_context
        .sql(timestamp_sql)
        .await
        .expect("SQL planning should succeed")
        .create_physical_plan()
        .await
        .expect("Physical plan creation should succeed");
    let plan_text = format_physical_plan(&plan);
    let scan_lines = paimon_scan_lines(&plan_text);

    assert!(
        !scan_lines.is_empty(),
        "plan should contain a PaimonTableScan, plan:\n{plan_text}"
    );
    assert!(
        scan_lines
            .iter()
            .any(|line| line.contains("predicate=ts = TS(")),
        "Temporal predicate should be pushed into PaimonTableScan, plan:\n{plan_text}"
    );
    assert!(
        !plan_text.contains("FilterExec"),
        "PaimonTableScan should evaluate accepted runtime filters exactly, plan:\n{plan_text}"
    );
    assert!(
        scan_lines
            .iter()
            .any(|line| line.contains("runtime_filters=[")),
        "Residual filter should be retained by PaimonTableScan, plan:\n{plan_text}"
    );

    let rows = common::collect_id_name(&sql_context, timestamp_sql).await;
    assert!(
        rows.is_empty(),
        "Residual filter should remove the row matched by the pushed temporal predicate"
    );

    let local_zoned_sql = "SELECT id, name FROM paimon.test_db.temporal_filter_pushdown \
        WHERE lzts = TIMESTAMP '2024-01-01 00:00:00.654321+00:00' AND id + 1 > 3";
    let plan = sql_context
        .sql(local_zoned_sql)
        .await
        .expect("SQL planning should succeed")
        .create_physical_plan()
        .await
        .expect("Physical plan creation should succeed");
    let plan_text = format_physical_plan(&plan);
    let scan_lines = paimon_scan_lines(&plan_text);

    assert!(
        !scan_lines.is_empty(),
        "plan should contain a PaimonTableScan, plan:\n{plan_text}"
    );
    assert!(
        scan_lines
            .iter()
            .any(|line| line.contains("predicate=lzts = LZTS(")),
        "Local zoned timestamp predicate should be pushed into PaimonTableScan, plan:\n{plan_text}"
    );
    assert!(
        !plan_text.contains("FilterExec"),
        "PaimonTableScan should evaluate accepted runtime filters exactly, plan:\n{plan_text}"
    );
    assert!(
        scan_lines
            .iter()
            .any(|line| line.contains("runtime_filters=[")),
        "Residual filter should be retained by PaimonTableScan, plan:\n{plan_text}"
    );

    let rows = common::collect_id_name(&sql_context, local_zoned_sql).await;
    assert!(
        rows.is_empty(),
        "Residual filter should remove the row matched by the pushed local zoned timestamp predicate"
    );
}

/// Regression test for pre-epoch, sub-millisecond timestamps.
///
/// The BinaryRow write path and the DataFusion literal pushdown path must
/// agree on the `(millis, nanos)` representation of a timestamp. Previously the
/// write path used truncating division while pushdown used euclidean division,
/// so `1969-12-31 23:59:59.999999` (`-1us`) was stored as `(0, -1000)` but the
/// pushed literal became `(-1, 999_000)`. A pushed equality/range predicate on
/// such a value was therefore false-negative and dropped matching rows.
#[tokio::test]
async fn test_negative_temporal_filter_pushdown_via_datafusion_scan() {
    let (_tmp, sql_context) = common::setup_sql_context().await;
    sql_context
        .sql(
            "CREATE TABLE paimon.test_db.negative_temporal_pushdown (
                id INT,
                name STRING,
                ts TIMESTAMP(6)
            )",
        )
        .await
        .expect("CREATE TABLE should succeed")
        .collect()
        .await
        .expect("CREATE TABLE should collect");
    // id=1 sits 1us before the epoch (negative micros), id=2 sits at the epoch,
    // id=3 is a normal post-epoch value. Each row is written in its own INSERT so
    // it lands in a separate data file, forcing per-file min/max stats pruning to
    // rely on the pushed predicate rather than a range that happens to overlap.
    for values in [
        "(1, 'neg_us', TIMESTAMP '1969-12-31 23:59:59.999999')",
        "(2, 'epoch', TIMESTAMP '1970-01-01 00:00:00.000000')",
        "(3, 'post', TIMESTAMP '1970-01-01 00:00:00.000001')",
    ] {
        sql_context
            .sql(&format!(
                "INSERT INTO paimon.test_db.negative_temporal_pushdown VALUES {values}"
            ))
            .await
            .expect("INSERT should succeed")
            .collect()
            .await
            .expect("INSERT should collect");
    }

    // Equality on the pre-epoch fractional timestamp must return exactly id=1.
    let eq_sql = "SELECT id, name FROM paimon.test_db.negative_temporal_pushdown \
        WHERE ts = TIMESTAMP '1969-12-31 23:59:59.999999'";
    let plan = sql_context
        .sql(eq_sql)
        .await
        .expect("SQL planning should succeed")
        .create_physical_plan()
        .await
        .expect("Physical plan creation should succeed");
    let plan_text = format_physical_plan(&plan);
    let scan_lines = paimon_scan_lines(&plan_text);
    assert!(
        scan_lines
            .iter()
            .any(|line| line.contains("predicate=ts = TS(")),
        "Temporal predicate should be pushed into PaimonTableScan, plan:\n{plan_text}"
    );

    let rows = common::collect_id_name(&sql_context, eq_sql).await;
    assert_eq!(
        rows,
        vec![(1, "neg_us".to_string())],
        "Pushed equality on a pre-epoch sub-millisecond timestamp must match the stored row"
    );

    // Range predicate across the epoch boundary must keep only the pre-epoch row.
    let range_sql = "SELECT id, name FROM paimon.test_db.negative_temporal_pushdown \
        WHERE ts < TIMESTAMP '1970-01-01 00:00:00.000000'";
    let rows = common::collect_id_name(&sql_context, range_sql).await;
    assert_eq!(
        rows,
        vec![(1, "neg_us".to_string())],
        "Pushed range predicate must order pre-epoch fractional timestamps correctly"
    );
}

#[tokio::test]
async fn test_limit_pushdown_on_data_evolution_table_returns_merged_rows() {
    let batches = collect_query("SELECT id, name FROM paimon.default.data_evolution_table LIMIT 3")
        .await
        .expect("Limit query on data evolution table should succeed");

    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
    assert_eq!(
        total_rows, 3,
        "LIMIT 3 should return exactly 3 rows for data evolution table"
    );

    let mut rows = extract_id_name_rows(&batches);
    rows.sort_by_key(|(id, _)| *id);

    assert_eq!(
        rows,
        vec![
            (1, "alice-v2".to_string()),
            (2, "bob".to_string()),
            (3, "carol-v2".to_string()),
        ],
        "Data evolution table LIMIT 3 should return merged rows"
    );
}

#[tokio::test]
async fn test_limit_pushdown_marks_safe_scan_limit_hint_and_keeps_correctness() {
    let sql = "SELECT id, name FROM paimon.default.simple_log_table LIMIT 2";
    let plan = create_physical_plan(sql)
        .await
        .expect("Physical plan creation should succeed");
    let plan_text = format_physical_plan(&plan);
    let scan_lines = paimon_scan_lines(&plan_text);

    assert!(
        scan_lines
            .iter()
            .any(|line| line.contains("limit=2") && !line.contains("fetch=")),
        "Safe LIMIT query should push a scan limit hint into PaimonTableScan, plan:\n{plan_text}"
    );

    let batches = collect_query(sql)
        .await
        .expect("LIMIT query should succeed");
    let total_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum();
    assert_eq!(total_rows, 2, "LIMIT 2 should still return exactly 2 rows");
}

#[tokio::test]
async fn test_offset_limit_pushdown_keeps_correctness_without_fetch_contract() {
    let sql = "SELECT id, name FROM paimon.default.partitioned_log_table OFFSET 1 LIMIT 1";
    let plan = create_physical_plan(sql)
        .await
        .expect("Physical plan creation should succeed");
    let plan_text = format_physical_plan(&plan);
    let scan_lines = paimon_scan_lines(&plan_text);

    assert!(
        plan_text.contains("GlobalLimitExec"),
        "OFFSET queries should keep a GlobalLimitExec in DataFusion, plan:\n{plan_text}"
    );
    assert!(
        scan_lines.iter().all(|line| !line.contains("fetch=")),
        "OFFSET + LIMIT should not rely on the removed DataFusion fetch contract in PaimonTableScan, plan:\n{plan_text}"
    );

    let batches = collect_query(sql)
        .await
        .expect("OFFSET + LIMIT query should succeed");

    let total_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum();
    assert_eq!(
        total_rows, 1,
        "OFFSET 1 LIMIT 1 should still return exactly 1 row"
    );
}

#[tokio::test]
async fn test_inexact_filter_limit_keeps_correctness_without_fetch_contract() {
    let sql = "SELECT id, name FROM paimon.default.partitioned_log_table WHERE id > 1 LIMIT 1";
    let plan = create_physical_plan(sql)
        .await
        .expect("Physical plan creation should succeed");
    let plan_text = format_physical_plan(&plan);
    let scan_lines = paimon_scan_lines(&plan_text);

    assert!(
        !scan_lines.is_empty(),
        "plan should contain a PaimonTableScan, plan:\n{plan_text}"
    );
    assert!(
        scan_lines.iter().all(|line| !line.contains("fetch=")),
        "Inexact filter queries should not revive the removed fetch contract, plan:\n{plan_text}"
    );

    let batches = collect_query(sql)
        .await
        .expect("Inexact filter + LIMIT query should succeed");
    let total_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum();
    assert_eq!(
        total_rows, 1,
        "Inexact filter + LIMIT should still return exactly 1 row"
    );
}

#[tokio::test]
async fn test_residual_filter_limit_keeps_connector_limit_and_correctness() {
    let sql = "SELECT id, name FROM paimon.default.simple_log_table WHERE id + 1 > 3 LIMIT 1";
    let plan = create_physical_plan(sql)
        .await
        .expect("Physical plan creation should succeed");
    let plan_text = format_physical_plan(&plan);
    let scan_lines = paimon_scan_lines(&plan_text);

    assert!(
        !scan_lines.is_empty(),
        "plan should contain a PaimonTableScan, plan:\n{plan_text}"
    );
    assert!(
        scan_lines.iter().all(|line| !line.contains("fetch=")),
        "Residual filter queries should not revive the removed fetch contract, plan:\n{plan_text}"
    );
    assert!(
        scan_lines
            .iter()
            .all(|line| !line.contains(", limit=")),
        "Residual filter queries should not push a scan limit hint when residual filters stay above the scan, plan:\n{plan_text}"
    );

    let batches = collect_query(sql)
        .await
        .expect("Residual filter + LIMIT query should succeed");
    let rows = extract_id_name_rows(&batches);

    assert_eq!(
        rows,
        vec![(3, "carol".to_string())],
        "Residual filter + LIMIT should still return the matching row"
    );
}

// ======================= Catalog Provider Tests =======================
#[tokio::test]
async fn test_query_via_catalog_provider() {
    let catalog = create_catalog();
    let catalog: Arc<dyn Catalog> = Arc::new(catalog);
    let mut ctx = SQLContext::new();
    ctx.register_catalog("paimon", catalog)
        .await
        .expect("Failed to register catalog");

    let df = ctx
        .sql("SELECT id, name FROM paimon.default.simple_log_table")
        .await
        .expect("Failed to execute query");

    let batches = df.collect().await.expect("Failed to collect results");
    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
    assert_eq!(total_rows, 3, "Expected 3 rows from simple_log_table");
}

#[tokio::test]
async fn test_missing_database_returns_no_schema() {
    let catalog = create_catalog();
    let provider = PaimonCatalogProvider::new(
        None,
        Arc::new(catalog),
        Default::default(),
        Default::default(),
        None,
    );

    assert!(
        provider.schema("definitely_missing_database").is_none(),
        "missing databases should not resolve to a schema provider"
    );
}

// ======================= Time Travel Tests =======================

/// Helper: create a SQLContext with catalog + relation planner for time travel.
async fn create_time_travel_context() -> SQLContext {
    let catalog = create_catalog();
    let catalog: Arc<dyn Catalog> = Arc::new(catalog);
    let mut ctx = SQLContext::new();
    ctx.register_catalog("paimon", catalog)
        .await
        .expect("Failed to register catalog");
    ctx
}

#[tokio::test]
async fn test_time_travel_by_snapshot_id() {
    let ctx = create_time_travel_context().await;

    // Snapshot 1: should contain only the first insert (alice, bob)
    let batches = ctx
        .sql("SELECT id, name FROM paimon.default.time_travel_table VERSION AS OF 1")
        .await
        .expect("time travel query should parse")
        .collect()
        .await
        .expect("time travel query should execute");

    let mut rows = extract_id_name_rows(&batches);
    rows.sort_by_key(|(id, _)| *id);
    assert_eq!(
        rows,
        vec![(1, "alice".to_string()), (2, "bob".to_string())],
        "Snapshot 1 should contain only the first batch of rows"
    );

    // Snapshot 2 (latest): should contain all rows
    let batches = ctx
        .sql("SELECT id, name FROM paimon.default.time_travel_table VERSION AS OF 2")
        .await
        .expect("time travel query should parse")
        .collect()
        .await
        .expect("time travel query should execute");

    let mut rows = extract_id_name_rows(&batches);
    rows.sort_by_key(|(id, _)| *id);
    assert_eq!(
        rows,
        vec![
            (1, "alice".to_string()),
            (2, "bob".to_string()),
            (3, "carol".to_string()),
            (4, "dave".to_string()),
        ],
        "Snapshot 2 should contain all rows"
    );
}

#[tokio::test]
async fn test_time_travel_by_tag_name() {
    let ctx = create_time_travel_context().await;

    // Tag 'snapshot1' points to snapshot 1: should contain only (alice, bob)
    let batches = ctx
        .sql("SELECT id, name FROM paimon.default.time_travel_table VERSION AS OF 'snapshot1'")
        .await
        .expect("tag time travel query should parse")
        .collect()
        .await
        .expect("tag time travel query should execute");

    let mut rows = extract_id_name_rows(&batches);
    rows.sort_by_key(|(id, _)| *id);
    assert_eq!(
        rows,
        vec![(1, "alice".to_string()), (2, "bob".to_string())],
        "Tag 'snapshot1' should contain only the first batch of rows"
    );

    // Tag 'snapshot2' points to snapshot 2: should contain all rows
    let batches = ctx
        .sql("SELECT id, name FROM paimon.default.time_travel_table VERSION AS OF 'snapshot2'")
        .await
        .expect("tag time travel query should parse")
        .collect()
        .await
        .expect("tag time travel query should execute");

    let mut rows = extract_id_name_rows(&batches);
    rows.sort_by_key(|(id, _)| *id);
    assert_eq!(
        rows,
        vec![
            (1, "alice".to_string()),
            (2, "bob".to_string()),
            (3, "carol".to_string()),
            (4, "dave".to_string()),
        ],
        "Tag 'snapshot2' should contain all rows"
    );
}

#[tokio::test]
async fn time_travel_schema_evolution() {
    assert_time_travel_schema_evolution_rows(
        "before_add_column",
        "id, name",
        None,
        &[&["1", "alice"], &["2", "bob"]],
    )
    .await;
    assert_time_travel_schema_evolution_rows(
        "before_add_column",
        "id, name",
        Some("id = 2"),
        &[&["2", "bob"]],
    )
    .await;

    assert_time_travel_schema_evolution_rows(
        "after_add_column",
        "id, name, age",
        None,
        &[
            &["1", "alice", "NULL"],
            &["2", "bob", "NULL"],
            &["3", "carol", "30"],
            &["4", "dave", "40"],
        ],
    )
    .await;
    assert_time_travel_schema_evolution_rows(
        "after_add_column",
        "id, name, age",
        Some("age IS NOT NULL"),
        &[&["3", "carol", "30"], &["4", "dave", "40"]],
    )
    .await;

    assert_time_travel_schema_evolution_rows(
        "before_rename",
        "id, name, age",
        Some("id >= 3"),
        &[&["3", "carol", "30"], &["4", "dave", "40"]],
    )
    .await;

    assert_time_travel_schema_evolution_rows(
        "after_rename",
        "id, full_name, age",
        None,
        &[
            &["1", "alice", "NULL"],
            &["2", "bob", "NULL"],
            &["3", "carol", "30"],
            &["4", "dave", "40"],
            &["5", "erin", "50"],
            &["6", "frank", "60"],
        ],
    )
    .await;
    assert_time_travel_schema_evolution_rows(
        "after_rename",
        "id, full_name, age",
        Some("full_name LIKE 'f%'"),
        &[&["6", "frank", "60"]],
    )
    .await;

    assert_time_travel_schema_evolution_rows(
        "before_drop",
        "id, full_name, age",
        Some("age >= 50"),
        &[&["5", "erin", "50"], &["6", "frank", "60"]],
    )
    .await;

    assert_time_travel_schema_evolution_rows(
        "after_drop",
        "id, full_name",
        None,
        &[
            &["1", "alice"],
            &["2", "bob"],
            &["3", "carol"],
            &["4", "dave"],
            &["5", "erin"],
            &["6", "frank"],
            &["7", "grace"],
            &["8", "hank"],
        ],
    )
    .await;
    assert_time_travel_schema_evolution_rows(
        "after_drop",
        "id, full_name",
        Some("id > 6"),
        &[&["7", "grace"], &["8", "hank"]],
    )
    .await;

    assert_time_travel_schema_evolution_rows(
        "before_reorder",
        "id, full_name",
        Some("id = 8"),
        &[&["8", "hank"]],
    )
    .await;

    assert_time_travel_schema_evolution_rows(
        "after_reorder",
        "id, full_name",
        None,
        &[
            &["1", "alice"],
            &["2", "bob"],
            &["3", "carol"],
            &["4", "dave"],
            &["5", "erin"],
            &["6", "frank"],
            &["7", "grace"],
            &["8", "hank"],
            &["9", "ivy"],
            &["10", "jane"],
        ],
    )
    .await;
    assert_time_travel_schema_evolution_rows(
        "after_reorder",
        "id, full_name",
        Some("id >= 9"),
        &[&["9", "ivy"], &["10", "jane"]],
    )
    .await;
}

#[tokio::test]
async fn test_time_travel_conflicting_selectors_fail() {
    // When both scan.version and scan.timestamp-millis are set on the same
    // provider, Paimon rejects the combination at scan time.
    let provider = create_provider_with_options(
        "time_travel_table",
        HashMap::from([
            ("scan.version".to_string(), "1".to_string()),
            ("scan.timestamp-millis".to_string(), "1234".to_string()),
        ]),
    )
    .await;

    let ctx = create_context().await;
    ctx.register_temp_table("paimon.default.time_travel_table", Arc::new(provider))
        .expect("Failed to register temp table");

    let err = ctx
        .sql("SELECT id, name FROM paimon.default.time_travel_table")
        .await
        .expect("query should parse")
        .collect()
        .await
        .expect_err("conflicting time-travel selectors should fail");

    let message = err.to_string();
    assert!(
        message.contains("Only one time-travel selector may be set"),
        "unexpected conflict error: {message}"
    );
    assert!(
        message.contains("scan.version"),
        "conflict error should mention scan.version: {message}"
    );
}

#[tokio::test]
async fn test_time_travel_invalid_version_fails() {
    let provider = create_provider_with_options(
        "time_travel_table",
        HashMap::from([("scan.version".to_string(), "nonexistent-tag".to_string())]),
    )
    .await;

    let ctx = create_context().await;
    ctx.register_temp_table("paimon.default.time_travel_table", Arc::new(provider))
        .expect("Failed to register temp table");

    let err = ctx
        .sql("SELECT id, name FROM paimon.default.time_travel_table")
        .await
        .expect("query should parse")
        .collect()
        .await
        .expect_err("invalid version should fail");

    let message = err.to_string();
    assert!(
        message.contains("is not a valid tag name or snapshot id"),
        "unexpected invalid version error: {message}"
    );
}

/// Verifies that data evolution merge correctly NULL-fills columns that no file in a
/// merge group provides (e.g. a newly added column after MERGE INTO on old rows).
/// Without the fix, `active_file_indices` would be empty and rows would be silently lost.
#[tokio::test]
async fn test_data_evolution_drop_column_null_fill() {
    let batches =
        collect_query("SELECT id, name, extra FROM paimon.default.data_evolution_drop_column")
            .await
            .expect("data_evolution_drop_column query should succeed");

    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
    assert_eq!(
        total_rows, 3,
        "Should return 3 rows (not silently drop rows from merge groups missing the new column)"
    );

    let mut rows: Vec<(i32, String, Option<String>)> = Vec::new();
    for batch in &batches {
        let id_array = batch
            .column_by_name("id")
            .and_then(|c| c.as_any().downcast_ref::<Int32Array>())
            .expect("Expected Int32Array for id");
        let name_array = batch
            .column_by_name("name")
            .expect("Expected string array for name");
        let extra_array = batch
            .column_by_name("extra")
            .expect("Expected string array for extra");

        for i in 0..batch.num_rows() {
            let extra = if extra_array.is_null(i) {
                None
            } else {
                Some(string_value(extra_array.as_ref(), i).to_string())
            };
            rows.push((
                id_array.value(i),
                string_value(name_array.as_ref(), i).to_string(),
                extra,
            ));
        }
    }
    rows.sort_by_key(|(id, _, _)| *id);

    assert_eq!(
        rows,
        vec![
            (1, "alice-v2".to_string(), None),
            (2, "bob".to_string(), None),
            (3, "carol".to_string(), Some("new".to_string())),
        ],
        "Old rows should have extra=NULL, new row should have extra='new'"
    );
}

#[tokio::test]
async fn test_sql_read_format_schema_evolution_add_column() {
    assert_sql_rows(
        "SELECT id, name, age FROM paimon.default.format_schema_evolution_add_column ORDER BY id",
        &[
            &["1", "alice", "NULL"],
            &["2", "bob", "NULL"],
            &["3", "carol", "30"],
            &["4", "dave", "40"],
            &["5", "eve", "50"],
            &["6", "frank", "60"],
        ],
    )
    .await;

    assert_sql_rows(
        "SELECT id, age FROM paimon.default.format_schema_evolution_add_column WHERE age IS NULL ORDER BY id",
        &[&["1", "NULL"], &["2", "NULL"]],
    )
    .await;
}

#[tokio::test]
async fn test_sql_read_partitioned_format_schema_evolution_add_column() {
    assert_sql_rows(
        "SELECT id, name, extra FROM paimon.default.partitioned_format_schema_evolution_add_column WHERE dt = '2024-01-02' ORDER BY id",
        &[
            &["2", "bob", "NULL"],
            &["5", "eve", "avro-extra-1"],
        ],
    )
    .await;

    assert_sql_rows(
        "SELECT id, name, extra FROM paimon.default.partitioned_format_schema_evolution_add_column WHERE dt = '2024-01-01' AND extra IS NULL ORDER BY id",
        &[&["1", "alice", "NULL"]],
    )
    .await;

    assert_sql_rows(
        "SELECT dt, id, extra FROM paimon.default.partitioned_format_schema_evolution_add_column ORDER BY id",
        &[
            &["2024-01-01", "1", "NULL"],
            &["2024-01-02", "2", "NULL"],
            &["2024-01-01", "3", "orc-extra-1"],
            &["2024-01-03", "4", "orc-extra-2"],
            &["2024-01-02", "5", "avro-extra-1"],
            &["2024-01-03", "6", "avro-extra-2"],
        ],
    )
    .await;
}

#[tokio::test]
async fn test_sql_read_format_schema_evolution_type_promotion() {
    assert_sql_rows(
        "SELECT id, value FROM paimon.default.format_schema_evolution_type_promotion ORDER BY id",
        &[
            &["1", "100"],
            &["2", "200"],
            &["3", "3000000000"],
            &["4", "4000000000"],
            &["5", "5000000000"],
            &["6", "6000000000"],
        ],
    )
    .await;

    assert_sql_rows(
        "SELECT id, value FROM paimon.default.format_schema_evolution_type_promotion WHERE value > 3000000000 ORDER BY id",
        &[&["4", "4000000000"], &["5", "5000000000"], &["6", "6000000000"]],
    )
    .await;
}

#[tokio::test]
async fn test_sql_read_schema_evolution_rename_column() {
    assert_sql_rows(
        "SELECT id, renamed_payload FROM paimon.default.schema_evolution_rename_column ORDER BY id",
        &[
            &["1", "parquet-old"],
            &["2", "parquet-old-2"],
            &["3", "orc-new"],
            &["4", "avro-new"],
        ],
    )
    .await;

    assert_sql_rows(
        "SELECT id, renamed_payload FROM paimon.default.schema_evolution_rename_column WHERE renamed_payload LIKE '%new' ORDER BY id",
        &[&["3", "orc-new"], &["4", "avro-new"]],
    )
    .await;
}

#[tokio::test]
async fn test_sql_read_mixed_format_schema_evolution_drop_column() {
    assert_sql_rows(
        "SELECT id, name FROM paimon.default.mixed_format_schema_evolution_drop_column ORDER BY id",
        &[
            &["1", "parquet-alice"],
            &["2", "parquet-bob"],
            &["3", "orc-carol"],
            &["4", "orc-dave"],
            &["5", "avro-eve"],
            &["6", "avro-frank"],
        ],
    )
    .await;

    assert_sql_rows(
        "SELECT id, name FROM paimon.default.mixed_format_schema_evolution_drop_column WHERE name LIKE 'avro-%' ORDER BY id",
        &[&["5", "avro-eve"], &["6", "avro-frank"]],
    )
    .await;
}

#[tokio::test]
async fn test_sql_read_mixed_format_schema_evolution_reorder_move_column() {
    assert_sql_rows(
        "SELECT right_value, left_value, id FROM paimon.default.mixed_format_schema_evolution_reorder_move_column ORDER BY id",
        &[
            &["parquet-right-1", "parquet-left-1", "1"],
            &["parquet-right-2", "parquet-left-2", "2"],
            &["orc-right-3", "orc-left-3", "3"],
            &["orc-right-4", "orc-left-4", "4"],
            &["avro-right-5", "avro-left-5", "5"],
            &["avro-right-6", "avro-left-6", "6"],
        ],
    )
    .await;

    assert_sql_rows(
        "SELECT id, right_value FROM paimon.default.mixed_format_schema_evolution_reorder_move_column WHERE right_value LIKE 'orc-%' ORDER BY id",
        &[&["3", "orc-right-3"], &["4", "orc-right-4"]],
    )
    .await;
}

// ======================= Complex Type Tests =======================

#[tokio::test]
async fn test_read_complex_type_table_via_datafusion() {
    let batches = collect_query(
        "SELECT id, int_array, string_map, row_field FROM paimon.default.complex_type_table ORDER BY id",
    )
    .await
    .expect("Complex type query should succeed");

    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
    assert_eq!(total_rows, 3, "Expected 3 rows from complex_type_table");

    // Verify column types exist and are correct
    for batch in &batches {
        let schema = batch.schema();
        assert!(
            schema.field_with_name("int_array").is_ok(),
            "int_array column should exist"
        );
        assert!(
            schema.field_with_name("string_map").is_ok(),
            "string_map column should exist"
        );
        assert!(
            schema.field_with_name("row_field").is_ok(),
            "row_field column should exist"
        );
    }

    // Extract and verify data using Arrow arrays
    let mut rows: Vec<(i32, String, String, String)> = Vec::new();
    for batch in &batches {
        let id_array = batch
            .column_by_name("id")
            .and_then(|c| c.as_any().downcast_ref::<Int32Array>())
            .expect("Expected Int32Array for id");
        let int_array_col = batch.column_by_name("int_array").expect("int_array");
        let string_map_col = batch.column_by_name("string_map").expect("string_map");
        let row_field_col = batch.column_by_name("row_field").expect("row_field");

        for i in 0..batch.num_rows() {
            use datafusion::arrow::util::display::ArrayFormatter;
            let fmt_opts = datafusion::arrow::util::display::FormatOptions::default();

            let arr_fmt = ArrayFormatter::try_new(int_array_col.as_ref(), &fmt_opts).unwrap();
            let map_fmt = ArrayFormatter::try_new(string_map_col.as_ref(), &fmt_opts).unwrap();
            let row_fmt = ArrayFormatter::try_new(row_field_col.as_ref(), &fmt_opts).unwrap();

            rows.push((
                id_array.value(i),
                arr_fmt.value(i).to_string(),
                map_fmt.value(i).to_string(),
                row_fmt.value(i).to_string(),
            ));
        }
    }
    rows.sort_by_key(|(id, _, _, _)| *id);

    assert_eq!(rows[0].0, 1);
    assert_eq!(rows[0].1, "[1, 2, 3]");
    assert_eq!(rows[0].2, "{a: 10, b: 20}");
    assert_eq!(rows[0].3, "{name: alice, value: 100}");

    assert_eq!(rows[1].0, 2);
    assert_eq!(rows[1].1, "[4, 5]");
    assert_eq!(rows[1].2, "{c: 30}");
    assert_eq!(rows[1].3, "{name: bob, value: 200}");

    assert_eq!(rows[2].0, 3);
    assert_eq!(rows[2].1, "[]");
    assert_eq!(rows[2].2, "{}");
    assert_eq!(rows[2].3, "{name: carol, value: 300}");
}

#[tokio::test]
async fn test_select_row_id_from_data_evolution_table() {
    use datafusion::arrow::array::Int64Array;

    let ctx = create_context().await;

    let batches = ctx
        .sql(r#"SELECT "_ROW_ID", id, name FROM paimon.default.data_evolution_table"#)
        .await
        .expect("SQL should parse")
        .collect()
        .await
        .expect("query should execute");

    assert!(!batches.is_empty());
    let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
    assert!(total_rows > 0);

    for batch in &batches {
        let row_id_col = batch
            .column_by_name("_ROW_ID")
            .expect("_ROW_ID column should exist");
        let row_id_array = row_id_col
            .as_any()
            .downcast_ref::<Int64Array>()
            .expect("_ROW_ID should be Int64");
        for i in 0..batch.num_rows() {
            assert!(
                row_id_array.is_valid(i),
                "_ROW_ID should not be null for data evolution table"
            );
            assert!(row_id_array.value(i) >= 0);
        }
    }
}

#[tokio::test]
async fn test_filter_row_id_from_data_evolution_table() {
    use datafusion::arrow::array::Int64Array;

    let ctx = create_context().await;

    let all_batches = ctx
        .sql(r#"SELECT "_ROW_ID" FROM paimon.default.data_evolution_table"#)
        .await
        .expect("SQL")
        .collect()
        .await
        .expect("collect");
    let all_count: usize = all_batches.iter().map(|b| b.num_rows()).sum();

    let filtered_batches = ctx
        .sql(r#"SELECT "_ROW_ID", id FROM paimon.default.data_evolution_table WHERE "_ROW_ID" = 0"#)
        .await
        .expect("SQL")
        .collect()
        .await
        .expect("collect");
    let filtered_count: usize = filtered_batches.iter().map(|b| b.num_rows()).sum();

    assert!(filtered_count <= all_count);
    for batch in &filtered_batches {
        let row_id_array = batch
            .column_by_name("_ROW_ID")
            .expect("_ROW_ID")
            .as_any()
            .downcast_ref::<Int64Array>()
            .expect("Int64");
        for i in 0..batch.num_rows() {
            assert_eq!(row_id_array.value(i), 0);
        }
    }
}

/// Pins that the DataFusion SQL path is always case-sensitive for column names.
///
/// Paimon exposes case-insensitive column matching only through the direct
/// ReadBuilder API (core / C / Python). It cannot be offered via SQL: DataFusion
/// resolves projection/filter columns against the provider schema during logical
/// planning, *before* `TableProvider::scan` runs. `enable_ident_normalization`
/// only lowercases unquoted identifiers at parse time; it does not make schema
/// resolution case-insensitive. So a differently-cased reference fails at
/// planning and never reaches the scan-level resolution. This test documents
/// that boundary so any future attempt to claim SQL-level case-insensitivity has
/// to confront the planner-side resolution first.
#[tokio::test]
async fn test_case_insensitive_column_not_supported_via_sql() {
    let (_tmp, sql_context) = common::setup_sql_context().await;
    // Quote the identifier so DDL parsing preserves the mixed-case field name.
    sql_context
        .sql("CREATE TABLE paimon.test_db.mixed_case_cols (id INT, \"Name\" STRING)")
        .await
        .expect("CREATE TABLE should succeed")
        .collect()
        .await
        .expect("CREATE TABLE should collect");

    // The exact-case (quoted) column resolves and the query plans fine.
    sql_context
        .sql("SELECT \"Name\" FROM paimon.test_db.mixed_case_cols")
        .await
        .expect("exact-case column reference should resolve");

    // An unquoted `name` is normalized to lowercase and fails during logical
    // planning against the `Name` schema field — the scan-level case-insensitive
    // resolution is never reached, confirming SQL reads stay case-sensitive.
    let err = sql_context
        .sql("SELECT name FROM paimon.test_db.mixed_case_cols")
        .await
        .expect_err("case-mismatched column must fail at planning, not resolve case-insensitively");
    let msg = err.to_string().to_lowercase();
    assert!(
        msg.contains("name"),
        "expected a column-resolution error referencing `name`, got: {err}"
    );
}

// ======================= Full-Text Search Tests =======================

#[cfg(feature = "fulltext")]
mod fulltext_tests {
    use std::sync::Arc;

    use bytes::Bytes;
    use datafusion::arrow::array::Int32Array;
    use paimon::catalog::Identifier;
    use paimon::spec::{GlobalIndexMeta, IndexFileMeta};
    use paimon::table::BranchManager;
    use paimon::{Catalog, CatalogOptions, CommitMessage, FileSystemCatalog, Options, TableCommit};
    use paimon_datafusion::{register_full_text_search, SQLContext};

    use super::common::string_value;
    use paimon_ftindex_core::io::PosWriter;
    use paimon_ftindex_core::{FullTextIndexConfig, FullTextIndexWriter};

    async fn create_fulltext_context() -> (SQLContext, Arc<FileSystemCatalog>, tempfile::TempDir) {
        let tmp = tempfile::tempdir().expect("Failed to create temp dir");
        let warehouse = format!("file://{}", tmp.path().display());
        let mut options = Options::new();
        options.set(CatalogOptions::WAREHOUSE, warehouse);
        let catalog = Arc::new(FileSystemCatalog::new(options).expect("Failed to create catalog"));

        let mut ctx = SQLContext::new();
        ctx.register_catalog("paimon", catalog.clone())
            .await
            .expect("Failed to register catalog");
        register_full_text_search(ctx.ctx(), catalog.clone(), "default");
        run_fulltext_sql(
            &ctx,
            "CREATE TABLE paimon.default.test_fulltext (
                id INT,
                content STRING
            ) WITH (
                'row-tracking.enabled' = 'true',
                'global-index.search-mode' = 'full'
            )",
        )
        .await;
        for insert in [
            "INSERT INTO paimon.default.test_fulltext VALUES (0, 'alphaft lakehouse storage')",
            "INSERT INTO paimon.default.test_fulltext VALUES (1, 'betafft gammaft full-text search')",
            "INSERT INTO paimon.default.test_fulltext VALUES (2, 'alphaft supports indexed queries')",
            "INSERT INTO paimon.default.test_fulltext VALUES (3, 'gammaft over raw fallback rows')",
            "INSERT INTO paimon.default.test_fulltext VALUES (4, 'alphaft tables can be queried')",
        ] {
            run_fulltext_sql(&ctx, insert).await;
        }
        build_fulltext_index(catalog.as_ref()).await;
        (ctx, catalog, tmp)
    }

    async fn build_fulltext_index(catalog: &dyn Catalog) {
        let table = catalog
            .get_table(&Identifier::new("default", "test_fulltext"))
            .await
            .expect("table should exist");
        let content_field_id = table
            .schema()
            .fields()
            .iter()
            .find(|field| field.name() == "content")
            .expect("content field")
            .id();

        let mut writer =
            FullTextIndexWriter::new(FullTextIndexConfig::new().with_text_fields(["content"]))
                .expect("full-text writer");
        for (row_id, content) in [
            (0, "alphaft lakehouse storage"),
            (1, "betafft gammaft full-text search"),
            (2, "alphaft supports indexed queries"),
            (3, "gammaft over raw fallback rows"),
            (4, "alphaft tables can be queried"),
        ] {
            writer
                .add_document_fields(row_id, [("content", content)])
                .expect("add full-text document");
        }
        let mut index_bytes = Vec::new();
        writer
            .write(&mut PosWriter::new(&mut index_bytes))
            .expect("write full-text index");

        let index_file_name = "ft-datafusion-fulltext.index";
        table
            .file_io()
            .mkdirs(&format!(
                "{}/index/",
                table.location().trim_end_matches('/')
            ))
            .await
            .expect("create index dir");
        table
            .file_io()
            .new_output(&format!(
                "{}/index/{index_file_name}",
                table.location().trim_end_matches('/')
            ))
            .expect("index output")
            .write(Bytes::from(index_bytes.clone()))
            .await
            .expect("write index file");

        let mut message = CommitMessage::new(Vec::new(), 0, Vec::new());
        message.new_index_files = vec![IndexFileMeta {
            index_type: "full-text".to_string(),
            file_name: index_file_name.to_string(),
            file_size: i64::try_from(index_bytes.len()).unwrap(),
            row_count: 5,
            deletion_vectors_ranges: None,
            global_index_meta: Some(GlobalIndexMeta {
                row_range_start: 0,
                row_range_end: 4,
                index_field_id: content_field_id,
                extra_field_ids: None,
                index_meta: None,
                source_meta: None,
            }),
        }];
        TableCommit::new(table, "test-user".to_string())
            .commit(vec![message])
            .await
            .expect("commit full-text index manifest");
    }

    async fn run_fulltext_sql(ctx: &SQLContext, sql: &str) {
        ctx.sql(sql)
            .await
            .unwrap_or_else(|e| panic!("Failed to plan `{sql}`: {e}"))
            .collect()
            .await
            .unwrap_or_else(|e| panic!("Failed to execute `{sql}`: {e}"));
    }

    fn extract_id_content_rows(
        batches: &[datafusion::arrow::record_batch::RecordBatch],
    ) -> Vec<(i32, String)> {
        let mut rows = Vec::new();
        for batch in batches {
            let id_array = batch
                .column_by_name("id")
                .and_then(|c| c.as_any().downcast_ref::<Int32Array>())
                .expect("Expected Int32Array for id");
            let content_array = batch
                .column_by_name("content")
                .expect("Expected string array for content");
            for i in 0..batch.num_rows() {
                rows.push((
                    id_array.value(i),
                    string_value(content_array.as_ref(), i).to_string(),
                ));
            }
        }
        rows.sort_by_key(|(id, _)| *id);
        rows
    }

    /// Search for 'alphaft' — rows 0, 2, 4 contain the token.
    #[tokio::test]
    async fn test_full_text_search_paimon() {
        let (ctx, _catalog, _tmp) = create_fulltext_context().await;
        let batches = ctx
            .sql("SELECT id, content FROM full_text_search('paimon.default.test_fulltext', 'content', 'alphaft', 10)")
            .await
            .expect("SQL should parse")
            .collect()
            .await
            .expect("query should execute");

        let rows = extract_id_content_rows(&batches);
        let ids: Vec<i32> = rows.iter().map(|(id, _)| *id).collect();
        assert_eq!(
            ids,
            vec![0, 2, 4],
            "Searching 'alphaft' should match rows 0, 2, 4"
        );
    }

    #[tokio::test]
    async fn test_full_text_search_branch() {
        let (ctx, catalog, _tmp) = create_fulltext_context().await;
        let identifier = Identifier::new("default", "test_fulltext");
        let table = catalog.get_table(&identifier).await.expect("load table");
        let snapshot_manager = table.snapshot_manager();
        let snapshot = snapshot_manager
            .get_latest_snapshot()
            .await
            .expect("load latest snapshot")
            .expect("latest snapshot");
        table
            .tag_manager()
            .create("branch-source", &snapshot)
            .await
            .expect("create tag");
        BranchManager::new(table.file_io().clone(), table.location().to_string())
            .create_branch_from_tag("b1", "branch-source")
            .await
            .expect("create branch");
        table
            .file_io()
            .delete_dir(&snapshot_manager.snapshot_dir())
            .await
            .expect("delete main snapshots");

        let batches = ctx
            .sql("SELECT id, content FROM full_text_search('paimon.default.test_fulltext$branch_b1', 'content', 'alphaft', 10)")
            .await
            .expect("SQL should parse")
            .collect()
            .await
            .expect("branch query should execute");

        let rows = extract_id_content_rows(&batches);
        let ids: Vec<i32> = rows.iter().map(|(id, _)| *id).collect();
        assert_eq!(ids, vec![0, 2, 4]);
    }

    /// Search for 'betafft' — only row 1.
    #[tokio::test]
    async fn test_full_text_search_betafft() {
        let (ctx, _catalog, _tmp) = create_fulltext_context().await;
        let batches = ctx
            .sql("SELECT id, content FROM full_text_search('paimon.default.test_fulltext', 'content', 'betafft', 10)")
            .await
            .expect("SQL should parse")
            .collect()
            .await
            .expect("query should execute");

        let rows = extract_id_content_rows(&batches);
        let ids: Vec<i32> = rows.iter().map(|(id, _)| *id).collect();
        assert_eq!(ids, vec![1], "Searching 'betafft' should match row 1");
    }

    /// Search for 'gammaft' — rows 1, 3 contain the token.
    #[tokio::test]
    async fn test_full_text_search_search() {
        let (ctx, _catalog, _tmp) = create_fulltext_context().await;
        let batches = ctx
            .sql("SELECT id, content FROM full_text_search('paimon.default.test_fulltext', 'content', 'gammaft', 10)")
            .await
            .expect("SQL should parse")
            .collect()
            .await
            .expect("query should execute");

        let rows = extract_id_content_rows(&batches);
        let ids: Vec<i32> = rows.iter().map(|(id, _)| *id).collect();
        assert_eq!(
            ids,
            vec![1, 3],
            "Searching 'gammaft' should match rows 1, 3"
        );
    }
}

// ======================= Vector Search Tests =======================

mod vector_search_tests {
    use std::sync::Arc;

    use datafusion::arrow::array::{ArrayRef, Float32Builder, Int32Array, ListBuilder};
    use datafusion::arrow::datatypes::{
        DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema,
    };
    use datafusion::arrow::record_batch::RecordBatch;
    use datafusion::datasource::MemTable;
    use paimon::catalog::Identifier;
    use paimon::spec::{ArrayType, DataType, FloatType, IntType, Schema};
    use paimon::table::BranchManager;
    use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options};
    use paimon_datafusion::{register_vector_search, SQLContext};

    use super::common::string_value;

    fn extract_test_warehouse(archive_name: &str) -> (tempfile::TempDir, String) {
        let archive_path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
            .join("testdata")
            .join(archive_name);
        let file = std::fs::File::open(&archive_path)
            .unwrap_or_else(|e| panic!("Failed to open {}: {e}", archive_path.display()));
        let decoder = flate2::read::GzDecoder::new(file);
        let mut archive = tar::Archive::new(decoder);

        let tmp = tempfile::tempdir().expect("Failed to create temp dir");
        let db_dir = tmp.path().join("default.db");
        std::fs::create_dir_all(&db_dir).unwrap();
        archive.unpack(&db_dir).unwrap();

        let warehouse = format!("file://{}", tmp.path().display());
        (tmp, warehouse)
    }

    async fn create_vector_search_context(
        archive_name: &str,
    ) -> (SQLContext, Arc<FileSystemCatalog>, tempfile::TempDir) {
        let (tmp, warehouse) = extract_test_warehouse(archive_name);
        let mut options = Options::new();
        options.set(CatalogOptions::WAREHOUSE, warehouse);
        let catalog = Arc::new(FileSystemCatalog::new(options).expect("Failed to create catalog"));

        let mut ctx = SQLContext::new();
        ctx.register_catalog("paimon", catalog.clone())
            .await
            .expect("Failed to register catalog");
        register_vector_search(ctx.ctx(), catalog.clone(), "default");
        (ctx, catalog, tmp)
    }

    async fn create_lumina_vector_search_context(
    ) -> (SQLContext, Arc<FileSystemCatalog>, tempfile::TempDir) {
        create_vector_search_context("test_lumina_vector.tar.gz").await
    }

    async fn create_java_vindex_vector_search_context(
    ) -> (SQLContext, Arc<FileSystemCatalog>, tempfile::TempDir) {
        create_vector_search_context("test_java_vindex_vector.tar.gz").await
    }

    async fn create_empty_vector_search_context(
    ) -> (SQLContext, Arc<FileSystemCatalog>, tempfile::TempDir) {
        let tmp = tempfile::tempdir().expect("Failed to create temp dir");
        let warehouse = format!("file://{}", tmp.path().display());
        let mut options = Options::new();
        options.set(CatalogOptions::WAREHOUSE, warehouse);
        let catalog = Arc::new(FileSystemCatalog::new(options).expect("Failed to create catalog"));

        let mut ctx = SQLContext::new();
        ctx.register_catalog("paimon", catalog.clone())
            .await
            .expect("Failed to register catalog");
        register_vector_search(ctx.ctx(), catalog.clone(), "default");
        (ctx, catalog, tmp)
    }

    fn build_lumina_table_schema() -> Schema {
        let mut options = std::collections::HashMap::new();
        options.insert("row-tracking.enabled".to_string(), "true".to_string());
        options.insert("data-evolution.enabled".to_string(), "true".to_string());
        options.insert("global-index.enabled".to_string(), "true".to_string());
        options.insert(
            "global-index.row-count-per-shard".to_string(),
            "3".to_string(),
        );
        options.insert("lumina.index.dimension".to_string(), "2".to_string());
        options.insert("lumina.encoding.type".to_string(), "rawf32".to_string());

        Schema::builder()
            .column("id", DataType::Int(IntType::new()))
            .column(
                "embedding",
                DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))),
            )
            .options(options)
            .build()
            .expect("Failed to build table schema")
    }

    fn build_vindex_table_schema() -> Schema {
        let mut options = std::collections::HashMap::new();
        options.insert("row-tracking.enabled".to_string(), "true".to_string());
        options.insert("data-evolution.enabled".to_string(), "true".to_string());
        options.insert("global-index.enabled".to_string(), "true".to_string());
        options.insert(
            "global-index.row-count-per-shard".to_string(),
            "3".to_string(),
        );

        Schema::builder()
            .column("id", DataType::Int(IntType::new()))
            .column(
                "embedding",
                DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))),
            )
            .options(options)
            .build()
            .expect("Failed to build table schema")
    }

    fn build_vector_batch(ids: Vec<i32>, vectors: Vec<Vec<f32>>) -> RecordBatch {
        let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true));
        let mut vector_builder =
            ListBuilder::new(Float32Builder::new()).with_field(element_field.clone());
        for vector in vectors {
            for value in vector {
                vector_builder.values().append_value(value);
            }
            vector_builder.append(true);
        }

        let schema = Arc::new(ArrowSchema::new(vec![
            ArrowField::new("id", ArrowDataType::Int32, false),
            ArrowField::new("embedding", ArrowDataType::List(element_field), true),
        ]));
        RecordBatch::try_new(
            schema,
            vec![
                Arc::new(Int32Array::from(ids)) as ArrayRef,
                Arc::new(vector_builder.finish()) as ArrayRef,
            ],
        )
        .expect("Failed to build vector batch")
    }

    fn extract_ids(batches: &[datafusion::arrow::record_batch::RecordBatch]) -> Vec<i32> {
        let mut ids = Vec::new();
        for batch in batches {
            let id_array = batch
                .column_by_name("id")
                .and_then(|c| c.as_any().downcast_ref::<Int32Array>())
                .expect("Expected Int32Array for id");
            for i in 0..batch.num_rows() {
                ids.push(id_array.value(i));
            }
        }
        ids.sort();
        ids
    }

    fn extract_query_result_ids(
        batches: &[datafusion::arrow::record_batch::RecordBatch],
    ) -> Vec<(i32, i32)> {
        let mut rows = Vec::new();
        for batch in batches {
            let query_id_array = batch
                .column_by_name("query_id")
                .and_then(|c| c.as_any().downcast_ref::<Int32Array>())
                .expect("Expected Int32Array for query_id");
            let result_id_array = batch
                .column_by_name("result_id")
                .and_then(|c| c.as_any().downcast_ref::<Int32Array>())
                .expect("Expected Int32Array for result_id");
            for i in 0..batch.num_rows() {
                rows.push((query_id_array.value(i), result_id_array.value(i)));
            }
        }
        rows.sort();
        rows
    }

    fn extract_index_rows(
        batches: &[datafusion::arrow::record_batch::RecordBatch],
    ) -> Vec<(String, i64, i64, i64, String)> {
        let mut rows = Vec::new();
        for batch in batches {
            let index_type_array = batch
                .column_by_name("index_type")
                .expect("Expected string array for index_type");
            let row_count_array = batch
                .column_by_name("row_count")
                .and_then(|c| {
                    c.as_any()
                        .downcast_ref::<datafusion::arrow::array::Int64Array>()
                })
                .expect("Expected Int64Array for row_count");
            let row_range_start_array = batch
                .column_by_name("row_range_start")
                .and_then(|c| {
                    c.as_any()
                        .downcast_ref::<datafusion::arrow::array::Int64Array>()
                })
                .expect("Expected Int64Array for row_range_start");
            let row_range_end_array = batch
                .column_by_name("row_range_end")
                .and_then(|c| {
                    c.as_any()
                        .downcast_ref::<datafusion::arrow::array::Int64Array>()
                })
                .expect("Expected Int64Array for row_range_end");
            let index_field_name_array = batch
                .column_by_name("index_field_name")
                .expect("Expected string array for index_field_name");

            for row_index in 0..batch.num_rows() {
                rows.push((
                    string_value(index_type_array.as_ref(), row_index).to_string(),
                    row_count_array.value(row_index),
                    row_range_start_array.value(row_index),
                    row_range_end_array.value(row_index),
                    string_value(index_field_name_array.as_ref(), row_index).to_string(),
                ));
            }
        }
        rows.sort_by_key(|row| row.2);
        rows
    }

    #[tokio::test]
    async fn test_vector_search_top3() {
        let (ctx, _catalog, _tmp) = create_lumina_vector_search_context().await;
        let batches = ctx
            .sql("SELECT id FROM vector_search('paimon.default.test_lumina_vector', 'embedding', '[1.0, 0.0, 0.0, 0.0]', 3)")
            .await
            .expect("SQL should parse")
            .collect()
            .await
            .expect("query should execute");

        let ids = extract_ids(&batches);
        assert_eq!(ids.len(), 3);
        assert!(ids.contains(&0), "exact match [1,0,0,0] should be in top 3");
    }

    #[tokio::test]
    async fn test_vector_search_branch() {
        let (ctx, catalog, _tmp) = create_java_vindex_vector_search_context().await;
        let identifier = Identifier::new("default", "test_java_vindex_vector");
        let table = catalog.get_table(&identifier).await.expect("load table");
        let snapshot_manager = table.snapshot_manager();
        let snapshot = snapshot_manager
            .get_latest_snapshot()
            .await
            .expect("load latest snapshot")
            .expect("latest snapshot");
        table
            .tag_manager()
            .create("branch-source", &snapshot)
            .await
            .expect("create tag");
        BranchManager::new(table.file_io().clone(), table.location().to_string())
            .create_branch_from_tag("b1", "branch-source")
            .await
            .expect("create branch");
        table
            .file_io()
            .delete_dir(&snapshot_manager.snapshot_dir())
            .await
            .expect("delete main snapshots");

        let batches = ctx
            .sql("SELECT id FROM vector_search('paimon.default.test_java_vindex_vector$branch_b1', 'embedding', '[1.0, 0.0, 0.0, 0.0]', 3)")
            .await
            .expect("SQL should parse")
            .collect()
            .await
            .expect("branch query should execute");

        let ids = extract_ids(&batches);
        assert_eq!(ids.len(), 3);
        assert!(ids.contains(&0));
    }

    #[tokio::test]
    async fn test_vector_search_top6_returns_all() {
        let (ctx, _catalog, _tmp) = create_lumina_vector_search_context().await;
        let batches = ctx
            .sql("SELECT id FROM vector_search('paimon.default.test_lumina_vector', 'embedding', '[1.0, 0.0, 0.0, 0.0]', 6)")
            .await
            .expect("SQL should parse")
            .collect()
            .await
            .expect("query should execute");

        let ids = extract_ids(&batches);
        assert_eq!(ids, vec![0, 1, 2, 3, 4, 5]);
    }

    #[tokio::test]
    async fn test_vector_search_without_matching_index_returns_empty() {
        let (ctx, _catalog, _tmp) = create_lumina_vector_search_context().await;
        let batches = ctx
            .sql("SELECT id FROM vector_search('paimon.default.test_lumina_vector', 'missing_embedding', '[1.0]', 10)")
            .await
            .expect("SQL should parse")
            .collect()
            .await
            .expect("query should execute");

        let total_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum();
        assert_eq!(
            total_rows, 0,
            "vector_search without a matching Lumina index should not fall back to a full table scan"
        );
    }

    #[tokio::test]
    async fn test_vector_search_java_vindex_table() {
        let (ctx, _catalog, _tmp) = create_java_vindex_vector_search_context().await;
        let batches = ctx
            .sql("SELECT id FROM vector_search('paimon.default.test_java_vindex_vector', 'embedding', '[1.0, 0.0, 0.0, 0.0]', 3)")
            .await
            .expect("SQL should parse")
            .collect()
            .await
            .expect("query should execute");

        let ids = extract_ids(&batches);
        assert_eq!(ids, vec![0, 1, 2]);
    }

    #[tokio::test]
    async fn test_vector_search_lateral_join_uses_query_vectors() {
        let (ctx, _catalog, _tmp) = create_java_vindex_vector_search_context().await;
        let query_batch = build_vector_batch(
            vec![10, 20],
            vec![vec![1.0, 0.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]],
        );
        let query_table = MemTable::try_new(query_batch.schema(), vec![vec![query_batch]])
            .expect("Failed to create query vector table");
        ctx.register_temp_table("paimon.default.queries", Arc::new(query_table))
            .expect("Failed to register query vector table");

        let batches = ctx
            .sql(
                "SELECT q.id AS query_id, r.id AS result_id \
                 FROM paimon.default.queries q \
                 CROSS JOIN LATERAL vector_search('paimon.default.test_java_vindex_vector', 'embedding', q.embedding, 2) AS r \
                 ORDER BY query_id, result_id",
            )
            .await
            .expect("lateral vector_search SQL should parse")
            .collect()
            .await
            .expect("lateral vector_search query should execute");

        let rows = extract_query_result_ids(&batches);
        assert_eq!(rows, vec![(10, 0), (10, 1), (20, 1), (20, 2)]);
    }

    // Manual run with a local Lumina native library:
    // LUMINA_LIB_PATH=/path/to/liblumina_py.so cargo test -p paimon-datafusion \
    //     vector_search_tests::test_lumina_build_then_vector_search_query \
    //     -- --ignored --exact
    #[tokio::test]
    #[ignore = "requires LUMINA_LIB_PATH"]
    async fn test_lumina_build_then_vector_search_query() {
        let (ctx, catalog, _tmp) = create_empty_vector_search_context().await;
        let identifier = Identifier::new("default", "lumina_build_query_e2e");
        catalog
            .create_table(&identifier, build_lumina_table_schema(), false)
            .await
            .expect("Failed to create table");
        let table = catalog
            .get_table(&identifier)
            .await
            .expect("Failed to load table");

        let write_builder = table
            .new_write_builder()
            .with_commit_user("test-user")
            .expect("Failed to configure write builder");
        let mut table_write = write_builder
            .new_write()
            .expect("Failed to create table write");
        table_write
            .write_arrow_batch(&build_vector_batch(
                vec![0, 1, 2, 3, 4, 5],
                vec![
                    vec![1.0, 0.0],
                    vec![0.9, 0.1],
                    vec![0.0, 1.0],
                    vec![-1.0, 0.0],
                    vec![0.0, -1.0],
                    vec![0.7, 0.3],
                ],
            ))
            .await
            .expect("Failed to write vector batch");
        let messages = table_write
            .prepare_commit()
            .await
            .expect("Failed to prepare commit");
        write_builder
            .new_commit()
            .commit(messages)
            .await
            .expect("Failed to commit vector data");

        ctx.sql("CALL sys.create_lumina_index(table => 'default.lumina_build_query_e2e', index_column => 'embedding')")
            .await
            .expect("Lumina index build SQL should parse")
            .collect()
            .await
            .expect("Lumina index build SQL should execute");

        let index_batches = ctx
            .sql("SELECT index_type, row_count, row_range_start, row_range_end, index_field_name FROM paimon.default.`lumina_build_query_e2e$table_indexes` WHERE index_type = 'lumina'")
            .await
            .expect("index metadata SQL should parse")
            .collect()
            .await
            .expect("index metadata query should execute");
        let index_rows = extract_index_rows(&index_batches);
        assert_eq!(
            index_rows,
            vec![
                ("lumina".to_string(), 3, 0, 2, "embedding".to_string()),
                ("lumina".to_string(), 3, 3, 5, "embedding".to_string()),
            ]
        );

        let search_batches = ctx
            .sql("SELECT id FROM vector_search('paimon.default.lumina_build_query_e2e', 'embedding', '[1.0, 0.0]', 2)")
            .await
            .expect("vector_search SQL should parse")
            .collect()
            .await
            .expect("vector_search query should execute");
        let ids = extract_ids(&search_batches);
        assert_eq!(ids.len(), 2);
        assert!(ids.contains(&0), "exact vector match should be returned");
        assert!(
            ids.iter().any(|id| matches!(id, 1 | 5)),
            "one same-direction neighbor should be returned, got {ids:?}"
        );
    }

    #[tokio::test]
    async fn test_vindex_build_then_vector_search_query() {
        let (ctx, catalog, _tmp) = create_empty_vector_search_context().await;
        let identifier = Identifier::new("default", "vindex_build_query_e2e");
        catalog
            .create_table(&identifier, build_vindex_table_schema(), false)
            .await
            .expect("Failed to create table");
        let table = catalog
            .get_table(&identifier)
            .await
            .expect("Failed to load table");

        let write_builder = table
            .new_write_builder()
            .with_commit_user("test-user")
            .expect("Failed to configure write builder");
        let mut table_write = write_builder
            .new_write()
            .expect("Failed to create table write");
        table_write
            .write_arrow_batch(&build_vector_batch(
                vec![0, 1, 2, 3, 4, 5],
                vec![
                    vec![1.0, 0.0],
                    vec![0.9, 0.1],
                    vec![0.0, 1.0],
                    vec![-1.0, 0.0],
                    vec![0.0, -1.0],
                    vec![0.7, 0.3],
                ],
            ))
            .await
            .expect("Failed to write vector batch");
        let messages = table_write
            .prepare_commit()
            .await
            .expect("Failed to prepare commit");
        write_builder
            .new_commit()
            .commit(messages)
            .await
            .expect("Failed to commit vector data");

        ctx.sql(
            "CALL sys.create_global_index( \
             table => 'default.vindex_build_query_e2e', \
             index_column => 'embedding', \
             index_type => 'ivf-flat', \
             options => 'ivf-flat.dimension=2,ivf-flat.nlist=1,ivf-flat.distance.metric=l2')",
        )
        .await
        .expect("vindex index build SQL should parse")
        .collect()
        .await
        .expect("vindex index build SQL should execute");

        let index_batches = ctx
            .sql("SELECT index_type, row_count, row_range_start, row_range_end, index_field_name FROM paimon.default.`vindex_build_query_e2e$table_indexes` WHERE index_type = 'ivf-flat'")
            .await
            .expect("index metadata SQL should parse")
            .collect()
            .await
            .expect("index metadata query should execute");
        let index_rows = extract_index_rows(&index_batches);
        assert_eq!(
            index_rows,
            vec![
                ("ivf-flat".to_string(), 3, 0, 2, "embedding".to_string()),
                ("ivf-flat".to_string(), 3, 3, 5, "embedding".to_string()),
            ]
        );

        let search_batches = ctx
            .sql("SELECT id FROM vector_search('paimon.default.vindex_build_query_e2e', 'embedding', '[1.0, 0.0]', 2)")
            .await
            .expect("vector_search SQL should parse")
            .collect()
            .await
            .expect("vector_search query should execute");
        let ids = extract_ids(&search_batches);
        assert_eq!(ids, vec![0, 1]);
    }
}

// ======================= Hybrid Search Tests =======================

mod hybrid_search_tests {
    use std::sync::Arc;

    #[cfg(feature = "fulltext")]
    use datafusion::arrow::array::{Array, Int64Array};
    use datafusion::arrow::array::{Float32Array, Int32Array};
    use paimon::catalog::Identifier;
    use paimon::table::BranchManager;
    use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options};
    use paimon_datafusion::SQLContext;

    fn extract_test_warehouse(archive_name: &str) -> (tempfile::TempDir, String) {
        let archive_path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
            .join("testdata")
            .join(archive_name);
        let file = std::fs::File::open(&archive_path)
            .unwrap_or_else(|e| panic!("Failed to open {}: {e}", archive_path.display()));
        let decoder = flate2::read::GzDecoder::new(file);
        let mut archive = tar::Archive::new(decoder);

        let tmp = tempfile::tempdir().expect("Failed to create temp dir");
        let db_dir = tmp.path().join("default.db");
        std::fs::create_dir_all(&db_dir).unwrap();
        archive.unpack(&db_dir).unwrap();

        let warehouse = format!("file://{}", tmp.path().display());
        (tmp, warehouse)
    }

    async fn create_hybrid_search_context(
    ) -> (SQLContext, Arc<FileSystemCatalog>, tempfile::TempDir) {
        let (tmp, warehouse) = extract_test_warehouse("test_java_vindex_vector.tar.gz");
        let mut options = Options::new();
        options.set(CatalogOptions::WAREHOUSE, warehouse);
        let catalog = Arc::new(FileSystemCatalog::new(options).expect("Failed to create catalog"));

        let mut ctx = SQLContext::new();
        ctx.register_catalog("paimon", catalog.clone())
            .await
            .expect("Failed to register catalog");
        (ctx, catalog, tmp)
    }

    fn extract_ids(batches: &[datafusion::arrow::record_batch::RecordBatch]) -> Vec<i32> {
        let mut ids = Vec::new();
        for batch in batches {
            let id_array = batch
                .column_by_name("id")
                .and_then(|c| c.as_any().downcast_ref::<Int32Array>())
                .expect("Expected Int32Array for id");
            for i in 0..batch.num_rows() {
                ids.push(id_array.value(i));
            }
        }
        ids.sort();
        ids
    }

    fn extract_scored_ids(
        batches: &[datafusion::arrow::record_batch::RecordBatch],
    ) -> Vec<(i32, f32)> {
        let mut rows = Vec::new();
        for batch in batches {
            let id_array = batch
                .column_by_name("id")
                .and_then(|column| column.as_any().downcast_ref::<Int32Array>())
                .expect("Expected Int32Array for id");
            let score_array = batch
                .column_by_name("__paimon_search_score")
                .and_then(|column| column.as_any().downcast_ref::<Float32Array>())
                .expect("Expected Float32Array for __paimon_search_score");
            for row in 0..batch.num_rows() {
                rows.push((id_array.value(row), score_array.value(row)));
            }
        }
        rows
    }

    fn extract_scores(batches: &[datafusion::arrow::record_batch::RecordBatch]) -> Vec<f32> {
        let mut scores = Vec::new();
        for batch in batches {
            assert_eq!(batch.num_columns(), 1, "only the score column was selected");
            let score_array = batch
                .column_by_name("__paimon_search_score")
                .and_then(|column| column.as_any().downcast_ref::<Float32Array>())
                .expect("Expected Float32Array for __paimon_search_score");
            scores.extend((0..batch.num_rows()).map(|row| score_array.value(row)));
        }
        scores
    }

    #[tokio::test]
    async fn test_hybrid_search_exposes_ranked_scores() {
        let (ctx, _catalog, _tmp) = create_hybrid_search_context().await;
        let batches = ctx
            .sql(
                "SELECT __paimon_search_score, id FROM hybrid_search( \
                 'paimon.default.test_java_vindex_vector', \
                 array(named_struct( \
                   'field', 'embedding', \
                   'query_vector', array(1.0, 0.0, 0.0, 0.0), \
                   'limit', 3, \
                   'weight', 1.0)), \
                 array(), \
                 3, \
                 'rrf') \
                 ORDER BY __paimon_search_score DESC",
            )
            .await
            .expect("hybrid_search score SQL should parse")
            .collect()
            .await
            .expect("hybrid_search score query should execute");

        let rows = extract_scored_ids(&batches);
        assert_eq!(rows.len(), 3);
        for ((id, score), (expected_id, expected_score)) in
            rows.into_iter()
                .zip([(0, 1.0 / 61.0), (1, 1.0 / 62.0), (2, 1.0 / 63.0)])
        {
            assert_eq!(id, expected_id);
            assert!(
                (score - expected_score).abs() < 1e-6,
                "score for id {id}: expected {expected_score}, got {score}"
            );
        }
    }

    #[tokio::test]
    async fn test_hybrid_search_score_only_and_empty_results() {
        let (ctx, _catalog, _tmp) = create_hybrid_search_context().await;
        let score_batches = ctx
            .sql(
                "SELECT __paimon_search_score FROM hybrid_search( \
                 'paimon.default.test_java_vindex_vector', \
                 array(named_struct( \
                   'field', 'embedding', \
                   'query_vector', array(1.0, 0.0, 0.0, 0.0), \
                   'limit', 3, \
                   'weight', 1.0)), \
                 array(), \
                 3, \
                 'rrf') \
                 ORDER BY __paimon_search_score DESC",
            )
            .await
            .expect("score-only hybrid_search SQL should parse")
            .collect()
            .await
            .expect("score-only hybrid_search query should execute");

        let scores = extract_scores(&score_batches);
        assert_eq!(scores.len(), 3);
        for (score, expected) in scores.into_iter().zip([1.0 / 61.0, 1.0 / 62.0, 1.0 / 63.0]) {
            assert!((score - expected).abs() < 1e-6);
        }

        let empty_batches = ctx
            .sql(
                "SELECT __paimon_search_score FROM hybrid_search( \
                 'paimon.default.test_java_vindex_vector', \
                 array(named_struct( \
                   'field', 'missing_embedding', \
                   'query_vector', array(1.0), \
                   'limit', 3, \
                   'weight', 1.0)), \
                 array(), \
                 3, \
                 'rrf')",
            )
            .await
            .expect("empty hybrid_search score SQL should parse")
            .collect()
            .await
            .expect("empty hybrid_search score query should execute");
        assert_eq!(
            empty_batches
                .iter()
                .map(|batch| batch.num_rows())
                .sum::<usize>(),
            0
        );
    }

    #[cfg(feature = "fulltext")]
    #[tokio::test]
    async fn test_hybrid_search_score_with_row_tracking_only_fulltext() {
        let (_tmp, ctx) = super::common::setup_sql_context().await;
        super::common::exec(
            &ctx,
            "CREATE TABLE paimon.test_db.hybrid_raw_fulltext (id INT, content STRING) \
             WITH ( \
               'row-tracking.enabled' = 'true', \
               'global-index.search-mode' = 'full')",
        )
        .await;
        super::common::exec(
            &ctx,
            "INSERT INTO paimon.test_db.hybrid_raw_fulltext VALUES \
             (1, 'paimon search'), (2, 'other'), (3, 'paimon table')",
        )
        .await;

        const SEARCH: &str = "hybrid_search( \
             'paimon.test_db.hybrid_raw_fulltext', \
             array(), \
             array(named_struct( \
               'column', 'content', \
               'query', 'paimon', \
               'limit', 10, \
               'weight', 1.0)), \
             10, \
             'rrf')";
        let explicit_sql = format!("SELECT id, __paimon_search_score FROM {SEARCH} ORDER BY id");
        let explicit_batches = ctx
            .sql(&explicit_sql)
            .await
            .expect("row-tracking-only score SQL should parse")
            .collect()
            .await
            .expect("row-tracking-only score query should execute");

        let mut ids = Vec::new();
        for batch in &explicit_batches {
            let id_array = batch
                .column_by_name("id")
                .and_then(|column| column.as_any().downcast_ref::<Int32Array>())
                .expect("id column");
            let score_array = batch
                .column_by_name("__paimon_search_score")
                .and_then(|column| column.as_any().downcast_ref::<Float32Array>())
                .expect("score column");
            assert_eq!(score_array.null_count(), 0);
            for row in 0..batch.num_rows() {
                ids.push(id_array.value(row));
                assert!(score_array.value(row) > 0.0);
            }
        }
        assert_eq!(ids, vec![1, 3]);

        let limited_sql = format!("SELECT id, __paimon_search_score FROM {SEARCH} LIMIT 2");
        let limited_batches = ctx
            .sql(&limited_sql)
            .await
            .expect("row-tracking-only limited SQL should parse")
            .collect()
            .await
            .expect("row-tracking-only limited query should execute");
        assert_eq!(extract_scored_ids(&limited_batches).len(), 2);

        let projected_sql = format!("SELECT id FROM {SEARCH} ORDER BY id");
        let projected_batches = ctx
            .sql(&projected_sql)
            .await
            .expect("row-tracking-only projected SQL should parse")
            .collect()
            .await
            .expect("row-tracking-only projected query should execute");
        assert_eq!(extract_ids(&projected_batches), vec![1, 3]);

        let count_sql = format!("SELECT COUNT(*) AS matches FROM {SEARCH}");
        let count_batches = ctx
            .sql(&count_sql)
            .await
            .expect("row-tracking-only aggregate SQL should parse")
            .collect()
            .await
            .expect("row-tracking-only aggregate query should execute");
        let matches = count_batches
            .first()
            .and_then(|batch| batch.column_by_name("matches"))
            .and_then(|column| column.as_any().downcast_ref::<Int64Array>())
            .expect("matches column");
        assert_eq!(matches.value(0), 2);

        let star_sql = format!("SELECT * FROM {SEARCH}");
        let star_batches = ctx
            .sql(&star_sql)
            .await
            .expect("row-tracking-only SELECT * SQL should parse")
            .collect()
            .await
            .expect("row-tracking-only SELECT * query should execute");
        assert_eq!(
            star_batches
                .iter()
                .map(|batch| batch.num_rows())
                .sum::<usize>(),
            2
        );
        for batch in &star_batches {
            assert_eq!(batch.num_columns(), 3);
            assert!(batch.column_by_name("__paimon_search_score").is_some());
            assert!(batch.column_by_name("_ROW_ID").is_none());
        }
    }

    #[tokio::test]
    async fn test_hybrid_search_multiple_vector_routes_spark_shape() {
        let (ctx, _catalog, _tmp) = create_hybrid_search_context().await;
        let batches = ctx
            .sql(
                "SELECT id FROM hybrid_search( \
                 'paimon.default.test_java_vindex_vector', \
                 array(named_struct( \
                   'field', 'embedding', \
                   'query_vector', array(1.0, 0.0, 0.0, 0.0), \
                   'limit', 3, \
                   'weight', 1.0), \
                 named_struct( \
                   'field', 'embedding', \
                   'query_vector', array(0.9, 0.1, 0.0, 0.0), \
                   'limit', 3, \
                   'weight', 1.0)), \
                 array(), \
                 3, \
                 'rrf')",
            )
            .await
            .expect("hybrid_search SQL should parse")
            .collect()
            .await
            .expect("hybrid_search query should execute");

        assert_eq!(extract_ids(&batches), vec![0, 1, 2]);
    }

    #[tokio::test]
    async fn test_hybrid_search_branch() {
        let (ctx, catalog, _tmp) = create_hybrid_search_context().await;
        let identifier = Identifier::new("default", "test_java_vindex_vector");
        let table = catalog.get_table(&identifier).await.expect("load table");
        let snapshot = table
            .snapshot_manager()
            .get_latest_snapshot()
            .await
            .expect("load latest snapshot")
            .expect("latest snapshot");
        table
            .tag_manager()
            .create("branch-source", &snapshot)
            .await
            .expect("create tag");
        BranchManager::new(table.file_io().clone(), table.location().to_string())
            .create_branch_from_tag("b1", "branch-source")
            .await
            .expect("create branch");

        let batches = ctx
            .sql(
                "SELECT id FROM hybrid_search( \
                 'paimon.default.test_java_vindex_vector$branch_b1', \
                 array(named_struct( \
                   'field', 'embedding', \
                   'query_vector', array(1.0, 0.0, 0.0, 0.0), \
                   'limit', 3, \
                   'weight', 1.0)), \
                 array(), \
                 3, \
                 'rrf')",
            )
            .await
            .expect("hybrid_search SQL should parse")
            .collect()
            .await
            .expect("branch query should execute");

        assert_eq!(extract_ids(&batches), vec![0, 1, 2]);
    }
}