partitionByRandom seed and new function mapWithIndex

Author: cesarpomar

ignis/driver/api/IDataFrame.py

@@ -89,10 +89,10 @@ def repartition(self, numPartitions, preserveOrdering, global_):
         except ignis.rpc.driver.exception.ttypes.IDriverException as ex:
             raise IDriverException(ex.message, ex.cause_)
 
-    def partitionByRandom(self, numPartitions):
+    def partitionByRandom(self, numPartitions, seed):
         try:
             with Ignis._clientPool().getClient() as client:
-                self._id = client.getDataFrameService().partitionByRandom(self._id, numPartitions)
+                self._id = client.getDataFrameService().partitionByRandom(self._id, numPartitions, seed)
         except ignis.rpc.driver.exception.ttypes.IDriverException as ex:
             raise IDriverException(ex.message, ex.cause_)
 
@@ -138,6 +138,13 @@ def keyBy(self, src):
         except ignis.rpc.driver.exception.ttypes.IDriverException as ex:
             raise IDriverException(ex.message, ex.cause_)
 
+    def mapWithIndex(self, src):
+        try:
+            with Ignis._clientPool().getClient() as client:
+                return IDataFrame(client.getDataFrameService().mapWithIndex(self._id, ISource.wrap(src).rpc()))
+        except ignis.rpc.driver.exception.ttypes.IDriverException as ex:
+            raise IDriverException(ex.message, ex.cause_)
+
     def mapPartitions(self, src):
         try:
             with Ignis._clientPool().getClient() as client:

ignis/executor/core/modules/IGeneralModule.py

@@ -49,6 +49,12 @@ def keyBy(self, src):
         except Exception as ex:
             self._pack_exception(ex)
 
+    def mapWithIndex(self, src):
+        try:
+            self.__pipe_impl.mapWithIndex(self._executor_data.loadLibrary(src))
+        except Exception as ex:
+            self._pack_exception(ex)
+
     def mapPartitions(self, src):
         try:
             self.__pipe_impl.mapPartitions(self._executor_data.loadLibrary(src))
@@ -149,9 +155,9 @@ def repartition(self, numPartitions, preserveOrdering, global_):
         except Exception as ex:
             self._pack_exception(ex)
 
-    def partitionByRandom(self, numPartitions):
+    def partitionByRandom(self, numPartitions, seed):
         try:
-            self.__repartition_impl.partitionByRandom(numPartitions)
+            self.__repartition_impl.partitionByRandom(numPartitions, seed)
         except Exception as ex:
             self._pack_exception(ex)
 

ignis/executor/core/modules/impl/IPipeImpl.py

@@ -99,6 +99,26 @@ def keyBy(self, f):
         f.after(context)
         self._executor_data.setPartitions(output)
 
+    def mapWithIndex(self, f):
+        context = self._executor_data.getContext()
+        input = self._executor_data.getAndDeletePartitions()
+        f.before(context)
+        output = self._executor_data.getPartitionTools().newPartitionGroup(input)
+        logger.info("General: mapWithIndex " + str(len(input)) + " partitions")
+
+        elems = sum([len(p) for p in input])
+        indices = self._executor_data.mpi().native().allgather(elems)
+        id = sum(indices[0:self._executor_data.mpi().rank()])
+
+        for i in range(len(input)):
+            it = output[i].writeIterator()
+            for elem in input[i]:
+                it.write(f.call(id, elem, context))
+                id += 1
+            input[i] = None
+        f.after(context)
+        self._executor_data.setPartitions(output)
+
     def mapPartitions(self, f):
         context = self._executor_data.getContext()
         input = self._executor_data.getAndDeletePartitions()
@@ -119,9 +139,13 @@ def mapPartitionsWithIndex(self, f):
         f.before(context)
         output = self._executor_data.getPartitionTools().newPartitionGroup(input)
         logger.info("General: mapPartitionsWithIndex " + str(len(input)) + " partitions")
