WISK: A Workload-aware Learned Index for Spatial Keyword Queries

We proceed to show that the optimal partitioning problem is NP-hard by reducing from the MaxSkip partitioning problem, which has been shown to be NP-hard (Sun et al., 2014; Yang et al., 2020).

A naïve method to find the value to make a split uses a brute-force search. Let $V_{d}$ ($d\in\{x,y\}$) be a sorted list of distinct object coordinate values along dimension $d$ in the current (sub-)space to be partitioned. Except for the first and the last values, every value in $V_{d}$ can be used to split the space into two sub-spaces. Examining all $|V_{d}|$ values takes $\mathcal{O}\big{(}f\cdot(|V_{x}|+|V_{y}|)\big{)}$ where $f$ denotes the time cost to split on a value and run queries based on such a splitting This approach becomes impractical for large datasets with a large value of $|V_{x}|+|V_{y}|$.

Motivated by the recent success of machine learning in solving complex problems (Prates et al., 2019; Bengio et al., 2021), we propose a learning-based method to predict the query costs given a split dimension and a split value, such that the optimal split can be approximated by minimizing the predicted query cost with high efficiency. At the core of the query cost prediction problem of a split is to (1) predict the number of resultant sub-spaces overlapping with the query, and (2) predict the number of objects that contain any of the query keywords and reside in the resultant sub-spaces.

To address the first prediction problem, we use the indicator function (Kleene, 1952) to denote whether a sub-space overlaps with the query region. For example, let $[q_{x_{b}},q_{x_{u}}]$ be the x range of query q and $p_{x}$ be a split value along dimension x. The indicator functions $\vmathbb{1}(p_{x}\geq q_{x_{b}})$ and $\vmathbb{1}(p_{x}<q_{x_{u}})$ are used to decide whether $q$ intersects with the resultant left and right sub-spaces, respectively. If a sub-space has an indicator function value of 1, we need to further predict the number of query result objects within the sub-space. Otherwise, we can ignore the sub-space when computing the query cost. The indicator function is not differentiable, and machine learning methods such as gradient descent cannot be applied to solve a split value optimization problem formulated by such functions. As such, we use the sigmoid function (Han and Moraga, 1995), $\sigma(\beta x)$ with $\beta=3$, to approximate the indicator function as does in prior work (Chen et al., 2013b; Cao and Zavala, 2020), e.g., $\vmathbb{1}(p_{x}\geq q_{x_{b}})=\vmathbb{1}(p_{x}-q_{x_{b}}\geq 0)\approx\sigma(3(p_{x}-q_{x_{b}}))$.

To address the second prediction problem, we follow the idea in recent studies (Kraska et al., 2018; Nathan et al., 2020; Li et al., 2021) that learn the Cumulative Distribution Function (CDF) to estimate the density of objects in a data space. Our goal is to learn the joint CDF $F_{X,Y}(x,y)$ of two variables X and Y, corresponding to the spatial coordinates in two dimensions. The learned CDF can quickly estimate the number of objects in a rectangular region, i.e., a sub-space. To accelerate the CDF learning, we assume that X and Y are independent, following a previous study (Nathan et al., 2020). Thus, we can decompose the joint CDF into the product of two marginal CDFs, $F_{X}(x)$ and $F_{Y}(y)$, as shown in Eq. 3.

For ease of presentation, we use F(x) and F(y) to denote the marginal CDFs of X and Y in the rest of the paper, respectively.

The CDF in Eq. 3 only estimates the spatial density of objects without considering the keyword distribution. To solve this issue, we learn the marginal CDFs, i.e., $F_{k}(x)$ and $F_{k}(y)$, for each keyword $k$. The choice of CDF models will be detailed in Section 6.

With the CDF models and the sigmoid functions, we formulate the cost for processing a query $q$ with region $[(x_{b},y_{b}),(x_{u},y_{u})]$ after splitting on dimensions $x$ or $y$ in Eq. 4.

