src/process_spider.py

"""
Spider Data Processing Script v5
Build schema-aware inputs for Spider to fine-tune T5/CodeT5/BART.

Input:  spider_data/{train_spider,dev,test}.json + tables.json + test_tables.json
Output: processed/spider/{train,dev,test}.json

Usage:
    python src/process_spider.py
    python src/process_spider.py --no_foreign_keys
    python src/process_spider.py --no_column_types
    python src/process_spider.py --no_primary_keys
    python src/process_spider.py --no_lowercase_identifiers
    python src/process_spider.py --oracle_schema_link --output_dir processed/spider_oracle
        (RQ6 ablation: filter schema to keep only tables used in gold SQL — ceiling-effect
        study for WRONG_TABLE error; not applied to test split which has no gold SQL)
"""

import re
import json
import argparse
import sys
from pathlib import Path

try:
    sys.stdout.reconfigure(encoding='utf-8')
except Exception:
    pass

try:
    import sqlglot
    import sqlglot.expressions as exp
    HAS_SQLGLOT = True
except ImportError:
    HAS_SQLGLOT = False
    print("WARNING: sqlglot not installed. Run: pip install sqlglot")


# ── SQL utilities ──────────────────────────────────────────────────────────────

def normalize_sql(sql):
    """Normalize SQL for Exact Match — preserve string literals (Spider DB is mixed-case)."""
    if not sql:
        return ""
    if HAS_SQLGLOT:
        try:
            return sqlglot.parse_one(sql, read='sqlite').sql(
                dialect='sqlite', pretty=False
            ).strip()
        except Exception:
            pass
    sql = sql.strip()
    sql = re.sub(r'\s+', ' ', sql)
    sql = re.sub(r'\s*,\s*', ' , ', sql)
    sql = re.sub(r'\s*(=|!=|<>|>=|<=|>|<)\s*', r' \1 ', sql)
    return sql.strip()


def normalize_question(question):
    return re.sub(r'\s+', ' ', question.strip()).rstrip('.?!').strip()


def normalize_identifier(name, lowercase=True):
    return name.strip().lower() if lowercase else name.strip()


def normalize_primary_keys(primary_keys):
    pk_indices = set()
    for pk in primary_keys:
        if isinstance(pk, list):
            pk_indices.update(pk)
        else:
            pk_indices.add(pk)
    return pk_indices


def check_sql_validity(sql):
    if not sql or not sql.strip():
        return False
    if HAS_SQLGLOT:
        try:
            sqlglot.parse_one(sql, read='sqlite')
            return True
        except Exception:
            return False
    return sql.strip().upper().startswith('SELECT')


# ── Schema builder ─────────────────────────────────────────────────────────────

def build_schema_string(db, use_foreign_keys=True, use_column_types=True,
                        use_primary_keys=True, lowercase_identifiers=True):
    table_names = [normalize_identifier(n, lowercase_identifiers)
                   for n in db['table_names_original']]
    col_names   = [(ti, normalize_identifier(cn, lowercase_identifiers))
                   for ti, cn in db['column_names_original']]
    col_types    = db.get('column_types', [])
    primary_keys = normalize_primary_keys(db.get('primary_keys', []))

    table_cols = {i: [] for i in range(len(table_names))}
    for col_idx, (ti, cn) in enumerate(col_names):
        if ti == -1:
            continue
        col = cn
        if use_primary_keys and col_idx in primary_keys:
            col += "*"
        if use_column_types and col_idx < len(col_types):
            col += f":{col_types[col_idx]}"
        table_cols[ti].append(col)

    parts = []
    for i, tname in enumerate(table_names):
        cols = table_cols.get(i, [])
        parts.append(f"{tname} ( {' , '.join(cols)} )" if cols else tname)
    schema = " | ".join(parts)

    if use_foreign_keys and db.get('foreign_keys'):
        fk_parts = []
        for c1, c2 in db['foreign_keys']:
            t1 = table_names[col_names[c1][0]]
            t2 = table_names[col_names[c2][0]]
            fk_parts.append(f"{t1}.{col_names[c1][1]} = {t2}.{col_names[c2][1]}")
        if fk_parts:
            schema += " | foreign_keys: " + " , ".join(fk_parts)

    return schema


