Integrate ScannerBuilder into index lookups

warehouse/core/src/main/java/datawave/query/scan/ExecutionHintHelper.java

@@ -0,0 +1,136 @@
+package datawave.query.scan;
+
+import java.util.Map;
+
+import org.apache.accumulo.core.client.ScannerBase.ConsistencyLevel;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A helper class that bridges the execution hints and consistency levels configured in the ShardQueryLogic with a specific context.
+ * <p>
+ * The shard query logic is configured with multiple execution hint maps that are keyed to a context. The context is often a table name, however the context
+ * could also be based on a stage of query planning, or even an entire class of query logic.
+ * <p>
+ * For example, scan resources may be reserved for a specific query logic.
+ * <p>
+ * In addition to execution hints there is also support for selecting a consistency level from a configuration map.
+ */
+public class ExecutionHintHelper {
+
+    private static final Logger log = LoggerFactory.getLogger(ExecutionHintHelper.class);
+
+    public static final String SCAN_TYPE = "scan_type";
+    public static final String PRIORITY = "priority";
+
+    private ExecutionHintHelper() {
+        // enforce static access
+    }
+
+    /**
+     * Get the execution hint map based on a key. If no execution hints are configured for the primary key, the secondary key is used.
+     * <p>
+     * The key is usually a table name, but in some cases may be tied to a particular query execution stage (e.g., index expansion).
+     *
+     * @param primary
+     *            the primary key
+     * @param secondary
+     *            the secondary key, used if the primary key is not found
+     * @param hintMap
+     *            the map of execution hints, organized by key
+     * @return the map of execution hints, or null no hint matches the primary or secondary key
+     */
+    public static Map<String,String> getExecutionHints(String primary, String secondary, Map<String,Map<String,String>> hintMap) {
+        Map<String,String> hints = getExecutionHints(primary, hintMap);
+
+        if (hints == null) {
+            hints = getExecutionHints(secondary, hintMap);
+        }
+
+        return hints;
+    }
+
+    /**
+     * Get the execution hint map for the provided key. If no execution hints are configured for the specified key a warning is logged.
+     *
+     * @param key
+     *            the key
+     * @param hintMap
+     *            a map of execution hints
+     * @return a map of execution hints, or null if no hint matches the provided key
+     */
+    public static Map<String,String> getExecutionHints(String key, Map<String,Map<String,String>> hintMap) {
+        if (hintMap == null) {
+            log.warn("Execution hints is null");
+            return null;
+        }
+        Map<String,String> hints = hintMap.get(key);
+        if (hints == null) {
+            log.warn("No execution hints found for key {}", key);
+        }
+        return hints;
+    }
+
+    /**
+     * Get the <code>scan_type</code> hint from the map of execution hints. A warning is logged if no scan type is found.
+     *
+     * @param executionHints
+     *            the map of execution hints
+     * @return the scan type, or null if no scan type is configured
+     */
+    public static String getScanType(Map<String,String> executionHints) {
+        if (executionHints == null) {
+            log.warn("Execution hints is null");
+            return null;
+        }
+
+        String scanType = executionHints.get(SCAN_TYPE);
+        if (scanType == null) {
+            log.warn("No scan type found");
+        }
+        return scanType;
+    }
+
+    /**
+     * Get the scan priority from the map of execution hints. A warning is logged if no priority is found
+     *
+     * @param executionHints
+     *            the map of execution hints
+     * @return the scan priority, or {@link Integer#MAX_VALUE} if no priority is configured
+     */
+    public static int getPriority(Map<String,String> executionHints) {
+        if (executionHints == null) {
+            log.warn("Execution hints is null");
+            return Integer.MAX_VALUE;
+        }
+
+        String value = executionHints.get(PRIORITY);
+        if (value == null) {
+            log.warn("No priority found");
+            return Integer.MAX_VALUE;
+        }
+        return Integer.parseInt(value);
+    }
+
+    /**
+     * Get the consistency level for the provided key
+     *
+     * @param key
+     *            the key
+     * @param levels
+     *            a map of consistency levels
+     * @return the consistency level, or null if no configured consistency level exists
+     */
+    public static ConsistencyLevel getConsistencyLevel(String key, Map<String,ConsistencyLevel> levels) {
+        if (levels == null) {
+            log.warn("ConsistencyLevels is null");
+            return null;
+        }
+
+        ConsistencyLevel consistencyLevel = levels.get(key);
+        if (consistencyLevel == null) {
+            log.warn("No consistency level found for key {}", key);
+        }
+        return consistencyLevel;
+    }
+}

warehouse/core/src/main/java/datawave/query/scan/ScanBuilder.java

@@ -1,5 +1,7 @@
 package datawave.query.scan;
 
+import static org.apache.accumulo.core.client.ScannerBase.ConsistencyLevel;
+
 import java.util.HashMap;
 import java.util.Map;
 
@@ -8,6 +10,8 @@
 import org.apache.accumulo.core.client.Scanner;
 import org.apache.accumulo.core.client.ScannerBase;
 import org.apache.accumulo.core.security.Authorizations;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import com.google.common.base.Preconditions;
 
