Bringing feature/splits-interface branch current - adding in unit test fixes

warehouse/ingest-core/src/main/java/datawave/ingest/mapreduce/job/IngestJob.java

@@ -362,8 +362,6 @@ public int run(String[] args) throws Exception {
         log.info("InputFormat: " + job.getInputFormatClass().getName());
         log.info("Mapper: " + job.getMapperClass().getName());
         log.info("Reduce tasks: " + (useMapOnly ? 0 : reduceTasks));
-        log.info("Split File: " + conf.get(TableSplitsCache.SPLITS_CACHE_DIR) + "/"
-                        + conf.get(TableSplitsCache.SPLITS_CACHE_FILE, TableSplitsCache.DEFAULT_SPLITS_CACHE_FILE));
 
         // Note that if we run any other jobs in the same vm (such as a sampler), then we may
         // need to catch and throw away an exception here
@@ -665,7 +663,7 @@ protected Configuration parseArguments(String[] args, Configuration conf) throws
             } else if (args[i].equals("-disableRefreshSplits")) {
                 conf.setBoolean(TableSplitsCache.REFRESH_SPLITS, false);
             } else if (args[i].equals("-splitsCacheDir")) {
-                conf.set(TableSplitsCache.SPLITS_CACHE_DIR, args[++i]);
+                conf.set(SplitsConstants.SPLITS_CACHE_DIR, args[++i]);
             } else if (args[i].equals("-multipleNumShardsCacheDir")) {
                 conf.set(NumShards.MULTIPLE_NUMSHARDS_CACHE_PATH, args[++i]);
             } else if (args[i].equals("-enableAccumuloConfigCache")) {
@@ -849,7 +847,7 @@ protected void configureBulkPartitionerAndOutputFormatter(Job job, AccumuloHelpe
         conf.setInt("splits.num.reduce", this.reduceTasks);
         // used by the output formatter and the sharded partitioner
         long before = System.currentTimeMillis();
-        SplitsFile.setupFile(job, conf);
+        SplitsCache.getInstance(conf).setupJob(job);
         long after = System.currentTimeMillis();
 
         log.info("Sharded splits files setup time: " + (after - before) + "ms");

warehouse/ingest-core/src/main/java/datawave/ingest/mapreduce/job/MultiRFileOutputFormatter.java

@@ -112,6 +112,7 @@ public class MultiRFileOutputFormatter extends FileOutputFormat<BulkIngestKey,Va
     protected Configuration conf;
     protected Map<String,ConfigurationCopy> tableConfigs;
     protected Set<String> tableIds = null;
+    protected SplitsCache splitsCache;
     protected long maxRFileSize = 0;
     protected int maxRFileEntries = 0;
     protected boolean generateMapFileRowKeys = false;
@@ -274,9 +275,10 @@ protected SizeTrackingWriter openWriter(String filename, AccumuloConfiguration t
 
         var builder = org.apache.accumulo.core.client.rfile.RFile.newWriter().to(filename).withFileSystem(fs).withTableProperties(tableConf);
         if(this.loadPlanningEnabled) {
-            var splits = SplitsFile.getSplits(conf, table);
-            if(splits != null && !splits.isEmpty()) {
-                LoadPlan.SplitResolver splitResolver = row->findContainingSplits(row, splits);
+            SplitsCache splits = SplitsCache.getInstance(conf);
+            List<Text> splitsFile = splits.getSplits(conf, table);
+            if(splitsFile != null && !splitsFile.isEmpty()) {
+                LoadPlan.SplitResolver splitResolver = row->SplitsFile.findContainingSplits(row, splitsFile);
                 builder = builder.withSplitResolver(splitResolver);
             }
         }
@@ -363,27 +365,6 @@ private void writeLoadPlans(TaskAttemptContext context) {
         log.debug("Finished writing bulk load plans to disk");
     }
 
-    /**
-     * Finds two contiguous table splits that contain the specified row should reside
-     *
-     * @param lookupRow
-     *            Row value to be mapped
-     * @param tableSplits
-     *            Splits for the table in question
-     * @return KeyExtent mapping for the given row
-     */
-    static LoadPlan.TableSplits findContainingSplits(Text lookupRow, List<Text> tableSplits) {
-        int position = Collections.binarySearch(tableSplits, lookupRow);
-        if (position < 0) {
-            position = -1 * (position + 1);
-        }
-
-        Text prevRow = position == 0 ? null : tableSplits.get(position - 1);
-        Text endRow = position == tableSplits.size() ? null : tableSplits.get(position);
-
-        return new LoadPlan.TableSplits(prevRow, endRow);
-    }
-
     public static class SizeTrackingWriter {
         private RFileWriter delegate;
         long size = 0;
@@ -550,6 +531,7 @@ public RecordWriter<BulkIngestKey,Value> getRecordWriter(final TaskAttemptContex
         FileOutputCommitter committer = (FileOutputCommitter) getOutputCommitter(context);
         workDir = committer.getWorkPath();
         conf = context.getConfiguration();
+        splitsCache = SplitsCache.getInstance(conf);
 
         setTableIdsAndConfigs();
 
@@ -743,15 +725,10 @@ private SizeTrackingWriter getOrCreateWriter(TaskAttemptContext context, String
                     if (generateMapFilePerShardLocation) {
                         // Look up the shard location (tablet server serving shard ID rowKey)
                         // If we don't have a location, then just use the rowKey itself.
-                        Map<Text,String> shardLocs = getShardLocations(tableName);
-                        shardLocation = shardLocs.containsKey(rowKey) ? shardLocs.get(rowKey) : null;
-                        if (shardLocation == null) {
-                            // in this case we have a shard id that has no split. Lets put this in one "extra" file
-                            shardLocation = "extra";
-                        } else {
-                            // Ensure there's no colon
-                            shardLocation = shardLocation.replace(":", "_");
-                        }
+                        String fallback = conf.get("shard.fallback.name." + rowKey.toString(), "extra");
+                        shardLocation = splitsCache.getExactLocation(tableName, rowKey, () -> fallback);
+                        // in the case we have a shard id that has no split. Lets put this in one "extra" file
+                        shardLocation = shardLocation.replace(":", "_");
                     }
                     // Combine table name with shard location so that we end up
                     // with all of the shard map files under directories that can be
@@ -777,27 +754,4 @@ private SizeTrackingWriter getOrCreateWriter(TaskAttemptContext context, String
             }
         };
     }
-
-    /**
-     * Read in the sequence file (that was created at job startup) for the given table that contains a list of shard IDs and the corresponding tablet server to
-     * which that shard is assigned.
-     *
-     * @param tableName
-     *            the table name
-     * @return a mapping of the shard ids and tablet server
-     * @throws IOException
-     *             if there is an issue with read or write
-     */
-    protected Map<Text,String> getShardLocations(String tableName) throws IOException {
-        // Create the Map of sharded table name to [shardId -> server]
-        if (this.tableShardLocations == null) {
-            this.tableShardLocations = new HashMap<>();
-        }
-
-        if (null == this.tableShardLocations.get(tableName)) {
-            this.tableShardLocations.put(tableName, SplitsFile.getSplitsAndLocations(conf, tableName));
-        }
-
-        return tableShardLocations.get(tableName);
-    }
 }

warehouse/ingest-core/src/main/java/datawave/ingest/mapreduce/job/SplitsCache.java

@@ -0,0 +1,33 @@
+package datawave.ingest.mapreduce.job;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.function.Supplier;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.io.Text;
+import org.apache.hadoop.mapreduce.Job;
+
+public interface SplitsCache extends AutoCloseable {
+    static SplitsCache getInstance(final Configuration conf) {
+        return SplitsCacheFactory.getSplitsCache(conf);
+    }
+
+    void init(Configuration conf);
+
+    void setupJob(final Job job) throws IOException;
+
+    boolean hasSplits();
+
+    int getSplitsCount(String table);
+
+    int getExactIndex(String table, Text key);
+
+    int getExactPartition(String table, Text key);
+
+    int getNearestPartition(String table, Text key);
+
+    String getExactLocation(String table, Text key, Supplier<String> defaultFn);
+
+    List<Text> getSplits(Configuration conf, String tableName) throws IOException;
+}

warehouse/ingest-core/src/main/java/datawave/ingest/mapreduce/job/SplitsCacheFactory.java

@@ -0,0 +1,35 @@
+package datawave.ingest.mapreduce.job;
+
+import java.lang.reflect.InvocationTargetException;
+
+import org.apache.hadoop.conf.Configuration;
+
+public class SplitsCacheFactory {
+    public static final String SPLITS_CACHE_IMPL = "datawave.ingest.splits.cache.impl";
+
+    static volatile SplitsCache INSTANCE;
+
+    public static SplitsCache getSplitsCache(final Configuration conf) {
+        if (INSTANCE == null) {
+            synchronized (SplitsCacheFactory.class) {
+                if (INSTANCE == null) {
+                    try {
+                        String splitsCacheImpl = conf.get(SPLITS_CACHE_IMPL);
+                        // noinspection unchecked
+                        Class<? extends SplitsCache> clazz = splitsCacheImpl != null ? (Class<? extends SplitsCache>) Class.forName(splitsCacheImpl)
+                                        : SplitsFile.class;
+                        INSTANCE = clazz.getDeclaredConstructor().newInstance();
+                        INSTANCE.init(conf);
+                    } catch (ClassNotFoundException | NoSuchMethodException | InstantiationException | IllegalAccessException | InvocationTargetException ex) {
+                        throw new RuntimeException(ex);
+                    }
+                }
+            }
+        }
+        return INSTANCE;
+    }
+
+    public static void clearInstance() {
+        INSTANCE = null;
+    }
+}

warehouse/ingest-core/src/main/java/datawave/ingest/mapreduce/job/SplitsConstants.java

@@ -0,0 +1,7 @@
+package datawave.ingest.mapreduce.job;
+
+public class SplitsConstants {
+    public static final String SPLITS_CACHE_DIR = "datawave.ingest.splits.cache.dir";
+    public static final String SPLITS_CACHE_FILE = "datawave.ingest.splits.cache.file";
+    public static final String DEFAULT_SPLITS_CACHE_DIR = "/data/splitsCache";
+}