uqa-engine 0.2.3

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! CREATE INDEX execution and index option validation.

use super::{
    ddl_storage_error, ColumnType, CreateIndex, Engine, HNSWIndexParams, IVFIndexParams, SQLError,
    SQLResult, VectorIndexSpec,
};

mod expressions;
mod unique;

pub(in crate::sql) fn run_create_index(
    engine: &Engine,
    mut c: CreateIndex,
) -> Result<SQLResult, SQLError> {
    c.table = engine
        .try_resolve_index_table_name(&c.table)?
        .ok_or_else(|| SQLError::UnknownTable(c.table.clone()))?;
    engine.ensure_table_owner(&c.table)?;
    engine.ensure_existing_relation_creation_privilege(&c.table)?;
    // Every accepted access method has a matching physical implementation.
    // Reject unknown methods before allocating a name or mutating any table,
    // index, analyzer, or catalog state.
    let am = c.access_method.to_ascii_lowercase();
    if !matches!(am.as_str(), "" | "btree" | "gin" | "ivf" | "hnsw") {
        return Err(SQLError::Unsupported(format!(
            "CREATE INDEX access method `{}` is not supported",
            c.access_method
        )));
    }

    let table_relation = crate::RelationIdentity::from_legacy_name(&c.table)
        .map_err(|error| SQLError::Internal(format!("resolve index table: {error}")))?;
    let name = if let Some(name) = c.name.as_ref() {
        name.clone()
    } else {
        allocate_default_index_name(engine, &table_relation, &c.columns)?
    };
    let relation = crate::RelationIdentity::new(&table_relation.schema, &name);
    if matches!(
        engine.resolve_bound_relation_kind(&relation.qualified_name())?,
        crate::engine_capabilities::RelationResolution::Found(_, _)
    ) {
        if c.if_not_exists {
            engine.push_sql_notice(
                "NOTICE",
                &format!("relation \"{name}\" already exists, skipping"),
            );
            return Ok(SQLResult::empty());
        }
        return Err(SQLError::Routine {
            sqlstate: "42P07".into(),
            message: format!("relation \"{name}\" already exists"),
        });
    }

    let attribute_keys = c
        .columns
        .iter()
        .cloned()
        .chain(
            c.included_columns
                .iter()
                .cloned()
                .map(uqa_sql::ast::IndexKey::Column),
        )
        .collect::<Vec<_>>();
    let key_names = expressions::key_names(&attribute_keys);
    let key_types = expressions::prepare_index_keys(engine, &mut c)?;
    if let Some(predicate) = c.predicate.as_deref_mut() {
        crate::sql::generated::prepare_index_predicate(engine, &c.table, predicate)?;
    }
    unique::validate_unique_index(engine, &c, &name)?;

    match am.as_str() {
        "gin" => {
            for col in &c.columns {
                let analyzer = c
                    .options
                    .iter()
                    .find(|(k, _)| k.eq_ignore_ascii_case("analyzer"))
                    .map(|(_, v)| v.as_str());
                let column = expressions::require_column_key(col, "gin")?;
                if let Err(e) =
                    engine.add_fts_field_with_analyzer(&c.table, column.to_string(), analyzer)
                {
                    return Err(SQLError::Internal(format!("add_fts_field: {e}")));
                }
            }
        }
        "" | "btree" => {}
        "ivf" | "hnsw" => create_vector_index(engine, &c, &am)?,
        _ => unreachable!("access method was validated above"),
    }
    // Persist the CREATE INDEX statement itself so reopen sees the
    // same set of registered indexes. The engine layer parses
    // `parameters_json` back into `(key, value)` pairs and re-runs
    // any access-method-specific side effects (e.g. add_fts_field
    // for `gin`) on restore.
    let catalog_index_type = if am.is_empty() { "btree" } else { &am };
    engine
        .register_catalog_index_definition(
            &relation.qualified_name(),
            catalog_index_type,
            &c.table,
            &c.columns,
            &c.options,
            &crate::engine_catalog_indexes::IndexDefinition {
                key_names,
                key_types,
                included_columns: c.included_columns.clone(),
                column_order: c.column_order.clone(),
                predicate: c.predicate.clone(),
                unique: c.unique,
                nulls_not_distinct: c.nulls_not_distinct,
            },
        )
        .map_err(|e| ddl_storage_error("CREATE INDEX", e))?;
    Ok(SQLResult::empty())
}

