You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@ignite.apache.org by is...@apache.org on 2017/10/31 12:38:15 UTC

ignite git commit: IGNITE-6761: Change block numeration in SparseBlockMatrixStorage.

Repository: ignite
Updated Branches:
  refs/heads/master 035dcb7f7 -> 94dd19d36


IGNITE-6761: Change block numeration in SparseBlockMatrixStorage.


Project: http://git-wip-us.apache.org/repos/asf/ignite/repo
Commit: http://git-wip-us.apache.org/repos/asf/ignite/commit/94dd19d3
Tree: http://git-wip-us.apache.org/repos/asf/ignite/tree/94dd19d3
Diff: http://git-wip-us.apache.org/repos/asf/ignite/diff/94dd19d3

Branch: refs/heads/master
Commit: 94dd19d36812d0bea3a1652ff745fbefc10eba70
Parents: 035dcb7
Author: Yury Babak <yb...@gridgain.com>
Authored: Tue Oct 31 15:35:28 2017 +0300
Committer: Igor Sapego <is...@gridgain.com>
Committed: Tue Oct 31 15:37:40 2017 +0300

----------------------------------------------------------------------
 .../math/distributed/keys/BlockMatrixKey.java   |  12 +-
 .../distributed/keys/impl/BlockMatrixKey.java   |  49 +++--
 .../matrix/SparseBlockDistributedMatrix.java    |  15 +-
 .../impls/matrix/SparseDistributedMatrix.java   |   2 +-
 .../storage/matrix/BlockMatrixStorage.java      | 213 +++++++++----------
 .../SparseDistributedBlockMatrixTest.java       |  15 +-
 6 files changed, 157 insertions(+), 149 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/ignite/blob/94dd19d3/modules/ml/src/main/java/org/apache/ignite/ml/math/distributed/keys/BlockMatrixKey.java
----------------------------------------------------------------------
diff --git a/modules/ml/src/main/java/org/apache/ignite/ml/math/distributed/keys/BlockMatrixKey.java b/modules/ml/src/main/java/org/apache/ignite/ml/math/distributed/keys/BlockMatrixKey.java
index c55b950..091b325 100644
--- a/modules/ml/src/main/java/org/apache/ignite/ml/math/distributed/keys/BlockMatrixKey.java
+++ b/modules/ml/src/main/java/org/apache/ignite/ml/math/distributed/keys/BlockMatrixKey.java
@@ -17,14 +17,22 @@
 
 package org.apache.ignite.ml.math.distributed.keys;
 
+import org.apache.ignite.internal.util.lang.IgnitePair;
 import org.apache.ignite.ml.math.impls.matrix.SparseBlockDistributedMatrix;
 
 /**
  * Cache key for blocks in {@link SparseBlockDistributedMatrix}.
+ *
+ * TODO: check if using {@link IgnitePair} will be better for block id.
  */
 public interface BlockMatrixKey extends MatrixCacheKey {
     /**
-     * @return block id.
+     * @return block row id.
+     */
+    public long blockRowId();
+
+    /**
+     * @return block col id.
      */
-    public long blockId();
+    public long blockColId();
 }