def build_model_input(question, db, use_foreign_keys=True, use_column_types=True,
                      use_primary_keys=True, lowercase_identifiers=True):
    schema = build_schema_string(db, use_foreign_keys, use_column_types,
                                  use_primary_keys, lowercase_identifiers)
    return f"question: {normalize_question(question)} | schema: {schema}"


# ── Oracle schema-linking (RQ6 ablation) ────────────────────────────────────────
# Ceiling-effect study: use gold SQL to filter schema to only tables actually used.
# Measures WRONG_TABLE error reduction under perfect schema-linking.
# NOT a real inference system (requires gold SQL in advance) — train/dev only.

def extract_gold_tables(sql, table_names_original):
    """Extract table names used in gold SQL via sqlglot AST (handles T1/T2 aliases
    and nested subqueries). Falls back to all original tables on parse error or no
    match — safe, never produces an empty schema."""
    valid = {t.lower() for t in table_names_original}
    if not HAS_SQLGLOT:
        return valid
    try:
        tree = sqlglot.parse_one(sql, read='sqlite')
        used = {t.name.lower() for t in tree.find_all(exp.Table)}
    except Exception:
        return valid
    matched = used & valid
    return matched if matched else valid


def filter_db_to_tables(db, keep_tables_lower):
    """Filter db dict to keep only tables in keep_tables_lower, remapping all
    column/table indices so build_schema_string() works unchanged.
    Foreign keys are kept only when both columns belong to retained tables."""
    table_names = db['table_names_original']
    keep_indices = [i for i, t in enumerate(table_names) if t.lower() in keep_tables_lower]
    if not keep_indices:
        return db  # an toan: khong loc duoc gi thi giu nguyen schema goc

    index_map = {old: new for new, old in enumerate(keep_indices)}
    new_table_names = [table_names[i] for i in keep_indices]

    col_names    = db['column_names_original']
    col_types    = db.get('column_types', [])
    new_col_names, new_col_types, col_index_map = [], [], {}
    for old_idx, (ti, cn) in enumerate(col_names):
        if ti == -1 or ti in index_map:
            new_ti = -1 if ti == -1 else index_map[ti]
            col_index_map[old_idx] = len(new_col_names)
            new_col_names.append([new_ti, cn])
            if old_idx < len(col_types):
                new_col_types.append(col_types[old_idx])

    new_primary_keys = []
    for pk in db.get('primary_keys', []):
        if isinstance(pk, list):
            filtered = [col_index_map[c] for c in pk if c in col_index_map]
            if filtered:
                new_primary_keys.append(filtered if len(filtered) > 1 else filtered[0])
        elif pk in col_index_map:
            new_primary_keys.append(col_index_map[pk])

    new_foreign_keys = [
        [col_index_map[c1], col_index_map[c2]]
        for c1, c2 in db.get('foreign_keys', [])
        if c1 in col_index_map and c2 in col_index_map
    ]

    return {
        **db,
        'table_names_original':  new_table_names,
        'column_names_original': new_col_names,
        'column_types':          new_col_types,
        'primary_keys':          new_primary_keys,
        'foreign_keys':          new_foreign_keys,
    }


# ── Data loading ───────────────────────────────────────────────────────────────

def load_schemas(tables_path):
    with open(tables_path, 'r', encoding='utf-8') as f:
        tables = json.load(f)
    schemas = {db['db_id']: db for db in tables}
    print(f"  Loaded {len(schemas):,} schemas from {Path(tables_path).name}")
    return schemas


# ── Process split ──────────────────────────────────────────────────────────────