where $\left|O_{1}\right|$ and $\left|O_{2}\right|$ denote the number of objects containing the query keywords in the two resulting sub-spaces, respectively, which are estimated through the learned keyword-based marginal CDF models. The sigmoid functions (e.g., $\sigma\big{(}3(x-x_{b})\big{)}$ and $\sigma\big{(}3(x_{u}-x)\big{)}$) predicts whether the query intersects the two resultant sub-spaces, respectively. We apply stochastic gradient descent (SGD) to minimize $L_{q}(x)$ and $L_{q}(y)$ using the query workload as the training data.

When splitting a data space, there are both profit and loss in the query costs. The profit is gained by the reduced number of objects to be checked while the loss reflects an increased number of sub-spaces to be checked. In Example 4.1, the profit and loss are equal to $2w_{2}$ and $2w_{1}$, respectively. The difference between the profit and the loss determines whether a split is needed, and where the split should be made.

Algorithm 2 summarizes our bottom cluster generation algorithm. The algorithm takes the query workload $W$ and the data space $S$ enclosing all geo-textual objects as the input, and it aims to return a set of clusters that minimize the cost of executing all the queries in $W$.

The algorithm maintains a priority queue $Q$ of sub-spaces to the examined, which are prioritized by their numbers of intersecting queries. At the start, $Q$ contains only the input data space $S$ (lines 1 and 2). Then, we iterate through the sub-spaces in $Q$. Let the current sub-space to be split be $s$. We set the initial object checking the cost of $s$ to be $|O_{s}|\cdot|W_{s}|\cdot w_{2}$ where $|O_{s}|$ and $|W_{s}|$ denote the number of objects in $s$ and the number of queries intersecting with $s$, respectively (lines 5 and 6). Then, we find the optimal split along both $x$- and $y$-dimensions, respectively (lines 7 and 8), and we use the one with a smaller object checking cost as our candidate split (line 9). If the reduction in the object checking cost from $C_{s}$ outweighs the increase in cluster checking cost, i.e., $w_{1}\cdot|W|$ (every split adds a cluster to be checked against $|W|$ queries), we execute the split and enqueue the resultant sub-spaces (lines 10 to 13). Otherwise, $s$ is finalized, and we generate the MBR for the data objects in $s$ and use it as a bottom cluster (lines 14 to 16). The process terminates when $Q$ becomes empty (line 4).

When finding the optimal splitting value along a dimension (lines 18 to 24), we apply SGD (Saad, 1998) to minimize Eq. 4 (line 21). Here, we use a map structure $opt$ to record the new object checking cost, the dimension, and the value of a learned optimal split.

The time complexity of each iteration in Algorithm 2 is $\mathcal{O}\big{(}h\cdot(E_{x}+E_{y})\big{)}$ where $h$ and $E_{d},d\in\{x,y\}$ denote the time complexity of SGD per iteration and the number of epoches respectively. Recall that the time complexity of the brute-force algorithm is $\mathcal{O}\big{(}f\cdot(|V_{x}|+|V_{y}|)\big{)}$. We note that $f$ is larger than $h$ because our heuristic algorithm does not need to run a split to calculate a query cost (while the brute-force algorithm does). $E_{x}$ and $E_{y}$ depend on the algorithm configurations, such as the learning rate and the number of model parameters. They are usually much smaller than $V_{x}$ and $V_{y}$, respectively. Therefore, the time complexity of our heuristic algorithm is lower than that of the brute-force algorithm.

The query time spent on node pruning can directly reflect the pruning capability of a hierarchical index. Measuring the query time, however, needs to run all queries in the query workload on an existing index, which is not suitable to be used as an optimization metric of our bottom-up packing problem. We observe that the pruning time cost is proportional to the number of accessed nodes for the workload such that it can be used to evaluate the pruning capability. To adopt this criterion, we associate each bottom cluster $c_{i}$ with a query label set denoted by $c_{i}.l$. If a cluster $c_{i}$ intersects with a training query $q_{j}$ and its textual document includes any keyword of $q_{j}$, we add $q_{j}$ to the query label set of this cluster, that is, $c_{i}.l=\{q_{j}\}$. During packing, the labels of a node in an upper level (an “upper node” for short hereafter) can be easily generated by merging all labels of its sub-tree.

Next, we define the bottom-up packing problem to minimize the number of accessed nodes.