fn create_vector_index(
    engine: &Engine,
    statement: &CreateIndex,
    access_method: &str,
) -> Result<(), SQLError> {
    let spec = match access_method {
        "ivf" => VectorIndexSpec::IVF(parse_ivf_index_params(&statement.options)?),
        "hnsw" => VectorIndexSpec::HNSW(parse_hnsw_index_params(&statement.options)?),
        _ => unreachable!("vector access method was validated above"),
    };
    let table = engine
        .try_resolve_table_name(&statement.table)
        .map_err(|err| ddl_storage_error("CREATE INDEX", err))?
        .ok_or_else(|| {
            SQLError::Unsupported(format!(
                "CREATE INDEX USING {access_method}: relation `{}` does not exist",
                statement.table
            ))
        })?;
    let mut fields = Vec::with_capacity(statement.columns.len());
    for key in &statement.columns {
        let column = expressions::require_column_key(key, access_method)?;
        let dimensions = match engine
            .column_type(&table, column)
            .map_err(|err| ddl_storage_error("CREATE INDEX", err))?
        {
            Some(ColumnType::Vector(dim) | ColumnType::Tensor(dim)) => dim,
            Some(other) => {
                return Err(SQLError::Unsupported(format!(
                    "CREATE INDEX USING {access_method} requires VECTOR or TENSOR column `{column}`, got {other:?}"
                )));
            }
            None => {
                return Err(SQLError::Unsupported(format!(
                    "CREATE INDEX USING {access_method}: column `{table}`.`{column}` does not exist"
                )));
            }
        };
        let existing = engine
            .vector_catalog_index_names_for_column(&table, column)
            .map_err(|err| ddl_storage_error("CREATE INDEX", err))?;
        if !existing.is_empty() {
            return Err(SQLError::Unsupported(format!(
                "CREATE INDEX USING {access_method}: `{table}`.`{column}` already has physical vector index `{}`",
                existing.join("`, `")
            )));
        }
        fields.push((column, dimensions));
    }
    for (column, dimensions) in fields {
        if !engine
            .rebuild_vector_field_with_spec(&table, column, dimensions, spec)
            .map_err(|err| ddl_storage_error("CREATE INDEX vector field", err))?
        {
            return Err(SQLError::Unsupported(format!(
                "CREATE INDEX USING {access_method}: relation `{table}` does not exist"
            )));
        }
    }
    Ok(())
}

fn allocate_default_index_name(
    engine: &Engine,
    table: &crate::RelationIdentity,
    columns: &[uqa_sql::ast::IndexKey],
) -> Result<String, SQLError> {
    fn component(raw: &str) -> String {
        let mut out = String::with_capacity(raw.len());
        let mut previous_was_separator = false;
        for ch in raw.chars() {
            if ch.is_alphanumeric() || ch == '_' {
                out.extend(ch.to_lowercase());
                previous_was_separator = false;
            } else if !previous_was_separator && !out.is_empty() {
                out.push('_');
                previous_was_separator = true;
            }
        }
        while out.ends_with('_') {
            out.pop();
        }
        out
    }

    let mut parts = std::iter::once(component(&table.name))
        .chain(
            expressions::key_names(columns)
                .iter()
                .map(|column| component(column)),
        )
        .filter(|part| !part.is_empty())
        .collect::<Vec<_>>();
    if parts.is_empty() {
        parts.push("index".to_string());
    }
    let base = format!("{}_idx", parts.join("_"));
    let available = |name: &str| -> Result<bool, SQLError> {
        let candidate = crate::RelationIdentity::new(&table.schema, name).qualified_name();
        Ok(matches!(
            engine.resolve_bound_relation_kind(&candidate)?,
            crate::engine_capabilities::RelationResolution::MissingRelation
        ))
    };
    if available(&base)? {
        return Ok(base);
    }
    for suffix in 1_u64.. {
        let candidate = format!("{base}{suffix}");
        if available(&candidate)? {
            return Ok(candidate);
        }
    }
    unreachable!("u64 index-name suffix space is non-empty")
}

