Representing Schema Structure with Graph Neural Networks for Text-to-SQL Parsing

Introduction

Semantic parsing Zelle and Mooney (1996); Zettlemoyer and Collins (2005) has recently taken increased interest in parsing questions into SQL queries, due to the popularity of SQL as a query language for relational databases (DBs).

Work on parsing to SQL Zhong et al. (2017); Iyer et al. (2017); Finegan-Dollak et al. (2018); Yu et al. (2018a) has either involved simple DBs that contain just one table, or had a single DB that is observed at both training and test time. Consequently, modeling the schema structure received little attention. Recently, Yu et al. (2018b) presented Spider, a text-to-SQL dataset, where at test time questions are executed against unseen and complex DBs. In this zero-shot setup, an informative representation of the schema structure is important. Consider the questions in Figure 1: while their language structure is similar, in the first query a ‘join’ operation is necessary because the information is distributed across three tables, while in the other query no ‘join’ is needed.

In this work, we propose a semantic parser that strongly uses the schema structure. We represent the structure of the schema as a graph, and use graph neural networks (GNNs) to provide a global representation for each node Li et al. (2016); De Cao et al. (2019); Sorokin and Gurevych (2018). We incorporate our schema representation into the encoder-decoder parser of Krishnamurthy et al. (2017), which was designed to parse questions into queries against unseen semi-structured tables. At encoding time we enrich each question word with a representation of the subgraph it is related to, and at decoding time we emit symbols from the schema that are related through the graph to previously decoded symbols.

We evaluate our parser on Spider, and show that encoding the schema structure improves accuracy from 33.8% to 39.4% (and from 14.6% to 26.8% on questions that involve multiple tables), well beyond 19.7%, the current state-of-the-art. We make our code publicly available at https://github.com/benbogin/spider-schema-gnn.

Problem Setup

We are given a training set $\{(x^{(k)},y^{(k)},S^{(k)})\}_{k=1}^{N}$, where $x^{(k)}$ is a natural language question, $y^{(k)}$ is its translation to a SQL query, and $S^{(k)}$ is the schema of the DB where $y^{(k)}$ is executed. Our goal is to learn a function that maps an unseen question-schema pair $(x,S)$ to its correct SQL query. Importantly, the schema $S$ was not seen at training time, that is, $S\neq S^{(k)}$ for all $k$.

A DB schema $S$ includes: (a) The set of DB tables $\mathcal{T}$ (e.g., singer), (b) a set of columns $\mathcal{C}_{t}$ for each $t\in\mathcal{T}$ (e.g., singer_name), and (c) a set of foreign key-primary key column pairs $\mathcal{F}$, where each $(c_{f},c_{p})\in\mathcal{F}$ is a relation from a foreign-key $c_{f}$ in one table to a primary-key $c_{p}$ in another. We term all schema tables and columns as schema items and denote them by $\mathcal{V}=\mathcal{T}\cup\{\mathcal{C}_{t}\}_{t\in\mathcal{T}}$.

A Neural Semantic Parser for SQL

We base our model on the parser of Krishnamurthy et al. (2017), along with a grammar for SQL provided by AllenNLP Gardner et al. (2018); Lin et al. (2019), which covers 98.3% of the examples in Spider. This parser uses a linking mechanism for handling unobserved DB constants at test time. We review this model in the context of text-to-SQL parsing, focusing on components we expand upon in §4.

To handle unseen schema items, Krishnamurthy et al. (2017) learn a similarity score $s_{\text{link}}(v,x_{i})$ between a word $x_{i}$ and a schema item $v$ that has type $\tau$.Types are tables, string columns, number columns, etc. The score is based on learned word embeddings and a few manually-crafted features.

where $\mathcal{V}_{\tau}$ are all schema items of type $\tau$ and $s_{\text{link}}(\varnothing,\cdot)=0$ for words that do not link to any schema item. The functions $p_{\text{link}}(\cdot)$ and $s_{\text{link}}(\cdot)$ will be used to decode unseen schema items.

Encoder

A Bidirectional LSTM Hochreiter and Schmidhuber (1997) provides a contextualized representation $h_{i}$ for each question word $x_{i}$. Importantly, the encoder input at time step $i$ is $[w_{x_{i}};l_{i}]$: the concatenation of the word embedding for $x_{i}$ and $l_{i}=\sum_{\tau}\sum_{v\in\mathcal{V}_{\tau}}p_{\text{link}}(v\mid x_{i})\cdot r_{v}$, where $r_{v}$ is a learned embedding for the schema item $v$, based on the type of $v$ and its schema neighbors. Thus, $p_{\text{link}}(v\mid x_{i})$ augments every word $x_{i}$ with information on the schema items it should link to.

