CDM-57 min/max partition defaults based on partitioner rather than co… (#150)

mieslep · web-flow · commit dd8843fb223d · 2023-05-26T08:43:30.000-04:00
* CDM-57 min/max partition defaults based on partitioner rather than config setting

* CDM-57 fixing feature post-merge

---------

Co-authored-by: Phil Miesle &lt;phil.miesle@datastax.com&gt;
diff --git a/src/main/java/com/datastax/cdm/cql/statement/OriginSelectByPartitionRangeStatement.java b/src/main/java/com/datastax/cdm/cql/statement/OriginSelectByPartitionRangeStatement.java
@@ -16,26 +16,20 @@
 public class OriginSelectByPartitionRangeStatement extends OriginSelectStatement {
     public Logger logger = LoggerFactory.getLogger(this.getClass().getName());
 
-    private final Long defaultMinPartition;
-    private final Long defaultMaxPartition;
-
     public OriginSelectByPartitionRangeStatement(IPropertyHelper propertyHelper, EnhancedSession session) {
         super(propertyHelper, session);
-
-        defaultMinPartition = propertyHelper.getLong(KnownProperties.PARTITION_MIN);
-        defaultMaxPartition = propertyHelper.getLong(KnownProperties.PARTITION_MAX);
     }
 
     @Override
     public BoundStatement bind(Object... binds) {
         if (null==binds
                 || binds.length != 2
-                || !(null==binds[0] || binds[0] instanceof BigInteger)
-                || !(null==binds[1] || binds[1] instanceof BigInteger))
-            throw new RuntimeException("Expected 2 nullable bind of type BigInteger, got " + binds.length);
+                || !(binds[0] instanceof BigInteger)
+                || !(binds[1] instanceof BigInteger))
+            throw new RuntimeException("Expected 2 not-null binds of type BigInteger, got " + binds.length);
 
-        BigInteger min = (null==binds[0]) ? BigInteger.valueOf(defaultMinPartition) : (BigInteger) binds[0];
-        BigInteger max = (null==binds[1]) ? BigInteger.valueOf(defaultMaxPartition) : (BigInteger) binds[1];
+        BigInteger min = (BigInteger) binds[0];
+        BigInteger max = (BigInteger) binds[1];
 
         PreparedStatement preparedStatement = prepareStatement();
         // random partitioner uses BigInteger, the normal partitioner uses long
diff --git a/src/main/java/com/datastax/cdm/properties/KnownProperties.java b/src/main/java/com/datastax/cdm/properties/KnownProperties.java
@@ -172,10 +172,8 @@ public enum PropertyType {
     public static final String PARTITION_MAX               = "spark.cdm.filter.cassandra.partition.max";
     public static final String FILTER_CQL_WHERE_CONDITION  = "spark.cdm.filter.cassandra.where.condition";
     static {
-           types.put(PARTITION_MIN, PropertyType.NUMBER);
-        defaults.put(PARTITION_MIN, "-9223372036854775808");
-           types.put(PARTITION_MAX, PropertyType.NUMBER);
-        defaults.put(PARTITION_MAX, "9223372036854775807");
+           types.put(PARTITION_MIN, PropertyType.STRING);
+           types.put(PARTITION_MAX, PropertyType.STRING);
            types.put(FILTER_CQL_WHERE_CONDITION, PropertyType.STRING);
     }
 
diff --git a/src/main/scala/com/datastax/cdm/job/BaseJob.scala b/src/main/scala/com/datastax/cdm/job/BaseJob.scala
@@ -57,8 +57,16 @@ abstract class BaseJob[T: ClassTag] extends App {
     propertyHelper = PropertyHelper.getInstance(sc);
 
     consistencyLevel = propertyHelper.getString(KnownProperties.READ_CL)
-    minPartition = new BigInteger(propertyHelper.getAsString(KnownProperties.PARTITION_MIN))
-    maxPartition = new BigInteger(propertyHelper.getAsString(KnownProperties.PARTITION_MAX))
+    val connectionFetcher = new ConnectionFetcher(sContext, propertyHelper)
+    originConnection = connectionFetcher.getConnection("ORIGIN", consistencyLevel)
+    targetConnection = connectionFetcher.getConnection("TARGET", consistencyLevel)
+
+    val hasRandomPartitioner: Boolean = {
+      val partitionerName = originConnection.withSessionDo(_.getMetadata.getTokenMap.get().getPartitionerName)
+      partitionerName.endsWith("RandomPartitioner")
+    }
+    minPartition = getMinPartition(propertyHelper.getString(KnownProperties.PARTITION_MIN), hasRandomPartitioner)
+    maxPartition = getMaxPartition(propertyHelper.getString(KnownProperties.PARTITION_MAX), hasRandomPartitioner)
     coveragePercent = propertyHelper.getInteger(KnownProperties.TOKEN_COVERAGE_PERCENT)
     numSplits = propertyHelper.getInteger(KnownProperties.PERF_NUM_PARTS)
     abstractLogger.info("PARAM -- Min Partition: " + minPartition)
@@ -71,9 +79,6 @@ abstract class BaseJob[T: ClassTag] extends App {
     abstractLogger.info("PARAM Calculated -- Total Partitions: " + parts.size())
     abstractLogger.info("Spark parallelize created : " + slices.count() + " slices!");
 
-    val connectionFetcher = new ConnectionFetcher(sContext, propertyHelper)
-    originConnection = connectionFetcher.getConnection("ORIGIN", consistencyLevel)
-    targetConnection = connectionFetcher.getConnection("TARGET", consistencyLevel)
   }
 
   def getParts(pieces: Int): util.Collection[T]
@@ -125,4 +130,15 @@ abstract class BaseJob[T: ClassTag] extends App {
     abstractLogger.info(bannerFill)
   }
 
+  def getMinPartition(minPartition: String, hasRandomPartitioner: Boolean): BigInteger = {
+    if (minPartition != null && minPartition.nonEmpty) new BigInteger(minPartition)
+    else if (hasRandomPartitioner) BigInteger.ZERO
+    else BigInteger.valueOf(Long.MinValue)
+  }
+
+  def getMaxPartition(maxPartition: String, hasRandomPartitioner: Boolean): BigInteger = {
+    if (maxPartition != null && maxPartition.nonEmpty) new BigInteger(maxPartition)
+    else if (hasRandomPartitioner) new BigInteger("2").pow(127).subtract(BigInteger.ONE)
+    else BigInteger.valueOf(Long.MaxValue)
+  }
 }
diff --git a/src/resources/sparkConf.properties b/src/resources/sparkConf.properties
@@ -235,8 +235,10 @@ spark.cdm.perfops.writeRateLimit                  40000
 #
 #  spark.cdm.filter.cassandra
 #    .partition
-#       .min            : Default is -9223372036854775808 (-2^63). Lower partition bound (inclusive).
-#       .min            : Default is 9223372036854775807 (2^63-1). Upper partition bound (inclusive).
+#       .min            : Default is 0 (when using RandomPartitioner) and -9223372036854775808 (-2^63)
+#                         otherwise. Lower partition bound (inclusive).
+#       .min            : Default is 2^127-1 (when using RandomPartitioner) and 9223372036854775807
+#                         (2^63-1) otherwise. Upper partition bound (inclusive).
 #    .where.condition   : CQL added to the WHERE clause of SELECTs from origin
 #-----------------------------------------------------------------------------------------------------------
 #spark.cdm.filter.cassandra.partition.min          -9223372036854775808
diff --git a/src/test/java/com/datastax/cdm/cql/CommonMocks.java b/src/test/java/com/datastax/cdm/cql/CommonMocks.java
@@ -109,8 +109,8 @@ public class CommonMocks {
     public List<String> targetColumnNames;
     public List<DataType> targetColumnTypes;
 
-    public Long minPartition;
-    public Long maxPartition;
+    public String minPartition;
+    public String maxPartition;
     public ConsistencyLevel readCL;
     public Integer fetchSizeInRows;
 
@@ -180,8 +180,8 @@ public void defaultClassVariables() {
         targetKeyspaceName = "target_ks";
         targetTableName = "table_name";
 
-        minPartition = -9876543L;
-        maxPartition =  1234567L;
+        minPartition = "-9876543";
+        maxPartition = "1234567";
         readCL = ConsistencyLevel.LOCAL_QUORUM;
         fetchSizeInRows = 999;
 
@@ -239,8 +239,8 @@ public void setPropertyHelperWhens() {
         when(propertyHelper.getAsString(KnownProperties.FILTER_COLUMN_NAME)).thenReturn("");
         when(propertyHelper.getString(KnownProperties.FILTER_COLUMN_VALUE)).thenReturn(null);
 
-        when(propertyHelper.getLong(KnownProperties.PARTITION_MIN)).thenReturn(minPartition);
-        when(propertyHelper.getLong(KnownProperties.PARTITION_MAX)).thenReturn(maxPartition);
+        when(propertyHelper.getString(KnownProperties.PARTITION_MIN)).thenReturn(minPartition);
+        when(propertyHelper.getString(KnownProperties.PARTITION_MAX)).thenReturn(maxPartition);
 
         when(propertyHelper.getStringList(KnownProperties.ORIGIN_COLUMN_NAMES_TO_TARGET)).thenReturn(originToTargetNameList);
     }
diff --git a/src/test/java/com/datastax/cdm/cql/statement/OriginSelectByPartitionRangeStatementTest.java b/src/test/java/com/datastax/cdm/cql/statement/OriginSelectByPartitionRangeStatementTest.java
@@ -55,15 +55,14 @@ public void originFilterCondition() {
         String cql = originSelectByPartitionRangeStatement.getCQL();
         assertEquals(sb.toString(),cql);
     }
+
     @Test
-    public void bind_withNullBinds_usesDefaultPartitions() {
-        originSelectByPartitionRangeStatement.bind(null, null);
+    public void bind_withNullBinds() {
         assertAll(
-                () -> verify(preparedStatement).bind(
-                        BigInteger.valueOf(minPartition).longValueExact(),
-                        BigInteger.valueOf(maxPartition).longValueExact()),
-                () -> verify(boundStatement).setConsistencyLevel(readCL),
-                () -> verify(boundStatement).setPageSize(fetchSizeInRows)
+                () -> assertThrows(RuntimeException.class, () -> originSelectByPartitionRangeStatement.bind(null,null), "two null"),
+                () -> assertThrows(RuntimeException.class, () -> originSelectByPartitionRangeStatement.bind(BigInteger.valueOf(20)), "missing second"),
+                () -> assertThrows(RuntimeException.class, () -> originSelectByPartitionRangeStatement.bind(BigInteger.valueOf(20),null), "null second"),
+                () -> assertThrows(RuntimeException.class, () -> originSelectByPartitionRangeStatement.bind(null,BigInteger.valueOf(20)), "null first")
         );
     }
 
@@ -74,41 +73,12 @@ public void bind_withNonNullBinds_usesProvidedPartitions() {
 
         originSelectByPartitionRangeStatement.bind(providedMin, providedMax);
         assertAll(
-                () -> verify(preparedStatement).bind(
-                        providedMin.longValueExact(),
-                        providedMax.longValueExact()),
+                () -> verify(preparedStatement).bind(providedMin.longValueExact(), providedMax.longValueExact()),
                 () -> verify(boundStatement).setConsistencyLevel(readCL),
                 () -> verify(boundStatement).setPageSize(fetchSizeInRows)
         );
     }
 
-    @Test
-    public void bind_withMixedBinds_usesDefaultAndProvidedPartitions() {
-        BigInteger providedMin = BigInteger.valueOf(12345L);
-
-        originSelectByPartitionRangeStatement.bind(providedMin, null);
-        assertAll(
-                () -> verify(preparedStatement).bind(
-                        providedMin.longValueExact(),
-                        BigInteger.valueOf(maxPartition).longValueExact()),
-                () -> verify(boundStatement).setConsistencyLevel(readCL),
-                () -> verify(boundStatement).setPageSize(fetchSizeInRows)
-        );
-    }
-
-    @Test
-    public void bind_withNullBinds_usesDefaultPartitions_whenRandomPartitioner() {
-        when(originTable.hasRandomPartitioner()).thenReturn(true);
-
-        originSelectByPartitionRangeStatement.bind(null, null);
-        assertAll(
-                () -> verify(preparedStatement).bind(
-                        BigInteger.valueOf(minPartition),
-                        BigInteger.valueOf(maxPartition)),
-                () -> verify(boundStatement).setConsistencyLevel(readCL),
-                () -> verify(boundStatement).setPageSize(fetchSizeInRows)
-        );
-    }
 
     @Test
     public void bind_withNonNullBinds_usesProvidedPartitions_whenRandomPartitioner() {
@@ -128,29 +98,11 @@ public void bind_withNonNullBinds_usesProvidedPartitions_whenRandomPartitioner()
     }
 
     @Test
-    public void bind_withMixedBinds_usesDefaultAndProvidedPartitions_whenRandomPartitioner() {
-        when(originTable.hasRandomPartitioner()).thenReturn(true);
-
-        BigInteger providedMax = BigInteger.valueOf(999999999L);
-
-        originSelectByPartitionRangeStatement.bind(null, providedMax);
+    public void bind_withInvalidBindType_throwsException() {
         assertAll(
-                () -> verify(preparedStatement).bind(
-                        BigInteger.valueOf(minPartition),
-                        providedMax),
-                () -> verify(boundStatement).setConsistencyLevel(readCL),
-                () -> verify(boundStatement).setPageSize(fetchSizeInRows)
+                () -> assertThrows(RuntimeException.class, () -> originSelectByPartitionRangeStatement.bind("invalidType", BigInteger.valueOf(20)), "invalid first"),
+                () -> assertThrows(RuntimeException.class, () -> originSelectByPartitionRangeStatement.bind(BigInteger.valueOf(20),"invalidType"), "invalid second")
         );
     }
 
-    @Test
-    public void bind_withInvalidBindLength_throwsException() {
-        assertThrows(RuntimeException.class, () -> originSelectByPartitionRangeStatement.bind(null));
-    }
-
-    @Test
-    public void bind_withInvalidBindType_throwsException() {
-        assertThrows(RuntimeException.class, () -> originSelectByPartitionRangeStatement.bind("invalidType", BigInteger.valueOf(20)));
-    }
-
 }
