/** * Sets partitions as input for a PartitionedFileSet. If both a PartitionFilter and Partition(s) are specified, the * PartitionFilter takes precedence and the specified Partition(s) will be ignored. * * @param arguments the runtime arguments for a partitioned dataset * @param partitionIterator the iterator of partitions to add as input */ public static void addInputPartitions(Map<String, String> arguments, Iterator<? extends Partition> partitionIterator) { Collection<PartitionKey> inputPartitionKeys = new ArrayList<>(getInputPartitionKeys(arguments)); while (partitionIterator.hasNext()) { inputPartitionKeys.add(partitionIterator.next().getPartitionKey()); } arguments.put(INPUT_PARTITION_LIST, GSON.toJson(inputPartitionKeys)); }
/** * Computes and returns the input partition keys given by the partition filter - if present. Otherwise, get the list * of partition keys explicitly specified in the runtime arguments. * Stores the result in a cache and returns it. */ private Collection<PartitionKey> getInputKeys() { if (inputKeysCache != null) { return inputKeysCache.get(); } Collection<PartitionKey> inputKeys = computeInputKeys(); if (inputKeys == null) { inputKeys = PartitionedFileSetArguments.getInputPartitionKeys(runtimeArguments); } inputKeysCache = new AtomicReference<>(inputKeys); return inputKeys; }
/** * Computes and returns the input partition keys given by the partition filter - if present. Otherwise, get the list * of partition keys explicitly specified in the runtime arguments. * Stores the result in a cache and returns it. */ private Collection<PartitionKey> getInputKeys() { if (inputKeysCache != null) { return inputKeysCache.get(); } Collection<PartitionKey> inputKeys = computeInputKeys(); if (inputKeys == null) { inputKeys = PartitionedFileSetArguments.getInputPartitionKeys(runtimeArguments); } inputKeysCache = new AtomicReference<>(inputKeys); return inputKeys; }
@Test public void testGetInputPartitionKeys() throws Exception { Map<String, String> arguments = new HashMap<>(); Assert.assertEquals(0, PartitionedFileSetArguments.getInputPartitionKeys(arguments).size()); List<? extends Partition> partitions = Lists.newArrayList(new BasicPartition(null, "path/doesn't/matter/1", generateUniqueKey()), new BasicPartition(null, "path/doesn't/matter/2", generateUniqueKey()), new BasicPartition(null, "path/doesn't/matter/3", generateUniqueKey())); for (Partition partition : partitions) { PartitionedFileSetArguments.addInputPartition(arguments, partition); } List<PartitionKey> inputPartitionKeys = Lists.transform(partitions, new Function<Partition, PartitionKey>() { @Nullable @Override public PartitionKey apply(Partition input) { return input.getPartitionKey(); } }); Assert.assertEquals(inputPartitionKeys, PartitionedFileSetArguments.getInputPartitionKeys(arguments)); arguments.clear(); PartitionedFileSetArguments.addInputPartitions(arguments, partitions.iterator()); Assert.assertEquals(inputPartitionKeys, PartitionedFileSetArguments.getInputPartitionKeys(arguments)); }