gaturchenko commented on code in PR #2542:
URL: https://github.com/apache/systemds/pull/2542#discussion_r3613929110


##########
src/main/java/org/apache/sysds/runtime/matrix/data/LibMatrixSketch.java:
##########
@@ -19,89 +19,695 @@
 
 package org.apache.sysds.runtime.matrix.data;
 
-import org.apache.sysds.common.Types;
-
+import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+
+import org.apache.sysds.common.Types;
+import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.util.CommonThreadPool;
+import org.apache.sysds.runtime.util.UtilFunctions;
 
 public class LibMatrixSketch {
+       private static final long PAR_UNIQUE_NUMCELL_THRESHOLD = 1024 * 16;
+       private static final long PAR_UNIQUE_MAX_LOCAL_BYTES_FRACTION = 4;
+       private static final long PAR_UNIQUE_LOCAL_BYTES_OVERHEAD = 8;
 
+       /**
+        * Computes unique values with the original single-threaded behavior.
+        * The overload with a parallelism argument keeps this path as the k=1 
baseline.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return matrix block containing unique values
+        */
        public static MatrixBlock getUniqueValues(MatrixBlock blkIn, 
Types.Direction dir) {
-               //similar to R's unique, this operation takes a matrix and 
computes the
-               //unique values (or rows in case of multiple column inputs)
-               
+               return getUniqueValues(blkIn, dir, 1);

Review Comment:
   The degree of parallelism is hard-coded here, so 
`AggregateUnaryCPInstruction` will always call the single-threaded 
implementation.



##########
src/main/java/org/apache/sysds/runtime/matrix/data/LibMatrixSketch.java:
##########
@@ -19,89 +19,695 @@
 
 package org.apache.sysds.runtime.matrix.data;
 
-import org.apache.sysds.common.Types;
-
+import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+
+import org.apache.sysds.common.Types;
+import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.util.CommonThreadPool;
+import org.apache.sysds.runtime.util.UtilFunctions;
 
 public class LibMatrixSketch {
+       private static final long PAR_UNIQUE_NUMCELL_THRESHOLD = 1024 * 16;
+       private static final long PAR_UNIQUE_MAX_LOCAL_BYTES_FRACTION = 4;
+       private static final long PAR_UNIQUE_LOCAL_BYTES_OVERHEAD = 8;
 
+       /**
+        * Computes unique values with the original single-threaded behavior.
+        * The overload with a parallelism argument keeps this path as the k=1 
baseline.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return matrix block containing unique values
+        */
        public static MatrixBlock getUniqueValues(MatrixBlock blkIn, 
Types.Direction dir) {
-               //similar to R's unique, this operation takes a matrix and 
computes the
-               //unique values (or rows in case of multiple column inputs)
-               
+               return getUniqueValues(blkIn, dir, 1);
+       }
+
+       /**
+        * Computes unique values. For sufficiently large inputs and k > 1, 
this uses
+        * parallel local deduplication or its batched variant.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing unique values
+        */
+       public static MatrixBlock getUniqueValues(MatrixBlock blkIn, 
Types.Direction dir, int k) {
+               // Similar to R's unique, this operation computes unique values 
according
+               // to the requested direction.
+               if( !satisfiesMultiThreadingConstraints(blkIn, dir, k) )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               boolean localDedupMemorySafe = 
isLocalDedupMemoryBudgetSafe(blkIn, dir);
+               switch(dir) {
+                       case RowCol:
+                               return localDedupMemorySafe ?
+                                       getUniqueValuesRowColParallel(blkIn, k) 
:
+                                       
getUniqueValuesRowColBatchedParallel(blkIn, dir, k);
+                       case Row:
+                               return localDedupMemorySafe ?
+                                       getUniqueRowValuesParallel(blkIn, k) :
+                                       
getUniqueRowValuesBatchedParallel(blkIn, dir, k);
+                       case Col:
+                               return localDedupMemorySafe ?
+                                       getUniqueColumnValuesParallel(blkIn, k) 
:
+                                       
getUniqueColumnValuesBatchedParallel(blkIn, dir, k);
+                       default:
+                               throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
+               }
+       }
+
+       /**
+        * Single-threaded baseline implementation for all unique directions.
+        * This preserves the original row-wise and column-wise unique behavior.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return matrix block containing unique values
+        */
+       private static MatrixBlock getUniqueValuesSequential(MatrixBlock blkIn, 
Types.Direction dir) {
                int rlen = blkIn.getNumRows();
                int clen = blkIn.getNumColumns();
 
                MatrixBlock blkOut = null;
-               // TODO optimize for dense/sparse/compressed (once multi-column 
support added)
-               
                switch (dir) {
-                       case RowCol: {
-                               // obtain set of unique items (dense input 
vector)
+                       case RowCol:
+                               // TODO optimize for dense/sparse/compressed 
(once multi-column support added)
+
+                               // obtain set of unique items
                                HashSet<Double> hashSet = new HashSet<>();
                                for( int i=0; i<rlen; i++ ) {
                                        for( int j=0; j<clen; j++ )
                                                hashSet.add(blkIn.get(i, j));
                                }
-                               
+
                                // allocate output block and place values
-                               int rlen2 = hashSet.size();
-                               blkOut = new MatrixBlock(rlen2, 1, 
false).allocateBlock();
-                               int pos = 0;
-                               for( Double val : hashSet )
-                                       blkOut.set(pos++, 0, val);
+                               blkOut = createRowColOutput(hashSet);
                                break;
-                       }
-                       case Row: {
-                               //2-pass algorithm to avoid unnecessarily large 
mem requirements
-                               HashSet<Double> hashSet = new HashSet<>();
+
+                       case Row:
+                               // 2-pass algorithm to avoid unnecessarily 
large mem requirements
+                               HashSet<Double> rowSet = new HashSet<>();
                                int clen2 = 0;
                                for( int i=0; i<rlen; i++ ) {
-                                       hashSet.clear();
+                                       rowSet.clear();
                                        for( int j=0; j<clen; j++ )
-                                               hashSet.add(blkIn.get(i, j));
-                                       clen2 = Math.max(clen2, hashSet.size());
+                                               rowSet.add(blkIn.get(i, j));
+                                       clen2 = Math.max(clen2, rowSet.size());
                                }
-                               
-                               //actual 
-                               blkOut = new MatrixBlock(rlen, clen2, 
false).allocateBlock();
+
+                               // actual
+                               blkOut = allocateOutputBlock(rlen, clen2);
                                for( int i=0; i<rlen; i++ ) {
-                                       hashSet.clear();
+                                       rowSet.clear();
                                        for( int j=0; j<clen; j++ )
-                                               hashSet.add(blkIn.get(i, j));
+                                               rowSet.add(blkIn.get(i, j));
                                        int pos = 0;
-                                       for( Double val : hashSet )
+                                       for( Double val : rowSet )
                                                blkOut.set(i, pos++, val);
                                }
                                break;
-                       }
-                       case Col: {
-                               //2-pass algorithm to avoid unnecessarily large 
mem requirements
-                               HashSet<Double> hashSet = new HashSet<>();
+
+                       case Col:
+                               // 2-pass algorithm to avoid unnecessarily 
large mem requirements
+                               HashSet<Double> colSet = new HashSet<>();
                                int rlen2 = 0;
                                for( int j=0; j<clen; j++ ) {
-                                       hashSet.clear();
+                                       colSet.clear();
                                        for( int i=0; i<rlen; i++ )
-                                               hashSet.add(blkIn.get(i, j));
-                                       rlen2 = Math.max(rlen2, hashSet.size());
+                                               colSet.add(blkIn.get(i, j));
+                                       rlen2 = Math.max(rlen2, colSet.size());
                                }
-                               
-                               //actual 
-                               blkOut = new MatrixBlock(rlen2, clen, 
false).allocateBlock();
+
+                               // actual
+                               blkOut = allocateOutputBlock(rlen2, clen);
                                for( int j=0; j<clen; j++ ) {
-                                       hashSet.clear();
+                                       colSet.clear();
                                        for( int i=0; i<rlen; i++ )
-                                               hashSet.add(blkIn.get(i, j));
+                                               colSet.add(blkIn.get(i, j));
                                        int pos = 0;
-                                       for( Double val : hashSet )
+                                       for( Double val : colSet )
                                                blkOut.set(pos++, j, val);
                                }
                                break;
+
+                       default:
+                               throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
+               }
+
+               blkOut.recomputeNonZeros();
+               blkOut.examSparsity();
+               return blkOut;
+       }
+
+       /**
+        * Parallel unique for all matrix values. Rows are split into balanced 
partitions,
+        * each task builds a local set, and the caller merges all local sets 
afterwards.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return one-column matrix block containing the unique values
+        */
+       private static MatrixBlock getUniqueValuesRowColParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumRows());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<UniqueValueTask> tasks = new ArrayList<>();
+                       for( int[] range : 
getBalancedRanges(blkIn.getNumRows(), numThreads) )
+                               tasks.add(new UniqueValueTask(blkIn, range[0], 
range[1]));
+
+                       // Merge local sets after the workers complete.
+                       HashSet<Double> hashSet = new HashSet<>();
+                       List<Future<HashSet<Double>>> rtasks = 
pool.invokeAll(tasks);
+                       for( int i = 0; i < rtasks.size(); i++ ) {
+                               HashSet<Double> localSet = rtasks.get(i).get();
+                               hashSet.addAll(localSet);
+                               localSet.clear();
+                               rtasks.set(i, null);
                        }
+
+                       return createRowColOutput(hashSet);
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Parallel row-wise unique values. A first pass computes the output 
width,
+        * and a second pass materializes the row-local unique values.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return matrix block containing row-wise unique values
+        */
+       private static MatrixBlock getUniqueRowValuesParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumRows());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<int[]> ranges = 
getBalancedRanges(blkIn.getNumRows(), numThreads);
+                       int clen2 = getMaxUniqueValues(pool, blkIn, 
Types.Direction.Row, ranges);
+                       MatrixBlock blkOut = 
allocateOutputBlock(blkIn.getNumRows(), clen2);
+                       fillUniqueValues(pool, blkIn, blkOut, 
Types.Direction.Row, ranges);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Parallel column-wise unique values. A first pass computes the output 
height,
+        * and a second pass materializes the column-local unique values.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return matrix block containing column-wise unique values
+        */
+       private static MatrixBlock getUniqueColumnValuesParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumColumns());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<int[]> ranges = 
getBalancedRanges(blkIn.getNumColumns(), numThreads);
+                       int rlen2 = getMaxUniqueValues(pool, blkIn, 
Types.Direction.Col, ranges);
+                       MatrixBlock blkOut = allocateOutputBlock(rlen2, 
blkIn.getNumColumns());
+                       fillUniqueValues(pool, blkIn, blkOut, 
Types.Direction.Col, ranges);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel unique for all matrix values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return one-column matrix block containing the unique values
+        */
+       private static MatrixBlock 
getUniqueValuesRowColBatchedParallel(MatrixBlock blkIn, Types.Direction dir, 
int k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       HashSet<Double> hashSet = new HashSet<>();

Review Comment:
   This `HashSet` is unbounded. When almost all values are unique, the set 
grows to `rlen x clen`, so it in principle matches the sequential case. 



##########
src/test/java/org/apache/sysds/runtime/matrix/data/LibMatrixSketchUniqueParallelTest.java:
##########
@@ -0,0 +1,209 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.sysds.runtime.matrix.data;

Review Comment:
   Please, move the tests to respective files in 
`src/test/java/org/apache/sysds/test/functions/unique`, feel free to extend for 
the multi-threaded case if necessary



##########
src/main/java/org/apache/sysds/runtime/matrix/data/LibMatrixSketch.java:
##########
@@ -19,89 +19,695 @@
 
 package org.apache.sysds.runtime.matrix.data;
 
-import org.apache.sysds.common.Types;
-
+import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+
+import org.apache.sysds.common.Types;
+import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.util.CommonThreadPool;
+import org.apache.sysds.runtime.util.UtilFunctions;
 
 public class LibMatrixSketch {
+       private static final long PAR_UNIQUE_NUMCELL_THRESHOLD = 1024 * 16;
+       private static final long PAR_UNIQUE_MAX_LOCAL_BYTES_FRACTION = 4;
+       private static final long PAR_UNIQUE_LOCAL_BYTES_OVERHEAD = 8;
 
+       /**
+        * Computes unique values with the original single-threaded behavior.
+        * The overload with a parallelism argument keeps this path as the k=1 
baseline.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return matrix block containing unique values
+        */
        public static MatrixBlock getUniqueValues(MatrixBlock blkIn, 
Types.Direction dir) {
-               //similar to R's unique, this operation takes a matrix and 
computes the
-               //unique values (or rows in case of multiple column inputs)
-               
+               return getUniqueValues(blkIn, dir, 1);
+       }
+
+       /**
+        * Computes unique values. For sufficiently large inputs and k > 1, 
this uses
+        * parallel local deduplication or its batched variant.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing unique values
+        */
+       public static MatrixBlock getUniqueValues(MatrixBlock blkIn, 
Types.Direction dir, int k) {
+               // Similar to R's unique, this operation computes unique values 
according
+               // to the requested direction.
+               if( !satisfiesMultiThreadingConstraints(blkIn, dir, k) )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               boolean localDedupMemorySafe = 
isLocalDedupMemoryBudgetSafe(blkIn, dir);
+               switch(dir) {
+                       case RowCol:
+                               return localDedupMemorySafe ?
+                                       getUniqueValuesRowColParallel(blkIn, k) 
:
+                                       
getUniqueValuesRowColBatchedParallel(blkIn, dir, k);
+                       case Row:
+                               return localDedupMemorySafe ?
+                                       getUniqueRowValuesParallel(blkIn, k) :
+                                       
getUniqueRowValuesBatchedParallel(blkIn, dir, k);
+                       case Col:
+                               return localDedupMemorySafe ?
+                                       getUniqueColumnValuesParallel(blkIn, k) 
:
+                                       
getUniqueColumnValuesBatchedParallel(blkIn, dir, k);
+                       default:
+                               throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
+               }
+       }
+
+       /**
+        * Single-threaded baseline implementation for all unique directions.
+        * This preserves the original row-wise and column-wise unique behavior.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return matrix block containing unique values
+        */
+       private static MatrixBlock getUniqueValuesSequential(MatrixBlock blkIn, 
Types.Direction dir) {
                int rlen = blkIn.getNumRows();
                int clen = blkIn.getNumColumns();
 
                MatrixBlock blkOut = null;
-               // TODO optimize for dense/sparse/compressed (once multi-column 
support added)
-               
                switch (dir) {
-                       case RowCol: {
-                               // obtain set of unique items (dense input 
vector)
+                       case RowCol:
+                               // TODO optimize for dense/sparse/compressed 
(once multi-column support added)
+
+                               // obtain set of unique items
                                HashSet<Double> hashSet = new HashSet<>();
                                for( int i=0; i<rlen; i++ ) {
                                        for( int j=0; j<clen; j++ )
                                                hashSet.add(blkIn.get(i, j));
                                }
-                               
+
                                // allocate output block and place values
-                               int rlen2 = hashSet.size();
-                               blkOut = new MatrixBlock(rlen2, 1, 
false).allocateBlock();
-                               int pos = 0;
-                               for( Double val : hashSet )
-                                       blkOut.set(pos++, 0, val);
+                               blkOut = createRowColOutput(hashSet);
                                break;
-                       }
-                       case Row: {
-                               //2-pass algorithm to avoid unnecessarily large 
mem requirements
-                               HashSet<Double> hashSet = new HashSet<>();
+
+                       case Row:
+                               // 2-pass algorithm to avoid unnecessarily 
large mem requirements
+                               HashSet<Double> rowSet = new HashSet<>();
                                int clen2 = 0;
                                for( int i=0; i<rlen; i++ ) {
-                                       hashSet.clear();
+                                       rowSet.clear();
                                        for( int j=0; j<clen; j++ )
-                                               hashSet.add(blkIn.get(i, j));
-                                       clen2 = Math.max(clen2, hashSet.size());
+                                               rowSet.add(blkIn.get(i, j));
+                                       clen2 = Math.max(clen2, rowSet.size());
                                }
-                               
-                               //actual 
-                               blkOut = new MatrixBlock(rlen, clen2, 
false).allocateBlock();
+
+                               // actual
+                               blkOut = allocateOutputBlock(rlen, clen2);
                                for( int i=0; i<rlen; i++ ) {
-                                       hashSet.clear();
+                                       rowSet.clear();
                                        for( int j=0; j<clen; j++ )
-                                               hashSet.add(blkIn.get(i, j));
+                                               rowSet.add(blkIn.get(i, j));
                                        int pos = 0;
-                                       for( Double val : hashSet )
+                                       for( Double val : rowSet )
                                                blkOut.set(i, pos++, val);
                                }
                                break;
-                       }
-                       case Col: {
-                               //2-pass algorithm to avoid unnecessarily large 
mem requirements
-                               HashSet<Double> hashSet = new HashSet<>();
+
+                       case Col:
+                               // 2-pass algorithm to avoid unnecessarily 
large mem requirements
+                               HashSet<Double> colSet = new HashSet<>();
                                int rlen2 = 0;
                                for( int j=0; j<clen; j++ ) {
-                                       hashSet.clear();
+                                       colSet.clear();
                                        for( int i=0; i<rlen; i++ )
-                                               hashSet.add(blkIn.get(i, j));
-                                       rlen2 = Math.max(rlen2, hashSet.size());
+                                               colSet.add(blkIn.get(i, j));
+                                       rlen2 = Math.max(rlen2, colSet.size());
                                }
-                               
-                               //actual 
-                               blkOut = new MatrixBlock(rlen2, clen, 
false).allocateBlock();
+
+                               // actual
+                               blkOut = allocateOutputBlock(rlen2, clen);
                                for( int j=0; j<clen; j++ ) {
-                                       hashSet.clear();
+                                       colSet.clear();
                                        for( int i=0; i<rlen; i++ )
-                                               hashSet.add(blkIn.get(i, j));
+                                               colSet.add(blkIn.get(i, j));
                                        int pos = 0;
-                                       for( Double val : hashSet )
+                                       for( Double val : colSet )
                                                blkOut.set(pos++, j, val);
                                }
                                break;
+
+                       default:
+                               throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
+               }
+
+               blkOut.recomputeNonZeros();
+               blkOut.examSparsity();
+               return blkOut;
+       }
+
+       /**
+        * Parallel unique for all matrix values. Rows are split into balanced 
partitions,
+        * each task builds a local set, and the caller merges all local sets 
afterwards.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return one-column matrix block containing the unique values
+        */
+       private static MatrixBlock getUniqueValuesRowColParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumRows());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<UniqueValueTask> tasks = new ArrayList<>();
+                       for( int[] range : 
getBalancedRanges(blkIn.getNumRows(), numThreads) )
+                               tasks.add(new UniqueValueTask(blkIn, range[0], 
range[1]));
+
+                       // Merge local sets after the workers complete.
+                       HashSet<Double> hashSet = new HashSet<>();
+                       List<Future<HashSet<Double>>> rtasks = 
pool.invokeAll(tasks);
+                       for( int i = 0; i < rtasks.size(); i++ ) {
+                               HashSet<Double> localSet = rtasks.get(i).get();
+                               hashSet.addAll(localSet);
+                               localSet.clear();
+                               rtasks.set(i, null);
                        }