Given the leaf nodes, i.e., bottom clusters, we can build a hierarchical index using techniques from traditional indexes such as the CDIR-Tree. However, those techniques only consider the underlying data distribution, which might lead to worse performance as shown in the later experiment.

To address these issues, and motivated by the strong performance of query-aware structures learned by reinforcement learning (RL), we propose an RL-based algorithm to learn a packing. We construct our index level by level with a bottom-up packing process, and we model the packing problem at each level as a sequential decision-making process, i.e., a Markov decision process, which makes it solvable by RL. To pack each level, the nodes from a lower level to be packed (“bottom nodes” hereafter) are processed sequentially, and we find an upper node to host each bottom node until there are no more bottom nodes. After the packing process of one level stops, the non-empty upper nodes become the new bottom nodes to be packed for the next level.

In this experiment, we fix other settings except for the query distribution and show the results on all datasets in Figure 8. Clearly, conventional indexes (SFC-Quad, ST2I, and CDIR-Tree) perform worse on the skewed workload since they do not use the query characteristics when constructing the index. For the learned indexes, TFI performs even worse than the conventional indexes since it only loosely combines a learned spatial index with a textual index. Flood-T shows a slight fluctuation in its performance since it learns from the underlying data and the query workload simultaneously, but in the geo-textual scenario, it only splits along one dimension, making it incompatible with the skewed workload. Our WISK improves the partitioning and adopts RL to build a tree, so it is less sensitive to this alteration.

We show the performance of all indexes, by varying the query region size varies from $0.005\%$ to $1\%$ of the whole region in Figure 9. Again, WISK performs the best on four datasets. Besides, Flood-T performs slightly worse than ST2I on SP, even though it optimizes its layout by learning from the data and query workload. It is because it only splits the whole region along one dimension. Thus, we improve this process to generate the leaf nodes of WISK and also pack the bottom clusters into a hierarchical structure. The two techniques simultaneously result in the superiority of WISK over the other indexes.

We evaluate the query sets with different numbers of keywords. Figure 10 shows that the query time of all indexes grows with the query keyword set size. The reason is that with the increase in the number of query keywords, more candidates need to be verified after the filtering step. Besides, WISK consistently outperforms other baseline indexes, and its cost grows much slower than those of others, e.g., the increased time of WISK on BPD is around 100 $\mu$s while those of Flood-T and ST2I are both over 250 $\mu$s. Hence, compared to other indexes, WISK is less sensitive to the number of query keywords.

We generate five sub-datasets of OSM containing from 1 to 100 million objects and run experiments on these sub-datasets. We choose ST2I, LSTI, and Flood-T as our baselines. As shown in Figure 12, the query processing time increases with the size of the dataset, but WISK performs more stable.

We evaluate the performance of ST2I, Flood-T, and WISK when the query distribution changes on FS. We initially train Flood-T and WISK based on the query workload with UNI distribution. Then, we keep the index consistent and adjust the ratio of queries with LAP distribution from 0.2 to 1.0 in the testing query set. As shown in Figure 12, the performance of query-aware indexes becomes worse when query distribution is more different from the training one. However, it can be seen that WISK is more robust than Flood-T due to its improved partitioning algorithm and the bottom-up packing process. Additionally, the query time of ST2I also increases since it ignores the query knowledge when building the index, but the fluctuation is less than the one of Flood-T.

Table 3 reports the index sizes. Overall, WISK costs less space than conventional indexes but is comparable to that of the best adapted learned indexes. In particular, the size of CDIR-Tree is larger than those of the others since each of its nodes has an inverted file. For query efficiency, we have not compressed SFC-Quad, which leads to a larger size. ST2I has the smallest size among the conventional indexes. Among the learned indexes, the sizes of TFI are the largest, as it uses inverted files. WISK takes more space than Flood-T on a small dataset since the number of its bottom clusters is similar to the number of the columns of Flood-T, but we build a hierarchical index. However, on larger datasets, WISK needs less space cost, since Flood-T splits more columns for better performance and builds inverted files for them.

We compare the efficiency of index construction algorithms and report the results in Table 4. It takes the minimum time to build SFC-Quad and ST2I on small datasets. However, the time cost of ST2I significantly increases when the dataset becomes larger, since ST2I is built based on the set of converted points, and its time cost is positively correlated to the total number of keywords. CDIR-Tree takes the highest time cost because it inserts the objects sequentially.

