SafeQL: Search-based Refinement for Safe and Efficient LLM-based Text-to-SQL

SafeQL: Search-based Refinement for Safe and Efficient LLM-based Text-to-SQL

Abstract Summary:

Large language models (LLMs) have advanced Text-to-SQL by enabling natural language interfaces to databases without task-specific fine-tuning. However, existing LLM-based systems remain unreliable, often generating SQL queries that are invalid under the database schema, referencing non-existent tables, attributes, functions, or values. Such errors persist because interactions with the database management system (DBMS) are typically limited to error messages, leaving it in a largely passive role during query refinement. This paper proposes SafeQL, a search-based refinement paradigm that redefines the role of the DBMS as an active guide in the refinement process. Instead of regenerating entire queries after execution failure, SafeQL interprets DBMS feedback to incrementally repair only the erroneous components. Each refinement step is formulated as a guided search within a safe query space, where candidate queries are progressively validated through DBMS execution, thereby converging to an executable query and preventing repeated regeneration of errors. Experiments on the Bird and Spider benchmarks show that SafeQL significantly improves execution accuracy and efficiency compared to regeneration-based methods.

Introduction Summary:

Natural language interfaces to databases (NLIDBs) have long been a key research goal. Text-to-SQL—the task of translating a natural language question into an executable SQL query—represents the most concrete realization of this vision. Research on Text-to-SQL has evolved through several generations: early symbolic systems (PRECISE, NaLIR, SQLizer, ATHENA) relied on rule-based parsing; subsequent neural semantic parsers (Seq2SQL, SQLNet, IRNet, RAT-SQL, SmBoP) replaced handcrafted rules with data-driven models; the emergence of LLMs (GPT) transformed the landscape with prompt-based approaches (DAIL-SQL, MCS-SQL) and multi-agent systems (MAC-SQL, CHESS-SQL, CHASE-SQL, Alpha-SQL, OpenSearch-SQL). Despite these advances, reliability remains a critical challenge due to the inherent hallucination tendencies of LLMs. Recent studies have attempted to mitigate this issue through regeneration-based refinement: executing the generated SQL, capturing DBMS error messages, and prompting the LLM to regenerate a corrected query. However, this paradigm treats the DBMS primarily as a passive checker that provides error messages, and repeatedly regenerates entire queries based on the errors. As a result, the regeneration loop is both inefficient—requiring excessive tokens and computation—and unreliable, often reintroducing previous errors with no guarantee of convergence.

In contrast, SafeQL proposes a refinement framework that achieves safe and efficient correction through a search-based refinement paradigm. Rather than discarding the initial query, SafeQL interprets and incrementally refines it. When a query fails, SafeQL does not regenerate a new query. Instead, it analyzes the DBMS feedback to precisely locate the faulty component—whether a relation, attribute, or function—and applies a minimal, structure-preserving correction. This transforms the initial query directly into refined and executable queries, while preserving all valid fragments of the original logic. SafeQL replaces the traditional trial-and-error regeneration loop with a guided search over a safe query space, where candidates are incrementally repaired to satisfy schema-level constraints enforced by the DBMS. Since this safe query space may contain multiple executable refinements for a single input, SafeQL selects the candidate that is semantically the most aligned with the initial query. This paradigm shift redefines the DBMS from a passive checker into an active reasoning engine that incrementally narrows the search to the most faithful executable query, improving both safety and efficiency.

The effectiveness of SafeQL critically depends on how search and optimization processes are conducted. The safe query space can grow exponentially, as each detected error may yield multiple candidate substitutions for relations, attributes, or values. To manage this combinatorial complexity, SafeQL employs a best-first search strategy guided by a semantic distance metric that combines structural similarity and embedding-based relevance. SafeQL incorporates type-based pruning, which eliminates type-inconsistent refinements, and top-K pruning, which restricts candidate substitutions to the most semantically relevant ones. SafeQL is integrated directly into the DBMS, enabling refinement through the database's own parser, binder, and type analyzer instead of treating SQL as plain text. This tight integration allows for precise error localization, while system-level optimizations—such as caching and vector indexing—make semantic similarity computations efficient at scale. Moreover, SafeQL incorporates a hybrid refinement strategy that combines its search-based refinement with selective regeneration to handle cases where the initial query is fundamentally misgenerated.

