diff --git a/client/src/main/java/com/metamx/druid/shard/NumberedShardSpec.java b/client/src/main/java/com/metamx/druid/shard/NumberedShardSpec.java new file mode 100644 index 00000000000..fa65a67d7d4 --- /dev/null +++ b/client/src/main/java/com/metamx/druid/shard/NumberedShardSpec.java @@ -0,0 +1,63 @@ +package com.metamx.druid.shard; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.base.Preconditions; +import com.metamx.druid.input.InputRow; +import com.metamx.druid.partition.NumberedPartitionChunk; +import com.metamx.druid.partition.PartitionChunk; + +import java.util.Map; + +public class NumberedShardSpec implements ShardSpec +{ + @JsonIgnore + final private int partitionNum; + + @JsonIgnore + final private int partitions; + + @JsonCreator + public NumberedShardSpec( + @JsonProperty("partitionNum") int partitionNum, + @JsonProperty("partitions") int partitions + ) + { + Preconditions.checkArgument(partitionNum >= 0, "partitionNum >= 0"); + Preconditions.checkArgument(partitionNum < partitions, "partitionNum < partitions"); + this.partitionNum = partitionNum; + this.partitions = partitions; + } + + @JsonProperty("partitionNum") + @Override + public int getPartitionNum() + { + return partitionNum; + } + + @JsonProperty("partitions") + public int getPartitions() + { + return partitions; + } + + @Override + public PartitionChunk createChunk(T obj) + { + return NumberedPartitionChunk.make(partitionNum, partitions, obj); + } + + @Override + public boolean isInChunk(Map dimensions) + { + return true; + } + + @Override + public boolean isInChunk(InputRow inputRow) + { + return true; + } +} diff --git a/client/src/main/java/com/metamx/druid/shard/ShardSpec.java b/client/src/main/java/com/metamx/druid/shard/ShardSpec.java index e2191c90e3f..b5b778283b5 100644 --- a/client/src/main/java/com/metamx/druid/shard/ShardSpec.java +++ b/client/src/main/java/com/metamx/druid/shard/ShardSpec.java @@ -33,7 +33,8 @@ import java.util.Map; @JsonSubTypes({ @JsonSubTypes.Type(name="single", value=SingleDimensionShardSpec.class), @JsonSubTypes.Type(name="none", value=NoneShardSpec.class), - @JsonSubTypes.Type(name="linear", value=LinearShardSpec.class) + @JsonSubTypes.Type(name="linear", value=LinearShardSpec.class), + @JsonSubTypes.Type(name="numbered", value=NumberedShardSpec.class) }) public interface ShardSpec { diff --git a/client/src/test/java/com/metamx/druid/shard/NumberedShardSpecTest.java b/client/src/test/java/com/metamx/druid/shard/NumberedShardSpecTest.java new file mode 100644 index 00000000000..20150cd4783 --- /dev/null +++ b/client/src/test/java/com/metamx/druid/shard/NumberedShardSpecTest.java @@ -0,0 +1,70 @@ +package com.metamx.druid.shard; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.base.Function; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Lists; +import com.metamx.druid.jackson.DefaultObjectMapper; +import com.metamx.druid.partition.PartitionChunk; +import junit.framework.Assert; +import org.junit.Test; + +import java.util.List; + +public class NumberedShardSpecTest +{ + @Test + public void testSerde() throws Exception + { + final ObjectMapper jsonMapper = new DefaultObjectMapper(); + final ShardSpec spec = new NumberedShardSpec(1, 2); + final ShardSpec spec2 = jsonMapper.readValue(jsonMapper.writeValueAsBytes(spec), ShardSpec.class); + Assert.assertEquals(1, spec2.getPartitionNum()); + Assert.assertEquals(2, ((NumberedShardSpec) spec).getPartitions()); + } + + @Test + public void testPartitionChunks() + { + final List specs = ImmutableList.of( + new NumberedShardSpec(0, 3), + new NumberedShardSpec(1, 3), + new NumberedShardSpec(2, 3) + ); + + final List> chunks = Lists.transform( + specs, + new Function>() + { + @Override + public PartitionChunk apply(ShardSpec shardSpec) + { + return shardSpec.createChunk("rofl"); + } + } + ); + + Assert.assertEquals(0, chunks.get(0).getChunkNumber()); + Assert.assertEquals(1, chunks.get(1).getChunkNumber()); + Assert.assertEquals(2, chunks.get(2).getChunkNumber()); + + Assert.assertTrue(chunks.get(0).isStart()); + Assert.assertFalse(chunks.get(1).isStart()); + Assert.assertFalse(chunks.get(2).isStart()); + + Assert.assertFalse(chunks.get(0).isEnd()); + Assert.assertFalse(chunks.get(1).isEnd()); + Assert.assertTrue(chunks.get(2).isEnd()); + + Assert.assertTrue(chunks.get(0).abuts(chunks.get(1))); + Assert.assertTrue(chunks.get(1).abuts(chunks.get(2))); + + Assert.assertFalse(chunks.get(0).abuts(chunks.get(0))); + Assert.assertFalse(chunks.get(0).abuts(chunks.get(2))); + Assert.assertFalse(chunks.get(1).abuts(chunks.get(0))); + Assert.assertFalse(chunks.get(1).abuts(chunks.get(1))); + Assert.assertFalse(chunks.get(2).abuts(chunks.get(0))); + Assert.assertFalse(chunks.get(2).abuts(chunks.get(1))); + Assert.assertFalse(chunks.get(2).abuts(chunks.get(2))); + } +} diff --git a/common/src/main/java/com/metamx/druid/partition/IntegerPartitionChunk.java b/common/src/main/java/com/metamx/druid/partition/IntegerPartitionChunk.java index a49313711ec..1faa3f1dee3 100644 --- a/common/src/main/java/com/metamx/druid/partition/IntegerPartitionChunk.java +++ b/common/src/main/java/com/metamx/druid/partition/IntegerPartitionChunk.java @@ -103,8 +103,9 @@ public class IntegerPartitionChunk implements PartitionChunk } return retVal; + } else { + throw new IllegalArgumentException("Cannot compare against something that is not an IntegerPartitionChunk."); } - throw new IllegalArgumentException("Cannot compare against something that is not a StringPartitionChunk."); } @Override @@ -139,3 +140,4 @@ public class IntegerPartitionChunk implements PartitionChunk '}'; } } + diff --git a/common/src/main/java/com/metamx/druid/partition/NumberedPartitionChunk.java b/common/src/main/java/com/metamx/druid/partition/NumberedPartitionChunk.java new file mode 100644 index 00000000000..242b7a1a25c --- /dev/null +++ b/common/src/main/java/com/metamx/druid/partition/NumberedPartitionChunk.java @@ -0,0 +1,98 @@ +package com.metamx.druid.partition; + +import com.google.common.base.Objects; +import com.google.common.base.Preconditions; +import com.google.common.collect.ComparisonChain; + +public class NumberedPartitionChunk implements PartitionChunk +{ + private final int chunkNumber; + private final int chunks; + private final T object; + + public static NumberedPartitionChunk make( + int chunkNumber, + int chunks, + T obj + ) + { + return new NumberedPartitionChunk(chunkNumber, chunks, obj); + } + + public NumberedPartitionChunk( + int chunkNumber, + int chunks, + T object + ) + { + Preconditions.checkArgument(chunkNumber >= 0, "chunkNumber >= 0"); + Preconditions.checkArgument(chunkNumber < chunks, "chunkNumber < chunks"); + this.chunkNumber = chunkNumber; + this.chunks = chunks; + this.object = object; + } + + @Override + public T getObject() + { + return object; + } + + @Override + public boolean abuts(final PartitionChunk other) + { + return other instanceof NumberedPartitionChunk && other.getChunkNumber() == chunkNumber + 1; + } + + @Override + public boolean isStart() + { + return chunkNumber == 0; + } + + @Override + public boolean isEnd() + { + return chunkNumber == chunks - 1; + } + + @Override + public int getChunkNumber() + { + return chunkNumber; + } + + @Override + public int compareTo(PartitionChunk other) + { + if (other instanceof NumberedPartitionChunk) { + final NumberedPartitionChunk castedOther = (NumberedPartitionChunk) other; + return ComparisonChain.start() + .compare(chunks, castedOther.chunks) + .compare(chunkNumber, castedOther.chunkNumber) + .result(); + } else { + throw new IllegalArgumentException("Cannot compare against something that is not a NumberedPartitionChunk."); + } + } + + @Override + @SuppressWarnings("unchecked") + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + + return compareTo((NumberedPartitionChunk) o) == 0; + } + + @Override + public int hashCode() + { + return Objects.hashCode(chunks, chunkNumber); + } +}