For the learned indexes, we report the training time. LSTI takes the least time to build because it only needs to scan the whole dataset twice. The time cost of TFI increases significantly when there are more different keywords. For Flood-T and WISK, we report the average time as the time costs of query-aware learned indexes usually increase with more query keywords.

We designed two training time acceleration techniques as presented in Section 6. We report the training and query times of WISK with different sampling ratios in Figure 13(a). The result of each sampling ratio is an average of 10 runs. While the training time decreases by 72%, we do not observe a large drop in query performance with a sample of only 30% of the full query workload. We also observe that the standard deviation (represented by the width of the bands) of the training and query times of WISK is consistently small for all sampling ratios. This demonstrates that WISK has a stable performance using stratified sampling. We vary the clustering ratio, i.e., the number of groups obtained over the number of bottom clusters, to balance the training and querying time. Figure 13(b) shows that even when the number of bottom clusters decreases by 80%, the query time of WISK still only changes slightly. We set the sampling ratio to 30% and the clustering ratio to 20% and generate the Accelerated WISK. As shown in Table 4, WISK has longer training times than the other learned indexes, but the acceleration techniques can reduce index training time up to 4 times while the query time is only affected marginally.

To update the index when query distribution changes, we can retrain WISK periodically following the former study (Nathan et al., 2020). We generate six workloads for FS. For each workload, we randomly select the query region size and the number of query keywords, and the query distribution adopts the default settings (MIX) and we randomly select the proportions of UNI and LAP. Each workload runs for 30 minutes and consists of 100 queries. As Figure 14 shows, at the start of each 30-minute period, i.e., a new query workload starts, retraining WISK is triggered, which happens in a separate thread and does not interrupt the query processing. While the index is being rebuilt, WISK runs the new queries on its old layout, which explains the jumps in the figure. The retraining lasts about 3 minutes, and then WISK switches to the new layout adapted to the new query workload. Thus, the query time drops back again.

To capture minor changes in query distribution when retraining, we propose to apply incremental updates to the original index, in parallel to the retraining process. We locate the bottom clusters that are affected by the new queries, re-partition these clusters if the query costs can be reduced using new queries, and then insert the new clusters back into the non-leaf nodes they previously belonged to. The incremental updates may also help reduce the query times, which explains the multiple drops (e.g., at 00:30 and 01:30).

Figure 14 also indicates the necessity of learning from the query workload. When a new workload arrives, the performance of WISK drops due to its outdated layout. Re-learning the layouts based on the new workload mitigates the impact of the changing query distribution. The other two indexes do not utilize the queries during construction, which leads to much worse performance than WISK (e.g. at 01:00, 02:00, and 02:30) on the skewed query workloads (i.e. high proportion of LAP).

WISK also handles data insertion well. Given a new object o, we can traverse WISK to find the bottom cluster where o falls. Next, we update the inverted file or the bitmap of the affected nodes to obtain an updated index. This simple process, however, cannot guarantee an optimal layout because the bottom clusters might need to be split after the insertion. Thus, we buffer the inserted objects and retrain our index when the buffer is full.

We set the buffer size at 100,000 (around 20MB) and run experiments to examine the impact of data insertions. We randomly select 500,0000 objects from FS for the insertions. We insert 100,000 objects every 30 minutes. Figure 15 shows the performance of ST2I, LSTI, and WISK. We compare with WISK using the simple insertion process without retraining. It can be seen that the query time of all indexes increases when more objects are inserted. Between the two WISK variants, we see that the query time of WISK without retraining increases faster with more insertions, thus verifying the importance of retraining in improving the query time of WISK under dynamic data settings. We also observe that the retraining process takes only 1 to 2 minutes each time, since only the affected bottom clusters need to be split, and the RL-based packing can inherit knowledge from the previous training process, i.e., the unaffected bottom clusters are initially packed into the previous corresponding upper nodes.