Decoder

We use a grammar-based Xiao et al. (2016); Cheng et al. (2017); Yin and Neubig (2017); Rabinovich et al. (2017) LSTM decoder with attention on the input question (Figure 2). At each decoding step, a non-terminal of type $\tau$ is expanded using one of the grammar rules. Rules are either schema-independent and generate non-terminals or SQL keywords, or schema-specific and generate schema items.

At each decoding step $j$, the decoding LSTM takes a vector $g_{j}$ as input, which is an embedding of the grammar rule decoded in the previous step, and outputs a vector $o_{j}$. If this rule is schema-independent, $g_{j}$ is a learned global embedding. If it is schema-specific, i.e., a schema item $v$ was generated, $g_{j}$ is a learned embedding $\tau(v)$ of its type. An attention distribution $a_{j}$ over the input words is computed in a standard manner Bahdanau et al. (2015), where the attention score for every word is $h_{i}^{\top}o_{j}$. It is then used to compute the weighted average of the input $c_{j}=\sum_{i}a_{j}h_{j}$. Now a distribution over grammar rules is computed by:

Modeling Schemas with GNNs

Schema structure is informative for predicting the SQL query. Consider a table with two columns, where each is a foreign key to two other tables (student_semester table in Figure 3). Such a table is commonly used for describing a many-to-many relation between two other tables, which affects the output query. We now show how we represent this information in a neural parser and use it to improve predictions.

At a high-level our model has the following parts (Figure 3). (a) The schema is converted to a graph. (b) The graph is softly pruned conditioned on the input question. (c) A Graph neural network generates a representation for nodes that is aware of the global schema structure. (d) The encoder and decoder use the schema representation. We will now elaborate on each part.

To convert the schema $S$ to a graph (Figure 3, left), we define the graph nodes as the schema items $\mathcal{V}$. We add three types of edges: for each column $c_{t}$ in a table $t$, we add edges $(c_{t},t)$ and $(t,c_{t})$ to the edge set $\mathcal{E}_{\leftrightarrow}$ (green edges). For each foreign-primary key column pair $(c_{t_{1}},c_{t_{2}})\in\mathcal{F}$, we add edges $(c_{t_{1}},c_{t_{2}})$ and $(t_{1},t_{2})$ to the edge set $\mathcal{E}_{\rightarrow}$ and edges $(c_{t_{2}},c_{t_{1}})$ and $(t_{2},t_{1})$ to $\mathcal{E}_{\leftarrow}$ (dashed edges). Edge types are used by the graph neural network to capture different ways in which columns and tables relate to one another.

Question-conditioned relevance

Each question refers to different parts of the schema, and thus, our representation should change conditioned on the question. For example, in Figure 3, the relation between the tables student_semester and program is irrelevant. To model that, we re-use the distribution $p_{\text{link}}(\cdot)$ from §3, and define a relevance score for a schema item $v$: $\rho_{v}=\max_{i}{p_{\text{link}}(v\mid x_{i})}$ — the maximum probability of $v$ for any word $x_{i}$. We use this score next to create a question-conditioned graph representation. Figure 3 shows relevant schema items in dark orange, and irrelevant items in light orange.

Neural graph representation

To learn a node representation that considers its relevance score and the global schema structure, we use gated GNNs Li et al. (2016). Each node $v$ is given an initial embedding conditioned on the relevance score: $h^{(0)}_{v}=r_{v}\cdot\rho_{v}$. We then apply the GNN recurrence for $L$ steps. At each step, each node re-computes its representation based on the representation of its neighbors in the previous step:

and then $h_{v}^{(l)}$ is computed as following, using a standard GRU Cho et al. (2014) update:

(see Li et al. (2016) for further details).

We denote the final representation of each schema item after $L$ steps by $\varphi_{v}=h_{v}^{(L)}$. We now show how this representation is used by the parser.

Encoder

In §3, a weighted average over schema items $l_{i}$ was concatenated to every word $x_{i}$. To enjoy the schema-aware representations, we compute $l^{\varphi}_{i}=\sum_{\tau}\sum_{v\in\mathcal{V}_{\tau}}\varphi_{v}p_{\text{link}}(v\mid x_{i})$, which is identical to $l_{i}$, except $\varphi_{v}$ is used instead of $r_{v}$. We concatenate $l^{\varphi}_{i}$ to the output of the encoder $h_{i}$, so that each word is augmented with the graph structure around the schema items it is linked to.

