| """Spider dataset loader and validator.""" |
|
|
| import json |
| import sqlite3 |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any, Dict, List, Optional, Set, Tuple |
|
|
|
|
| @dataclass |
| class SpiderExample: |
| question: str |
| db_id: str |
| query: str |
| hardness: Optional[str] = None |
|
|
|
|
| @dataclass |
| class SpiderTableSchema: |
| db_id: str |
| table_names: List[str] |
| column_names: List[Tuple[int, str]] |
| column_types: List[str] |
| primary_keys: List[int] |
| foreign_keys: List[Tuple[int, int]] |
|
|
|
|
| def load_spider_examples(json_path: str) -> List[SpiderExample]: |
| path = Path(json_path) |
| if not path.exists(): |
| raise FileNotFoundError(f"Spider json file not found: {json_path}") |
|
|
| with open(path, "r", encoding="utf-8") as f: |
| data = json.load(f) |
|
|
| examples: List[SpiderExample] = [] |
| for item in data: |
| ex = SpiderExample( |
| question=item["question"], |
| db_id=item["db_id"], |
| query=item["query"], |
| hardness=item.get("hardness"), |
| ) |
| examples.append(ex) |
|
|
| return examples |
|
|
|
|
| def load_tables_json(tables_path: str) -> Dict[str, SpiderTableSchema]: |
| path = Path(tables_path) |
| if not path.exists(): |
| raise FileNotFoundError(f"Tables json file not found: {tables_path}") |
|
|
| with open(path, "r", encoding="utf-8") as f: |
| data = json.load(f) |
|
|
| tables: Dict[str, SpiderTableSchema] = {} |
| for item in data: |
| schema = SpiderTableSchema( |
| db_id=item["db_id"], |
| table_names=item["table_names_original"], |
| column_names=[(col[0], col[1]) for col in item["column_names_original"]], |
| column_types=item["column_types"], |
| primary_keys=item["primary_keys"], |
| foreign_keys=[(fk[0], fk[1]) for fk in item["foreign_keys"]], |
| ) |
| tables[schema.db_id] = schema |
|
|
| return tables |
|
|
|
|
| def validate_dataset( |
| data_dir: str = "data/raw/spider", reports_dir: str = "reports" |
| ) -> Dict[str, Any]: |
| raw_path = Path(data_dir) |
| rep_path = Path(reports_dir) |
| rep_path.mkdir(parents=True, exist_ok=True) |
|
|
| spider_train_path = raw_path / "train_spider.json" |
| others_train_path = raw_path / "train_others.json" |
| dev_path = raw_path / "dev.json" |
| db_dir = raw_path / "database" |
|
|
| |
| train_spider_exs = load_spider_examples(str(spider_train_path)) |
| train_others_exs = ( |
| load_spider_examples(str(others_train_path)) if others_train_path.exists() else [] |
| ) |
| full_train_exs = train_spider_exs + train_others_exs |
| dev_exs = load_spider_examples(str(dev_path)) |
|
|
| |
| train_dbs: Set[str] = {ex.db_id for ex in full_train_exs} |
| dev_dbs: Set[str] = {ex.db_id for ex in dev_exs} |
| db_overlap = train_dbs.intersection(dev_dbs) |
|
|
| |
| missing_dbs: List[str] = [] |
| all_dbs = train_dbs.union(dev_dbs) |
| for db_id in all_dbs: |
| db_file = db_dir / db_id / f"{db_id}.sqlite" |
| if not db_file.exists(): |
| missing_dbs.append(db_id) |
|
|
| |
| def test_sql_execution(examples: List[SpiderExample]) -> Tuple[int, int, List[Dict[str, str]]]: |
| passed = 0 |
| failed = 0 |
| failures: List[Dict[str, str]] = [] |
| for ex in examples: |
| db_file = db_dir / ex.db_id / f"{ex.db_id}.sqlite" |
| if not db_file.exists(): |
| failed += 1 |
| failures.append( |
| {"db_id": ex.db_id, "query": ex.query, "error": "Database file missing"} |
| ) |
| continue |
|
|
| try: |
| |
| conn = sqlite3.connect(f"file:{db_file}?mode=ro", uri=True, timeout=5.0) |
| cursor = conn.cursor() |
| cursor.execute(ex.query) |
| _ = cursor.fetchall() |
| conn.close() |
| passed += 1 |
| except Exception as e: |
| failed += 1 |
| failures.append({"db_id": ex.db_id, "query": ex.query, "error": str(e)}) |
|
|
| return passed, failed, failures |
|
|
| train_passed, train_failed, train_failures = test_sql_execution(full_train_exs) |
| dev_passed, dev_failed, dev_failures = test_sql_execution(dev_exs) |
|
|
| |
| dev_hardness: Dict[str, int] = {} |
| for ex in dev_exs: |
| if ex.hardness: |
| dev_hardness[ex.hardness] = dev_hardness.get(ex.hardness, 0) + 1 |
|
|
| profile = { |
| "train_spider_questions": len(train_spider_exs), |
| "train_others_questions": len(train_others_exs), |
| "train_total_questions": len(full_train_exs), |
| "dev_questions": len(dev_exs), |
| "train_databases_count": len(train_dbs), |
| "dev_databases_count": len(dev_dbs), |
| "db_overlap_count": len(db_overlap), |
| "db_overlap_list": list(db_overlap), |
| "missing_database_files": missing_dbs, |
| "train_gold_sql_passed": train_passed, |
| "train_gold_sql_failed": train_failed, |
| "dev_gold_sql_passed": dev_passed, |
| "dev_gold_sql_failed": dev_failed, |
| "dev_hardness_distribution": dev_hardness, |
| } |
|
|
| output_json = rep_path / "data_profile.json" |
| with open(output_json, "w", encoding="utf-8") as f: |
| json.dump(profile, f, indent=2) |
|
|
| return profile |
|
|
|
|
| if __name__ == "__main__": |
| profile = validate_dataset() |
| print("Data Profile:") |
| print(json.dumps(profile, indent=2)) |
|
|
