Some examples in the text2sql datasets use ID as a column reference to the
    column of a table which has a primary key. This causes problems if you are trying
    to constrain a grammar to only produce the column names directly, because you don't
    know what ID refers to. So instead of dealing with that, we just replace it.

def resolve_primary_keys_in_schema(sql_tokens: List[str],
                                   schema: Dict[str, List[TableColumn]]) -> List[str]:
    """
    Some examples in the text2sql datasets use ID as a column reference to the
    column of a table which has a primary key. This causes problems if you are trying
    to constrain a grammar to only produce the column names directly, because you don't
    know what ID refers to. So instead of dealing with that, we just replace it.
    """
    primary_keys_for_tables = {name: max(columns, key=lambda x: x.is_primary_key).name
                               for name, columns in schema.items()}
    resolved_tokens = []
    for i, token in enumerate(sql_tokens):
        if i > 2:
            table_name = sql_tokens[i - 2]
            if token == "ID" and table_name in primary_keys_for_tables.keys():
                token = primary_keys_for_tables[table_name]
        resolved_tokens.append(token)
    return resolved_tokens