Vector index building in distributed computing system
A data processing service receives a data table which includes a plurality of rows of data. The service assigns the plurality of rows of data to a plurality of computing nodes. The computing nodes belongs to a distributed computing system. At each assigned computing node, the service applies a machine learning (ML) model to the assigned rows of data for indexing each row in the data table. The ML model is trained to determine a set of centroids, and each centroid is associated with a centroid ID. For each row of data in the data table, the service receives a centroid ID of one of the plurality of centroids that the respective row of data is assigned to. The service generates an index of the data table with at least the received centroid ID of each of the plurality of row of data in the data table.
1 . A method comprising:
receiving a data table, the data table comprising a plurality of rows of data;
assigning the plurality of rows of data to a plurality of computing nodes;
applying, at each assigned computing node, a machine learning model to the assigned rows of data for indexing each row in the data table, wherein training the machine learning model comprises:
accessing a plurality of samples to be processed over one or more iterations of training the machine learning model;
in each iteration,
shuffling the plurality of samples;
partitioning the plurality of samples into one or more sets, each set comprising a set of the samples;
assigning the one or more sets of samples to the plurality of computing nodes;
applying, at each of the plurality of computing nodes, the machine learning model to the assigned sets to cluster the plurality of samples; and
receiving an output from the machine learning model, the output comprising one or more centroids, each centroid clustered with a cluster of samples;
performing the one or more iterations of training the machine learning model to stabilize the plurality of centroids; and
determining a set of stabilized centroids, each centroid having a centroid ID;
receiving, for each row of data in the data table, a centroid ID of one of the plurality of centroids that the respective row of data is assigned to; and
generating an index of the data table with at least the received centroid ID of each of the plurality of rows of data in the data table.
2 . The method of claim 1 , wherein accessing a plurality of samples to be processed over one or more iterations of training the machine learning model comprises:
generating a plurality of embedding vectors, each embedding vector representing one of the plurality of samples.
3 . The method of claim 1 , wherein applying, at each of the plurality of computing nodes, the machine learning model to the assigned sets to cluster the plurality of samples comprises:
applying K-Means algorithm to each of the assigned set of samples; and
assigning each of the assigned set of samples to one of the one or more centroids, the one of the one or more centroids is a nearest centroid to the respective sample.
4 . The method of claim 1 , wherein performing the one or more iterations of training the machine learning model to stabilize the plurality of centroids comprises:
determining that a change in the plurality of centroids between two successive iterations meets a predetermined condition.
5 . The method of claim 1 , wherein generating an index of the data table comprises:
computing a product quantization (PQ) code for each row of data; and
generating the index of the data table comprising at least the centroid ID and the PQ code for each of the plurality of rows of data in the data table.
6 . The method of claim 1 , further comprising:
partitioning the plurality of rows of data in inverted file index (IVF) partitions based on a plurality of centroid IDs, each IVF partition comprising a set of rows of data that are assigned to a same centroid.
7 . The method of claim 1 , wherein generating an index of the data table comprises:
generating a secondary index comprising metadata indexing, the metadata comprising information of the respective row of data.
8 . A non-transitory computer readable storage medium comprising stored program code, the program code comprising instructions, the instructions when executed cause a processor system to:
receive a data table, the data table comprising a plurality of rows of data;
assign the plurality of rows of data to a plurality of computing nodes;
apply, at each assigned computing node, a machine learning model to the assigned rows of data for indexing each row in the data table, wherein training the machine learning model comprises:
accessing a plurality of samples to be processed over one or more iterations of training the machine learning model;
in each iteration,
shuffling the plurality of samples;
partitioning the plurality of samples into one or more sets, each set comprising a set of the samples;
assigning the one or more sets of samples to the plurality of computing nodes;
applying, at each of the plurality of computing nodes, the machine learning model to the assigned sets to cluster the plurality of samples; and
receiving an output from the machine learning model, the output comprising one or more centroids, each centroid clustered with a cluster of samples;
performing the one or more iterations of training the machine learning model to stabilize the plurality of centroids; and
determining a set of stabilized centroids, each centroid having a centroid ID;
receive, for each row of data in the data table, a centroid ID of one of the plurality of centroids that the respective row of data is assigned to; and
generate an index of the data table with at least the received centroid ID of each of the plurality of rows of data in the data table.
9 . The non-transitory computer readable storage medium of claim 8 , wherein the instructions to access a plurality of samples to be processed over one or more iterations of training the machine learning model cause the processor system to:
generate a plurality of embedding vectors, each embedding vector representing one of the plurality of samples.
10 . The non-transitory computer readable storage medium of claim 8 , wherein the instructions to apply, at each of the plurality of computing nodes, the machine learning model to the assigned sets to cluster the plurality of samples cause the processor system to:
apply K-Means algorithm to each of the assigned set of samples; and
assign each of the assigned set of samples to one of the one or more centroids, the one of the one or more centroids is a nearest centroid to the respective sample.
11 . The non-transitory computer readable storage medium of claim 8 , wherein the instructions to perform the one or more iterations of training the machine learning model to stabilize the plurality of centroids cause the processor system to:
determine that a change in the plurality of centroids between two successive iterations meets a predetermined condition.
12 . The non-transitory computer readable storage medium of claim 8 , wherein the instructions to generate an index of the data table cause the processor system to:
compute a product quantization (PQ) code for each row of data; and
generate the index of the data table comprising at least the centroid ID and the PQ code for each of the plurality of rows of data in the data table.
13 . The non-transitory computer readable storage medium of claim 8 , wherein the instructions cause a processor system to:
partition the plurality of rows of data in inverted file index (IVF) partitions based on a plurality of centroid IDs, each IVF partition comprising a set of rows of data that are assigned to a same centroid.
14 . The non-transitory computer readable storage medium of claim 8 , wherein the instructions to generate an index of the data table cause the processor system to:
generate a secondary index comprising metadata indexing, the metadata comprising information of the respective row of data.
15 . A system comprising:
one or more computer processors; and
one or more computer-readable mediums storing instructions that, when executed by the one or more computer processors, cause the system to:
receive a data table, the data table comprising a plurality of rows of data;
assign the plurality of rows of data to a plurality of computing nodes;
apply, at each assigned computing node, a machine learning model to the assigned rows of data for indexing each row in the data table, wherein training the machine learning model comprises:
accessing a plurality of samples to be processed over one or more iterations of training the machine learning model;
in each iteration,
shuffling the plurality of samples;
partitioning the plurality of samples into one or more sets, each set comprising a set of the samples;
assigning the one or more sets of samples to the plurality of computing nodes;
applying, at each of the plurality of computing nodes, the machine learning model to the assigned sets to cluster the plurality of samples; and
receiving an output from the machine learning model, the output comprising one or more centroids, each centroid clustered with a cluster of samples;
performing the one or more iterations of training the machine learning model to stabilize the plurality of centroids; and
determining a set of stabilized centroids, each centroid having a centroid ID;
receive, for each row of data in the data table, a centroid ID of one of the plurality of centroids that the respective row of data is assigned to; and
generate an index of the data table with at least the received centroid ID of each of the plurality of rows of data in the data table.
16 . The system of claim 15 , wherein the instructions to access a plurality of samples to be processed over one or more iterations of training the machine learning model cause the system to:
generate a plurality of embedding vectors, each embedding vector representing one of the plurality of samples.
17 . The system of claim 15 , wherein the instructions to apply, at each of the plurality of computing nodes, the machine learning model to the assigned sets to cluster the plurality of samples cause the system to:
apply K-Means algorithm to each of the assigned set of samples; and
assign each of the assigned set of samples to one of the one or more centroids, the one of the one or more centroids is a nearest centroid to the respective sample.
18 . The system of claim 15 , wherein the instructions to perform the one or more iterations of training the machine learning model to stabilize the plurality of centroids cause the system to:
determine that a change in the plurality of centroids between two successive iterations meets a predetermined condition.
19 . The system of claim 15 , wherein the instructions to generate an index of the data table cause the system to:
compute a product quantization (PQ) code for each row of data; and
generate the index of the data table comprising at least the centroid ID and the PQ code for each of the plurality of rows of data in the data table.
20 . The system of claim 15 , wherein the instructions cause the system to:
partition the plurality of rows of data in inverted file index (IVF) partitions based on a plurality of centroid IDs, each IVF partition comprising a set of rows of data that are assigned to a same centroid.