http://git-wip-us.apache.org/repos/asf/ignite/blob/94dd19d3/modules/ml/src/main/java/org/apache/ignite/ml/math/distributed/keys/impl/BlockMatrixKey.java
----------------------------------------------------------------------
diff --git a/modules/ml/src/main/java/org/apache/ignite/ml/math/distributed/keys/impl/BlockMatrixKey.java b/modules/ml/src/main/java/org/apache/ignite/ml/math/distributed/keys/impl/BlockMatrixKey.java
index 5fd1a16..2edd9cb 100644
--- a/modules/ml/src/main/java/org/apache/ignite/ml/math/distributed/keys/impl/BlockMatrixKey.java
+++ b/modules/ml/src/main/java/org/apache/ignite/ml/math/distributed/keys/impl/BlockMatrixKey.java
@@ -42,8 +42,10 @@ import org.jetbrains.annotations.Nullable;
 public class BlockMatrixKey implements org.apache.ignite.ml.math.distributed.keys.BlockMatrixKey, Externalizable, Binarylizable {
     /** */
     private static final long serialVersionUID = 0L;
-    /** Block ID */
-    private long blockId;
+    /** Block row ID */
+    private long blockIdRow;
+    /** Block col ID */
+    private long blockIdCol;
     /** Matrix ID */
     private IgniteUuid matrixUuid;
     /** Block affinity key. */
@@ -59,22 +61,28 @@ public class BlockMatrixKey implements org.apache.ignite.ml.math.distributed.key
     /**
      * Construct matrix block key.
      *
-     * @param blockId Block id.
      * @param matrixUuid Matrix uuid.
      * @param affinityKey Affinity key.
      */
-    public BlockMatrixKey(long blockId, IgniteUuid matrixUuid, @Nullable IgniteUuid affinityKey) {
-        assert blockId >= 0;
+    public BlockMatrixKey(long rowId, long colId,  IgniteUuid matrixUuid, @Nullable IgniteUuid affinityKey) {
+        assert rowId >= 0;
+        assert colId >= 0;
         assert matrixUuid != null;
 
-        this.blockId = blockId;
+        this.blockIdRow = rowId;
+        this.blockIdCol = colId;
         this.matrixUuid = matrixUuid;
         this.affinityKey = affinityKey;
     }
 
     /** {@inheritDoc} */
-    @Override public long blockId() {
-        return blockId;
+    @Override public long blockRowId() {
+        return blockIdRow;
+    }
+
+    /** {@inheritDoc} */
+    @Override public long blockColId() {
+        return blockIdCol;
     }
 
     /** {@inheritDoc} */
@@ -91,14 +99,16 @@ public class BlockMatrixKey implements org.apache.ignite.ml.math.distributed.key
     @Override public void writeExternal(ObjectOutput out) throws IOException {
         U.writeGridUuid(out, matrixUuid);
         U.writeGridUuid(out, affinityKey);
-        out.writeLong(blockId);
+        out.writeLong(blockIdRow);
+        out.writeLong(blockIdCol);
     }
 
     /** {@inheritDoc} */
     @Override public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException {
         matrixUuid = U.readGridUuid(in);
         affinityKey = U.readGridUuid(in);
-        blockId = in.readLong();
+        blockIdRow = in.readLong();
+        blockIdCol = in.readLong();
     }
 
     /** {@inheritDoc} */
@@ -107,7 +117,8 @@ public class BlockMatrixKey implements org.apache.ignite.ml.math.distributed.key
 
         BinaryUtils.writeIgniteUuid(out, matrixUuid);
         BinaryUtils.writeIgniteUuid(out, affinityKey);
-        out.writeLong(blockId);
+        out.writeLong(blockIdRow);
+        out.writeLong(blockIdCol);
     }
 
     /** {@inheritDoc} */
@@ -116,12 +127,19 @@ public class BlockMatrixKey implements org.apache.ignite.ml.math.distributed.key
 
         matrixUuid = BinaryUtils.readIgniteUuid(in);
         affinityKey = BinaryUtils.readIgniteUuid(in);
-        blockId = in.readLong();
+        blockIdRow = in.readLong();
+        blockIdCol = in.readLong();
     }
 
     /** {@inheritDoc} */
     @Override public int hashCode() {
-        return matrixUuid.hashCode() + (int)(blockId ^ (blockId >>> 32));
+        int res = 37;
+
+        res += res * 37 + blockIdCol;
+        res += res * 37 + blockIdRow;
+        res += res * 37 + matrixUuid.hashCode();
+
+        return res;
     }
 
     /** {@inheritDoc} */
@@ -134,11 +152,14 @@ public class BlockMatrixKey implements org.apache.ignite.ml.math.distributed.key
 
         BlockMatrixKey that = (BlockMatrixKey)obj;
 
-        return blockId == that.blockId && matrixUuid.equals(that.matrixUuid) && F.eq(affinityKey, that.affinityKey);
+        return blockIdRow == that.blockIdRow && blockIdCol == that.blockIdCol && matrixUuid.equals(that.matrixUuid)
+            && F.eq(affinityKey, that.affinityKey);
     }
 
     /** {@inheritDoc} */
     @Override public String toString() {
         return S.toString(BlockMatrixKey.class, this);
     }
+
+
 }