def process_split(data_path, schemas, split_name,
                  use_foreign_keys=True, use_column_types=True,
                  use_primary_keys=True, lowercase_identifiers=True,
                  oracle_schema_link=False):
    print(f"\n[{split_name.upper()}] Processing...")
    with open(data_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    dataset, errors, not_found, invalid_sql = [], 0, 0, 0
    is_test = (split_name == 'test')
    # Oracle filtering only applies when gold SQL is available (train/dev).
    # Test split has no gold SQL — skip.
    apply_oracle = oracle_schema_link and not is_test
    tables_before, tables_after = [], []

    for idx, item in enumerate(data):
        try:
            db_id    = item['db_id']
            question = normalize_question(item['question'])
            if db_id not in schemas:
                not_found += 1
                continue

            db  = schemas[db_id]
            sql = item.get('query', item.get('SQL', '')).strip() if not is_test else ''

            db_for_schema = db
            if apply_oracle and sql:
                keep_tables = extract_gold_tables(sql, db['table_names_original'])
                db_for_schema = filter_db_to_tables(db, keep_tables)
                tables_before.append(len(db['table_names_original']))
                tables_after.append(len(db_for_schema['table_names_original']))

            record = {
                'input':    build_model_input(question, db_for_schema, use_foreign_keys,
                                               use_column_types, use_primary_keys,
                                               lowercase_identifiers),
                'question': question,
                'db_id':    db_id,
            }

            if not is_test:
                normalized_sql = normalize_sql(sql)
                valid          = check_sql_validity(sql)
                if not valid:
                    invalid_sql += 1
                record['output']            = normalized_sql
                record['raw_output']        = sql
                record['normalized_output'] = normalized_sql
                record['is_valid_sql']      = valid
                if 'query_toks' in item:
                    record['query_toks'] = item['query_toks']

            dataset.append(record)
        except Exception as e:
            errors += 1
            if errors <= 5:
                print(f"  WARNING idx {idx}: {e}")

    print(f"  Processed: {len(dataset):,} | Not found: {not_found} | "
          f"Errors: {errors} | Invalid SQL: {invalid_sql}")
    if apply_oracle and tables_before:
        avg_before = sum(tables_before) / len(tables_before)
        avg_after  = sum(tables_after) / len(tables_after)
        reduction  = (1 - avg_after / avg_before) * 100 if avg_before else 0
        print(f"  [oracle_schema_link] Avg tables/schema: {avg_before:.2f} -> "
              f"{avg_after:.2f}  (-{reduction:.1f}%)")
    return dataset


# ── Main ───────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description='Process Spider v5')
    parser.add_argument('--data_dir',                  default=r'C:\Users\ADMIN\Documents\Text  to SQL\spider_data')
    parser.add_argument('--output_dir',                default=r'C:\Users\ADMIN\Documents\Text  to SQL\processed\spider')
    parser.add_argument('--no_foreign_keys',           action='store_true', default=False)
    parser.add_argument('--no_column_types',           action='store_true', default=False)
    parser.add_argument('--no_primary_keys',           action='store_true', default=False)
    parser.add_argument('--no_lowercase_identifiers',  action='store_true', default=False)
    parser.add_argument('--oracle_schema_link',        action='store_true', default=False,
                        help='RQ6 ablation: filter schema to keep only tables in gold SQL '
                             '(ceiling-effect study for WRONG_TABLE error). Not applied to '
                             'test split (no gold SQL available).')
    args = parser.parse_args()

    data_dir   = Path(args.data_dir)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    use_fk                = not args.no_foreign_keys
    use_column_types      = not args.no_column_types
    use_primary_keys      = not args.no_primary_keys
    lowercase_identifiers = not args.no_lowercase_identifiers
    oracle_schema_link    = args.oracle_schema_link

    schemas      = load_schemas(str(data_dir / 'tables.json'))
    test_schemas = (load_schemas(str(data_dir / 'test_tables.json'))
                    if (data_dir / 'test_tables.json').exists() else schemas)

    splits = [
        ('train', data_dir / 'train_spider.json', schemas),
        ('dev',   data_dir / 'dev.json',          schemas),
        ('test',  data_dir / 'test.json',          test_schemas),
    ]
    for split_name, fpath, split_schemas in splits:
        if not fpath.exists():
            print(f"WARNING: {fpath} not found, skipping")
            continue
        dataset = process_split(str(fpath), split_schemas, split_name,
                                use_fk, use_column_types,
                                use_primary_keys, lowercase_identifiers,
                                oracle_schema_link)
        out = output_dir / f"{split_name}.json"
        with open(out, 'w', encoding='utf-8') as f:
            json.dump(dataset, f, ensure_ascii=False, indent=2)
        size = out.stat().st_size / 1e6
        print(f"  Saved -> {out}  ({len(dataset):,} samples, {size:.1f} MB)")


if __name__ == '__main__':
    main()