Main Contributions:

• A search-based refinement paradigm that reformulates Text-to-SQL refinement as a guided exploration of a safe query space, where execution feedback serves as actionable semantic guidance rather than passive error reporting.

• A best-first search guided by a semantic distance metric, augmented with type-based and top-K pruning, enabling scalable and accurate refinement.

• SafeQL implemented as a DBMS-integrated refinement framework, incorporating caching, vector indexing, and hybrid refinement to achieve efficient and robust execution.

• Comprehensive experiments on Bird and Spider benchmarks, demonstrating up to 5.8% improvements in execution accuracy and 15× reductions in token usage compared to regeneration-based baselines.

Preliminaries Summary:

The Text-to-SQL task translates a natural language question into an executable SQL query over a given database. The process is decomposed into three stages: (1) extracting database information, (2) generating candidate SQL queries, and (3) refining them using DBMS execution. The refinement stage is crucial to ensure reliability. Execution errors provide essential feedback signals that enable the Text-to-SQL system to identify and correct invalid parts of a query. The most common errors include: unknown relation (query references a non-existent table), unknown attribute (query refers to an undefined column), unknown function or operator (query invokes a function not supported by the DBMS), and empty result (query executes successfully but returns no tuples due to unsatisfiable predicates). Regeneration-based Text-to-SQL methods attempt to correct these errors through a regeneration-based refinement process, which consists of execution (DBMS executes the current query and returns an error if failed) and regeneration (the refiner updates the prompt with error information and regenerates a new query). This repeats until a valid query is found. However, identical error messages may correspond to different underlying causes, each requiring a distinct refinement strategy, and DBMS error messages often provide little or no actionable guidance.

Safe Query Space Summary:

SafeQL systematically explores a structured space of refinements—syntactic transformations that incrementally repair an erroneous query based on its observed error. SafeQL focuses on error-directed transformations, i.e., targeted adjustments that resolve the specific cause of failure while retaining the original query's structure and intent. A refinement step is the application of one of several canonical operations: Relation Refinement (replace a relation with another), Join Refinement (add a join with another relation and a suitable condition), Attribute Refinement (replace an attribute with another in either SELECT or WHERE), Value Refinement (replace a constant value with another), and Function Refinement (replace a function symbol with another while preserving its argument structure).

The safe refinement tree is a directed tree that represents error-resolving refinement paths from the initial query. The root is the initial query, an edge is added if the parent query produces an execution error, and each edge corresponds to a refinement step. The set of admissible refinement operations allowed for each error type is: Relation Refinement for unknown relation; Relation, Join, and Attribute Refinement for unknown attribute; Attribute and Function Refinement for unknown function; and Attribute and Value Refinement for empty result. For unknown attribute and unknown function errors, multiple refinement operations are admissible because their causes are often indirect—an unknown attribute may not only indicate a misspelled column name but also reveal that the necessary relation has not been joined at all, and an unknown function may arise either from an undefined function name or from invalid arguments within its call.

The paper proves Lemma 7 (Reachability of Error-resolving Refinements): Given a database and the safe refinement tree, every execution error at a node has a corresponding refinement edge such that the refinement operation resolves the error under the database. The proof analyzes each error type by making its operational error condition explicit and showing that at least one refinement always restores the violated execution premise. For unknown relation, a Relation Refinement replaces the invalid relation with a valid one. For unknown attribute, the error can be repaired by Relation Refinement (replace a relation with one containing the attribute), Join Refinement (add a relation containing the attribute), or Attribute Refinement (replace the unresolved attribute with a valid one). For unknown function, the error can be repaired by Function Refinement (replace the unsupported function with a supported one) or Attribute Refinement (replace arguments with attributes of the required types). For empty result, a Value Refinement adjusts the predicate to produce a non-empty result.

The safe query space is defined as the set of all leaf nodes in the refinement tree that satisfy the condition that their execution does not produce an error. By construction, all leaf nodes in the safe refinement tree are error-free, and their set constitutes the safe query space.

Search Algorithm Summary:

SafeQL navigates the safe query space to identify the most semantically faithful executable query. The search objective is based on semantic distance, which jointly measures structural change and embedding-level drift. The semantic distance between two queries is defined as a weighted combination of structural deviation (using normalized tree edit distance between abstract syntax trees) and semantic deviation (based on embeddings of replaced tokens). This hybrid formulation is conceptually analogous to sparse-dense hybrid retrieval in information retrieval, combining complementary lexical and semantic signals to improve ranking robustness.

The refinement process finds the safe query nearest to the input query within a distance metric. The overall refinement process follows a best-first exploration of the safe query space. The search initializes a priority queue ordered by semantic distance from the LLM-generated initial query, expanding the closest candidates first. This design takes the initial query as the reference point under the assumption that it reflects the LLM's interpretation of the user's intent. During each iteration, the algorithm dequeues the nearest candidate query and attempts to execute it. If execution succeeds, the algorithm immediately returns the query as the nearest safe query, thus guaranteeing that the first valid query discovered is the one most semantically similar to the original input. When execution fails, SafeQL captures the corresponding error type and uses it to guide subsequent refinements. The GenerateRefinements procedure selects the appropriate refinement operations based on the error type, while PruneRefinements filters out redundant or semantically invalid refinements. Each surviving refinement is then applied to produce a new candidate query, which is enqueued for future exploration only if it has not been visited previously.

Pruning Methods Summary:

Although all queries in the safe space are executable, the space can grow rapidly as each refinement generates many candidates. For each error, all schema-level elements are considered, causing a single step to branch into dozens of candidates. Value refinements can expand even further, since every possible value in the database becomes a potential substitution. The total size of the search space is exponential with branching factor and depth. To mitigate this issue, SafeQL introduces two pruning mechanisms:

Type-based Refinement Pruning: Type information provides a strong pruning signal because any type mismatch immediately yields an invalid query. In SQL, each operator and function has a fixed input signature; thus, if a substituted attribute or value violates this signature, execution inevitably results in an error. SafeQL leverages the typing environment to discard such candidates before semantic evaluation. It validates refinements under the typing environment and prunes type-inconsistent candidates.

Top-K Refinement Pruning: Even after type-based pruning, many refinements remain type-consistent but semantically distant. To improve efficiency, SafeQL ranks candidates by embedding similarity and retains only the top-K per category. This is particularly beneficial for value refinements, where the out-degree is large and most database values are semantically unrelated to the query. A smaller K yields higher efficiency but risks over-pruning, while a larger K improves accuracy at increased search cost.

System Architecture Design Summary:

Realizing SafeQL as a refinement framework for building practical text-to-SQL systems requires addressing several system-level challenges. SafeQL's architecture is built around five key design considerations:

In-DBMS Integration: External Text-to-SQL systems rely on error strings, which may mention identifiers but do not reliably indicate error locations; some DBMSs provide no positions, and even when they do, they refer only to character offsets and cannot distinguish multiple occurrences. SafeQL instead performs refinement inside the DBMS over the Abstract Syntax Tree (AST). When a query fails, the analyzer directly identifies the corresponding AST node. Modifications are restricted to these nodes; since the AST encodes attribute-relation links, updates can be applied without affecting unrelated occurrences, and each refined query is re-analyzed and validated. This enables precise and safer structurally grounded refinement beyond text-based approaches.

Systematic Optimization: Executing refinement inside the DBMS introduces a new bottleneck: repeated embedding and similarity computations across overlapping query fragments. SafeQL incorporates two in-database optimization layers: Embedding Cache (embeddings of relations, attributes, and values are precomputed and stored on disk, while frequently accessed items are cached in memory) and Vector Index (an HNSW-based vector similarity index over cached embeddings for efficient nearest-neighbor search).

Hybrid Refinement Strategy: Even with precise in-database refinement, some queries are fundamentally misgenerated. SafeQL adopts a hybrid strategy combining search-based and regeneration-based refinement. When the number of refinement steps reaches a threshold (default 100) without producing a valid query, SafeQL triggers a regeneration loop. Such cases typically indicate that the initial query is too far from a valid structure for localized refinement to succeed, in which case an LLM generates a new candidate conditioned on the original input, the current query, and the DBMS error feedback.