@@ -27,6 +31,8 @@
  */
 public abstract class ScanBuilder<B> {
 
+    private static final Logger log = LoggerFactory.getLogger(ScanBuilder.class);
+
     // required variables
     protected final AccumuloClient client;
     protected String tableName;
@@ -96,6 +102,10 @@ public B setAuthorizations(Authorizations authorizations) {
      * @return this builder
      */
     public B setConsistencyLevel(ScannerBase.ConsistencyLevel consistencyLevel) {
+        if (consistencyLevel == null) {
+            log.warn("Null consistencyLevel");
+            throw new IllegalArgumentException("consistencyLevel cannot be null");
+        }
         this.consistencyLevel = consistencyLevel;
         return self();
     }
@@ -114,7 +124,11 @@ public B setConsistencyLevel(ScannerBase.ConsistencyLevel consistencyLevel) {
      * @return this builder
      */
     public B setConsistencyLevel(String consistencyLevel) {
-        this.consistencyLevel = ScannerBase.ConsistencyLevel.valueOf(consistencyLevel);
+        if (consistencyLevel == null) {
+            log.warn("Null consistencyLevel");
+            throw new IllegalArgumentException("consistencyLevel cannot be null");
+        }
+        this.consistencyLevel = ConsistencyLevel.valueOf(consistencyLevel);
         return self();
     }
 
@@ -128,6 +142,10 @@ public B setConsistencyLevel(String consistencyLevel) {
      * @return this builder
      */
     public B setScanType(String scanType) {
+        if (scanType == null) {
+            log.warn("Null scanType");
+            throw new IllegalArgumentException("scanType cannot be null");
+        }
         executionHints.put(SCAN_TYPE_KEY, scanType);
         return self();
     }
@@ -142,6 +160,10 @@ public B setScanType(String scanType) {
      * @return this builder
      */
     public B setScanPriority(int scanPriority) {
+        if (scanPriority < 0) {
+            log.warn("Negative scanPriority");
+            throw new IllegalArgumentException("scanPriority cannot be negative");
+        }
         executionHints.put(PRIORITY_KEY, Integer.toString(scanPriority));
         return self();
     }

warehouse/core/src/test/java/datawave/query/scan/ExecutionHintHelperTest.java

@@ -0,0 +1,190 @@
+package datawave.query.scan;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import org.apache.accumulo.core.client.ScannerBase.ConsistencyLevel;
+import org.junit.jupiter.api.Test;
+
+public class ExecutionHintHelperTest {
+
+    private final Map<String,String> tableBasedHints = createTableBasedHints();
+    private final Map<String,String> stageBasedHints = createStageBasedHints();
+    private final Map<String,String> defaultHints = createDefaultHints();
+    private final Map<String,Map<String,String>> executionHints = createExecutionHints();
+
+    private final Map<String,ConsistencyLevel> consistencyLevels = createConsistencyLevels();
+
+    private Map<String,String> createTableBasedHints() {
+        Map<String,String> hints = new HashMap<>();
+        hints.put(ExecutionHintHelper.SCAN_TYPE, "executor-pool-a");
+        hints.put(ExecutionHintHelper.PRIORITY, "1");
+        return hints;
+    }
+
+    private Map<String,String> createStageBasedHints() {
+        Map<String,String> hints = new HashMap<>();
+        hints.put(ExecutionHintHelper.SCAN_TYPE, "executor-pool-b");
+        hints.put(ExecutionHintHelper.PRIORITY, "2");
+        return hints;
+    }
+
+    private Map<String,String> createDefaultHints() {
+        Map<String,String> hints = new HashMap<>();
+        hints.put(ExecutionHintHelper.SCAN_TYPE, "executor-pool-c");
+        hints.put(ExecutionHintHelper.PRIORITY, "3");
+        return hints;
+    }
+
+    private Map<String,Map<String,String>> createExecutionHints() {
+        Map<String,Map<String,String>> hints = new HashMap<>();
+        hints.put("tableBased", tableBasedHints);
+        hints.put("stageBased", stageBasedHints);
+        hints.put("default", defaultHints);
+        return hints;
+    }
+
+    private Map<String,ConsistencyLevel> createConsistencyLevels() {
+        Map<String,ConsistencyLevel> levels = new HashMap<>();
+        levels.put("tableBased", ConsistencyLevel.EVENTUAL);
+        levels.put("stageBased", ConsistencyLevel.IMMEDIATE);
+        levels.put("default", ConsistencyLevel.IMMEDIATE);
+        return levels;
+    }
+
+    @Test
+    public void testGetExecutionHintsSingleKeyNullMap() {
+        Map<String,String> hints = ExecutionHintHelper.getExecutionHints("tableBased", null);
+        assertNull(hints);
+    }
+
+    @Test
+    public void testGetExecutionHintsSingleKeyNullKey() {
+        Map<String,String> hints = ExecutionHintHelper.getExecutionHints(null, executionHints);
+        assertNull(hints);
+    }
+
+    @Test
+    public void testGetExecutionHintsSingleKeyNoMatch() {
+        Map<String,String> hints = ExecutionHintHelper.getExecutionHints("dateBased", executionHints);
+        assertNull(hints);
+    }
+
+    @Test
+    public void testGetExecutionHintsSingleKeyGoodMatch() {
+        Map<String,String> hints = ExecutionHintHelper.getExecutionHints("tableBased", executionHints);
+        assertNotNull(hints);
+        assertScanType("executor-pool-a", hints);
+        assertScanPriority("1", hints);
+    }
+
+    @Test
+    public void testGetExecutionHintsPrimaryKeySecondaryKeyNullMap() {
+        Map<String,String> hints = ExecutionHintHelper.getExecutionHints("tableBased", "default", null);
+        assertNull(hints);
+    }
+
+    @Test
+    public void testGetExecutionHintsPrimaryNullKeySecondaryNullKey() {
+        Map<String,String> hints = ExecutionHintHelper.getExecutionHints(null, null, executionHints);
+        assertNull(hints);
+    }
+
+    @Test
+    public void testGetExecutionHintsPrimaryNullKeySecondaryKeyMatches() {
+        Map<String,String> hints = ExecutionHintHelper.getExecutionHints(null, "default", executionHints);
+        assertEquals(defaultHints, hints);
+    }
+
+    @Test
+    public void testGetExecutionHintsPrimaryKeyNotFoundSecondaryKeyMatches() {
+        Map<String,String> hints = ExecutionHintHelper.getExecutionHints("timeBased", "stageBased", executionHints);
+        assertEquals(stageBasedHints, hints);
+    }
+
+    @Test
+    public void testGetExecutionHintsPrimaryKeyMatchesSecondaryKeyNotNeeded() {
+        Map<String,String> hints = ExecutionHintHelper.getExecutionHints("stageBased", "default", executionHints);
+        assertEquals(stageBasedHints, hints);
+    }
+
+    @Test
+    public void testGetScanTypeNullMap() {
+        String scanType = ExecutionHintHelper.getScanType(null);
+        assertNull(scanType);
+    }
+
+    @Test
+    public void testGetScanTypeNoMatch() {
+        Map<String,String> hints = Map.of("priority", "7");
+        String scanType = ExecutionHintHelper.getScanType(hints);
+        assertNull(scanType);
+    }
+
+    @Test
+    public void testGetScanTypeMatches() {
+        Map<String,String> hints = Map.of("scan_type", "executor-pool-a");
+        String scanType = ExecutionHintHelper.getScanType(hints);
+        assertEquals("executor-pool-a", scanType);
+    }
+
+    @Test
+    public void testGetPriorityNullMap() {
+        int priority = ExecutionHintHelper.getPriority(null);
+        assertEquals(Integer.MAX_VALUE, priority);
+    }
+
+    @Test
+    public void testGetPriorityNoMatch() {
+        Map<String,String> hints = Map.of("scan_type", "executor-pool-c");
+        int priority = ExecutionHintHelper.getPriority(hints);
+        assertEquals(Integer.MAX_VALUE, priority);
+    }
+
+    @Test
+    public void testGetPriorityMatches() {
+        Map<String,String> hints = Map.of("priority", "2");
+        int priority = ExecutionHintHelper.getPriority(hints);
+        assertEquals(2, priority);
+    }
+
+    @Test
+    public void testGetConsistencyLevelNullMap() {
+        ConsistencyLevel level = ExecutionHintHelper.getConsistencyLevel("tableBased", null);
+        assertNull(level);
+    }
+
+    @Test
+    public void testGetConsistencyLevelNullKey() {
+        ConsistencyLevel level = ExecutionHintHelper.getConsistencyLevel(null, consistencyLevels);
+        assertNull(level);
+    }
+
+    @Test
+    public void testGetConsistencyLevelKeyNotFound() {
+        ConsistencyLevel level = ExecutionHintHelper.getConsistencyLevel("timeBased", consistencyLevels);
+        assertNull(level);
+    }
+
+    @Test
+    public void testGetConsistencyLevel() {
+        assertEquals(ConsistencyLevel.EVENTUAL, ExecutionHintHelper.getConsistencyLevel("tableBased", consistencyLevels));
+        assertEquals(ConsistencyLevel.IMMEDIATE, ExecutionHintHelper.getConsistencyLevel("stageBased", consistencyLevels));
+        assertEquals(ConsistencyLevel.IMMEDIATE, ExecutionHintHelper.getConsistencyLevel("default", consistencyLevels));
+    }
+
+    private void assertScanType(String expected, Map<String,String> hints) {
+        assertTrue(hints.containsKey(ExecutionHintHelper.SCAN_TYPE));
+        assertEquals(expected, hints.get(ExecutionHintHelper.SCAN_TYPE));
+    }
+
+    private void assertScanPriority(String expected, Map<String,String> hints) {
+        assertTrue(hints.containsKey(ExecutionHintHelper.PRIORITY));
+        assertEquals(expected, hints.get(ExecutionHintHelper.PRIORITY));
+    }
+}