fn parse_ivf_index_params(options: &[(String, String)]) -> Result<IVFIndexParams, SQLError> {
    let mut params = IVFIndexParams::default();
    let mut seen = std::collections::BTreeSet::new();
    for (key, value) in options {
        if key.eq_ignore_ascii_case("lists") || key.eq_ignore_ascii_case("nlist") {
            claim_index_option(&mut seen, "nlist", "ivf", key)?;
            params.nlist = parse_positive_usize_option("ivf", key, value)?;
        } else if key.eq_ignore_ascii_case("probes") || key.eq_ignore_ascii_case("nprobe") {
            claim_index_option(&mut seen, "nprobe", "ivf", key)?;
            params.nprobe = parse_positive_usize_option("ivf", key, value)?;
        } else if key.eq_ignore_ascii_case("train_threshold")
            || key.eq_ignore_ascii_case("train-threshold")
            || key.eq_ignore_ascii_case("min_train")
        {
            claim_index_option(&mut seen, "train_threshold", "ivf", key)?;
            params.train_threshold = parse_positive_usize_option("ivf", key, value)?;
        } else {
            return Err(SQLError::Unsupported(format!(
                "CREATE INDEX USING ivf option `{key}` is not supported"
            )));
        }
    }
    Ok(params)
}

fn parse_hnsw_index_params(options: &[(String, String)]) -> Result<HNSWIndexParams, SQLError> {
    let mut params = HNSWIndexParams::default();
    let mut seen = std::collections::BTreeSet::new();
    for (key, value) in options {
        if key.eq_ignore_ascii_case("m") {
            claim_index_option(&mut seen, "m", "hnsw", key)?;
            params.m = parse_positive_usize_option("hnsw", key, value)?;
        } else if key.eq_ignore_ascii_case("ef_construction")
            || key.eq_ignore_ascii_case("ef-construction")
        {
            claim_index_option(&mut seen, "ef_construction", "hnsw", key)?;
            params.ef_construction = parse_positive_usize_option("hnsw", key, value)?;
        } else if key.eq_ignore_ascii_case("ef_search") || key.eq_ignore_ascii_case("ef-search") {
            claim_index_option(&mut seen, "ef_search", "hnsw", key)?;
            params.ef_search = parse_positive_usize_option("hnsw", key, value)?;
        } else if key.eq_ignore_ascii_case("rebuild_threshold")
            || key.eq_ignore_ascii_case("rebuild-threshold")
        {
            claim_index_option(&mut seen, "rebuild_threshold", "hnsw", key)?;
            params.rebuild_threshold = parse_positive_usize_option("hnsw", key, value)?;
        } else if key.eq_ignore_ascii_case("seed") {
            claim_index_option(&mut seen, "seed", "hnsw", key)?;
            params.seed = value.parse::<u64>().map_err(|_| {
                SQLError::TypeMismatch(format!(
                    "CREATE INDEX USING hnsw option `{key}` must be an unsigned integer"
                ))
            })?;
        } else {
            return Err(SQLError::Unsupported(format!(
                "CREATE INDEX USING hnsw option `{key}` is not supported"
            )));
        }
    }
    params
        .validate()
        .map_err(|error| SQLError::TypeMismatch(format!("CREATE INDEX USING hnsw: {error}")))
}

fn claim_index_option(
    seen: &mut std::collections::BTreeSet<&'static str>,
    canonical: &'static str,
    access_method: &str,
    source: &str,
) -> Result<(), SQLError> {
    if !seen.insert(canonical) {
        return Err(SQLError::Unsupported(format!(
            "CREATE INDEX USING {access_method} option `{source}` duplicates `{canonical}`"
        )));
    }
    Ok(())
}

fn parse_positive_usize_option(
    access_method: &str,
    key: &str,
    value: &str,
) -> Result<usize, SQLError> {
    let parsed = value.parse::<usize>().map_err(|_| {
        SQLError::TypeMismatch(format!(
            "CREATE INDEX USING {access_method} option `{key}` must be a positive integer"
        ))
    })?;
    if parsed == 0 {
        return Err(SQLError::TypeMismatch(format!(
            "CREATE INDEX USING {access_method} option `{key}` must be a positive integer"
        )));
    }
    Ok(parsed)
}