We conduct an experiment to compare the cost at the leaf level and that at the non-leaf level. Figure 17 shows that the time at the leaf level dominates the query processing time, which occupies around 90% of the total cost, verifying the way to define cost function is reasonable. In Figure 17, we observe that packing our bottom clusters by directly using CDIR-Tree construction method might affect the query time because it may pack some leaf nodes intersecting with various queries.

We evaluate the effectiveness of the bottom-up construction process. As shown in Figure 18(a), the improvement of the different number of keywords is similar. This is because the number of query keywords has little effect on the number of bottom clusters. Thus, the improvement is stable using this RL-based grouping algorithm.

It can be seen from Figure 18(b) that the improvement of hierarchical indexes becomes more significant for queries of larger region sizes. This is because queries of larger region size correspond to a wider space covered by these queries and more bottom nodes such that the bottom-up packing process can reduce more filtering cost. However, we also observe that the improvement becomes stable as the region size continues to get larger, as the query region covers most of the data space.

When generating bottom clusters, we use CDF models to estimate the number of objects sharing the same keywords inside a region. To reduce the parameters, we propose to use Gaussian functions and NNs for keywords with different frequencies. In Figure 19(a), we compare our method with the settings in which only Gaussian models or NN models are used. Although the Gaussian-only method has the least training time, its estimation results are inaccurate, leading to much worse query time. In contrast, the NN-only method achieves the best query time, but it needs much more training time. In comparison, the proposed mixed method achieves similar query performance as the NN-only method without significantly increasing the training time.

To speed up pre-processing, we assume that the two spatial dimensions are independent following the existing work (Nathan et al., 2020). We next study the impact of such an assumption.

We observe that keywords with higher frequency have a stronger impact on the query time, and thus they need a more accurate CDF estimation. We run experiments with a randomly selected high-frequency keyword on FS. We train two marginal (1D) models and a joint (2D) model for the selected keyword on the whole dataset. The 1D models and the 2D model all employ a neural network with 2 hidden layers and 16 hidden units. We also compare with two variants of the 2D model, one uses more hidden units (32 units) and the other uses more layers (3 layers). We randomly sample 1,000 rectangular query regions within the data space as testing data. For each query region, we compute the proportion of objects that fall within it as the ground truth. The product of the two 1D models and the output of the 2D models are used as the estimation results corresponding to 1D and 2D models, respectively. For every 50 training rounds, we calculate the mean squared error (MSE) between the estimations and the ground truth. Figure 19(b) shows that the 1D model converges much faster while its final loss is comparable to those of the 2D models. These results justify the use of 1D CDFs in our method.

The FI mining extracts the frequent interrelation among keywords. In this experiment, the minimum support is set to $0.01\text{\textperthousand}$ and the maximum size of target itemsets is equal to the number of query keywords. Figure 20 shows the effect of FI on the index construction by query efficiency on FS and BPD. We see that the FI mining improves the performance of WISK consistently when there is more than one query keyword. Without the FI mining component, we learn models of each keyword separately where redundancies occur when there is more than one keyword. We also observe that this adaptation is more beneficial given more query keywords. In general, more query keywords lead to a higher possibility of resulting in redundancy since the probability of an object including more than one query keyword increases.

Besides, the performance improvement on BPD is more obvious than that on FS. This is because the number of distinct keywords on FS is much less, making the number of frequent itemsets less than others. Additionally, the improvement becomes consistent when the number of query keywords is larger than a threshold. The reason is that each object includes finite keywords, and thus we cannot generate frequent itemsets with more keywords.

When using the RL framework, the environment applies the action mask to reduce the action space. Here, we evaluate its effectiveness in two aspects. We use the SmoothL1Loss with the sum reduction as the loss function in our RL framework. Figure 21(a) shows that the RL framework with the action mask can speed up the model convergence and reach a smaller loss. In Figure 21(b), we sum the total rewards and the number of bottom nodes in each epoch to show the reduction in the average number of accessed nodes. The result shows that the pruning capability is always better when applying the action mask.

there are some fluctuations in the results as the RL agent learns from the feedback through trial-and-error interactions with the environment. RL balances the trade-off between exploration and exploitation, which reduces the fluctuation with more training epochs. These results confirm that the action mask helps to decrease the number of training epochs, leading to less training time.