Finding the Right Tables and Columns: A Benchmark and Corpus-Adaptive Embeddings for SQL Schema Retrieval

2026-07-14

AI Summary

This work addresses the underperformance of existing embedding models in table- and column-level schema retrieval for large-scale text-to-SQL tasks. It formulates schema linking as a standalone retrieval problem and introduces a novel unsupervised, corpus-adaptive fine-tuning approach that synthesizes queries based on target database schemas, mines granularity-aware hard negatives, and optimizes embeddings via contrastive learning. The method is lightweight and backbone-agnostic, achieving substantial gains on a newly constructed multi-granularity SQL schema retrieval benchmark: a 305M-parameter model improves Recall@10 from 60.4 to 75.6 and nDCG@10 from 51.9 to 68.0, surpassing models with billions of parameters; an 8B-parameter variant further attains a state-of-the-art Recall@10 of 78.4.

Abstract

Retrieval in the SQL setting has largely been studied as the task of finding, within a large collection of SQL statements, the statement that answers a natural-language question. At scale, however, a more fundamental retrieval problem precedes generation: schema retrieval, identifying the tables and columns a question requires in a database that may contain thousands of them, far more than fit in a model's context. We argue that this step warrants first-class evaluation. To this end, we recast five text-to-SQL datasets (Spider, BIRD, BEAVER, and two LiveSQLBench variants) as retrieval tasks at both table and column granularity, covering realistic and enterprise-scale schemas under two document representations, and we show that off-the-shelf text and code embedders transfer poorly to this setting. We then propose corpus-adaptive fine-tuning: natural-language queries are synthesized directly from the target schema corpus, granularity-aware hard negatives are mined, and a 305M-parameter embedder is fine-tuned contrastively. This procedure raises average recall@10 from 60.4 to 75.6 (nDCG@10 from 51.9 to 68.0), making the 305M model the strongest retriever under one billion parameters and competitive with state-of-the-art embedders of 4-8B parameters, more than an order of magnitude larger. The same recipe improves an 8B state-of-the-art embedder from 77.8 to 78.4 recall@10, matching the best result on the benchmark and indicating that the adaptation is backbone-agnostic. Leave-one-corpus-out experiments and a leakage audit show that these gains reflect a transferable schema-retrieval ability rather than memorization of the evaluation data. Our results establish schema linking as a standalone retrieval task and lightweight, label-free corpus adaptation as a practical route to deploying it at enterprise scale.

Problem

• schema retrieval
• text-to-SQL
• table and column linking
• database schema
• retrieval task

Innovation

• schema retrieval
• corpus-adaptive embedding
• text-to-SQL
• contrastive fine-tuning
• hard negative mining