CASSANALYTICS-32 expose number of rows which violated constraints to JobStats

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/BulkSparkConf.java

@@ -121,6 +121,7 @@ public class BulkSparkConf implements Serializable
     public final Integer sstableDataSizeInMiB;
     public final int commitBatchSize;
     public final boolean skipExtendedVerify;
+    public final boolean skipRowsViolatingConstraints;
     public final WriteMode writeMode;
     public final int commitThreadsPerInstance;
     public final double importCoordinatorTimeoutMultiplier;
@@ -166,6 +167,8 @@ public BulkSparkConf(SparkConf conf, Map<String, String> options)
         this.table = MapUtils.getOrThrow(options, WriterOptions.TABLE.name());
         this.skipExtendedVerify = MapUtils.getBoolean(options, WriterOptions.SKIP_EXTENDED_VERIFY.name(), true,
                                                       "skip extended verification of SSTables by Cassandra");
+        this.skipRowsViolatingConstraints = MapUtils.getBoolean(options, WriterOptions.SKIP_ROWS_VIOLATING_CONSTRAINTS.name(), false,
+                                                                "skip rows which failed to be written because of violated constraints");
         this.consistencyLevel = ConsistencyLevel.CL.valueOf(MapUtils.getOrDefault(options, WriterOptions.BULK_WRITER_CL.name(), "EACH_QUORUM"));
         String dc = MapUtils.getOrDefault(options, WriterOptions.LOCAL_DC.name(), null);
         if (!consistencyLevel.isLocal() && dc != null)

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraBulkSourceRelation.java

