File size: 5,547 Bytes
388137e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """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]] # list of (table_index, column_name)
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"
# Load train and dev examples
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))
# Databases in train and dev
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)
# Verify every db_id has a SQLite file
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)
# Test gold SQL execution
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:
# Read-only connection
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)
# Calculate hardness distribution for dev if available
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))
|