Decoder

As mentioned (§3), when a schema item $v$ is decoded, the input in the next time step is its type $\tau(v)$. A first change is to replace $\tau(v)$ by $\varphi_{v}$, which has knowledge of the structure around $v$. A second change is a self-attention mechanism that links to the schema, which we describe next.

When scoring a schema item, its score should depend on its relation to previously decoded schema items. E.g., in Figure 3, once the table semester has been decoded, it is likely to be joined to a related table. We capture this intuition with a self-attention mechanism.

Training

We maximize the log-likelihood of the gold sequence during training, and use beam-search (of size 10) at test time, similar to Krishnamurthy et al. 2017 and prior work. We run the GNN for $L=2$ steps.

Experiments and Results

We evaluate on Spider Yu et al. (2018b), which contains 7,000/1,034/2,147 train/development/test examples.

We pre-process examples to remove table aliases (AS T1/T2/...) from the queries and use the explicit table name instead (i.e. we replace T1.col with table1_name.col), as in the majority of the cases ($>$ 99% in Spider) these aliases are redundant. In addition, we add a table reference to all columns that do not have one (i.e. we replace col with table_name.col).

We use the official evaluation script from Spider to ≈compute accuracy, i.e., whether the predicted query is equivalent to the gold query.

Results

Our full model (GNN) obtains 39.4% accuracy on the test set, substantially higher than prior state-of-the-art (SyntaxSQLNet), which is at 19.7%. Removing the GNN from the parser (No GNN), which results in the parser of Krishnamurthy et al. (2017), augmented with a grammar for SQL, obtains an accuracy of 33.8%, showing the importance of encoding the schema structure.

Table 1 shows results on the development set for baselines and ablations. The first column describes accuracy on the entire dataset, and the next two columns show accuracy when partitioning examples to queries involving only one table (Single) vs. more than one table (Multi).

GNN dramatically outperforms previously published baselines SQLNet and SyntaxSQLNet, and improves the performance of No GNN from 34.9% to 40.7%. Importantly, using schema structure specifically improves performance on questions with multiple tables from 14.6% to 26.8%.

We ablate the major novel components of our model to assess their impact. First, we remove the self-attention component (No Self Attend). We observe that performance drops by 2 points, where Single slightly improves, and Multi drops by 6.5 points. Second, to verify that improvement is not only due to self-attention, we ablate all other uses of the GNN. Namely, We use a model identical to No GNN, except it can access the GNN representations through the self-attention (Only Self Attend). We observe a large drop in performance to 35.9%, showing that all components are important. Last, we ablate the relevance score by setting $\rho_{v}=1$ for all schema items (No Rel.). Indeed, accuracy drops to 37.0%.

To assess the ceiling performance possible with a perfect relevance score, we run an oracle experiment, where we set $\rho_{v}=1$ for all schema items that are in the gold query, and $\rho_{v}=0$ for all other schema items (GNN Oracle Rel.). We see that a perfect relevance score substantially improves performance to 54.3%, indicating substantial headroom for future research.

join analysis For any model, we can examine the proportion of predicted queries with a join, where the structure of the join is “bad”: (a) when the join condition clause uses the same table twice (ON t1.column1 = t1.column2), and (b) when the joined table are not connected through a primary-foreign key relation.

We find that No GNN predicts such joins in 83.4% of the cases, while GNN does so in only 15.6% of cases. When automatically omitting from the beam candidates where condition (a) occurs, No GNN predicts a “bad” join in 14.2% of the cases vs. 4.3% for GNN (total accuracy increases by 0.3% for both models). As an example, in Figure 3, $s^{\text{loc}}_{j}$ scores the table student the highest, although it is not related to the previously decoded table semester. Adding the self-attention score $s^{\text{att}}_{j}$ corrects this and leads to the correct student_semester, probably because the model learns to prefer connected tables.

Conclusion

We present a semantic parser that encodes the structure of the DB schema with a graph neural network, and uses this representation to make schema-aware decisions both at encoding and decoding time. We demonstrate the effectivness of this method on Spider, a dataset that contains complex schemas which are not seen at training time, and show substantial improvement over current state-of-the-art.

Acknowledgments

We thank Kevin Lin and Mark Neumann from Allen Institute for Artificial Intelligence for their help with the SQL grammar. This research was supported by Facebook. This work was completed in partial fulfillment for the Ph.D degree of the first author.

References