@@ -199,13 +199,14 @@ private void persist(@NotNull JavaPairRDD<DecoratedKey, Object[]> sortedRDD, Str
 
             long rowCount = streamResults.stream().mapToLong(res -> res.rowCount).sum();
             long totalBytesWritten = streamResults.stream().mapToLong(res -> res.bytesWritten).sum();
+            long rowsViolatedConstraints = streamResults.stream().mapToLong(res -> res.rowsViolatedConstraints).sum();
             boolean hasClusterTopologyChanged = writeResults.stream().anyMatch(WriteResult::isClusterResizeDetected);
 
             onCloudStorageTransport(context -> waitForImportCompletion(context, rowCount, totalBytesWritten, hasClusterTopologyChanged, streamResults));
 
             LOGGER.info("Bulk writer job complete. rowCount={} totalBytes={} hasClusterTopologyChanged={}",
                         rowCount, totalBytesWritten, hasClusterTopologyChanged);
-            publishSuccessfulJobStats(rowCount, totalBytesWritten, hasClusterTopologyChanged);
+            publishSuccessfulJobStats(rowCount, totalBytesWritten, rowsViolatedConstraints, hasClusterTopologyChanged);
         }
         catch (Throwable throwable)
         {
@@ -300,13 +301,14 @@ private void awaitImportCompletion(TransportContext.CloudStorageTransportContext
         importCoordinator.await();
     }
 
-    private void publishSuccessfulJobStats(long rowCount, long totalBytesWritten, boolean hasClusterTopologyChanged)
+    private void publishSuccessfulJobStats(long rowCount, long rowsViolatingConstraints, long totalBytesWritten, boolean hasClusterTopologyChanged)
     {
         writerContext.jobStats().publish(new HashMap<String, String>() // type declaration required to compile with java8
         {{
                 put("jobId", writerContext.job().getId().toString());
                 put("transportInfo", writerContext.job().transportInfo().toString());
                 put("rowsWritten", Long.toString(rowCount));
+                put("rowsViolatingConstraints", Long.toString(rowsViolatingConstraints));
                 put("bytesWritten", Long.toString(totalBytesWritten));
                 put("jobStatus", "Succeeded");
                 put("clusterResizeDetected", String.valueOf(hasClusterTopologyChanged));

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraJobInfo.java

@@ -76,6 +76,12 @@ public boolean skipExtendedVerify()
         return conf.skipExtendedVerify;
     }
 
+    @Override
+    public boolean skipRowsViolatingConstraints()
+    {
+        return conf.skipRowsViolatingConstraints;
+    }
+
     @Override
     public boolean getSkipClean()
     {

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/DirectStreamResult.java

@@ -32,9 +32,9 @@ public class DirectStreamResult extends StreamResult
 
     public DirectStreamResult(String sessionID, Range<BigInteger> tokenRange,
                               List<StreamError> failures, List<RingInstance> passed,
-                              long rowCount, long bytesWritten)
+                              long rowCount, long bytesWritten, long rowsViolatedConstraints)
     {
-        super(sessionID, tokenRange, failures, passed, rowCount, bytesWritten);
+        super(sessionID, tokenRange, failures, passed, rowCount, bytesWritten, rowsViolatedConstraints);
     }
 
     public void setCommitResults(List<CommitResult> commitResult)
@@ -53,6 +53,7 @@ public String toString()
                + ", commitResults=" + commitResults
                + ", passed=" + passed
                + ", bytesWritten=" + bytesWritten
+               + ", rowsViolatedConstraints=" + rowsViolatedConstraints
                + '}';
     }
 }

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/DirectStreamSession.java

@@ -123,7 +123,8 @@ protected StreamResult doFinalizeStream()
                                                                  errors,
                                                                  new ArrayList<>(replicas),
                                                                  sstableWriter.rowCount(),
-                                                                 sstableWriter.bytesWritten());
+                                                                 sstableWriter.bytesWritten(),
+                                                                 sstableWriter.rowsViolatedConstraints());
         List<CommitResult> cr;
         try
         {

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/JobInfo.java

@@ -70,6 +70,13 @@ default String getId()
 
     boolean skipExtendedVerify();
 
+    /**
+     * Defaults to false, in that case, when a row violates a constraint, a job will fail.
+     *
+     * @return true if rows violating constraints will not fail a job, false otherwise
+     */
+    boolean skipRowsViolatingConstraints();
+
     boolean getSkipClean();
 
     /**

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SortedSSTableWriter.java

@@ -77,12 +77,14 @@ public class SortedSSTableWriter
     private BigInteger maxToken = null;
     private final Map<Path, Digest> overallFileDigests = new HashMap<>();
     private final DigestAlgorithm digestAlgorithm;
+    private final boolean skipRowsViolatingConstraints;
 
     private volatile boolean isClosed = false;
 
     private int sstableCount = 0;
     private long rowCount = 0;
     private long bytesWritten = 0;
+    private long rowsViolatedConstraints = 0;
 
     public SortedSSTableWriter(org.apache.cassandra.bridge.SSTableWriter tableWriter, Path outDir,
                                DigestAlgorithm digestAlgorithm,
@@ -92,13 +94,15 @@ public SortedSSTableWriter(org.apache.cassandra.bridge.SSTableWriter tableWriter
         this.outDir = outDir;
         this.digestAlgorithm = digestAlgorithm;
         this.partitionId = partitionId;
+        this.skipRowsViolatingConstraints = false;
     }
 
     public SortedSSTableWriter(BulkWriterContext writerContext, Path outDir, DigestAlgorithm digestAlgorithm, int partitionId)
     {
         this.outDir = outDir;
         this.digestAlgorithm = digestAlgorithm;
         this.partitionId = partitionId;
+        this.skipRowsViolatingConstraints = writerContext.job().skipRowsViolatingConstraints();
 
         String lowestCassandraVersion = writerContext.cluster().getLowestCassandraVersion();
         String packageVersion = getPackageVersion(lowestCassandraVersion);
@@ -137,8 +141,24 @@ public void addRow(BigInteger token, Map<String, Object> boundValues) throws IOE
         }
         // rows are sorted. Therefore, only update the maxToken
         maxToken = token;
-        cqlSSTableWriter.addRow(boundValues);
-        rowCount += 1;
+        try
+        {
+            cqlSSTableWriter.addRow(boundValues);
+            rowCount += 1;
+        }
+        catch (Throwable t)
+        {
+            if (t.getCause() != null && "org.apache.cassandra.cql3.constraints.ConstraintViolationException".equals(t.getCause().getClass().getName()))
+            {
+                rowsViolatedConstraints += 1;
+                if (!skipRowsViolatingConstraints)
+                    throw t;
+            }
+            else
+            {
+                throw t;
+            }
+        }
     }
 
     public void setSSTablesProducedListener(Consumer<Set<SSTableDescriptor>> listener)
@@ -324,4 +344,9 @@ public Map<Path, Digest> fileDigestMap()
     {
         return Collections.unmodifiableMap(overallFileDigests);
     }
+
+    public long rowsViolatedConstraints()
+    {
+        return rowsViolatedConstraints;
+    }
 }

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/StreamResult.java

@@ -34,19 +34,22 @@ public abstract class StreamResult implements Serializable
     public final List<RingInstance> passed;
     public final long rowCount;
     public final long bytesWritten;
+    public final long rowsViolatedConstraints;
 
     protected StreamResult(String sessionID,
                            Range<BigInteger> tokenRange,
                            List<StreamError> failures,
                            List<RingInstance> passed,
                            long rowCount,
-                           long bytesWritten)
+                           long bytesWritten,
+                           long rowsViolatedConstraints)
     {
         this.sessionID = sessionID;
         this.tokenRange = tokenRange;
         this.failures = failures;
         this.passed = passed;
         this.rowCount = rowCount;
         this.bytesWritten = bytesWritten;
+        this.rowsViolatedConstraints = rowsViolatedConstraints;
     }
 }

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/WriterOptions.java

@@ -45,6 +45,7 @@ public enum WriterOptions implements WriterOption
     COMMIT_THREADS_PER_INSTANCE,
     COMMIT_BATCH_SIZE,
     SKIP_EXTENDED_VERIFY,
+    SKIP_ROWS_VIOLATING_CONSTRAINTS,
     WRITE_MODE,
     KEYSTORE_PASSWORD,
     KEYSTORE_PATH,

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/cloudstorage/CloudStorageStreamResult.java

@@ -44,7 +44,7 @@ public class CloudStorageStreamResult extends StreamResult
 
     public static CloudStorageStreamResult empty(String sessionID, Range<BigInteger> tokenRange)
     {
-        return new CloudStorageStreamResult(sessionID, tokenRange, new ArrayList<>(), new ArrayList<>(), new HashSet<>(), 0, 0, 0);
+        return new CloudStorageStreamResult(sessionID, tokenRange, new ArrayList<>(), new ArrayList<>(), new HashSet<>(), 0, 0, 0, 0);
     }
 
     public CloudStorageStreamResult(String sessionID,
@@ -54,9 +54,10 @@ public CloudStorageStreamResult(String sessionID,
                                     Set<CreatedRestoreSlice> createdRestoreSlices,
                                     int objectCount,
                                     long rowCount,
-                                    long bytesWritten)
+                                    long bytesWritten,
+                                    long rowsViolatedConstraints)
     {
-        super(sessionID, tokenRange, failures, passed, rowCount, bytesWritten);
+        super(sessionID, tokenRange, failures, passed, rowCount, bytesWritten, rowsViolatedConstraints);
         this.createdRestoreSlices = Collections.unmodifiableSet(createdRestoreSlices);
         this.objectCount = objectCount;
     }
@@ -70,6 +71,7 @@ public String toString()
                + ", objectCount=" + objectCount
                + ", rowCount=" + rowCount
                + ", bytesWritten=" + bytesWritten
+               + ", rowsViolatedConstraints=" + rowsViolatedConstraints
                + ", failures=" + failures
                + ", createdRestoreSlices=" + createdRestoreSlices
                + ", passed=" + passed

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/cloudstorage/CloudStorageStreamSession.java

@@ -189,7 +189,8 @@ protected StreamResult doFinalizeStream()
                                                                              createdRestoreSlices,
                                                                              sliceCount,
                                                                              sstableWriter.rowCount(),
-                                                                             sstableWriter.bytesWritten());
+                                                                             sstableWriter.bytesWritten(),
+                                                                             sstableWriter.rowsViolatedConstraints());
         LOGGER.info("StreamResult: {}", streamResult);
         // CL validation is only required for non-coordinated write.
         // For coordinated write, executors only need to make sure the bundles are uploaded and restore slice is created on sidecar.

cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/CommitCoordinatorTest.java

@@ -232,7 +232,7 @@ DirectStreamResult build()
             DirectStreamResult sr = new DirectStreamResult(UUID.randomUUID().toString(),
                                                            TEST_RANGE,
                                                            buildFailures(),
-                                                           buildPassed(), 0, 0);
+                                                           buildPassed(), 0, 0, 0);
             if (successfulCommits > 0 || failedCommits > 0)
             {
                 List<CommitResult> commitResults = new ArrayList<>();

cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/MockBulkWriterContext.java

@@ -259,6 +259,12 @@ public boolean skipExtendedVerify()
         return false;
     }
 
+    @Override
+    public boolean skipRowsViolatingConstraints()
+    {
+        return false;
+    }
+
     @Override
     public boolean getSkipClean()
     {

cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/bulkwriter/cloudstorage/ImportCompletionCoordinatorTest.java

@@ -481,7 +481,7 @@ else if (failedPerReplica-- > 0)
             }
             CloudStorageStreamResult result = new CloudStorageStreamResult("", mock(Range.class), Collections.emptyList(),
                                                                            passedReplicaSet, createdRestoreSlices,
-                                                                           createdRestoreSlices.size(), 0, 0);
+                                                                           createdRestoreSlices.size(), 0, 0, 0);
             resultList.add(result);
         }
         return resultList;