http://git-wip-us.apache.org/repos/asf/ignite/blob/94dd19d3/modules/ml/src/main/java/org/apache/ignite/ml/math/impls/matrix/SparseBlockDistributedMatrix.java
----------------------------------------------------------------------
diff --git a/modules/ml/src/main/java/org/apache/ignite/ml/math/impls/matrix/SparseBlockDistributedMatrix.java b/modules/ml/src/main/java/org/apache/ignite/ml/math/impls/matrix/SparseBlockDistributedMatrix.java
index 1dcf1d8..3d542bc 100644
--- a/modules/ml/src/main/java/org/apache/ignite/ml/math/impls/matrix/SparseBlockDistributedMatrix.java
+++ b/modules/ml/src/main/java/org/apache/ignite/ml/math/impls/matrix/SparseBlockDistributedMatrix.java
@@ -25,6 +25,7 @@ import org.apache.ignite.IgniteCache;
 import org.apache.ignite.Ignition;
 import org.apache.ignite.cache.affinity.Affinity;
 import org.apache.ignite.cluster.ClusterNode;
+import org.apache.ignite.internal.util.lang.IgnitePair;
 import org.apache.ignite.lang.IgniteUuid;
 import org.apache.ignite.ml.math.Matrix;
 import org.apache.ignite.ml.math.StorageConstants;
@@ -106,7 +107,7 @@ public class SparseBlockDistributedMatrix extends AbstractMatrix implements Stor
 
         CacheUtils.bcast(cacheName, () -> {
             Ignite ignite = Ignition.localIgnite();
-            Affinity affinity = ignite.affinity(cacheName);
+            Affinity<BlockMatrixKey> affinity = ignite.affinity(cacheName);
 
             IgniteCache<BlockMatrixKey, BlockEntry> cache = ignite.getOrCreateCache(cacheName);
             ClusterNode locNode = ignite.cluster().localNode();
@@ -122,11 +123,15 @@ public class SparseBlockDistributedMatrix extends AbstractMatrix implements Stor
             // compute Cij locally on each node
             // TODO: IGNITE:5114, exec in parallel
             locKeys.forEach(key -> {
-                long newBlockId = key.blockId();
+                long newBlockIdRow = key.blockRowId();
+                long newBlockIdCol = key.blockColId();
+
+                IgnitePair<Long> newBlockId = new IgnitePair<>(newBlockIdRow, newBlockIdCol);
+
                 BlockEntry blockC = null;
 
-                List<BlockEntry> aRow = matrixA.storage().getRowForBlock(newBlockId, storageC);
-                List<BlockEntry> bCol = matrixB.storage().getColForBlock(newBlockId, storageC);
+                List<BlockEntry> aRow = matrixA.storage().getRowForBlock(newBlockId);
+                List<BlockEntry> bCol = matrixB.storage().getColForBlock(newBlockId);
 
                 for (int i = 0; i < aRow.size(); i++) {
                     BlockEntry blockA = aRow.get(i);
@@ -137,7 +142,7 @@ public class SparseBlockDistributedMatrix extends AbstractMatrix implements Stor
                     blockC = blockC == null ? tmpBlock : new BlockEntry(blockC.plus(tmpBlock));
                 }
 
-                cache.put(storageC.getCacheKey(newBlockId), blockC);
+                cache.put(storageC.getCacheKey(newBlockIdRow, newBlockIdCol), blockC);
             });
         });
 

