""" 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()