Grammar Extensions: Practical Text-to-SQL systems must handle more complex SQL than the simplified grammar. SafeQL extends its support to standard SQL constructs appearing in benchmarks, including nested subqueries (refining subqueries in an inside-out manner), alias refinement (extending relation and attribute refinem…

Based on this paper, here are the specific improvements I can make to AI systems:

1. Error-Localized Repair Instead of Full Regeneration

• Improvement: When generating SQL (or similar structured outputs) fails, instead of regenerating the entire output, I can parse the error message, identify the exact faulty component (relation, attribute, function, or value), and apply a minimal, targeted correction while preserving all valid parts.

• What the improved system can do: Reduce token usage by up to 15× and avoid reintroducing previously fixed errors, leading to faster convergence on valid outputs.

2. Search-Based Refinement with Best-First Exploration

• Improvement: I can implement a best-first search algorithm that explores a structured space of possible corrections, guided by a semantic distance metric that combines structural similarity (tree edit distance) and embedding-based relevance.

• What the improved system can do: Systematically find the most semantically faithful executable output, rather than relying on stochastic regeneration that may not converge.

3. Type-Based Pruning for Constraint Satisfaction

• Improvement: Before evaluating candidate corrections, I can validate them against the type system of the target language or schema (e.g., SQL type signatures for operators and functions), discarding any candidate that would cause a type mismatch.

• What the improved system can do: Eliminate a large portion of invalid candidates upfront, dramatically reducing search space and improving efficiency without sacrificing accuracy.

4. Top-K Semantic Pruning with Embedding Similarity

• Improvement: After type-based filtering, I can rank remaining candidates by embedding similarity to the original query and retain only the top-K most semantically relevant options per error category.

• What the improved system can do: Handle cases with large candidate sets (e.g., value substitutions where every database value is a potential candidate) efficiently, maintaining near-optimal accuracy with K≥3 while keeping runtime manageable.

5. In-System Integration for Precise Error Localization

• Improvement: Instead of relying on text-based error messages, I can integrate directly with the target system's internal parser, binder, and type analyzer to work on the abstract syntax tree (AST), precisely identifying which AST node caused the error.

• What the improved system can do: Distinguish between multiple occurrences of the same identifier, apply corrections only to the faulty node, and ensure structural consistency through re-validation after each refinement.

6. Hybrid Refinement Strategy with Fallback Trigger

• Improvement: I can combine search-based refinement with selective regeneration, triggering LLM-based regeneration only when the number of refinement steps exceeds a threshold (e.g., 100), indicating the initial output is fundamentally misgenerated.

• What the improved system can do: Handle both minor errors (via fast, token-free search) and major structural failures (via regeneration), achieving the best of both approaches—high accuracy with minimal token overhead.

7. Caching and Vector Indexing for Repeated Computations

• Improvement: I can precompute and cache embeddings of all schema elements (relations, attributes, values) and use an HNSW-based vector index for efficient nearest-neighbor search during refinement.

• What the improved system can do: Reduce computational overhead by over 30× compared to naive approaches, making search-based refinement practical for real-time applications.

8. Inside-Out Refinement for Nested Structures

• Improvement: For outputs with nested subqueries or hierarchical structures, I can refine inner components first (inside-out), ensuring that outer structures are built on valid foundations.

• What the improved system can do: Handle complex, multi-level outputs (e.g., nested SQL subqueries) reliably, preventing cascading errors from propagating outward.

9. Semantic Distance as a Quality Metric for Output Selection

• Improvement: When multiple valid outputs exist, I can select the one with the minimum semantic distance to the original generated output, using a weighted combination of structural deviation and embedding-level drift.

• What the improved system can do: Ensure the final output most closely preserves the original intent, even when multiple valid alternatives exist, improving user satisfaction and correctness.

10. Model-Agnostic Error Correction

• Improvement: The refinement process can be applied as a post-processing layer regardless of the underlying generation model (prompt-based or agent-based, proprietary or open-source).

• What the improved system can do: Consistently improve execution accuracy by +4–12% across different model sizes and architectures, including narrowing the performance gap between open-source and proprietary models, without requiring model-specific fine-tuning.