http://git-wip-us.apache.org/repos/asf/ignite/blob/94dd19d3/modules/ml/src/main/java/org/apache/ignite/ml/math/impls/matrix/SparseDistributedMatrix.java
----------------------------------------------------------------------
diff --git a/modules/ml/src/main/java/org/apache/ignite/ml/math/impls/matrix/SparseDistributedMatrix.java b/modules/ml/src/main/java/org/apache/ignite/ml/math/impls/matrix/SparseDistributedMatrix.java
index 92d7c39..9a18f8b 100644
--- a/modules/ml/src/main/java/org/apache/ignite/ml/math/impls/matrix/SparseDistributedMatrix.java
+++ b/modules/ml/src/main/java/org/apache/ignite/ml/math/impls/matrix/SparseDistributedMatrix.java
@@ -119,7 +119,7 @@ public class SparseDistributedMatrix extends AbstractMatrix implements StorageCo
 
         CacheUtils.bcast(cacheName, () -> {
             Ignite ignite = Ignition.localIgnite();
-            Affinity affinity = ignite.affinity(cacheName);
+            Affinity<RowColMatrixKey> affinity = ignite.affinity(cacheName);
 
             IgniteCache<RowColMatrixKey, BlockEntry> cache = ignite.getOrCreateCache(cacheName);
             ClusterNode locNode = ignite.cluster().localNode();

http://git-wip-us.apache.org/repos/asf/ignite/blob/94dd19d3/modules/ml/src/main/java/org/apache/ignite/ml/math/impls/storage/matrix/BlockMatrixStorage.java
----------------------------------------------------------------------
diff --git a/modules/ml/src/main/java/org/apache/ignite/ml/math/impls/storage/matrix/BlockMatrixStorage.java b/modules/ml/src/main/java/org/apache/ignite/ml/math/impls/storage/matrix/BlockMatrixStorage.java
index 0f285c2..0978e5a 100644
--- a/modules/ml/src/main/java/org/apache/ignite/ml/math/impls/storage/matrix/BlockMatrixStorage.java
+++ b/modules/ml/src/main/java/org/apache/ignite/ml/math/impls/storage/matrix/BlockMatrixStorage.java
@@ -24,8 +24,6 @@ import java.util.HashSet;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.LongStream;
 import org.apache.ignite.IgniteCache;
 import org.apache.ignite.Ignition;
 import org.apache.ignite.cache.CacheAtomicityMode;
@@ -33,6 +31,7 @@ import org.apache.ignite.cache.CacheMode;
 import org.apache.ignite.cache.CachePeekMode;
 import org.apache.ignite.cache.CacheWriteSynchronizationMode;
 import org.apache.ignite.configuration.CacheConfiguration;
+import org.apache.ignite.internal.util.lang.IgnitePair;
 import org.apache.ignite.internal.util.typedef.internal.U;
 import org.apache.ignite.lang.IgniteUuid;
 import org.apache.ignite.ml.math.MatrixStorage;
@@ -88,10 +87,8 @@ public class BlockMatrixStorage extends CacheUtils implements MatrixStorage, Sto
         this.rows = rows;
         this.cols = cols;
 
-        //cols % maxBlockEdge > 0 ? 1 : 0
-
-        this.blocksInRow = cols % maxBlockEdge == 0 ? cols / maxBlockEdge : cols / maxBlockEdge + 1;
-        this.blocksInCol = rows % maxBlockEdge == 0 ? rows / maxBlockEdge : rows / maxBlockEdge + 1;
+        this.blocksInRow = rows % maxBlockEdge == 0 ? rows / maxBlockEdge : rows / maxBlockEdge + 1;
+        this.blocksInCol = cols % maxBlockEdge == 0 ? cols / maxBlockEdge : cols / maxBlockEdge + 1;
 
         cache = newCache();
 
@@ -135,6 +132,20 @@ public class BlockMatrixStorage extends CacheUtils implements MatrixStorage, Sto
         return RANDOM_ACCESS_MODE;
     }
 
+    /**
+     * @return Blocks in column.
+     */
+    public int blocksInCol() {
+        return blocksInCol;
+    }
+
+    /**
+     * @return Blocks in row.
+     */
+    public int blocksInRow() {
+        return blocksInRow;
+    }
+
     /** {@inheritDoc} */
     @Override public void writeExternal(ObjectOutput out) throws IOException {
         out.writeInt(rows);
@@ -182,37 +193,7 @@ public class BlockMatrixStorage extends CacheUtils implements MatrixStorage, Sto
 
     /** Delete all data from cache. */
     @Override public void destroy() {
-        long maxBlockId = getBlockId(cols, rows);
-
-        Set<BlockMatrixKey> keyset = LongStream.rangeClosed(0, maxBlockId).mapToObj(this::getCacheKey).collect(Collectors.toSet());
-
-        cache.clearAll(keyset);
-    }
-
-    /** {@inheritDoc} */
-    @Override public int hashCode() {
-        int res = 1;
-
-        res = res * 37 + cols;
-        res = res * 37 + rows;
-        res = res * 37 + uuid.hashCode();
-        res = res * 37 + cache.hashCode();
-
-        return res;
-    }
-
-    /** {@inheritDoc} */
-    @Override public boolean equals(Object obj) {
-        if (this == obj)
-            return true;
-
-        if (obj == null || getClass() != obj.getClass())
-            return false;
-
-        BlockMatrixStorage that = (BlockMatrixStorage)obj;
-
-        return rows == that.rows && cols == that.cols && uuid.equals(that.uuid)
-            && (cache != null ? cache.equals(that.cache) : that.cache == null);
+        cache.clearAll(getAllKeys());
     }
 
     /**
@@ -225,126 +206,123 @@ public class BlockMatrixStorage extends CacheUtils implements MatrixStorage, Sto
     }
 
     /**
-     * Build the cache key for the given block id
+     * Build the cache key for the given blocks id.
+     *
+     * NB: NOT cell indices.
      */
-    public BlockMatrixKey getCacheKey(long blockId) {
-        return new BlockMatrixKey(blockId, uuid, getAffinityKey(blockId));
+    public BlockMatrixKey getCacheKey(long blockIdRow, long blockIdCol) {
+        return new BlockMatrixKey(blockIdRow, blockIdCol, uuid, getAffinityKey(blockIdRow, blockIdCol));
     }
 
     /**
-     * Get rows for current block.
+     * Build the cache key for the given blocks id.
      *
-     * @param blockId block id.
-     * @param storageC result storage.
-     * @return The list of block entries.
+     * NB: NOT cell indices.
      */
-    public List<BlockEntry> getRowForBlock(long blockId, BlockMatrixStorage storageC) {
-        long blockRow = blockId / storageC.blocksInCol;
-        long blockCol = blockId % storageC.blocksInRow;
+    public BlockMatrixKey getCacheKey(IgnitePair<Long> blockId) {
+        return new BlockMatrixKey(blockId.get1(), blockId.get2(), uuid, getAffinityKey(blockId.get1(), blockId.get2()));
+    }
+
+    /** {@inheritDoc} */
+    @Override public Set<BlockMatrixKey> getAllKeys() {
+        IgnitePair<Long> maxBlockId = getBlockId(rows, cols);
+
+        Set<BlockMatrixKey> keyset = new HashSet<>();
 
-        long locBlock = this.blocksInRow * (blockRow) + (blockCol >= this.blocksInRow ? (blocksInRow - 1) : blockCol);
+        for(int i = 0; i <= maxBlockId.get1(); i++)
+            for(int j = 0; j <= maxBlockId.get2(); j++)
+                keyset.add(getCacheKey(i,j));
 
-        return getRowForBlock(locBlock);
+        return keyset;
+    }
+
+    /** {@inheritDoc} */
+    @Override public String cacheName() {
+        return CACHE_NAME;
     }
 
     /**
-     * Get cols for current block.
+     * Get rows for current block.
      *
      * @param blockId block id.
-     * @param storageC result storage.
      * @return The list of block entries.
      */
-    public List<BlockEntry> getColForBlock(long blockId, BlockMatrixStorage storageC) {
-        long blockRow = blockId / storageC.blocksInCol;
-        long blockCol = blockId % storageC.blocksInRow;
+    public List<BlockEntry> getRowForBlock(IgnitePair<Long> blockId) {
+        List<BlockEntry> res = new LinkedList<>();
 
-        long locBlock = this.blocksInRow * (blockRow) + (blockCol >= this.blocksInRow ? (blocksInRow - 1) : blockCol);
+        for (int i = 0; i < blocksInCol; i++)
+            res.add(getEntryById(new IgnitePair<>((long) i, blockId.get2())));
 
-        return getColForBlock(locBlock);
+        return res;
     }
 
-    /** {@inheritDoc} */
-    @Override public Set<BlockMatrixKey> getAllKeys() {
-        long maxBlockId = numberOfBlocks();
-        Set<BlockMatrixKey> keys = new HashSet<>();
+    /**
+     * Get cols for current block.
+     *
+     * @param blockId block id.
+     * @return The list of block entries.
+     */
+    public List<BlockEntry> getColForBlock(IgnitePair<Long> blockId) {
+        List<BlockEntry> res = new LinkedList<>();
 
-        for (long id = 0; id < maxBlockId; id++)
-            keys.add(getCacheKey(id));
+        for (int i = 0; i < blocksInRow; i++)
+            res.add(getEntryById(new IgnitePair<>(blockId.get1(), (long) i)));
 
-        return keys;
+        return res;
     }
 
     /** {@inheritDoc} */
-    @Override public String cacheName() {
-        return CACHE_NAME;
-    }
-
-    /** */
-    private List<BlockEntry> getRowForBlock(long blockId) {
-        List<BlockEntry> res = new LinkedList<>();
-
-        boolean isFirstRow = blockId < blocksInRow;
-
-        long startBlock = isFirstRow ? 0 : blockId - blockId % blocksInRow;
-        long endBlock = startBlock + blocksInRow - 1;
+    @Override public int hashCode() {
+        int res = 1;
 
-        for (long i = startBlock; i <= endBlock; i++)
-            res.add(getEntryById(i));
+        res = res * 37 + cols;
+        res = res * 37 + rows;
+        res = res * 37 + uuid.hashCode();
+        res = res * 37 + cache.hashCode();
 
         return res;
     }
 
-    /** */
-    private List<BlockEntry> getColForBlock(long blockId) {
-        List<BlockEntry> res = new LinkedList<>();
+    /** {@inheritDoc} */
+    @Override public boolean equals(Object obj) {
+        if (this == obj)
+            return true;
 
-        long startBlock = blockId % blocksInRow;
-        long endBlock = startBlock + blocksInRow * (blocksInCol - 1);
+        if (obj == null || getClass() != obj.getClass())
+            return false;
 
-        for (long i = startBlock; i <= endBlock; i += blocksInRow)
-            res.add(getEntryById(i));
+        BlockMatrixStorage that = (BlockMatrixStorage)obj;
 
-        return res;
+        return rows == that.rows && cols == that.cols && uuid.equals(that.uuid)
+            && (cache != null ? cache.equals(that.cache) : that.cache == null);
     }
 
     /**
      *
      */
-    private BlockEntry getEntryById(long blockId) {
-        BlockMatrixKey key = getCacheKey(blockId);
+    private BlockEntry getEntryById(IgnitePair<Long> blockId) {
+        BlockMatrixKey key = getCacheKey(blockId.get1(), blockId.get2());
 
         BlockEntry entry = cache.localPeek(key);
         entry = entry != null ? entry : cache.get(key);
 
         if (entry == null) {
-            long colId = blockId == 0 ? 0 : blockId + 1;
 
-            boolean isLastRow = (blockId) >= blocksInRow * (blocksInCol - 1);
-            boolean isLastCol = (colId) % blocksInRow == 0;
+            int colSize = blockId.get1() == (blocksInCol - 1) ? rows % maxBlockEdge : maxBlockEdge;
+            int rowSize = blockId.get2() == (blocksInRow -1 ) ? cols % maxBlockEdge : maxBlockEdge;
 
-            entry = new BlockEntry(isLastRow && rows % maxBlockEdge != 0 ? rows % maxBlockEdge : maxBlockEdge, isLastCol && cols % maxBlockEdge != 0 ? cols % maxBlockEdge : maxBlockEdge);
+            entry = new BlockEntry(rowSize, colSize);
         }
 
         return entry;
     }
 
     /**
-     *
-     */
-    private long numberOfBlocks() {
-        int rows = rowSize();
-        int cols = columnSize();
-
-        return ((rows / maxBlockEdge) + (((rows % maxBlockEdge) > 0) ? 1 : 0))
-            * ((cols / maxBlockEdge) + (((cols % maxBlockEdge) > 0) ? 1 : 0));
-    }
-
-    /**
      * TODO: IGNITE-5646, WIP
      *
      * Get affinity key for the given id.
      */
-    private IgniteUuid getAffinityKey(long id) {
+    private IgniteUuid getAffinityKey(long blockIdRow, long blockIdCol) {
         return null;
     }
 
@@ -356,12 +334,12 @@ public class BlockMatrixStorage extends CacheUtils implements MatrixStorage, Sto
      * @param v New value to set.
      */
     private void matrixSet(int a, int b, double v) {
-        long id = getBlockId(a, b);
+        IgnitePair<Long> blockId = getBlockId(a, b);
         // Remote set on the primary node (where given row or column is stored locally).
-        ignite().compute(groupForKey(CACHE_NAME, id)).run(() -> {
+        ignite().compute(groupForKey(CACHE_NAME, blockId)).run(() -> {
             IgniteCache<BlockMatrixKey, BlockEntry> cache = Ignition.localIgnite().getOrCreateCache(CACHE_NAME);
 
-            BlockMatrixKey key = getCacheKey(getBlockId(a, b));
+            BlockMatrixKey key = getCacheKey(blockId.get1(), blockId.get2());
 
             // Local get.
             BlockEntry block = cache.localPeek(key, CachePeekMode.PRIMARY);
@@ -380,21 +358,26 @@ public class BlockMatrixStorage extends CacheUtils implements MatrixStorage, Sto
     }
 
     /** */
-    private long getBlockId(int x, int y) {
-        return (y / maxBlockEdge) * blockShift(cols) + (x / maxBlockEdge);
+    public IgnitePair<Long> getBlockId(int x, int y) {
+        return new IgnitePair<>((long)x / maxBlockEdge, (long)y / maxBlockEdge);
     }
 
     /** */
     private BlockEntry initBlockFor(int x, int y) {
-        int blockRows = rows - x >= maxBlockEdge ? maxBlockEdge : rows - x;
-        int blockCols = cols - y >= maxBlockEdge ? maxBlockEdge : cols - y;
+        int blockRows = 0;
+        int blockCols = 0;
 
-        return new BlockEntry(blockRows, blockCols);
-    }
+        if(rows >= maxBlockEdge)
+            blockRows = rows - x >= maxBlockEdge ? maxBlockEdge : rows - x;
+        else
+            blockRows =  rows;
 
-    /** */
-    private int blockShift(int i) {
-        return (i) / maxBlockEdge + ((i) % maxBlockEdge > 0 ? 1 : 0);
+        if(cols >= maxBlockEdge)
+            blockCols = cols - y >= maxBlockEdge ? maxBlockEdge : cols - y;
+        else
+            blockCols = cols;
+
+        return new BlockEntry(blockRows, blockCols);
     }
 
     /**

http://git-wip-us.apache.org/repos/asf/ignite/blob/94dd19d3/modules/ml/src/test/java/org/apache/ignite/ml/math/impls/matrix/SparseDistributedBlockMatrixTest.java
----------------------------------------------------------------------
diff --git a/modules/ml/src/test/java/org/apache/ignite/ml/math/impls/matrix/SparseDistributedBlockMatrixTest.java b/modules/ml/src/test/java/org/apache/ignite/ml/math/impls/matrix/SparseDistributedBlockMatrixTest.java
index 56ff638..35329ae 100644
--- a/modules/ml/src/test/java/org/apache/ignite/ml/math/impls/matrix/SparseDistributedBlockMatrixTest.java
+++ b/modules/ml/src/test/java/org/apache/ignite/ml/math/impls/matrix/SparseDistributedBlockMatrixTest.java
@@ -23,12 +23,10 @@ import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
 import java.util.Collection;
-import java.util.HashSet;
 import java.util.Set;
 import org.apache.ignite.Ignite;
 import org.apache.ignite.IgniteCache;
 import org.apache.ignite.internal.util.IgniteUtils;
-import org.apache.ignite.lang.IgniteUuid;
 import org.apache.ignite.ml.math.Matrix;
 import org.apache.ignite.ml.math.distributed.DistributedStorage;
 import org.apache.ignite.ml.math.distributed.keys.impl.BlockMatrixKey;
@@ -293,6 +291,8 @@ public class SparseDistributedBlockMatrixTest extends GridCommonAbstractTest {
 
         Matrix res = cacheMatrix1.times(cacheMatrix2);
 
+        BlockMatrixStorage storage = (BlockMatrixStorage)res.getStorage();
+
         for(int i = 0; i < size; i++)
             for(int j = 0; j < size; j++)
                 if (i == j)
@@ -364,17 +364,8 @@ public class SparseDistributedBlockMatrixTest extends GridCommonAbstractTest {
 
     /** Build key set for SparseBlockDistributedMatrix. */
     private Set<BlockMatrixKey> buildKeySet(SparseBlockDistributedMatrix m){
-        Set<BlockMatrixKey> set = new HashSet<>();
-
         BlockMatrixStorage storage = (BlockMatrixStorage)m.getStorage();
 
-        IgniteUuid uuid = storage.getUUID();
-
-        long maxBlock = (rows / 32 + (rows % 32 > 0 ? 1 : 0)) * (cols / 32 + (cols % 32 > 0 ? 1 : 0));
-
-        for (long i = 0; i < maxBlock; i++)
-            set.add(new BlockMatrixKey(i,uuid,null));
-
-        return set;
+        return storage.getAllKeys();
     }
 }