Merge pull request #225 from metamx/num-shard-spec

NumberedShardSpec

client/src/main/java/com/metamx/druid/shard/NumberedShardSpec.java

+/*
+ * Druid - a distributed column store.
+ * Copyright (C) 2012  Metamarkets Group Inc.
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+
+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 <T> PartitionChunk<T> createChunk(T obj)
+  {
+    return NumberedPartitionChunk.make(partitionNum, partitions, obj);
+  }
+
+  @Override
+  public boolean isInChunk(Map<String, String> dimensions)
+  {
+    return true;
+  }
+
+  @Override
+  public boolean isInChunk(InputRow inputRow)
+  {
+    return true;
+  }
+}

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
 {

client/src/test/java/com/metamx/druid/shard/NumberedShardSpecTest.java

+/*
+ * Druid - a distributed column store.
+ * Copyright (C) 2012  Metamarkets Group Inc.
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+
+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 testSerdeRoundTrip() throws Exception
+  {
+    final ObjectMapper jsonMapper = new DefaultObjectMapper();
+    final ShardSpec spec = jsonMapper.readValue(
+        jsonMapper.writeValueAsBytes(new NumberedShardSpec(1, 2)),
+        ShardSpec.class
+    );
+    Assert.assertEquals(1, spec.getPartitionNum());
+    Assert.assertEquals(2, ((NumberedShardSpec) spec).getPartitions());
+  }
+
+  @Test
+  public void testSerdeBackwardsCompat() throws Exception
+  {
+    final ObjectMapper jsonMapper = new DefaultObjectMapper();
+    final ShardSpec spec = jsonMapper.readValue(
+        "{\"type\": \"numbered\", \"partitions\": 2, \"partitionNum\": 1}",
+        ShardSpec.class
+    );
+    Assert.assertEquals(1, spec.getPartitionNum());
+    Assert.assertEquals(2, ((NumberedShardSpec) spec).getPartitions());
+  }
+
+  @Test
+  public void testPartitionChunks()
+  {
+    final List<ShardSpec> specs = ImmutableList.<ShardSpec>of(
+        new NumberedShardSpec(0, 3),
+        new NumberedShardSpec(1, 3),
+        new NumberedShardSpec(2, 3)
+    );
+
+    final List<PartitionChunk<String>> chunks = Lists.transform(
+        specs,
+        new Function<ShardSpec, PartitionChunk<String>>()
+        {
+          @Override
+          public PartitionChunk<String> 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)));
+  }
+}

common/src/main/java/com/metamx/druid/partition/IntegerPartitionChunk.java

@@ -96,10 +96,10 @@ public class IntegerPartitionChunk<T> implements PartitionChunk<T>
   {
     if (chunk instanceof IntegerPartitionChunk) {
       IntegerPartitionChunk<T> intChunk = (IntegerPartitionChunk<T>) chunk;
-
       return comparator.compare(chunkNumber, intChunk.chunkNumber);
+    } 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
@@ -134,3 +134,4 @@ public class IntegerPartitionChunk<T> implements PartitionChunk<T>
            '}';
   }
 }
+

common/src/main/java/com/metamx/druid/partition/NumberedPartitionChunk.java

+/*
+ * Druid - a distributed column store.
+ * Copyright (C) 2012  Metamarkets Group Inc.
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+
+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<T> implements PartitionChunk<T>
+{
+  private final int chunkNumber;
+  private final int chunks;
+  private final T object;
+
+  public static <T> NumberedPartitionChunk<T> make(
+      int chunkNumber,
+      int chunks,
+      T obj
+  )
+  {
+    return new NumberedPartitionChunk<T>(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<T> 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<T> 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<T>) o) == 0;
+  }
+
+  @Override
+  public int hashCode()
+  {
+    return Objects.hashCode(chunks, chunkNumber);
+  }
+}