+
+                       return createRowColOutput(hashSet);
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Parallel row-wise unique values. A first pass computes the output 
width,
+        * and a second pass materializes the row-local unique values.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return matrix block containing row-wise unique values
+        */
+       private static MatrixBlock getUniqueRowValuesParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumRows());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<int[]> ranges = 
getBalancedRanges(blkIn.getNumRows(), numThreads);
+                       int clen2 = getMaxUniqueValues(pool, blkIn, 
Types.Direction.Row, ranges);
+                       MatrixBlock blkOut = 
allocateOutputBlock(blkIn.getNumRows(), clen2);
+                       fillUniqueValues(pool, blkIn, blkOut, 
Types.Direction.Row, ranges);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Parallel column-wise unique values. A first pass computes the output 
height,
+        * and a second pass materializes the column-local unique values.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return matrix block containing column-wise unique values
+        */
+       private static MatrixBlock getUniqueColumnValuesParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumColumns());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<int[]> ranges = 
getBalancedRanges(blkIn.getNumColumns(), numThreads);
+                       int rlen2 = getMaxUniqueValues(pool, blkIn, 
Types.Direction.Col, ranges);
+                       MatrixBlock blkOut = allocateOutputBlock(rlen2, 
blkIn.getNumColumns());
+                       fillUniqueValues(pool, blkIn, blkOut, 
Types.Direction.Col, ranges);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel unique for all matrix values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return one-column matrix block containing the unique values
+        */
+       private static MatrixBlock 
getUniqueValuesRowColBatchedParallel(MatrixBlock blkIn, Types.Direction dir, 
int k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       HashSet<Double> hashSet = new HashSet<>();
+                       for( int pos = 0; pos < config._len; ) {
+                               ArrayList<UniqueValueTask> tasks = new 
ArrayList<>();
+                               for( int i = 0; i < config._numThreads && pos < 
config._len; i++ ) {
+                                       int end = Math.min(pos + 
config._taskLen, config._len);
+                                       tasks.add(new UniqueValueTask(blkIn, 
pos, end));
+                                       pos = end;
+                               }
+
+                               List<Future<HashSet<Double>>> rtasks = 
pool.invokeAll(tasks);
+                               for( int i = 0; i < rtasks.size(); i++ ) {
+                                       HashSet<Double> localSet = 
rtasks.get(i).get();
+                                       hashSet.addAll(localSet);
+                                       localSet.clear();
+                                       rtasks.set(i, null);
+                               }
+                       }
+
+                       return createRowColOutput(hashSet);
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel row-wise unique values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing row-wise unique values
+        */
+       private static MatrixBlock 
getUniqueRowValuesBatchedParallel(MatrixBlock blkIn, Types.Direction dir, int 
k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       int clen2 = getMaxUniqueValuesBatched(pool, blkIn, dir, 
config);
+                       MatrixBlock blkOut = 
allocateOutputBlock(blkIn.getNumRows(), clen2);
+                       fillUniqueValuesBatched(pool, blkIn, blkOut, dir, 
config);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel column-wise unique values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing column-wise unique values
+        */
+       private static MatrixBlock 
getUniqueColumnValuesBatchedParallel(MatrixBlock blkIn, Types.Direction dir, 
int k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       int rlen2 = getMaxUniqueValuesBatched(pool, blkIn, dir, 
config);
+                       MatrixBlock blkOut = allocateOutputBlock(rlen2, 
blkIn.getNumColumns());
+                       fillUniqueValuesBatched(pool, blkIn, blkOut, dir, 
config);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Computes the maximum row-wise or column-wise unique count over 
balanced ranges.
+        */
+       private static int getMaxUniqueValues(ExecutorService pool, MatrixBlock 
blkIn, Types.Direction dir,
+               ArrayList<int[]> ranges) throws Exception {
+               ArrayList<UniqueCountTask> tasks = new ArrayList<>();
+               for( int[] range : ranges )
+                       tasks.add(new UniqueCountTask(blkIn, dir, range[0], 
range[1]));
+
+               int ret = 0;
+               List<Future<Integer>> rtasks = pool.invokeAll(tasks);
+               for( int i = 0; i < rtasks.size(); i++ ) {
+                       ret = Math.max(ret, rtasks.get(i).get());
+                       rtasks.set(i, null);
+               }
+               return ret;
+       }
+
+       /**
+        * Fills row-wise or column-wise unique values over balanced ranges.
+        */
+       private static void fillUniqueValues(ExecutorService pool, MatrixBlock 
blkIn, MatrixBlock blkOut,
+               Types.Direction dir, ArrayList<int[]> ranges) throws Exception {
+               ArrayList<UniqueOutputTask> tasks = new ArrayList<>();
+               for( int[] range : ranges )
+                       tasks.add(new UniqueOutputTask(blkIn, blkOut, dir, 
range[0], range[1]));
+               List<Future<Void>> rtasks = pool.invokeAll(tasks);
+               for( int i = 0; i < rtasks.size(); i++ ) {
+                       rtasks.get(i).get();
+                       rtasks.set(i, null);
+               }
+       }
+
+       /**
+        * Batched variant of getMaxUniqueValues.
+        */
+       private static int getMaxUniqueValuesBatched(ExecutorService pool, 
MatrixBlock blkIn, Types.Direction dir,
+               BatchConfig config) throws Exception {
+               int ret = 0;
+               for( int pos = 0; pos < config._len; ) {
+                       ArrayList<UniqueCountTask> tasks = new ArrayList<>();
+                       for( int i = 0; i < config._numThreads && pos < 
config._len; i++ ) {
+                               int end = Math.min(pos + config._taskLen, 
config._len);
+                               tasks.add(new UniqueCountTask(blkIn, dir, pos, 
end));
+                               pos = end;
+                       }
+
+                       List<Future<Integer>> rtasks = pool.invokeAll(tasks);
+                       for( int i = 0; i < rtasks.size(); i++ ) {
+                               ret = Math.max(ret, rtasks.get(i).get());
+                               rtasks.set(i, null);
+                       }
+               }
+               return ret;
+       }
+
+       /**
+        * Batched variant of fillUniqueValues.
+        */
+       private static void fillUniqueValuesBatched(ExecutorService pool, 
MatrixBlock blkIn, MatrixBlock blkOut,
+               Types.Direction dir, BatchConfig config) throws Exception {
+               for( int pos = 0; pos < config._len; ) {
+                       ArrayList<UniqueOutputTask> tasks = new ArrayList<>();
+                       for( int i = 0; i < config._numThreads && pos < 
config._len; i++ ) {
+                               int end = Math.min(pos + config._taskLen, 
config._len);
+                               tasks.add(new UniqueOutputTask(blkIn, blkOut, 
dir, pos, end));
+                               pos = end;
+                       }
+
+                       List<Future<Void>> rtasks = pool.invokeAll(tasks);
+                       for( int i = 0; i < rtasks.size(); i++ ) {
+                               rtasks.get(i).get();
+                               rtasks.set(i, null);
+                       }
+               }
+       }
+
+       /**
+        * Decides whether the input is large enough to justify local 
deduplication tasks.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return true if the parallel path should be used
+        */
+       private static boolean satisfiesMultiThreadingConstraints(MatrixBlock 
blkIn, Types.Direction dir, int k) {
+               if( k <= 1 || ((long) blkIn.getNumRows()) * 
blkIn.getNumColumns() < PAR_UNIQUE_NUMCELL_THRESHOLD )
+                       return false;
+
+               switch(dir) {
+                       case RowCol:
+                               return blkIn.getNumRows() > 1;
+                       case Row:
+                               return blkIn.getNumRows() > 1;
+                       case Col:
+                               return blkIn.getNumColumns() > 1;
                        default:
                                throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
                }
+       }
+
+       /**
+        * Creates balanced half-open ranges [start, end) using the same 
utility pattern as
+        * other SystemDS matrix libraries.
+        *
+        * @param len number of rows or columns to partition
+        * @param k requested degree of parallelism
+        * @return list of balanced index ranges
+        */
+       private static ArrayList<int[]> getBalancedRanges(int len, int k) {
+               ArrayList<int[]> ranges = new ArrayList<>();
+               ArrayList<Integer> blklens = 
UtilFunctions.getBalancedBlockSizesDefault(len, getNumThreads(k, len), false);
+               for( int i = 0, lb = 0; i < blklens.size(); lb += 
blklens.get(i), i++ )
+                       ranges.add(new int[] {lb, lb + blklens.get(i)});
+               return ranges;
+       }
+
+       /**
+        * Caps the number of workers by the number of row or column partitions 
available.
+        *
+        * @param k requested degree of parallelism
+        * @param len number of rows or columns to partition
+        * @return effective number of worker threads
+        */
+       private static int getNumThreads(int k, int len) {
+               return Math.max(1, Math.min(k, len));
+       }
+
+       /**
+        * Builds the execution plan for the batched parallel path.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return batch configuration, or null if batching would not use 
parallelism
+        */
+       private static BatchConfig getBatchConfig(MatrixBlock blkIn, 
Types.Direction dir, int k) {
+               int len = getPartitionLength(blkIn, dir);
+               long maxBatchIndexes = getMaxLocalDedupIndexes(blkIn, dir);
+               if( maxBatchIndexes < 2 )
+                       return null;
+
+               int numThreads = getNumThreads(k, (int) Math.min(len, 
maxBatchIndexes));
+               if( numThreads <= 1 )
+                       return null;
+
+               int taskLen = Math.max(1, (int) Math.min(Integer.MAX_VALUE, 
maxBatchIndexes / numThreads));
+               return new BatchConfig(numThreads, taskLen, len);
+       }
+
+       /**
+        * Returns the number of rows or columns that define the partition 
direction.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return number of partition indexes
+        */
+       private static int getPartitionLength(MatrixBlock blkIn, 
Types.Direction dir) {
+               return dir == Types.Direction.Col ? blkIn.getNumColumns() : 
blkIn.getNumRows();
+       }
+
+       /**
+        * Estimates how many partition indexes can be processed in one batch.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return maximum number of rows or columns per batch
+        */
+       private static long getMaxLocalDedupIndexes(MatrixBlock blkIn, 
Types.Direction dir) {
+               long cellsPerIndex = dir == Types.Direction.Col ? 
blkIn.getNumRows() : blkIn.getNumColumns();
+               if( cellsPerIndex <= 0 ||
+                       cellsPerIndex > Long.MAX_VALUE / Double.BYTES / 
PAR_UNIQUE_LOCAL_BYTES_OVERHEAD )
+                       return 0;
+
+               long bytesPerIndex = cellsPerIndex * Double.BYTES * 
PAR_UNIQUE_LOCAL_BYTES_OVERHEAD;
+               long maxLocalBytes = Runtime.getRuntime().maxMemory() / 
PAR_UNIQUE_MAX_LOCAL_BYTES_FRACTION;

Review Comment:
   Read from `InfrastructureAnalyzer.getLocalMaxMemory()` instead



##########
src/main/java/org/apache/sysds/runtime/matrix/data/LibMatrixSketch.java:
##########
@@ -19,89 +19,695 @@
 
 package org.apache.sysds.runtime.matrix.data;
 
-import org.apache.sysds.common.Types;
-
+import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+
+import org.apache.sysds.common.Types;
+import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.util.CommonThreadPool;
+import org.apache.sysds.runtime.util.UtilFunctions;
 
 public class LibMatrixSketch {
+       private static final long PAR_UNIQUE_NUMCELL_THRESHOLD = 1024 * 16;
+       private static final long PAR_UNIQUE_MAX_LOCAL_BYTES_FRACTION = 4;
+       private static final long PAR_UNIQUE_LOCAL_BYTES_OVERHEAD = 8;
 
+       /**
+        * Computes unique values with the original single-threaded behavior.
+        * The overload with a parallelism argument keeps this path as the k=1 
baseline.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return matrix block containing unique values
+        */
        public static MatrixBlock getUniqueValues(MatrixBlock blkIn, 
Types.Direction dir) {
-               //similar to R's unique, this operation takes a matrix and 
computes the
-               //unique values (or rows in case of multiple column inputs)
-               
+               return getUniqueValues(blkIn, dir, 1);
+       }
+
+       /**
+        * Computes unique values. For sufficiently large inputs and k > 1, 
this uses
+        * parallel local deduplication or its batched variant.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing unique values
+        */
+       public static MatrixBlock getUniqueValues(MatrixBlock blkIn, 
Types.Direction dir, int k) {
+               // Similar to R's unique, this operation computes unique values 
according
+               // to the requested direction.
+               if( !satisfiesMultiThreadingConstraints(blkIn, dir, k) )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               boolean localDedupMemorySafe = 
isLocalDedupMemoryBudgetSafe(blkIn, dir);
+               switch(dir) {
+                       case RowCol:
+                               return localDedupMemorySafe ?
+                                       getUniqueValuesRowColParallel(blkIn, k) 
:
+                                       
getUniqueValuesRowColBatchedParallel(blkIn, dir, k);
+                       case Row:
+                               return localDedupMemorySafe ?
+                                       getUniqueRowValuesParallel(blkIn, k) :
+                                       
getUniqueRowValuesBatchedParallel(blkIn, dir, k);
+                       case Col:
+                               return localDedupMemorySafe ?
+                                       getUniqueColumnValuesParallel(blkIn, k) 
:
+                                       
getUniqueColumnValuesBatchedParallel(blkIn, dir, k);
+                       default:
+                               throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
+               }
+       }
+
+       /**
+        * Single-threaded baseline implementation for all unique directions.
+        * This preserves the original row-wise and column-wise unique behavior.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return matrix block containing unique values
+        */
+       private static MatrixBlock getUniqueValuesSequential(MatrixBlock blkIn, 
Types.Direction dir) {
                int rlen = blkIn.getNumRows();
                int clen = blkIn.getNumColumns();
 
                MatrixBlock blkOut = null;
-               // TODO optimize for dense/sparse/compressed (once multi-column 
support added)
-               
                switch (dir) {
-                       case RowCol: {
-                               // obtain set of unique items (dense input 
vector)
+                       case RowCol:
+                               // TODO optimize for dense/sparse/compressed 
(once multi-column support added)
+
+                               // obtain set of unique items
                                HashSet<Double> hashSet = new HashSet<>();
                                for( int i=0; i<rlen; i++ ) {
                                        for( int j=0; j<clen; j++ )
                                                hashSet.add(blkIn.get(i, j));
                                }
-                               
+
                                // allocate output block and place values
-                               int rlen2 = hashSet.size();
-                               blkOut = new MatrixBlock(rlen2, 1, 
false).allocateBlock();
-                               int pos = 0;
-                               for( Double val : hashSet )
-                                       blkOut.set(pos++, 0, val);
+                               blkOut = createRowColOutput(hashSet);
                                break;
-                       }
-                       case Row: {
-                               //2-pass algorithm to avoid unnecessarily large 
mem requirements
-                               HashSet<Double> hashSet = new HashSet<>();
+
+                       case Row:
+                               // 2-pass algorithm to avoid unnecessarily 
large mem requirements
+                               HashSet<Double> rowSet = new HashSet<>();
                                int clen2 = 0;
                                for( int i=0; i<rlen; i++ ) {
-                                       hashSet.clear();
+                                       rowSet.clear();
                                        for( int j=0; j<clen; j++ )
-                                               hashSet.add(blkIn.get(i, j));
-                                       clen2 = Math.max(clen2, hashSet.size());
+                                               rowSet.add(blkIn.get(i, j));
+                                       clen2 = Math.max(clen2, rowSet.size());
                                }
-                               
-                               //actual 
-                               blkOut = new MatrixBlock(rlen, clen2, 
false).allocateBlock();
+
+                               // actual
+                               blkOut = allocateOutputBlock(rlen, clen2);
                                for( int i=0; i<rlen; i++ ) {
-                                       hashSet.clear();
+                                       rowSet.clear();
                                        for( int j=0; j<clen; j++ )
-                                               hashSet.add(blkIn.get(i, j));
+                                               rowSet.add(blkIn.get(i, j));
                                        int pos = 0;
-                                       for( Double val : hashSet )
+                                       for( Double val : rowSet )
                                                blkOut.set(i, pos++, val);
                                }
                                break;
-                       }
-                       case Col: {
-                               //2-pass algorithm to avoid unnecessarily large 
mem requirements
-                               HashSet<Double> hashSet = new HashSet<>();
+
+                       case Col:
+                               // 2-pass algorithm to avoid unnecessarily 
large mem requirements
+                               HashSet<Double> colSet = new HashSet<>();
                                int rlen2 = 0;
                                for( int j=0; j<clen; j++ ) {
-                                       hashSet.clear();
+                                       colSet.clear();
                                        for( int i=0; i<rlen; i++ )
-                                               hashSet.add(blkIn.get(i, j));
-                                       rlen2 = Math.max(rlen2, hashSet.size());
+                                               colSet.add(blkIn.get(i, j));
+                                       rlen2 = Math.max(rlen2, colSet.size());
                                }
-                               
-                               //actual 
-                               blkOut = new MatrixBlock(rlen2, clen, 
false).allocateBlock();
+
+                               // actual
+                               blkOut = allocateOutputBlock(rlen2, clen);
                                for( int j=0; j<clen; j++ ) {
-                                       hashSet.clear();
+                                       colSet.clear();
                                        for( int i=0; i<rlen; i++ )
-                                               hashSet.add(blkIn.get(i, j));
+                                               colSet.add(blkIn.get(i, j));
                                        int pos = 0;
-                                       for( Double val : hashSet )
+                                       for( Double val : colSet )
                                                blkOut.set(pos++, j, val);
                                }
                                break;
+
+                       default:
+                               throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
+               }
+
+               blkOut.recomputeNonZeros();
+               blkOut.examSparsity();
+               return blkOut;
+       }
+
+       /**
+        * Parallel unique for all matrix values. Rows are split into balanced 
partitions,
+        * each task builds a local set, and the caller merges all local sets 
afterwards.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return one-column matrix block containing the unique values
+        */
+       private static MatrixBlock getUniqueValuesRowColParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumRows());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<UniqueValueTask> tasks = new ArrayList<>();
+                       for( int[] range : 
getBalancedRanges(blkIn.getNumRows(), numThreads) )
+                               tasks.add(new UniqueValueTask(blkIn, range[0], 
range[1]));
+
+                       // Merge local sets after the workers complete.
+                       HashSet<Double> hashSet = new HashSet<>();
+                       List<Future<HashSet<Double>>> rtasks = 
pool.invokeAll(tasks);
+                       for( int i = 0; i < rtasks.size(); i++ ) {
+                               HashSet<Double> localSet = rtasks.get(i).get();
+                               hashSet.addAll(localSet);
+                               localSet.clear();
+                               rtasks.set(i, null);
                        }
+
+                       return createRowColOutput(hashSet);
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Parallel row-wise unique values. A first pass computes the output 
width,
+        * and a second pass materializes the row-local unique values.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return matrix block containing row-wise unique values
+        */
+       private static MatrixBlock getUniqueRowValuesParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumRows());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<int[]> ranges = 
getBalancedRanges(blkIn.getNumRows(), numThreads);
+                       int clen2 = getMaxUniqueValues(pool, blkIn, 
Types.Direction.Row, ranges);
+                       MatrixBlock blkOut = 
allocateOutputBlock(blkIn.getNumRows(), clen2);
+                       fillUniqueValues(pool, blkIn, blkOut, 
Types.Direction.Row, ranges);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Parallel column-wise unique values. A first pass computes the output 
height,
+        * and a second pass materializes the column-local unique values.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return matrix block containing column-wise unique values
+        */
+       private static MatrixBlock getUniqueColumnValuesParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumColumns());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<int[]> ranges = 
getBalancedRanges(blkIn.getNumColumns(), numThreads);
+                       int rlen2 = getMaxUniqueValues(pool, blkIn, 
Types.Direction.Col, ranges);
+                       MatrixBlock blkOut = allocateOutputBlock(rlen2, 
blkIn.getNumColumns());
+                       fillUniqueValues(pool, blkIn, blkOut, 
Types.Direction.Col, ranges);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel unique for all matrix values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return one-column matrix block containing the unique values
+        */
+       private static MatrixBlock 
getUniqueValuesRowColBatchedParallel(MatrixBlock blkIn, Types.Direction dir, 
int k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       HashSet<Double> hashSet = new HashSet<>();
+                       for( int pos = 0; pos < config._len; ) {
+                               ArrayList<UniqueValueTask> tasks = new 
ArrayList<>();
+                               for( int i = 0; i < config._numThreads && pos < 
config._len; i++ ) {
+                                       int end = Math.min(pos + 
config._taskLen, config._len);
+                                       tasks.add(new UniqueValueTask(blkIn, 
pos, end));
+                                       pos = end;
+                               }
+
+                               List<Future<HashSet<Double>>> rtasks = 
pool.invokeAll(tasks);
+                               for( int i = 0; i < rtasks.size(); i++ ) {
+                                       HashSet<Double> localSet = 
rtasks.get(i).get();
+                                       hashSet.addAll(localSet);
+                                       localSet.clear();
+                                       rtasks.set(i, null);
+                               }
+                       }
+
+                       return createRowColOutput(hashSet);
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel row-wise unique values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing row-wise unique values
+        */
+       private static MatrixBlock 
getUniqueRowValuesBatchedParallel(MatrixBlock blkIn, Types.Direction dir, int 
k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       int clen2 = getMaxUniqueValuesBatched(pool, blkIn, dir, 
config);
+                       MatrixBlock blkOut = 
allocateOutputBlock(blkIn.getNumRows(), clen2);
+                       fillUniqueValuesBatched(pool, blkIn, blkOut, dir, 
config);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel column-wise unique values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing column-wise unique values
+        */
+       private static MatrixBlock 
getUniqueColumnValuesBatchedParallel(MatrixBlock blkIn, Types.Direction dir, 
int k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       int rlen2 = getMaxUniqueValuesBatched(pool, blkIn, dir, 
config);
+                       MatrixBlock blkOut = allocateOutputBlock(rlen2, 
blkIn.getNumColumns());
+                       fillUniqueValuesBatched(pool, blkIn, blkOut, dir, 
config);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Computes the maximum row-wise or column-wise unique count over 
balanced ranges.
+        */
+       private static int getMaxUniqueValues(ExecutorService pool, MatrixBlock 
blkIn, Types.Direction dir,
+               ArrayList<int[]> ranges) throws Exception {
+               ArrayList<UniqueCountTask> tasks = new ArrayList<>();
+               for( int[] range : ranges )
+                       tasks.add(new UniqueCountTask(blkIn, dir, range[0], 
range[1]));
+
+               int ret = 0;
+               List<Future<Integer>> rtasks = pool.invokeAll(tasks);
+               for( int i = 0; i < rtasks.size(); i++ ) {
+                       ret = Math.max(ret, rtasks.get(i).get());
+                       rtasks.set(i, null);
+               }
+               return ret;
+       }
+
+       /**
+        * Fills row-wise or column-wise unique values over balanced ranges.
+        */
+       private static void fillUniqueValues(ExecutorService pool, MatrixBlock 
blkIn, MatrixBlock blkOut,
+               Types.Direction dir, ArrayList<int[]> ranges) throws Exception {
+               ArrayList<UniqueOutputTask> tasks = new ArrayList<>();
+               for( int[] range : ranges )
+                       tasks.add(new UniqueOutputTask(blkIn, blkOut, dir, 
range[0], range[1]));
+               List<Future<Void>> rtasks = pool.invokeAll(tasks);
+               for( int i = 0; i < rtasks.size(); i++ ) {
+                       rtasks.get(i).get();
+                       rtasks.set(i, null);
+               }
+       }
+
+       /**
+        * Batched variant of getMaxUniqueValues.
+        */
+       private static int getMaxUniqueValuesBatched(ExecutorService pool, 
MatrixBlock blkIn, Types.Direction dir,
+               BatchConfig config) throws Exception {
+               int ret = 0;
+               for( int pos = 0; pos < config._len; ) {
+                       ArrayList<UniqueCountTask> tasks = new ArrayList<>();
+                       for( int i = 0; i < config._numThreads && pos < 
config._len; i++ ) {
+                               int end = Math.min(pos + config._taskLen, 
config._len);

Review Comment:
   Same overflow issue



##########
src/main/java/org/apache/sysds/runtime/matrix/data/LibMatrixSketch.java:
##########
@@ -19,89 +19,695 @@
 
 package org.apache.sysds.runtime.matrix.data;
 
-import org.apache.sysds.common.Types;
-
+import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+
+import org.apache.sysds.common.Types;
+import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.util.CommonThreadPool;
+import org.apache.sysds.runtime.util.UtilFunctions;
 
 public class LibMatrixSketch {
+       private static final long PAR_UNIQUE_NUMCELL_THRESHOLD = 1024 * 16;
+       private static final long PAR_UNIQUE_MAX_LOCAL_BYTES_FRACTION = 4;
+       private static final long PAR_UNIQUE_LOCAL_BYTES_OVERHEAD = 8;
 
+       /**
+        * Computes unique values with the original single-threaded behavior.
+        * The overload with a parallelism argument keeps this path as the k=1 
baseline.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return matrix block containing unique values
+        */
        public static MatrixBlock getUniqueValues(MatrixBlock blkIn, 
Types.Direction dir) {
-               //similar to R's unique, this operation takes a matrix and 
computes the
-               //unique values (or rows in case of multiple column inputs)
-               
+               return getUniqueValues(blkIn, dir, 1);
+       }
+
+       /**
+        * Computes unique values. For sufficiently large inputs and k > 1, 
this uses
+        * parallel local deduplication or its batched variant.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing unique values
+        */
+       public static MatrixBlock getUniqueValues(MatrixBlock blkIn, 
Types.Direction dir, int k) {
+               // Similar to R's unique, this operation computes unique values 
according
+               // to the requested direction.
+               if( !satisfiesMultiThreadingConstraints(blkIn, dir, k) )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               boolean localDedupMemorySafe = 
isLocalDedupMemoryBudgetSafe(blkIn, dir);
+               switch(dir) {
+                       case RowCol:
+                               return localDedupMemorySafe ?
+                                       getUniqueValuesRowColParallel(blkIn, k) 
:
+                                       
getUniqueValuesRowColBatchedParallel(blkIn, dir, k);
+                       case Row:
+                               return localDedupMemorySafe ?
+                                       getUniqueRowValuesParallel(blkIn, k) :
+                                       
getUniqueRowValuesBatchedParallel(blkIn, dir, k);
+                       case Col:
+                               return localDedupMemorySafe ?
+                                       getUniqueColumnValuesParallel(blkIn, k) 
:
+                                       
getUniqueColumnValuesBatchedParallel(blkIn, dir, k);
+                       default:
+                               throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
+               }
+       }
+
+       /**
+        * Single-threaded baseline implementation for all unique directions.
+        * This preserves the original row-wise and column-wise unique behavior.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return matrix block containing unique values
+        */
+       private static MatrixBlock getUniqueValuesSequential(MatrixBlock blkIn, 
Types.Direction dir) {
                int rlen = blkIn.getNumRows();
                int clen = blkIn.getNumColumns();
 
                MatrixBlock blkOut = null;
-               // TODO optimize for dense/sparse/compressed (once multi-column 
support added)
-               
                switch (dir) {
-                       case RowCol: {
-                               // obtain set of unique items (dense input 
vector)
+                       case RowCol:
+                               // TODO optimize for dense/sparse/compressed 
(once multi-column support added)
+
+                               // obtain set of unique items
                                HashSet<Double> hashSet = new HashSet<>();
                                for( int i=0; i<rlen; i++ ) {
                                        for( int j=0; j<clen; j++ )
                                                hashSet.add(blkIn.get(i, j));
                                }
-                               
+
                                // allocate output block and place values
-                               int rlen2 = hashSet.size();
-                               blkOut = new MatrixBlock(rlen2, 1, 
false).allocateBlock();
-                               int pos = 0;
-                               for( Double val : hashSet )
-                                       blkOut.set(pos++, 0, val);
+                               blkOut = createRowColOutput(hashSet);
                                break;
-                       }
-                       case Row: {
-                               //2-pass algorithm to avoid unnecessarily large 
mem requirements
-                               HashSet<Double> hashSet = new HashSet<>();
+
+                       case Row:
+                               // 2-pass algorithm to avoid unnecessarily 
large mem requirements
+                               HashSet<Double> rowSet = new HashSet<>();
                                int clen2 = 0;
                                for( int i=0; i<rlen; i++ ) {
-                                       hashSet.clear();
+                                       rowSet.clear();
                                        for( int j=0; j<clen; j++ )
-                                               hashSet.add(blkIn.get(i, j));
-                                       clen2 = Math.max(clen2, hashSet.size());
+                                               rowSet.add(blkIn.get(i, j));
+                                       clen2 = Math.max(clen2, rowSet.size());
                                }
-                               
-                               //actual 
-                               blkOut = new MatrixBlock(rlen, clen2, 
false).allocateBlock();
+
+                               // actual
+                               blkOut = allocateOutputBlock(rlen, clen2);
                                for( int i=0; i<rlen; i++ ) {
-                                       hashSet.clear();
+                                       rowSet.clear();
                                        for( int j=0; j<clen; j++ )
-                                               hashSet.add(blkIn.get(i, j));
+                                               rowSet.add(blkIn.get(i, j));
                                        int pos = 0;
-                                       for( Double val : hashSet )
+                                       for( Double val : rowSet )
                                                blkOut.set(i, pos++, val);
                                }
                                break;
-                       }
-                       case Col: {
-                               //2-pass algorithm to avoid unnecessarily large 
mem requirements
-                               HashSet<Double> hashSet = new HashSet<>();
+
+                       case Col:
+                               // 2-pass algorithm to avoid unnecessarily 
large mem requirements
+                               HashSet<Double> colSet = new HashSet<>();
                                int rlen2 = 0;
                                for( int j=0; j<clen; j++ ) {
-                                       hashSet.clear();
+                                       colSet.clear();
                                        for( int i=0; i<rlen; i++ )
-                                               hashSet.add(blkIn.get(i, j));
-                                       rlen2 = Math.max(rlen2, hashSet.size());
+                                               colSet.add(blkIn.get(i, j));
+                                       rlen2 = Math.max(rlen2, colSet.size());
                                }
-                               
-                               //actual 
-                               blkOut = new MatrixBlock(rlen2, clen, 
false).allocateBlock();
+
+                               // actual
+                               blkOut = allocateOutputBlock(rlen2, clen);
                                for( int j=0; j<clen; j++ ) {
-                                       hashSet.clear();
+                                       colSet.clear();
                                        for( int i=0; i<rlen; i++ )
-                                               hashSet.add(blkIn.get(i, j));
+                                               colSet.add(blkIn.get(i, j));
                                        int pos = 0;
-                                       for( Double val : hashSet )
+                                       for( Double val : colSet )
                                                blkOut.set(pos++, j, val);
                                }
                                break;
+
+                       default:
+                               throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
+               }
+
+               blkOut.recomputeNonZeros();
+               blkOut.examSparsity();
+               return blkOut;
+       }
+
+       /**
+        * Parallel unique for all matrix values. Rows are split into balanced 
partitions,
+        * each task builds a local set, and the caller merges all local sets 
afterwards.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return one-column matrix block containing the unique values
+        */
+       private static MatrixBlock getUniqueValuesRowColParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumRows());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<UniqueValueTask> tasks = new ArrayList<>();
+                       for( int[] range : 
getBalancedRanges(blkIn.getNumRows(), numThreads) )
+                               tasks.add(new UniqueValueTask(blkIn, range[0], 
range[1]));
+
+                       // Merge local sets after the workers complete.
+                       HashSet<Double> hashSet = new HashSet<>();
+                       List<Future<HashSet<Double>>> rtasks = 
pool.invokeAll(tasks);
+                       for( int i = 0; i < rtasks.size(); i++ ) {
+                               HashSet<Double> localSet = rtasks.get(i).get();
+                               hashSet.addAll(localSet);
+                               localSet.clear();
+                               rtasks.set(i, null);
                        }
+
+                       return createRowColOutput(hashSet);
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Parallel row-wise unique values. A first pass computes the output 
width,
+        * and a second pass materializes the row-local unique values.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return matrix block containing row-wise unique values
+        */
+       private static MatrixBlock getUniqueRowValuesParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumRows());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<int[]> ranges = 
getBalancedRanges(blkIn.getNumRows(), numThreads);
+                       int clen2 = getMaxUniqueValues(pool, blkIn, 
Types.Direction.Row, ranges);
+                       MatrixBlock blkOut = 
allocateOutputBlock(blkIn.getNumRows(), clen2);
+                       fillUniqueValues(pool, blkIn, blkOut, 
Types.Direction.Row, ranges);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Parallel column-wise unique values. A first pass computes the output 
height,
+        * and a second pass materializes the column-local unique values.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return matrix block containing column-wise unique values
+        */
+       private static MatrixBlock getUniqueColumnValuesParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumColumns());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<int[]> ranges = 
getBalancedRanges(blkIn.getNumColumns(), numThreads);
+                       int rlen2 = getMaxUniqueValues(pool, blkIn, 
Types.Direction.Col, ranges);
+                       MatrixBlock blkOut = allocateOutputBlock(rlen2, 
blkIn.getNumColumns());
+                       fillUniqueValues(pool, blkIn, blkOut, 
Types.Direction.Col, ranges);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel unique for all matrix values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return one-column matrix block containing the unique values
+        */
+       private static MatrixBlock 
getUniqueValuesRowColBatchedParallel(MatrixBlock blkIn, Types.Direction dir, 
int k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       HashSet<Double> hashSet = new HashSet<>();
+                       for( int pos = 0; pos < config._len; ) {
+                               ArrayList<UniqueValueTask> tasks = new 
ArrayList<>();
+                               for( int i = 0; i < config._numThreads && pos < 
config._len; i++ ) {
+                                       int end = Math.min(pos + 
config._taskLen, config._len);
+                                       tasks.add(new UniqueValueTask(blkIn, 
pos, end));
+                                       pos = end;
+                               }
+
+                               List<Future<HashSet<Double>>> rtasks = 
pool.invokeAll(tasks);
+                               for( int i = 0; i < rtasks.size(); i++ ) {
+                                       HashSet<Double> localSet = 
rtasks.get(i).get();
+                                       hashSet.addAll(localSet);
+                                       localSet.clear();
+                                       rtasks.set(i, null);
+                               }
+                       }
+
+                       return createRowColOutput(hashSet);
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel row-wise unique values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing row-wise unique values
+        */
+       private static MatrixBlock 
getUniqueRowValuesBatchedParallel(MatrixBlock blkIn, Types.Direction dir, int 
k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       int clen2 = getMaxUniqueValuesBatched(pool, blkIn, dir, 
config);
+                       MatrixBlock blkOut = 
allocateOutputBlock(blkIn.getNumRows(), clen2);
+                       fillUniqueValuesBatched(pool, blkIn, blkOut, dir, 
config);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel column-wise unique values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing column-wise unique values
+        */
+       private static MatrixBlock 
getUniqueColumnValuesBatchedParallel(MatrixBlock blkIn, Types.Direction dir, 
int k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       int rlen2 = getMaxUniqueValuesBatched(pool, blkIn, dir, 
config);
+                       MatrixBlock blkOut = allocateOutputBlock(rlen2, 
blkIn.getNumColumns());
+                       fillUniqueValuesBatched(pool, blkIn, blkOut, dir, 
config);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Computes the maximum row-wise or column-wise unique count over 
balanced ranges.
+        */
+       private static int getMaxUniqueValues(ExecutorService pool, MatrixBlock 
blkIn, Types.Direction dir,
+               ArrayList<int[]> ranges) throws Exception {
+               ArrayList<UniqueCountTask> tasks = new ArrayList<>();
+               for( int[] range : ranges )
+                       tasks.add(new UniqueCountTask(blkIn, dir, range[0], 
range[1]));
+
+               int ret = 0;
+               List<Future<Integer>> rtasks = pool.invokeAll(tasks);
+               for( int i = 0; i < rtasks.size(); i++ ) {
+                       ret = Math.max(ret, rtasks.get(i).get());
+                       rtasks.set(i, null);
+               }
+               return ret;
+       }
+
+       /**
+        * Fills row-wise or column-wise unique values over balanced ranges.
+        */
+       private static void fillUniqueValues(ExecutorService pool, MatrixBlock 
blkIn, MatrixBlock blkOut,
+               Types.Direction dir, ArrayList<int[]> ranges) throws Exception {
+               ArrayList<UniqueOutputTask> tasks = new ArrayList<>();
+               for( int[] range : ranges )
+                       tasks.add(new UniqueOutputTask(blkIn, blkOut, dir, 
range[0], range[1]));
+               List<Future<Void>> rtasks = pool.invokeAll(tasks);
+               for( int i = 0; i < rtasks.size(); i++ ) {
+                       rtasks.get(i).get();
+                       rtasks.set(i, null);
+               }
+       }
+
+       /**
+        * Batched variant of getMaxUniqueValues.
+        */
+       private static int getMaxUniqueValuesBatched(ExecutorService pool, 
MatrixBlock blkIn, Types.Direction dir,
+               BatchConfig config) throws Exception {
+               int ret = 0;
+               for( int pos = 0; pos < config._len; ) {
+                       ArrayList<UniqueCountTask> tasks = new ArrayList<>();
+                       for( int i = 0; i < config._numThreads && pos < 
config._len; i++ ) {
+                               int end = Math.min(pos + config._taskLen, 
config._len);
+                               tasks.add(new UniqueCountTask(blkIn, dir, pos, 
end));
+                               pos = end;
+                       }
+
+                       List<Future<Integer>> rtasks = pool.invokeAll(tasks);
+                       for( int i = 0; i < rtasks.size(); i++ ) {
+                               ret = Math.max(ret, rtasks.get(i).get());
+                               rtasks.set(i, null);
+                       }
+               }
+               return ret;
+       }
+
+       /**
+        * Batched variant of fillUniqueValues.
+        */
+       private static void fillUniqueValuesBatched(ExecutorService pool, 
MatrixBlock blkIn, MatrixBlock blkOut,
+               Types.Direction dir, BatchConfig config) throws Exception {
+               for( int pos = 0; pos < config._len; ) {
+                       ArrayList<UniqueOutputTask> tasks = new ArrayList<>();
+                       for( int i = 0; i < config._numThreads && pos < 
config._len; i++ ) {
+                               int end = Math.min(pos + config._taskLen, 
config._len);
+                               tasks.add(new UniqueOutputTask(blkIn, blkOut, 
dir, pos, end));
+                               pos = end;
+                       }
+
+                       List<Future<Void>> rtasks = pool.invokeAll(tasks);
+                       for( int i = 0; i < rtasks.size(); i++ ) {
+                               rtasks.get(i).get();
+                               rtasks.set(i, null);
+                       }
+               }
+       }
+
+       /**
+        * Decides whether the input is large enough to justify local 
deduplication tasks.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return true if the parallel path should be used
+        */
+       private static boolean satisfiesMultiThreadingConstraints(MatrixBlock 
blkIn, Types.Direction dir, int k) {
+               if( k <= 1 || ((long) blkIn.getNumRows()) * 
blkIn.getNumColumns() < PAR_UNIQUE_NUMCELL_THRESHOLD )
+                       return false;
+
+               switch(dir) {
+                       case RowCol:
+                               return blkIn.getNumRows() > 1;
+                       case Row:
+                               return blkIn.getNumRows() > 1;
+                       case Col:
+                               return blkIn.getNumColumns() > 1;
                        default:
                                throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
                }
+       }
+
+       /**
+        * Creates balanced half-open ranges [start, end) using the same 
utility pattern as
+        * other SystemDS matrix libraries.
+        *
+        * @param len number of rows or columns to partition
+        * @param k requested degree of parallelism
+        * @return list of balanced index ranges
+        */
+       private static ArrayList<int[]> getBalancedRanges(int len, int k) {
+               ArrayList<int[]> ranges = new ArrayList<>();
+               ArrayList<Integer> blklens = 
UtilFunctions.getBalancedBlockSizesDefault(len, getNumThreads(k, len), false);
+               for( int i = 0, lb = 0; i < blklens.size(); lb += 
blklens.get(i), i++ )
+                       ranges.add(new int[] {lb, lb + blklens.get(i)});
+               return ranges;
+       }
+
+       /**
+        * Caps the number of workers by the number of row or column partitions 
available.
+        *
+        * @param k requested degree of parallelism
+        * @param len number of rows or columns to partition
+        * @return effective number of worker threads
+        */
+       private static int getNumThreads(int k, int len) {
+               return Math.max(1, Math.min(k, len));
+       }
+
+       /**
+        * Builds the execution plan for the batched parallel path.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return batch configuration, or null if batching would not use 
parallelism
+        */
+       private static BatchConfig getBatchConfig(MatrixBlock blkIn, 
Types.Direction dir, int k) {
+               int len = getPartitionLength(blkIn, dir);
+               long maxBatchIndexes = getMaxLocalDedupIndexes(blkIn, dir);
+               if( maxBatchIndexes < 2 )
+                       return null;
+
+               int numThreads = getNumThreads(k, (int) Math.min(len, 
maxBatchIndexes));
+               if( numThreads <= 1 )
+                       return null;
+
+               int taskLen = Math.max(1, (int) Math.min(Integer.MAX_VALUE, 
maxBatchIndexes / numThreads));

Review Comment:
   Then `numThreads x taskLen` consumes the entire 1/4 of the heap in the worst 
case



##########
src/main/java/org/apache/sysds/runtime/matrix/data/LibMatrixSketch.java:
##########
@@ -19,89 +19,695 @@
 
 package org.apache.sysds.runtime.matrix.data;
 
-import org.apache.sysds.common.Types;
-
+import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+
+import org.apache.sysds.common.Types;
+import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.util.CommonThreadPool;
+import org.apache.sysds.runtime.util.UtilFunctions;
 
 public class LibMatrixSketch {
+       private static final long PAR_UNIQUE_NUMCELL_THRESHOLD = 1024 * 16;
+       private static final long PAR_UNIQUE_MAX_LOCAL_BYTES_FRACTION = 4;
+       private static final long PAR_UNIQUE_LOCAL_BYTES_OVERHEAD = 8;
 
+       /**
+        * Computes unique values with the original single-threaded behavior.
+        * The overload with a parallelism argument keeps this path as the k=1 
baseline.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return matrix block containing unique values
+        */
        public static MatrixBlock getUniqueValues(MatrixBlock blkIn, 
Types.Direction dir) {
-               //similar to R's unique, this operation takes a matrix and 
computes the
-               //unique values (or rows in case of multiple column inputs)
-               
+               return getUniqueValues(blkIn, dir, 1);
+       }
+
+       /**
+        * Computes unique values. For sufficiently large inputs and k > 1, 
this uses
+        * parallel local deduplication or its batched variant.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing unique values
+        */
+       public static MatrixBlock getUniqueValues(MatrixBlock blkIn, 
Types.Direction dir, int k) {
+               // Similar to R's unique, this operation computes unique values 
according
+               // to the requested direction.
+               if( !satisfiesMultiThreadingConstraints(blkIn, dir, k) )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               boolean localDedupMemorySafe = 
isLocalDedupMemoryBudgetSafe(blkIn, dir);
+               switch(dir) {
+                       case RowCol:
+                               return localDedupMemorySafe ?
+                                       getUniqueValuesRowColParallel(blkIn, k) 
:
+                                       
getUniqueValuesRowColBatchedParallel(blkIn, dir, k);
+                       case Row:
+                               return localDedupMemorySafe ?
+                                       getUniqueRowValuesParallel(blkIn, k) :
+                                       
getUniqueRowValuesBatchedParallel(blkIn, dir, k);
+                       case Col:
+                               return localDedupMemorySafe ?
+                                       getUniqueColumnValuesParallel(blkIn, k) 
:
+                                       
getUniqueColumnValuesBatchedParallel(blkIn, dir, k);
+                       default:
+                               throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
+               }
+       }
+
+       /**
+        * Single-threaded baseline implementation for all unique directions.
+        * This preserves the original row-wise and column-wise unique behavior.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return matrix block containing unique values
+        */
+       private static MatrixBlock getUniqueValuesSequential(MatrixBlock blkIn, 
Types.Direction dir) {
                int rlen = blkIn.getNumRows();
                int clen = blkIn.getNumColumns();
 
                MatrixBlock blkOut = null;
-               // TODO optimize for dense/sparse/compressed (once multi-column 
support added)
-               
                switch (dir) {
-                       case RowCol: {
-                               // obtain set of unique items (dense input 
vector)
+                       case RowCol:
+                               // TODO optimize for dense/sparse/compressed 
(once multi-column support added)
+
+                               // obtain set of unique items
                                HashSet<Double> hashSet = new HashSet<>();
                                for( int i=0; i<rlen; i++ ) {
                                        for( int j=0; j<clen; j++ )
                                                hashSet.add(blkIn.get(i, j));
                                }
-                               
+
                                // allocate output block and place values
-                               int rlen2 = hashSet.size();
-                               blkOut = new MatrixBlock(rlen2, 1, 
false).allocateBlock();
-                               int pos = 0;
-                               for( Double val : hashSet )
-                                       blkOut.set(pos++, 0, val);
+                               blkOut = createRowColOutput(hashSet);
                                break;
-                       }
-                       case Row: {
-                               //2-pass algorithm to avoid unnecessarily large 
mem requirements
-                               HashSet<Double> hashSet = new HashSet<>();
+
+                       case Row:
+                               // 2-pass algorithm to avoid unnecessarily 
large mem requirements
+                               HashSet<Double> rowSet = new HashSet<>();
                                int clen2 = 0;
                                for( int i=0; i<rlen; i++ ) {
-                                       hashSet.clear();
+                                       rowSet.clear();
                                        for( int j=0; j<clen; j++ )
-                                               hashSet.add(blkIn.get(i, j));
-                                       clen2 = Math.max(clen2, hashSet.size());
+                                               rowSet.add(blkIn.get(i, j));
+                                       clen2 = Math.max(clen2, rowSet.size());
                                }
-                               
-                               //actual 
-                               blkOut = new MatrixBlock(rlen, clen2, 
false).allocateBlock();
+
+                               // actual
+                               blkOut = allocateOutputBlock(rlen, clen2);
                                for( int i=0; i<rlen; i++ ) {
-                                       hashSet.clear();
+                                       rowSet.clear();
                                        for( int j=0; j<clen; j++ )
-                                               hashSet.add(blkIn.get(i, j));
+                                               rowSet.add(blkIn.get(i, j));
                                        int pos = 0;
-                                       for( Double val : hashSet )
+                                       for( Double val : rowSet )
                                                blkOut.set(i, pos++, val);
                                }
                                break;
-                       }
-                       case Col: {
-                               //2-pass algorithm to avoid unnecessarily large 
mem requirements
-                               HashSet<Double> hashSet = new HashSet<>();
+
+                       case Col:
+                               // 2-pass algorithm to avoid unnecessarily 
large mem requirements
+                               HashSet<Double> colSet = new HashSet<>();
                                int rlen2 = 0;
                                for( int j=0; j<clen; j++ ) {
-                                       hashSet.clear();
+                                       colSet.clear();
                                        for( int i=0; i<rlen; i++ )
-                                               hashSet.add(blkIn.get(i, j));
-                                       rlen2 = Math.max(rlen2, hashSet.size());
+                                               colSet.add(blkIn.get(i, j));
+                                       rlen2 = Math.max(rlen2, colSet.size());
                                }
-                               
-                               //actual 
-                               blkOut = new MatrixBlock(rlen2, clen, 
false).allocateBlock();
+
+                               // actual
+                               blkOut = allocateOutputBlock(rlen2, clen);
                                for( int j=0; j<clen; j++ ) {
-                                       hashSet.clear();
+                                       colSet.clear();
                                        for( int i=0; i<rlen; i++ )
-                                               hashSet.add(blkIn.get(i, j));
+                                               colSet.add(blkIn.get(i, j));
                                        int pos = 0;
-                                       for( Double val : hashSet )
+                                       for( Double val : colSet )
                                                blkOut.set(pos++, j, val);
                                }
                                break;
+
+                       default:
+                               throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
+               }
+
+               blkOut.recomputeNonZeros();
+               blkOut.examSparsity();
+               return blkOut;
+       }
+
+       /**
+        * Parallel unique for all matrix values. Rows are split into balanced 
partitions,
+        * each task builds a local set, and the caller merges all local sets 
afterwards.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return one-column matrix block containing the unique values
+        */
+       private static MatrixBlock getUniqueValuesRowColParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumRows());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<UniqueValueTask> tasks = new ArrayList<>();
+                       for( int[] range : 
getBalancedRanges(blkIn.getNumRows(), numThreads) )
+                               tasks.add(new UniqueValueTask(blkIn, range[0], 
range[1]));
+
+                       // Merge local sets after the workers complete.
+                       HashSet<Double> hashSet = new HashSet<>();
+                       List<Future<HashSet<Double>>> rtasks = 
pool.invokeAll(tasks);
+                       for( int i = 0; i < rtasks.size(); i++ ) {
+                               HashSet<Double> localSet = rtasks.get(i).get();
+                               hashSet.addAll(localSet);
+                               localSet.clear();
+                               rtasks.set(i, null);
                        }
+
+                       return createRowColOutput(hashSet);
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Parallel row-wise unique values. A first pass computes the output 
width,
+        * and a second pass materializes the row-local unique values.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return matrix block containing row-wise unique values
+        */
+       private static MatrixBlock getUniqueRowValuesParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumRows());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<int[]> ranges = 
getBalancedRanges(blkIn.getNumRows(), numThreads);
+                       int clen2 = getMaxUniqueValues(pool, blkIn, 
Types.Direction.Row, ranges);
+                       MatrixBlock blkOut = 
allocateOutputBlock(blkIn.getNumRows(), clen2);
+                       fillUniqueValues(pool, blkIn, blkOut, 
Types.Direction.Row, ranges);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Parallel column-wise unique values. A first pass computes the output 
height,
+        * and a second pass materializes the column-local unique values.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return matrix block containing column-wise unique values
+        */
+       private static MatrixBlock getUniqueColumnValuesParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumColumns());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<int[]> ranges = 
getBalancedRanges(blkIn.getNumColumns(), numThreads);
+                       int rlen2 = getMaxUniqueValues(pool, blkIn, 
Types.Direction.Col, ranges);
+                       MatrixBlock blkOut = allocateOutputBlock(rlen2, 
blkIn.getNumColumns());
+                       fillUniqueValues(pool, blkIn, blkOut, 
Types.Direction.Col, ranges);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel unique for all matrix values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return one-column matrix block containing the unique values
+        */
+       private static MatrixBlock 
getUniqueValuesRowColBatchedParallel(MatrixBlock blkIn, Types.Direction dir, 
int k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       HashSet<Double> hashSet = new HashSet<>();
+                       for( int pos = 0; pos < config._len; ) {
+                               ArrayList<UniqueValueTask> tasks = new 
ArrayList<>();
+                               for( int i = 0; i < config._numThreads && pos < 
config._len; i++ ) {
+                                       int end = Math.min(pos + 
config._taskLen, config._len);

Review Comment:
   This will overflow when `pos` is close to `Integer.MAX_VALUE`, i.e., when 
the number of rows in the matrix is very large



##########
src/main/java/org/apache/sysds/runtime/matrix/data/LibMatrixSketch.java:
##########
@@ -19,89 +19,695 @@
 
 package org.apache.sysds.runtime.matrix.data;
 
-import org.apache.sysds.common.Types;
-
+import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+
+import org.apache.sysds.common.Types;
+import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.util.CommonThreadPool;
+import org.apache.sysds.runtime.util.UtilFunctions;
 
 public class LibMatrixSketch {
+       private static final long PAR_UNIQUE_NUMCELL_THRESHOLD = 1024 * 16;
+       private static final long PAR_UNIQUE_MAX_LOCAL_BYTES_FRACTION = 4;
+       private static final long PAR_UNIQUE_LOCAL_BYTES_OVERHEAD = 8;
 
+       /**
+        * Computes unique values with the original single-threaded behavior.
+        * The overload with a parallelism argument keeps this path as the k=1 
baseline.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return matrix block containing unique values
+        */
        public static MatrixBlock getUniqueValues(MatrixBlock blkIn, 
Types.Direction dir) {
-               //similar to R's unique, this operation takes a matrix and 
computes the
-               //unique values (or rows in case of multiple column inputs)
-               
+               return getUniqueValues(blkIn, dir, 1);
+       }
+
+       /**
+        * Computes unique values. For sufficiently large inputs and k > 1, 
this uses
+        * parallel local deduplication or its batched variant.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing unique values
+        */
+       public static MatrixBlock getUniqueValues(MatrixBlock blkIn, 
Types.Direction dir, int k) {
+               // Similar to R's unique, this operation computes unique values 
according
+               // to the requested direction.
+               if( !satisfiesMultiThreadingConstraints(blkIn, dir, k) )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               boolean localDedupMemorySafe = 
isLocalDedupMemoryBudgetSafe(blkIn, dir);
+               switch(dir) {
+                       case RowCol:
+                               return localDedupMemorySafe ?
+                                       getUniqueValuesRowColParallel(blkIn, k) 
:
+                                       
getUniqueValuesRowColBatchedParallel(blkIn, dir, k);
+                       case Row:
+                               return localDedupMemorySafe ?
+                                       getUniqueRowValuesParallel(blkIn, k) :
+                                       
getUniqueRowValuesBatchedParallel(blkIn, dir, k);
+                       case Col:
+                               return localDedupMemorySafe ?
+                                       getUniqueColumnValuesParallel(blkIn, k) 
:
+                                       
getUniqueColumnValuesBatchedParallel(blkIn, dir, k);
+                       default:
+                               throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
+               }
+       }
+
+       /**
+        * Single-threaded baseline implementation for all unique directions.
+        * This preserves the original row-wise and column-wise unique behavior.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return matrix block containing unique values
+        */
+       private static MatrixBlock getUniqueValuesSequential(MatrixBlock blkIn, 
Types.Direction dir) {
                int rlen = blkIn.getNumRows();
                int clen = blkIn.getNumColumns();
 
                MatrixBlock blkOut = null;
-               // TODO optimize for dense/sparse/compressed (once multi-column 
support added)
-               
                switch (dir) {
-                       case RowCol: {
-                               // obtain set of unique items (dense input 
vector)
+                       case RowCol:
+                               // TODO optimize for dense/sparse/compressed 
(once multi-column support added)
+
+                               // obtain set of unique items
                                HashSet<Double> hashSet = new HashSet<>();
                                for( int i=0; i<rlen; i++ ) {
                                        for( int j=0; j<clen; j++ )
                                                hashSet.add(blkIn.get(i, j));
                                }
-                               
+
                                // allocate output block and place values
-                               int rlen2 = hashSet.size();
-                               blkOut = new MatrixBlock(rlen2, 1, 
false).allocateBlock();
-                               int pos = 0;
-                               for( Double val : hashSet )
-                                       blkOut.set(pos++, 0, val);
+                               blkOut = createRowColOutput(hashSet);
                                break;
-                       }
-                       case Row: {
-                               //2-pass algorithm to avoid unnecessarily large 
mem requirements
-                               HashSet<Double> hashSet = new HashSet<>();
+
+                       case Row:
+                               // 2-pass algorithm to avoid unnecessarily 
large mem requirements
+                               HashSet<Double> rowSet = new HashSet<>();
                                int clen2 = 0;
                                for( int i=0; i<rlen; i++ ) {
-                                       hashSet.clear();
+                                       rowSet.clear();
                                        for( int j=0; j<clen; j++ )
-                                               hashSet.add(blkIn.get(i, j));
-                                       clen2 = Math.max(clen2, hashSet.size());
+                                               rowSet.add(blkIn.get(i, j));
+                                       clen2 = Math.max(clen2, rowSet.size());
                                }
-                               
-                               //actual 
-                               blkOut = new MatrixBlock(rlen, clen2, 
false).allocateBlock();
+
+                               // actual
+                               blkOut = allocateOutputBlock(rlen, clen2);
                                for( int i=0; i<rlen; i++ ) {
-                                       hashSet.clear();
+                                       rowSet.clear();
                                        for( int j=0; j<clen; j++ )
-                                               hashSet.add(blkIn.get(i, j));
+                                               rowSet.add(blkIn.get(i, j));
                                        int pos = 0;
-                                       for( Double val : hashSet )
+                                       for( Double val : rowSet )
                                                blkOut.set(i, pos++, val);
                                }
                                break;
-                       }
-                       case Col: {
-                               //2-pass algorithm to avoid unnecessarily large 
mem requirements
-                               HashSet<Double> hashSet = new HashSet<>();
+
+                       case Col:
+                               // 2-pass algorithm to avoid unnecessarily 
large mem requirements
+                               HashSet<Double> colSet = new HashSet<>();
                                int rlen2 = 0;
                                for( int j=0; j<clen; j++ ) {
-                                       hashSet.clear();
+                                       colSet.clear();
                                        for( int i=0; i<rlen; i++ )
-                                               hashSet.add(blkIn.get(i, j));
-                                       rlen2 = Math.max(rlen2, hashSet.size());
+                                               colSet.add(blkIn.get(i, j));
+                                       rlen2 = Math.max(rlen2, colSet.size());
                                }
-                               
-                               //actual 
-                               blkOut = new MatrixBlock(rlen2, clen, 
false).allocateBlock();
+
+                               // actual
+                               blkOut = allocateOutputBlock(rlen2, clen);
                                for( int j=0; j<clen; j++ ) {
-                                       hashSet.clear();
+                                       colSet.clear();
                                        for( int i=0; i<rlen; i++ )
-                                               hashSet.add(blkIn.get(i, j));
+                                               colSet.add(blkIn.get(i, j));
                                        int pos = 0;
-                                       for( Double val : hashSet )
+                                       for( Double val : colSet )
                                                blkOut.set(pos++, j, val);
                                }
                                break;
+
+                       default:
+                               throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
+               }
+
+               blkOut.recomputeNonZeros();
+               blkOut.examSparsity();
+               return blkOut;
+       }
+
+       /**
+        * Parallel unique for all matrix values. Rows are split into balanced 
partitions,
+        * each task builds a local set, and the caller merges all local sets 
afterwards.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return one-column matrix block containing the unique values
+        */
+       private static MatrixBlock getUniqueValuesRowColParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumRows());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<UniqueValueTask> tasks = new ArrayList<>();
+                       for( int[] range : 
getBalancedRanges(blkIn.getNumRows(), numThreads) )
+                               tasks.add(new UniqueValueTask(blkIn, range[0], 
range[1]));
+
+                       // Merge local sets after the workers complete.
+                       HashSet<Double> hashSet = new HashSet<>();
+                       List<Future<HashSet<Double>>> rtasks = 
pool.invokeAll(tasks);
+                       for( int i = 0; i < rtasks.size(); i++ ) {
+                               HashSet<Double> localSet = rtasks.get(i).get();
+                               hashSet.addAll(localSet);
+                               localSet.clear();
+                               rtasks.set(i, null);
                        }
+
+                       return createRowColOutput(hashSet);
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Parallel row-wise unique values. A first pass computes the output 
width,
+        * and a second pass materializes the row-local unique values.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return matrix block containing row-wise unique values
+        */
+       private static MatrixBlock getUniqueRowValuesParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumRows());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<int[]> ranges = 
getBalancedRanges(blkIn.getNumRows(), numThreads);
+                       int clen2 = getMaxUniqueValues(pool, blkIn, 
Types.Direction.Row, ranges);
+                       MatrixBlock blkOut = 
allocateOutputBlock(blkIn.getNumRows(), clen2);
+                       fillUniqueValues(pool, blkIn, blkOut, 
Types.Direction.Row, ranges);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Parallel column-wise unique values. A first pass computes the output 
height,
+        * and a second pass materializes the column-local unique values.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return matrix block containing column-wise unique values
+        */
+       private static MatrixBlock getUniqueColumnValuesParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumColumns());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<int[]> ranges = 
getBalancedRanges(blkIn.getNumColumns(), numThreads);
+                       int rlen2 = getMaxUniqueValues(pool, blkIn, 
Types.Direction.Col, ranges);
+                       MatrixBlock blkOut = allocateOutputBlock(rlen2, 
blkIn.getNumColumns());
+                       fillUniqueValues(pool, blkIn, blkOut, 
Types.Direction.Col, ranges);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel unique for all matrix values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return one-column matrix block containing the unique values
+        */
+       private static MatrixBlock 
getUniqueValuesRowColBatchedParallel(MatrixBlock blkIn, Types.Direction dir, 
int k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       HashSet<Double> hashSet = new HashSet<>();
+                       for( int pos = 0; pos < config._len; ) {
+                               ArrayList<UniqueValueTask> tasks = new 
ArrayList<>();
+                               for( int i = 0; i < config._numThreads && pos < 
config._len; i++ ) {
+                                       int end = Math.min(pos + 
config._taskLen, config._len);
+                                       tasks.add(new UniqueValueTask(blkIn, 
pos, end));
+                                       pos = end;
+                               }
+
+                               List<Future<HashSet<Double>>> rtasks = 
pool.invokeAll(tasks);
+                               for( int i = 0; i < rtasks.size(); i++ ) {
+                                       HashSet<Double> localSet = 
rtasks.get(i).get();
+                                       hashSet.addAll(localSet);
+                                       localSet.clear();
+                                       rtasks.set(i, null);
+                               }
+                       }
+
+                       return createRowColOutput(hashSet);
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel row-wise unique values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing row-wise unique values
+        */
+       private static MatrixBlock 
getUniqueRowValuesBatchedParallel(MatrixBlock blkIn, Types.Direction dir, int 
k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       int clen2 = getMaxUniqueValuesBatched(pool, blkIn, dir, 
config);
+                       MatrixBlock blkOut = 
allocateOutputBlock(blkIn.getNumRows(), clen2);
+                       fillUniqueValuesBatched(pool, blkIn, blkOut, dir, 
config);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel column-wise unique values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing column-wise unique values
+        */
+       private static MatrixBlock 
getUniqueColumnValuesBatchedParallel(MatrixBlock blkIn, Types.Direction dir, 
int k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       int rlen2 = getMaxUniqueValuesBatched(pool, blkIn, dir, 
config);
+                       MatrixBlock blkOut = allocateOutputBlock(rlen2, 
blkIn.getNumColumns());
+                       fillUniqueValuesBatched(pool, blkIn, blkOut, dir, 
config);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Computes the maximum row-wise or column-wise unique count over 
balanced ranges.
+        */
+       private static int getMaxUniqueValues(ExecutorService pool, MatrixBlock 
blkIn, Types.Direction dir,
+               ArrayList<int[]> ranges) throws Exception {
+               ArrayList<UniqueCountTask> tasks = new ArrayList<>();
+               for( int[] range : ranges )
+                       tasks.add(new UniqueCountTask(blkIn, dir, range[0], 
range[1]));
+
+               int ret = 0;
+               List<Future<Integer>> rtasks = pool.invokeAll(tasks);
+               for( int i = 0; i < rtasks.size(); i++ ) {
+                       ret = Math.max(ret, rtasks.get(i).get());
+                       rtasks.set(i, null);
+               }
+               return ret;
+       }
+
+       /**
+        * Fills row-wise or column-wise unique values over balanced ranges.
+        */
+       private static void fillUniqueValues(ExecutorService pool, MatrixBlock 
blkIn, MatrixBlock blkOut,
+               Types.Direction dir, ArrayList<int[]> ranges) throws Exception {
+               ArrayList<UniqueOutputTask> tasks = new ArrayList<>();
+               for( int[] range : ranges )
+                       tasks.add(new UniqueOutputTask(blkIn, blkOut, dir, 
range[0], range[1]));
+               List<Future<Void>> rtasks = pool.invokeAll(tasks);
+               for( int i = 0; i < rtasks.size(); i++ ) {
+                       rtasks.get(i).get();
+                       rtasks.set(i, null);
+               }
+       }
+
+       /**
+        * Batched variant of getMaxUniqueValues.
+        */
+       private static int getMaxUniqueValuesBatched(ExecutorService pool, 
MatrixBlock blkIn, Types.Direction dir,
+               BatchConfig config) throws Exception {
+               int ret = 0;
+               for( int pos = 0; pos < config._len; ) {
+                       ArrayList<UniqueCountTask> tasks = new ArrayList<>();
+                       for( int i = 0; i < config._numThreads && pos < 
config._len; i++ ) {
+                               int end = Math.min(pos + config._taskLen, 
config._len);
+                               tasks.add(new UniqueCountTask(blkIn, dir, pos, 
end));
+                               pos = end;
+                       }
+
+                       List<Future<Integer>> rtasks = pool.invokeAll(tasks);
+                       for( int i = 0; i < rtasks.size(); i++ ) {
+                               ret = Math.max(ret, rtasks.get(i).get());
+                               rtasks.set(i, null);
+                       }
+               }
+               return ret;
+       }
+
+       /**
+        * Batched variant of fillUniqueValues.
+        */
+       private static void fillUniqueValuesBatched(ExecutorService pool, 
MatrixBlock blkIn, MatrixBlock blkOut,
+               Types.Direction dir, BatchConfig config) throws Exception {
+               for( int pos = 0; pos < config._len; ) {
+                       ArrayList<UniqueOutputTask> tasks = new ArrayList<>();
+                       for( int i = 0; i < config._numThreads && pos < 
config._len; i++ ) {
+                               int end = Math.min(pos + config._taskLen, 
config._len);

Review Comment:
   Same overflow issue



##########
src/main/java/org/apache/sysds/runtime/matrix/data/LibMatrixSketch.java:
##########
@@ -19,89 +19,695 @@
 
 package org.apache.sysds.runtime.matrix.data;
 
-import org.apache.sysds.common.Types;
-
+import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+
+import org.apache.sysds.common.Types;
+import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.util.CommonThreadPool;
+import org.apache.sysds.runtime.util.UtilFunctions;
 
 public class LibMatrixSketch {
+       private static final long PAR_UNIQUE_NUMCELL_THRESHOLD = 1024 * 16;
+       private static final long PAR_UNIQUE_MAX_LOCAL_BYTES_FRACTION = 4;
+       private static final long PAR_UNIQUE_LOCAL_BYTES_OVERHEAD = 8;
 
+       /**
+        * Computes unique values with the original single-threaded behavior.
+        * The overload with a parallelism argument keeps this path as the k=1 
baseline.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return matrix block containing unique values
+        */
        public static MatrixBlock getUniqueValues(MatrixBlock blkIn, 
Types.Direction dir) {
-               //similar to R's unique, this operation takes a matrix and 
computes the
-               //unique values (or rows in case of multiple column inputs)
-               
+               return getUniqueValues(blkIn, dir, 1);
+       }
+
+       /**
+        * Computes unique values. For sufficiently large inputs and k > 1, 
this uses
+        * parallel local deduplication or its batched variant.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing unique values
+        */
+       public static MatrixBlock getUniqueValues(MatrixBlock blkIn, 
Types.Direction dir, int k) {
+               // Similar to R's unique, this operation computes unique values 
according
+               // to the requested direction.
+               if( !satisfiesMultiThreadingConstraints(blkIn, dir, k) )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               boolean localDedupMemorySafe = 
isLocalDedupMemoryBudgetSafe(blkIn, dir);
+               switch(dir) {
+                       case RowCol:
+                               return localDedupMemorySafe ?
+                                       getUniqueValuesRowColParallel(blkIn, k) 
:
+                                       
getUniqueValuesRowColBatchedParallel(blkIn, dir, k);
+                       case Row:
+                               return localDedupMemorySafe ?
+                                       getUniqueRowValuesParallel(blkIn, k) :
+                                       
getUniqueRowValuesBatchedParallel(blkIn, dir, k);
+                       case Col:
+                               return localDedupMemorySafe ?
+                                       getUniqueColumnValuesParallel(blkIn, k) 
:
+                                       
getUniqueColumnValuesBatchedParallel(blkIn, dir, k);
+                       default:
+                               throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
+               }
+       }
+
+       /**
+        * Single-threaded baseline implementation for all unique directions.
+        * This preserves the original row-wise and column-wise unique behavior.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return matrix block containing unique values
+        */
+       private static MatrixBlock getUniqueValuesSequential(MatrixBlock blkIn, 
Types.Direction dir) {
                int rlen = blkIn.getNumRows();
                int clen = blkIn.getNumColumns();
 
                MatrixBlock blkOut = null;
-               // TODO optimize for dense/sparse/compressed (once multi-column 
support added)
-               
                switch (dir) {
-                       case RowCol: {
-                               // obtain set of unique items (dense input 
vector)
+                       case RowCol:
+                               // TODO optimize for dense/sparse/compressed 
(once multi-column support added)
+
+                               // obtain set of unique items
                                HashSet<Double> hashSet = new HashSet<>();
                                for( int i=0; i<rlen; i++ ) {
                                        for( int j=0; j<clen; j++ )
                                                hashSet.add(blkIn.get(i, j));
                                }
-                               
+
                                // allocate output block and place values
-                               int rlen2 = hashSet.size();
-                               blkOut = new MatrixBlock(rlen2, 1, 
false).allocateBlock();
-                               int pos = 0;
-                               for( Double val : hashSet )
-                                       blkOut.set(pos++, 0, val);
+                               blkOut = createRowColOutput(hashSet);
                                break;
-                       }
-                       case Row: {
-                               //2-pass algorithm to avoid unnecessarily large 
mem requirements
-                               HashSet<Double> hashSet = new HashSet<>();
+
+                       case Row:
+                               // 2-pass algorithm to avoid unnecessarily 
large mem requirements
+                               HashSet<Double> rowSet = new HashSet<>();
                                int clen2 = 0;
                                for( int i=0; i<rlen; i++ ) {
-                                       hashSet.clear();
+                                       rowSet.clear();
                                        for( int j=0; j<clen; j++ )
-                                               hashSet.add(blkIn.get(i, j));
-                                       clen2 = Math.max(clen2, hashSet.size());
+                                               rowSet.add(blkIn.get(i, j));
+                                       clen2 = Math.max(clen2, rowSet.size());
                                }
-                               
-                               //actual 
-                               blkOut = new MatrixBlock(rlen, clen2, 
false).allocateBlock();
+
+                               // actual
+                               blkOut = allocateOutputBlock(rlen, clen2);
                                for( int i=0; i<rlen; i++ ) {
-                                       hashSet.clear();
+                                       rowSet.clear();
                                        for( int j=0; j<clen; j++ )
-                                               hashSet.add(blkIn.get(i, j));
+                                               rowSet.add(blkIn.get(i, j));
                                        int pos = 0;
-                                       for( Double val : hashSet )
+                                       for( Double val : rowSet )
                                                blkOut.set(i, pos++, val);
                                }
                                break;
-                       }
-                       case Col: {
-                               //2-pass algorithm to avoid unnecessarily large 
mem requirements
-                               HashSet<Double> hashSet = new HashSet<>();
+
+                       case Col:
+                               // 2-pass algorithm to avoid unnecessarily 
large mem requirements
+                               HashSet<Double> colSet = new HashSet<>();
                                int rlen2 = 0;
                                for( int j=0; j<clen; j++ ) {
-                                       hashSet.clear();
+                                       colSet.clear();
                                        for( int i=0; i<rlen; i++ )
-                                               hashSet.add(blkIn.get(i, j));
-                                       rlen2 = Math.max(rlen2, hashSet.size());
+                                               colSet.add(blkIn.get(i, j));
+                                       rlen2 = Math.max(rlen2, colSet.size());
                                }
-                               
-                               //actual 
-                               blkOut = new MatrixBlock(rlen2, clen, 
false).allocateBlock();
+
+                               // actual
+                               blkOut = allocateOutputBlock(rlen2, clen);
                                for( int j=0; j<clen; j++ ) {
-                                       hashSet.clear();
+                                       colSet.clear();
                                        for( int i=0; i<rlen; i++ )
-                                               hashSet.add(blkIn.get(i, j));
+                                               colSet.add(blkIn.get(i, j));
                                        int pos = 0;
-                                       for( Double val : hashSet )
+                                       for( Double val : colSet )
                                                blkOut.set(pos++, j, val);
                                }
                                break;
+
+                       default:
+                               throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
+               }
+
+               blkOut.recomputeNonZeros();
+               blkOut.examSparsity();
+               return blkOut;
+       }
+
+       /**
+        * Parallel unique for all matrix values. Rows are split into balanced 
partitions,
+        * each task builds a local set, and the caller merges all local sets 
afterwards.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return one-column matrix block containing the unique values
+        */
+       private static MatrixBlock getUniqueValuesRowColParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumRows());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<UniqueValueTask> tasks = new ArrayList<>();
+                       for( int[] range : 
getBalancedRanges(blkIn.getNumRows(), numThreads) )
+                               tasks.add(new UniqueValueTask(blkIn, range[0], 
range[1]));
+
+                       // Merge local sets after the workers complete.
+                       HashSet<Double> hashSet = new HashSet<>();
+                       List<Future<HashSet<Double>>> rtasks = 
pool.invokeAll(tasks);
+                       for( int i = 0; i < rtasks.size(); i++ ) {
+                               HashSet<Double> localSet = rtasks.get(i).get();
+                               hashSet.addAll(localSet);
+                               localSet.clear();
+                               rtasks.set(i, null);
                        }
+
+                       return createRowColOutput(hashSet);
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Parallel row-wise unique values. A first pass computes the output 
width,
+        * and a second pass materializes the row-local unique values.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return matrix block containing row-wise unique values
+        */
+       private static MatrixBlock getUniqueRowValuesParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumRows());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<int[]> ranges = 
getBalancedRanges(blkIn.getNumRows(), numThreads);
+                       int clen2 = getMaxUniqueValues(pool, blkIn, 
Types.Direction.Row, ranges);
+                       MatrixBlock blkOut = 
allocateOutputBlock(blkIn.getNumRows(), clen2);
+                       fillUniqueValues(pool, blkIn, blkOut, 
Types.Direction.Row, ranges);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Parallel column-wise unique values. A first pass computes the output 
height,
+        * and a second pass materializes the column-local unique values.
+        *
+        * @param blkIn input matrix block
+        * @param k requested degree of parallelism
+        * @return matrix block containing column-wise unique values
+        */
+       private static MatrixBlock getUniqueColumnValuesParallel(MatrixBlock 
blkIn, int k) {
+               int numThreads = getNumThreads(k, blkIn.getNumColumns());
+               ExecutorService pool = CommonThreadPool.get(numThreads);
+               try {
+                       ArrayList<int[]> ranges = 
getBalancedRanges(blkIn.getNumColumns(), numThreads);
+                       int rlen2 = getMaxUniqueValues(pool, blkIn, 
Types.Direction.Col, ranges);
+                       MatrixBlock blkOut = allocateOutputBlock(rlen2, 
blkIn.getNumColumns());
+                       fillUniqueValues(pool, blkIn, blkOut, 
Types.Direction.Col, ranges);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel unique for all matrix values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return one-column matrix block containing the unique values
+        */
+       private static MatrixBlock 
getUniqueValuesRowColBatchedParallel(MatrixBlock blkIn, Types.Direction dir, 
int k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       HashSet<Double> hashSet = new HashSet<>();
+                       for( int pos = 0; pos < config._len; ) {
+                               ArrayList<UniqueValueTask> tasks = new 
ArrayList<>();
+                               for( int i = 0; i < config._numThreads && pos < 
config._len; i++ ) {
+                                       int end = Math.min(pos + 
config._taskLen, config._len);
+                                       tasks.add(new UniqueValueTask(blkIn, 
pos, end));
+                                       pos = end;
+                               }
+
+                               List<Future<HashSet<Double>>> rtasks = 
pool.invokeAll(tasks);
+                               for( int i = 0; i < rtasks.size(); i++ ) {
+                                       HashSet<Double> localSet = 
rtasks.get(i).get();
+                                       hashSet.addAll(localSet);
+                                       localSet.clear();
+                                       rtasks.set(i, null);
+                               }
+                       }
+
+                       return createRowColOutput(hashSet);
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel row-wise unique values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing row-wise unique values
+        */
+       private static MatrixBlock 
getUniqueRowValuesBatchedParallel(MatrixBlock blkIn, Types.Direction dir, int 
k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       int clen2 = getMaxUniqueValuesBatched(pool, blkIn, dir, 
config);
+                       MatrixBlock blkOut = 
allocateOutputBlock(blkIn.getNumRows(), clen2);
+                       fillUniqueValuesBatched(pool, blkIn, blkOut, dir, 
config);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Batched parallel column-wise unique values.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return matrix block containing column-wise unique values
+        */
+       private static MatrixBlock 
getUniqueColumnValuesBatchedParallel(MatrixBlock blkIn, Types.Direction dir, 
int k) {
+               BatchConfig config = getBatchConfig(blkIn, dir, k);
+               if( config == null )
+                       return getUniqueValuesSequential(blkIn, dir);
+
+               ExecutorService pool = CommonThreadPool.get(config._numThreads);
+               try {
+                       int rlen2 = getMaxUniqueValuesBatched(pool, blkIn, dir, 
config);
+                       MatrixBlock blkOut = allocateOutputBlock(rlen2, 
blkIn.getNumColumns());
+                       fillUniqueValuesBatched(pool, blkIn, blkOut, dir, 
config);
+
+                       blkOut.recomputeNonZeros();
+                       blkOut.examSparsity();
+                       return blkOut;
+               }
+               catch(Exception ex) {
+                       throw new DMLRuntimeException(ex);
+               }
+               finally {
+                       pool.shutdown();
+               }
+       }
+
+       /**
+        * Computes the maximum row-wise or column-wise unique count over 
balanced ranges.
+        */
+       private static int getMaxUniqueValues(ExecutorService pool, MatrixBlock 
blkIn, Types.Direction dir,
+               ArrayList<int[]> ranges) throws Exception {
+               ArrayList<UniqueCountTask> tasks = new ArrayList<>();
+               for( int[] range : ranges )
+                       tasks.add(new UniqueCountTask(blkIn, dir, range[0], 
range[1]));
+
+               int ret = 0;
+               List<Future<Integer>> rtasks = pool.invokeAll(tasks);
+               for( int i = 0; i < rtasks.size(); i++ ) {
+                       ret = Math.max(ret, rtasks.get(i).get());
+                       rtasks.set(i, null);
+               }
+               return ret;
+       }
+
+       /**
+        * Fills row-wise or column-wise unique values over balanced ranges.
+        */
+       private static void fillUniqueValues(ExecutorService pool, MatrixBlock 
blkIn, MatrixBlock blkOut,
+               Types.Direction dir, ArrayList<int[]> ranges) throws Exception {
+               ArrayList<UniqueOutputTask> tasks = new ArrayList<>();
+               for( int[] range : ranges )
+                       tasks.add(new UniqueOutputTask(blkIn, blkOut, dir, 
range[0], range[1]));
+               List<Future<Void>> rtasks = pool.invokeAll(tasks);
+               for( int i = 0; i < rtasks.size(); i++ ) {
+                       rtasks.get(i).get();
+                       rtasks.set(i, null);
+               }
+       }
+
+       /**
+        * Batched variant of getMaxUniqueValues.
+        */
+       private static int getMaxUniqueValuesBatched(ExecutorService pool, 
MatrixBlock blkIn, Types.Direction dir,
+               BatchConfig config) throws Exception {
+               int ret = 0;
+               for( int pos = 0; pos < config._len; ) {
+                       ArrayList<UniqueCountTask> tasks = new ArrayList<>();
+                       for( int i = 0; i < config._numThreads && pos < 
config._len; i++ ) {
+                               int end = Math.min(pos + config._taskLen, 
config._len);
+                               tasks.add(new UniqueCountTask(blkIn, dir, pos, 
end));
+                               pos = end;
+                       }
+
+                       List<Future<Integer>> rtasks = pool.invokeAll(tasks);
+                       for( int i = 0; i < rtasks.size(); i++ ) {
+                               ret = Math.max(ret, rtasks.get(i).get());
+                               rtasks.set(i, null);
+                       }
+               }
+               return ret;
+       }
+
+       /**
+        * Batched variant of fillUniqueValues.
+        */
+       private static void fillUniqueValuesBatched(ExecutorService pool, 
MatrixBlock blkIn, MatrixBlock blkOut,
+               Types.Direction dir, BatchConfig config) throws Exception {
+               for( int pos = 0; pos < config._len; ) {
+                       ArrayList<UniqueOutputTask> tasks = new ArrayList<>();
+                       for( int i = 0; i < config._numThreads && pos < 
config._len; i++ ) {
+                               int end = Math.min(pos + config._taskLen, 
config._len);
+                               tasks.add(new UniqueOutputTask(blkIn, blkOut, 
dir, pos, end));
+                               pos = end;
+                       }
+
+                       List<Future<Void>> rtasks = pool.invokeAll(tasks);
+                       for( int i = 0; i < rtasks.size(); i++ ) {
+                               rtasks.get(i).get();
+                               rtasks.set(i, null);
+                       }
+               }
+       }
+
+       /**
+        * Decides whether the input is large enough to justify local 
deduplication tasks.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return true if the parallel path should be used
+        */
+       private static boolean satisfiesMultiThreadingConstraints(MatrixBlock 
blkIn, Types.Direction dir, int k) {
+               if( k <= 1 || ((long) blkIn.getNumRows()) * 
blkIn.getNumColumns() < PAR_UNIQUE_NUMCELL_THRESHOLD )
+                       return false;
+
+               switch(dir) {
+                       case RowCol:
+                               return blkIn.getNumRows() > 1;
+                       case Row:
+                               return blkIn.getNumRows() > 1;
+                       case Col:
+                               return blkIn.getNumColumns() > 1;
                        default:
                                throw new 
IllegalArgumentException("Unrecognized direction: " + dir);
                }
+       }
+
+       /**
+        * Creates balanced half-open ranges [start, end) using the same 
utility pattern as
+        * other SystemDS matrix libraries.
+        *
+        * @param len number of rows or columns to partition
+        * @param k requested degree of parallelism
+        * @return list of balanced index ranges
+        */
+       private static ArrayList<int[]> getBalancedRanges(int len, int k) {
+               ArrayList<int[]> ranges = new ArrayList<>();
+               ArrayList<Integer> blklens = 
UtilFunctions.getBalancedBlockSizesDefault(len, getNumThreads(k, len), false);
+               for( int i = 0, lb = 0; i < blklens.size(); lb += 
blklens.get(i), i++ )
+                       ranges.add(new int[] {lb, lb + blklens.get(i)});
+               return ranges;
+       }
+
+       /**
+        * Caps the number of workers by the number of row or column partitions 
available.
+        *
+        * @param k requested degree of parallelism
+        * @param len number of rows or columns to partition
+        * @return effective number of worker threads
+        */
+       private static int getNumThreads(int k, int len) {
+               return Math.max(1, Math.min(k, len));
+       }
+
+       /**
+        * Builds the execution plan for the batched parallel path.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @param k requested degree of parallelism
+        * @return batch configuration, or null if batching would not use 
parallelism
+        */
+       private static BatchConfig getBatchConfig(MatrixBlock blkIn, 
Types.Direction dir, int k) {
+               int len = getPartitionLength(blkIn, dir);
+               long maxBatchIndexes = getMaxLocalDedupIndexes(blkIn, dir);
+               if( maxBatchIndexes < 2 )
+                       return null;
+
+               int numThreads = getNumThreads(k, (int) Math.min(len, 
maxBatchIndexes));
+               if( numThreads <= 1 )
+                       return null;
+
+               int taskLen = Math.max(1, (int) Math.min(Integer.MAX_VALUE, 
maxBatchIndexes / numThreads));
+               return new BatchConfig(numThreads, taskLen, len);
+       }
+
+       /**
+        * Returns the number of rows or columns that define the partition 
direction.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return number of partition indexes
+        */
+       private static int getPartitionLength(MatrixBlock blkIn, 
Types.Direction dir) {
+               return dir == Types.Direction.Col ? blkIn.getNumColumns() : 
blkIn.getNumRows();
+       }
+
+       /**
+        * Estimates how many partition indexes can be processed in one batch.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return maximum number of rows or columns per batch
+        */
+       private static long getMaxLocalDedupIndexes(MatrixBlock blkIn, 
Types.Direction dir) {
+               long cellsPerIndex = dir == Types.Direction.Col ? 
blkIn.getNumRows() : blkIn.getNumColumns();
+               if( cellsPerIndex <= 0 ||
+                       cellsPerIndex > Long.MAX_VALUE / Double.BYTES / 
PAR_UNIQUE_LOCAL_BYTES_OVERHEAD )
+                       return 0;
+
+               long bytesPerIndex = cellsPerIndex * Double.BYTES * 
PAR_UNIQUE_LOCAL_BYTES_OVERHEAD;
+               long maxLocalBytes = Runtime.getRuntime().maxMemory() / 
PAR_UNIQUE_MAX_LOCAL_BYTES_FRACTION;
+               return maxLocalBytes / bytesPerIndex;
+       }
+
+       /**
+        * Conservative memory guard for full local deduplication. The estimate 
includes
+        * a small overhead factor for set objects.
+        *
+        * @param blkIn input matrix block
+        * @param dir unique direction
+        * @return true if local deduplication is small enough for the parallel 
path
+        */
+       private static boolean isLocalDedupMemoryBudgetSafe(MatrixBlock blkIn, 
Types.Direction dir) {
+               return getMaxLocalDedupIndexes(blkIn, dir) >= 
getPartitionLength(blkIn, dir);

Review Comment:
   If I understand correctly, you check here whether you can hold in memory 
each dedup set for every row at the same time. This is correct for `RowCol`, 
but due to `hashSet.clear()` it is unnecessarily conservative for `Row` and 
`Col`, because the peak memory is then `numThreads × clen × 64 bytes` which 
doesn't depend on `rlen`.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to