#32 add EntityPartitionMapper

us-address/src/main/java/io/github/mincongh/batch/EntityPartitionMapper.java

@@ -0,0 +1,65 @@
+package io.github.mincongh.batch;
+
+import java.util.HashSet;
+import java.util.Properties;
+import java.util.Set;
+
+import javax.batch.api.BatchProperty;
+import javax.batch.api.partition.PartitionMapper;
+import javax.batch.api.partition.PartitionPlan;
+import javax.batch.api.partition.PartitionPlanImpl;
+import javax.inject.Inject;
+import javax.inject.Named;
+
+@Named
+public class EntityPartitionMapper implements PartitionMapper {
+
+    @Inject @BatchProperty private String rootEntitiesStr;
+
+    @Override
+    public PartitionPlan mapPartitions() throws Exception {
+
+        String[] rootEntities = parse(rootEntitiesStr);
+
+        return new PartitionPlanImpl() {
+
+            @Override
+            public int getPartitions() {
+                System.out.printf("#mapPartitions(): %d partitions.%n", rootEntities.length);
+                return rootEntities.length;
+            }
+
+            @Override
+            public int getThreads() {
+                System.out.printf("#getThreads(): %d threads.%n", getPartitions());
+                return getPartitions();
+            }
+
+            @Override
+            public Properties[] getPartitionProperties() {
+                Properties[] props = new Properties[getPartitions()];
+                for (int i = 0; i < props.length; i++) {
+                    props[i] = new Properties();
+                    props[i].setProperty("entityType", rootEntities[i]);
+                }
+                return props;
+            }
+        };
+    }
+
+    /**
+     * Parse a set of entities in string into a set of entity-types.
+     * 
+     * @param raw a set of entities concatenated in string, with "[]" on both
+     *          ends, separated by ",".
+     * @return a set of entity-types
+     * @throws NullPointerException thrown if the entity-token is not found.
+     */
+    private String[] parse(String raw) throws NullPointerException {
+
+        if (raw == null) {
+            throw new NullPointerException("Not any target entity to index");
+        }
+        return raw.substring(1, raw.length() - 1).split(", ");
+    }
+}

us-address/src/main/java/io/github/mincongh/session/BatchSession.java

@@ -1,6 +1,8 @@
 package io.github.mincongh.session;
 
+import java.util.HashSet;
 import java.util.Properties;
+import java.util.Set;
 
 import javax.batch.operations.JobOperator;
 import javax.batch.runtime.BatchRuntime;
@@ -80,6 +82,7 @@ public void massIndex() {
         jobParams.setProperty("purgeAtStart", String.valueOf(true));
         jobParams.setProperty("optimizeAfterPurge", String.valueOf(true));
         jobParams.setProperty("optimizeAtEnd", String.valueOf(true));
+        jobParams.setProperty("rootEntitiesStr", getRootEntities().toString());
         Long executionId = jobOperator.start("mass-index", jobParams);
 
         // calculate the performance
@@ -111,4 +114,11 @@ public void printAddressesTop1000() throws InterruptedException {
         String msg = i <= 10 ? "Finished" : "Failed";
         System.out.println(msg);
     }
+
+    private Set<String> getRootEntities() {
+        Set<String> rootEntities = new HashSet<>();
+        rootEntities.add("io.github.mincongh.entity.Address");
+        rootEntities.add("io.github.mincongh.entity.Stock");
+        return rootEntities;
+    }
 }

us-address/src/main/resources/META-INF/batch-jobs/mass-index.xml

@@ -31,6 +31,12 @@
            the end !!! In order to debug the situation, the next step is set to
            the optimization directly.
        -->
+      <mapper ref="entityPartitionMapper">
+        <properties>
+          <property name="rootEntitiesStr" value="#{jobParameters['rootEntitiesStr']}"/>
+        </properties>
+      </mapper>
+<!-- 
       <plan partitions="2">
         <properties partition="0">
           <property name="entityType" value="io.github.mincongh.entity.Address"/>
@@ -39,6 +45,7 @@
           <property name="entityType" value="io.github.mincongh.entity.Stock"/>
         </properties>
       </plan>
+ -->
     </partition>
   </step>
   <decision id="purgeDecision" ref="purgeDecider">