[incubator-pinot] branch master updated: Add support for Long in Modulo partition function. (#6929)

[ma...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

mayanks pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-pinot.git


The following commit(s) were added to refs/heads/master by this push:
     new 5b7023a  Add support for Long in Modulo partition function. (#6929)
5b7023a is described below

commit 5b7023a4e75d91ea75d4f5f575d440b602bf3df6
Author: Mayank Shrivastava <ma...@apache.org>
AuthorDate: Mon May 17 14:02:05 2021 -0700

    Add support for Long in Modulo partition function. (#6929)
    
    - Currently, the `ModuloPartitionFunction` class only supports `int` and `String`,
      with this PR, we also add support for `long`.
    
    - Also fixed String to be parsed as Long, to be able to support both `long` and `int`.
---
 .../local/partition/ModuloPartitionFunction.java   |  6 ++++-
 .../local/partition/PartitionFunctionTest.java     | 29 ++++++++++++++++++----
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/partition/ModuloPartitionFunction.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/partition/ModuloPartitionFunction.java
index 52371ac..4d32f3c 100644
--- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/partition/ModuloPartitionFunction.java
+++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/partition/ModuloPartitionFunction.java
@@ -53,8 +53,12 @@ public class ModuloPartitionFunction implements PartitionFunction {
   public int getPartition(Object value) {
     if (value instanceof Integer) {
       return ((Integer) value) % _numPartitions;
+    } else if (value instanceof Long) {
+      // Since _numPartitions is int, the modulo should also be int.
+      return (int) (((Long) value) % _numPartitions);
     } else if (value instanceof String) {
-      return ((Integer.parseInt((String) value)) % _numPartitions);
+      // Parse String as Long, to support both Integer and Long.
+      return (int) ((Long.parseLong((String) value)) % _numPartitions);
     } else {
       throw new IllegalArgumentException(
           "Illegal argument for partitioning, expected Integer, got: " + value.getClass().getName());
diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/partition/PartitionFunctionTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/partition/PartitionFunctionTest.java
index ccda88c..c5fcb55 100644
--- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/partition/PartitionFunctionTest.java
+++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/partition/PartitionFunctionTest.java
@@ -57,10 +57,29 @@ public class PartitionFunctionTest {
       Assert.assertEquals(partitionFunction.toString().toLowerCase(), functionName.toLowerCase());
       Assert.assertEquals(partitionFunction.getNumPartitions(), expectedNumPartitions);
 
+      // Test int values.
       for (int j = 0; j < 1000; j++) {
         int value = random.nextInt();
         Assert.assertEquals(partitionFunction.getPartition(value), (value % expectedNumPartitions));
       }
+
+      // Test long values.
+      for (int j = 0; j < 1000; j++) {
+        long value = random.nextLong();
+        Assert.assertEquals(partitionFunction.getPartition(value), (value % expectedNumPartitions));
+      }
+
+      // Test Integer represented as String.
+      for (int j = 0; j < 1000; j++) {
+        int value = random.nextInt();
+        Assert.assertEquals(partitionFunction.getPartition(Integer.toString(value)), (value % expectedNumPartitions));
+      }
+
+      // Test Long represented as String.
+      for (int j = 0; j < 1000; j++) {
+        long value = random.nextLong();
+        Assert.assertEquals(partitionFunction.getPartition(Long.toString(value)), (value % expectedNumPartitions));
+      }
     }
   }
 
@@ -92,9 +111,10 @@ public class PartitionFunctionTest {
       Assert.assertEquals(partitionFunction.getNumPartitions(), expectedNumPartitions);
 
       for (int j = 0; j < 1000; j++) {
-        Integer value = random.nextInt();
-        Assert.assertTrue(partitionFunction.getPartition(value.toString()) < expectedNumPartitions,
-            "Illegal: " + partitionFunction.getPartition(value.toString()) + " " + expectedNumPartitions);
+        int value = random.nextInt();
+        String stringValue = Integer.toString(value);
+        Assert.assertTrue(partitionFunction.getPartition(stringValue) < expectedNumPartitions,
+            "Illegal: " + partitionFunction.getPartition(stringValue) + " " + expectedNumPartitions);
       }
     }
   }
@@ -169,8 +189,7 @@ public class PartitionFunctionTest {
     // 10 values of size 7, were randomly generated, using {@link Random::nextBytes} with seed 100
     // Applied org.apache.kafka.common.utils.Utils::murmur2 to those values and stored in expectedMurmurValues
     int[] expectedMurmurValues =
-        new int[]{-1044832774, -594851693, 1441878663, 1766739604, 1034724141, -296671913, 443511156, 1483601453,
-            1819695080, -931669296};
+        new int[]{-1044832774, -594851693, 1441878663, 1766739604, 1034724141, -296671913, 443511156, 1483601453, 1819695080, -931669296};
 
     long seed = 100;
     Random random = new Random(seed);

---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org