SplitGenerator { List<List<RingRange>> generateSplits(long totalSplitCount, List<BigInteger> ringTokens) { int tokenRangeCount = ringTokens.size(); List<RingRange> splits = new ArrayList<>(); for (int i = 0; i < tokenRangeCount; i++) { BigInteger start = ringTokens.get(i); BigInteger stop = ringTokens.get((i + 1) % tokenRangeCount); if (!inRange(start) || !inRange(stop)) { throw new RuntimeException( String.format("Tokens (%s,%s) not in range of %s", start, stop, partitioner)); } if (start.equals(stop) && tokenRangeCount != 1) { throw new RuntimeException( String.format("Tokens (%s,%s): two nodes have the same token", start, stop)); } BigInteger rs = stop.subtract(start); if (rs.compareTo(BigInteger.ZERO) <= 0) { rs = rs.add(rangeSize); } BigInteger[] splitCountAndRemainder = rs.multiply(BigInteger.valueOf(totalSplitCount)).divideAndRemainder(rangeSize); int splitCount = splitCountAndRemainder[0].intValue() + (splitCountAndRemainder[1].equals(BigInteger.ZERO) ? 0 : 1); LOG.debug("Dividing token range [{},{}) into {} splits", start, stop, splitCount); List<BigInteger> endpointTokens = new ArrayList<>(); for (int j = 0; j <= splitCount; j++) { BigInteger offset = rs.multiply(BigInteger.valueOf(j)).divide(BigInteger.valueOf(splitCount)); BigInteger token = start.add(offset); if (token.compareTo(rangeMax) > 0) { token = token.subtract(rangeSize); } endpointTokens.add( token.equals(BigInteger.valueOf(Long.MIN_VALUE)) ? token.add(BigInteger.ONE) : token); } for (int j = 0; j < splitCount; j++) { splits.add(new RingRange(endpointTokens.get(j), endpointTokens.get(j + 1))); LOG.debug("Split #{}: [{},{})", j + 1, endpointTokens.get(j), endpointTokens.get(j + 1)); } } BigInteger total = BigInteger.ZERO; for (RingRange split : splits) { BigInteger size = split.span(rangeSize); total = total.add(size); } if (!total.equals(rangeSize)) { throw new RuntimeException( "Some tokens are missing from the splits. " + "This should not happen."); } return coalesceSplits(getTargetSplitSize(totalSplitCount), splits); } SplitGenerator(String partitioner);   }
@Test public void testGenerateSegments() { List<BigInteger> tokens = Stream.of( "0", "1", "56713727820156410577229101238628035242", "56713727820156410577229101238628035243", "113427455640312821154458202477256070484", "113427455640312821154458202477256070485") .map(BigInteger::new) .collect(Collectors.toList()); SplitGenerator generator = new SplitGenerator("foo.bar.RandomPartitioner"); List<List<RingRange>> segments = generator.generateSplits(10, tokens); assertEquals(12, segments.size()); assertEquals("[(0,1], (1,14178431955039102644307275309657008811]]", segments.get(0).toString()); assertEquals( "[(14178431955039102644307275309657008811,28356863910078205288614550619314017621]]", segments.get(1).toString()); assertEquals( "[(70892159775195513221536376548285044053,85070591730234615865843651857942052863]]", segments.get(5).toString()); tokens = Stream.of( "5", "6", "56713727820156410577229101238628035242", "56713727820156410577229101238628035243", "113427455640312821154458202477256070484", "113427455640312821154458202477256070485") .map(BigInteger::new) .collect(Collectors.toList()); segments = generator.generateSplits(10, tokens); assertEquals(12, segments.size()); assertEquals("[(5,6], (6,14178431955039102644307275309657008815]]", segments.get(0).toString()); assertEquals( "[(70892159775195513221536376548285044053,85070591730234615865843651857942052863]]", segments.get(5).toString()); assertEquals( "[(141784319550391026443072753096570088109,155962751505430129087380028406227096921]]", segments.get(10).toString()); }
@Test(expected = RuntimeException.class) public void testZeroSizeRange() { List<String> tokenStrings = Arrays.asList( "0", "1", "56713727820156410577229101238628035242", "56713727820156410577229101238628035242", "113427455640312821154458202477256070484", "113427455640312821154458202477256070485"); List<BigInteger> tokens = tokenStrings.stream().map(BigInteger::new).collect(Collectors.toList()); SplitGenerator generator = new SplitGenerator("foo.bar.RandomPartitioner"); generator.generateSplits(10, tokens); }
@Test public void testRotatedRing() { List<String> tokenStrings = Arrays.asList( "56713727820156410577229101238628035243", "113427455640312821154458202477256070484", "113427455640312821154458202477256070485", "5", "6", "56713727820156410577229101238628035242"); List<BigInteger> tokens = tokenStrings.stream().map(BigInteger::new).collect(Collectors.toList()); SplitGenerator generator = new SplitGenerator("foo.bar.RandomPartitioner"); List<List<RingRange>> segments = generator.generateSplits(5, tokens); assertEquals(6, segments.size()); assertEquals( "[(85070591730234615865843651857942052863,113427455640312821154458202477256070484]," + " (113427455640312821154458202477256070484,113427455640312821154458202477256070485]]", segments.get(1).toString()); assertEquals( "[(113427455640312821154458202477256070485," + "141784319550391026443072753096570088109]]", segments.get(2).toString()); assertEquals( "[(141784319550391026443072753096570088109,5], (5,6]]", segments.get(3).toString()); }
@Test(expected = RuntimeException.class) public void testDisorderedRing() { List<String> tokenStrings = Arrays.asList( "0", "113427455640312821154458202477256070485", "1", "56713727820156410577229101238628035242", "56713727820156410577229101238628035243", "113427455640312821154458202477256070484"); List<BigInteger> tokens = tokenStrings.stream().map(BigInteger::new).collect(Collectors.toList()); SplitGenerator generator = new SplitGenerator("foo.bar.RandomPartitioner"); generator.generateSplits(10, tokens); }