+
+        indices = self._executor_data.mpi().native().allgather(len(input))
+        offset = sum(indices[0:self._executor_data.mpi().rank()])
+
         for i in range(len(input)):
             it = output[i].writeIterator()
-            for elem in f.call(i, input[i].readIterator(), context):
+            for elem in f.call(offset + i, input[i].readIterator(), context):
                 it.write(elem)
             input[i] = None
         f.after(context)

ignis/executor/core/modules/impl/IRepartitionImpl.py

@@ -226,8 +226,9 @@ def __local_repartition(self, numPartitions):
 
 		self._executor_data.setPartitions(output)
 
-	def partitionByRandom(self, numPartitions):
-		self.__partitionBy_impl(lambda elem, ctx: random.randint(0, numPartitions), numPartitions)
+	def partitionByRandom(self, numPartitions, seed):
+		r = random.Random(seed)
+		self.__partitionBy_impl(lambda elem, ctx: r.randint(0, numPartitions), numPartitions)
 
 	def partitionByHash(self, numPartitions):
 		self.__partitionBy_impl(lambda elem, ctx: hash(elem), numPartitions)

ignis/rpc/driver/dataframe/IDataFrameService-remote

@@ -34,13 +34,14 @@ if len(sys.argv) <= 1 or sys.argv[1] == '--help':
     print('  void saveAsTextFile(IDataFrameId id, string path)')
     print('  void saveAsJsonFile(IDataFrameId id, string path, bool pretty)')
     print('  IDataFrameId repartition(IDataFrameId id, i64 numPartitions, bool preserveOrdering, bool global_)')
-    print('  IDataFrameId partitionByRandom(IDataFrameId id, i64 numPartitions)')
+    print('  IDataFrameId partitionByRandom(IDataFrameId id, i64 numPartitions, i32 seed)')
     print('  IDataFrameId partitionByHash(IDataFrameId id, i64 numPartitions)')
     print('  IDataFrameId partitionBy(IDataFrameId id, ISource src, i64 numPartitions)')
     print('  IDataFrameId map_(IDataFrameId id, ISource src)')
     print('  IDataFrameId filter(IDataFrameId id, ISource src)')
     print('  IDataFrameId flatmap(IDataFrameId id, ISource src)')
     print('  IDataFrameId keyBy(IDataFrameId id, ISource src)')
+    print('  IDataFrameId mapWithIndex(IDataFrameId id, ISource src)')
     print('  IDataFrameId mapPartitions(IDataFrameId id, ISource src)')
     print('  IDataFrameId mapPartitionsWithIndex(IDataFrameId id, ISource src)')
     print('  IDataFrameId mapExecutor(IDataFrameId id, ISource src)')
@@ -246,10 +247,10 @@ elif cmd == 'repartition':
     pp.pprint(client.repartition(eval(args[0]), eval(args[1]), eval(args[2]), eval(args[3]),))
 
 elif cmd == 'partitionByRandom':
-    if len(args) != 2:
-        print('partitionByRandom requires 2 args')
+    if len(args) != 3:
+        print('partitionByRandom requires 3 args')
         sys.exit(1)
-    pp.pprint(client.partitionByRandom(eval(args[0]), eval(args[1]),))
+    pp.pprint(client.partitionByRandom(eval(args[0]), eval(args[1]), eval(args[2]),))
 
 elif cmd == 'partitionByHash':
     if len(args) != 2:
@@ -287,6 +288,12 @@ elif cmd == 'keyBy':
         sys.exit(1)
     pp.pprint(client.keyBy(eval(args[0]), eval(args[1]),))
 
+elif cmd == 'mapWithIndex':
+    if len(args) != 2:
+        print('mapWithIndex requires 2 args')
+        sys.exit(1)
+    pp.pprint(client.mapWithIndex(eval(args[0]), eval(args[1]),))
+
 elif cmd == 'mapPartitions':
     if len(args) != 2:
         print('mapPartitions requires 2 args')