Preserve Row Indices with Spark Matrix Multiplication
By: Shomit Goyal
Tags: Getting Started, Programming Tricks, Spark, Collaborative Filtering, Math Tricks

Preserve Row Indices with Spark Matrix Multiplication

Photo by Compare Fibre on Unsplash

1. Introduction

Matrix multiplications are quite common in machine learning. For example, in case of a fully connected neural network we can vectorise the forward prop and define it as a sequence of matrix multiplications. This vectorised implementation gives significant boost to the computations and makes training and scoring of deep neural nets feasible. Matrix multiplications are also very common in collaborative filtering implementations where we have a guest matrix consisting of every guest’s latent vectors and an item matrix consisting of every item’s latent vectors. The multiplication of these two matrices gives us the full guest-item rating matrix indicating every guest’s predicted rating for every item.

In this article we’ll start by briefly describing spark MLlib’s matrix API and how we can use it for large scale matrix multiplication. Then we’ll get into the problem of row index preservation and talk about a neat little trick where we’ll embed the row indices as dummy values in the original matrix to solve the problem. This will enable us to get back the original row indices from the matrix multiplication output. We’ll then conclude the discussion with the pyspark implementation of this solution.

2. Spark’s Matrix API

Spark’s MLlib provides a few convenient RDD based APIs to enable efficient matrix computations. There are 2 types of matrices that are available:

Local Matrix: These are by definition local to the system, i.e. they are stored on a single machine and not distributed over cluster nodes. Local Matrix API supports both dense and sparse matrix formats and we can choose one as per our data pattern. Distributed Matrix: As the name suggests, these matrices are stored distributively on one or more nodes, this enables us to store and manipulate very large matrices such as the ones that we generally deal with while building machine learning algorithms such as collaborative filtering. There are 4 types of distributed matrix formats that are available at the moment:

RowMatrix : Stores every row of the matrix as local vector and is defined on top of an rdd where every element of the rdd represents a single row of matrix. It does not store the row numbers (indices).

: Stores every row of the matrix as local vector and is defined on top of an rdd where every element of the rdd represents a single row of matrix. It does not store the row numbers (indices). IndexedRowMatrix : Just like a RowMatrix but it stores the row numbers (indices) along with the row vector. It is defined on top of an rdd of indexed rows where every IndexedRow contains the row index and a local vector.

: Just like a RowMatrix but it stores the row numbers (indices) along with the row vector. It is defined on top of an rdd of indexed rows where every IndexedRow contains the row index and a local vector. CoordinateMatrix : Stores the row and column indices along with the actual value. It’s defined on top of an rdd where every element of the rdd is a 3-tuple containing (row index, column index, entry value).

: Stores the row and column indices along with the actual value. It’s defined on top of an rdd where every element of the rdd is a 3-tuple containing (row index, column index, entry value). BlockMatrix : Splits the original matrix into smaller sub-matrices called blocks where every block is of size rowsPerBlock x colsPerBlock . These matrices can be used to distribute very large matrices wrt both rows and columns.

The matrix APIs discussed above provide functionalities to perform various matrix algebra operations such as addition, multiplication, transpose etc. In this article we’ll be focusing on matrix multiplication and a challenge associated with it, called row index preservation.

3. Matrix Multiplication

Say we want to multiply a matrix A of dimensions m x n with another matrix B of dimensions n x k which will result into a new matrix of dimensions m x k . In its most basic form, this multiplication can be thought of as the dot product between every row vector of matrix A with every column vector of matrix B resulting into a total of m*k vector dot products. If the matrix A and B were small and runtime wasn’t a concern we could simply store both the matrices A and B as separate RDDs and perform a cross join of two RDDs followed by a map operation to get an element-wise dot product and this will give us the values of the desired matrix. But in practise, while working on big data applications (think deep learning, collaborative filtering etc) we deal with matrices which are very large, say 100 million x 1 million . In such large settings, carrying out the cross join and vector dot product is very inefficient and suboptimal. Here we would like to utilise the vectorised algorithms for matrix multiplication which are much faster than the element wise multiplication or vector dot products. The matrix APIs provided by spark allows us to do exactly that. We can use either IndexedRowMatrix or BlockMatrix from spark’s MLlib to carry out distributed matrix multiplications. In practise I have found BlockMatrix multiplications to be rather inefficient with respect to the computation time, so we will focus our attention to matrix multiplications using IndexedRowMatrix but the content that follows is applicable to BlockMatrix as well.

4. Problem of row index preservation

Entry at (i, j) position in a matrix that is a result of a multiplication of two matrices represents the multiplication of ith row of the first matrix with jth row of the second matrix. When using the matrix multiplication function of spark’s IndexedRowMatrix API which lets us multiply an IndexedRowMatrix with a LocalMatrix , the row ordering is guaranteed to be preserved only within every partition. What this means is that if the underlying rdd of our IndexedRowMatrix is distributed across many partitions then after multiplying it with a local matrix the ith row in the resulting matrix is not guaranteed to correspond to the ith row of the original matrix. This poses a huge issue as we don’t have any way to link the results back to the original row identifiers.

Some related issues on apache spark’s Jira board and github (not yet resolved):

https://issues.apache.org/jira/browse/SPARK-8614

https://github.com/hl475/svd/pull/1

For example say we are working on a collaborative filtering algorithm and we have a guest matrix where we know which row index maps to which guest, we also have an item matrix. We multiply these two matrices and get the guest-item rating matrix, but if the row orders are not preserved, a particular row in the rating matrix will not necessarily correspond to the same guest that it was pointing to in the original guest matrix.

5. A quick solution to extract the original row indices after multiplying two matrices