Spaces:
Running on Zero
Running on Zero
File size: 4,062 Bytes
ac8ed28 | 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 | import re
from typing import List
_FROM_TABLE_RE = re.compile(r'\bFROM\s+table\b', flags=re.IGNORECASE)
_AGG_RE = re.compile(r'\b(COUNT|MAX|MIN|AVG|SUM)\s+([A-Za-z0-9_\s`"]+?)(?=\sFROM|\sWHERE|,|$)', flags=re.IGNORECASE)
_WHERE_RE = re.compile(r'\bWHERE\b\s+(.*)', flags=re.IGNORECASE)
_AND_OR_RE = re.compile(r'\s+\b(AND|OR)\b\s+', flags=re.IGNORECASE)
_COND_RE = re.compile(r'(.*?)\s*(=|>|<|>=|<=|!=)\s*(.*)')
def replace_table(sql: str, table_name: str = "data") -> str:
"""Replaces 'FROM table' (case-insensitive) with the actual table name."""
return _FROM_TABLE_RE.sub(f'FROM "{table_name}"', sql)
def add_agg_parentheses(sql: str) -> str:
"""
Wraps column names after SQL aggregations in parentheses.
E.g., 'SELECT COUNT Player' -> 'SELECT COUNT(Player)'
"""
def repl(match: re.Match) -> str:
func = match.group(1)
col = match.group(2).strip()
# If it already has parens, leave it alone
if col.startswith('(') and col.endswith(')'):
return f"{func}{col}"
return f"{func}({col})"
return _AGG_RE.sub(repl, sql)
def quote_multiword_columns(sql: str, columns: List[str]) -> str:
"""
Wraps multi-word column names in double quotes to prevent SQLite tokenization errors.
Only quotes columns that need it: those containing spaces or special characters.
Single-word, simple columns (e.g., 'Age', 'Team') are left unquoted.
"""
# Filter to columns that actually need quoting (contain spaces or non-alphanumeric/underscore chars)
needs_quoting = [col for col in columns if ' ' in col or not re.match(r'^[A-Za-z_][A-Za-z0-9_]*$', col)]
# Sort by length descending so we replace longer names first (e.g. 'Player Name' before 'Player')
sorted_cols = sorted(needs_quoting, key=len, reverse=True)
for col in sorted_cols:
# Find the column name NOT already surrounded by quotes/backticks.
# We can't use \b word boundaries because columns may contain special chars
# like parentheses or slashes (e.g., "Revenue ($)", "Win/Loss %").
# Instead, match at string boundaries or after/before SQL-meaningful characters.
pattern = rf'(?<![`"])(?<![A-Za-z0-9_]){re.escape(col)}(?![A-Za-z0-9_`"])'
sql = re.sub(pattern, f'"{col}"', sql, flags=re.IGNORECASE)
return sql
def quote_string_values(sql: str, columns: List[str], types: List[str]) -> str:
"""
Checks column schema; if text, converts `= Value` to `= 'Value'`.
Robust against AND/OR statements.
"""
where_match = _WHERE_RE.search(sql)
if not where_match:
return sql
where_clause = where_match.group(1)
parts = _AND_OR_RE.split(where_clause)
type_map = {c.lower(): t for c, t in zip(columns, types)}
processed_parts = []
for part in parts:
if part.upper() in ['AND', 'OR']:
processed_parts.append(part.upper())
continue
cond_match = _COND_RE.match(part)
if cond_match:
col_ref = cond_match.group(1).strip()
op = cond_match.group(2)
val = cond_match.group(3).strip()
clean_col = col_ref.replace('"', '').replace('`', '').lower()
if clean_col in type_map and type_map[clean_col] == 'text':
if not (val.startswith("'") and val.endswith("'")):
val = f"'{val}'"
processed_parts.append(f"{col_ref} {op} {val}")
else:
processed_parts.append(part)
new_where = " ".join(processed_parts)
return sql[:where_match.start()] + "WHERE " + new_where
def post_process_sql(sql: str, columns: List[str], types: List[str], table_name: str = "data") -> str:
"""Runs the full post-processing pipeline on the raw T5 output."""
sql = replace_table(sql, table_name)
sql = quote_multiword_columns(sql, columns)
sql = add_agg_parentheses(sql)
sql = quote_string_values(sql, columns, types)
return sql
|