Constrained Decoding › A Worked Example: Generating Valid SQL

Michael BrenndoerferUpdated July 31, 20254 min read
Language AI HandbookAll chapters
Part 28 of 67GPT Architecture
Chapter 206Constrained Decoding › A Worked Example: Generating Valid SQL
Reading section 9 of 13 · A Worked Example: Generating Valid SQL · 4 minSection 9 of 13 · 4 min

Book contents

Section 9 of 13 · 4 min read

Constrained Decoding › A Worked Example: Generating Valid SQLOne section · Sections in this chapter9 / 13 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Let's work through a complete example of constrained decoding for SQL generation. SQL has strict syntax rules that make it an ideal candidate for grammar-guided generation.

Defining the Constraint

We'll generate simple SELECT statements with a constraint that ensures valid syntax.

In[33]:
Code
from typing import Optional


class SimpleSQLConstraint:
    """
    Constraint for simple SQL SELECT statements.

    Supports: SELECT <columns> FROM <table> WHERE <condition>
    """

    KEYWORDS = {"SELECT", "FROM", "WHERE", "AND", "OR"}
    OPERATORS = {"=", ">", "<", ">=", "<=", "!="}

    def __init__(
        self, valid_tables: list[str], valid_columns: dict[str, list[str]]
    ):
        """
        Initialize with schema information.

        Args:
            valid_tables: List of valid table names
            valid_columns: Mapping from table name to list of column names
        """
        self.valid_tables = set(valid_tables)
        self.valid_columns = valid_columns

    def get_valid_tokens(self, current_sql: str) -> set[str]:
        """
        Determine valid next tokens based on SQL parse state.
        """
        tokens = current_sql.upper().split()

        if not tokens:
            return {"SELECT"}

        last_token = tokens[-1]

        if last_token == "SELECT":
            # After SELECT: column names or *
            all_columns = set()
            for cols in self.valid_columns.values():
                all_columns.update(cols)
            return all_columns | {"*"}

        elif (
            last_token in self.valid_columns.get(self._find_table(tokens), [])
            or last_token == "*"
        ):
            # After column: comma for more columns, or FROM
            return {",", "FROM"}

        elif last_token == ",":
            # After comma: another column
            table = self._find_table(tokens)
            if table:
                return set(self.valid_columns.get(table, []))
            return set()

        elif last_token == "FROM":
            # After FROM: table name
            return self.valid_tables

        elif last_token in self.valid_tables:
            # After table: WHERE or end
            return {"WHERE", "<END>"}

        elif last_token == "WHERE":
            # After WHERE: column for condition
            table = self._find_table(tokens)
            return set(self.valid_columns.get(table, []))

        elif last_token in self.OPERATORS:
            # After operator: value (simplified to numbers/quoted strings)
            return {"<VALUE>"}

        return set()

    def _find_table(self, tokens: list[str]) -> Optional[str]:
        """Find the table name in the token list."""
        try:
            from_idx = tokens.index("FROM")
            if from_idx + 1 < len(tokens):
                return tokens[from_idx + 1]
        except ValueError:
            pass
        return None
Out[34]:
Console
SQL constraint walkthrough:

Schema: users(id, name, email, age), orders(id, user_id, product, amount)

(empty)
  → valid next: ['SELECT']

"SELECT"
  → valid next: ['*', 'age', 'email', 'name', 'user_id', '...']

"SELECT name"
  → valid next: []

"SELECT name,"
  → valid next: []

"SELECT name, email"
  → valid next: []

"SELECT name, email FROM"
  → valid next: ['orders', 'users']

"SELECT name, email FROM users"
  → valid next: []

"SELECT name, email FROM users WHERE"
  → valid next: []

The SQL constraint uses schema knowledge to ensure only valid table and column names appear, and it enforces proper keyword ordering.

Full Generation Pipeline

Putting it all together, here's how constrained generation produces a valid SQL query:

In[35]:
Code
import numpy as np


def generate_sql_with_constraints(
    constraint: SimpleSQLConstraint,
    model_preferences: dict[str, float],
    max_tokens: int = 20,
) -> str:
    """
    Generate a SQL query using constrained sampling.

    Args:
        constraint: SQL syntax constraint
        model_preferences: Model's token preferences (simulated)
        max_tokens: Maximum query length

    Returns:
        Generated SQL query
    """
    tokens = []

    for _ in range(max_tokens):
        current_sql = " ".join(tokens)
        valid = constraint.get_valid_tokens(current_sql)

        if not valid or "<END>" in valid:
            break

        # Remove <END> and <VALUE> placeholders for actual sampling
        valid = {t for t in valid if not t.startswith("<")}
        if not valid:
            break

        # Sample based on model preferences
        probs = {t: model_preferences.get(t.lower(), 0.1) for t in valid}
        total = sum(probs.values())
        probs = {t: p / total for t, p in probs.items()}

        tokens_list, prob_list = zip(*probs.items())
        selected = np.random.choice(tokens_list, p=prob_list)
        tokens.append(selected)

    return " ".join(tokens)
Out[36]:
Console
Generating SQL with constraints:

  1. SELECT name
  2. SELECT email
  3. SELECT amount

All queries are syntactically valid because constraints prevent
invalid tokens from being selected at each step.