CASSANALYTICS-31 SAI index support in analytics

CHANGES.txt

@@ -1,5 +1,6 @@
 0.5.0
 -----
+ * Add Storage Attached Index (SAI) support to the bulk writer (CASSANALYTICS-31)
  * SSTable-version-based bridge determination (CASSANALYTICS-24)
  * Upgrade sidecar version to 0.4.0 (CASSANALYTICS-176)
  * Exclude IP address from RingInstance equality so node replacement does not fail bulk write jobs (CASSANALYTICS-175)

analytics-sidecar-client-common/src/main/java/org/apache/cassandra/sidecar/common/data/SSTableImportOptions.java

@@ -33,6 +33,8 @@ public class SSTableImportOptions extends HashMap<String, String>
     private static final String INVALIDATE_CACHES = "invalidateCaches";
     private static final String EXTENDED_VERIFY = "extendedVerify";
     private static final String COPY_DATA = "copyData";
+    private static final String FAIL_ON_MISSING_INDEX = "failOnMissingIndex";
+    private static final String VALIDATE_INDEX_CHECKSUM = "validateIndexChecksum";
 
     public static SSTableImportOptions defaults()
     {
@@ -126,4 +128,34 @@ public boolean copyData()
     {
         return Boolean.parseBoolean(get(COPY_DATA));
     }
+
+    /**
+     * When enabled, SSTable import fails if a table with Storage Attached Indexes (SAI) is missing its index
+     * components, instead of silently rebuilding them. Only meaningful for SAI tables (Cassandra 5.0+).
+     */
+    public SSTableImportOptions failOnMissingIndex(boolean enabled)
+    {
+        put(FAIL_ON_MISSING_INDEX, Boolean.toString(enabled));
+        return this;
+    }
+
+    public boolean failOnMissingIndex()
+    {
+        return Boolean.parseBoolean(get(FAIL_ON_MISSING_INDEX));
+    }
+
+    /**
+     * When enabled, SSTable import validates the checksums of Storage Attached Index (SAI) components.
+     * Only meaningful for SAI tables (Cassandra 5.0+).
+     */
+    public SSTableImportOptions validateIndexChecksum(boolean enabled)
+    {
+        put(VALIDATE_INDEX_CHECKSUM, Boolean.toString(enabled));
+        return this;
+    }
+
+    public boolean validateIndexChecksum()
+    {
+        return Boolean.parseBoolean(get(VALIDATE_INDEX_CHECKSUM));
+    }
 }

analytics-sidecar-client-common/src/main/java/org/apache/cassandra/sidecar/common/request/ImportSSTableRequest.java

@@ -65,6 +65,8 @@ public static class ImportOptions
         private Boolean invalidateCaches;
         private Boolean extendedVerify;
         private Boolean copyData;
+        private Boolean failOnMissingIndex;
+        private Boolean validateIndexChecksum;
 
         public ImportOptions()
         {
@@ -153,6 +155,32 @@ public ImportOptions copyData(boolean copyData)
             this.copyData = copyData;
             return this;
         }
+
+        /**
+         * Sets the {@code failOnMissingIndex} and returns a reference to this ImportOptions enabling method chaining.
+         * When enabled, Cassandra validates SAI index components during import.
+         *
+         * @param failOnMissingIndex the {@code failOnMissingIndex} to set
+         * @return a reference to this ImportOptions
+         */
+        public ImportOptions failOnMissingIndex(boolean failOnMissingIndex)
+        {
+            this.failOnMissingIndex = failOnMissingIndex;
+            return this;
+        }
+
+        /**
+         * Sets the {@code validateIndexChecksum} and returns a reference to this ImportOptions enabling method chaining.
+         * When enabled, Cassandra verifies SAI index component checksums during import.
+         *
+         * @param validateIndexChecksum the {@code validateIndexChecksum} to set
+         * @return a reference to this ImportOptions
+         */
+        public ImportOptions validateIndexChecksum(boolean validateIndexChecksum)
+        {
+            this.validateIndexChecksum = validateIndexChecksum;
+            return this;
+        }
     }
 
     static String requestURI(String keyspace, String tableName, String uploadId, ImportOptions importOptions)
@@ -205,6 +233,14 @@ private static List<String> selectedOptions(ImportOptions importOptions)
         {
             options.add("copyData=" + importOptions.copyData);
         }
+        if (importOptions.failOnMissingIndex != null)
+        {
+            options.add("failOnMissingIndex=" + importOptions.failOnMissingIndex);
+        }
+        if (importOptions.validateIndexChecksum != null)
+        {
+            options.add("validateIndexChecksum=" + importOptions.validateIndexChecksum);
+        }
 
         return options;
     }

cassandra-analytics-cdc/src/test/java/org/apache/cassandra/cdc/CdcTests.java

@@ -175,7 +175,7 @@ public void testMockedCdc(CassandraVersion version)
                                Partitioner.Murmur3Partitioner,
                                table.udtCreateStmts(bridge.cassandraTypes()),
                                null,
-                               0,
+                               Collections.emptySet(),
                                true);
             SchemaSupplier schemaSupplier = () -> CompletableFuture.completedFuture(ImmutableSet.of(table));
             AtomicReference<byte[]> state = new AtomicReference<>();
@@ -514,13 +514,13 @@ public void testMultiTable(CassandraVersion version)
                            ReplicationFactor.simpleStrategy(1),
                            Partitioner.Murmur3Partitioner,
                            Collections.emptySet(),
-                           null, 0, schema2.withCdc);
+                           null, Collections.emptySet(), schema2.withCdc);
         bridge.buildSchema(cqlTable3.createStatement(),
                            cqlTable3.keyspace(),
                            ReplicationFactor.simpleStrategy(1),
                            Partitioner.Murmur3Partitioner,
                            Collections.emptySet(),
-                           null, 0, schema3.withCdc);
+                           null, Collections.emptySet(), schema3.withCdc);
         int numRows = DEFAULT_NUM_ROWS;
 
         AtomicReference<TestSchema> schema1Holder = new AtomicReference<>();

cassandra-analytics-cdc/src/test/java/org/apache/cassandra/cdc/MicroBatchIteratorTests.java

@@ -390,7 +390,7 @@ private void runTest(CassandraBridge bridge,
                            ReplicationFactor.simpleStrategy(1),
                            Partitioner.Murmur3Partitioner,
                            Collections.emptySet(),
-                           null, 0, schema.withCdc);
+                           null, Collections.emptySet(), schema.withCdc);
         schema.setCassandraVersion(bridge.getVersion());
 
         try

cassandra-analytics-cdc/src/test/java/org/apache/cassandra/cdc/test/CdcTester.java

@@ -21,6 +21,7 @@
 
 import java.nio.file.Path;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -266,7 +267,7 @@ public void run()
         {
             LOGGER.info("Running CDC test testId={} schema='{}' thread={}", testId, cqlTable.fields(), Thread.currentThread().getName());
             Set<String> udtStmts = schema.udts.stream().map(e -> e.createStatement(bridge.cassandraTypes(), schema.keyspace)).collect(Collectors.toSet());
-            bridge.buildSchema(schema.createStatement, schema.keyspace, schema.rf, partitioner, udtStmts, null, 0, true);
+            bridge.buildSchema(schema.createStatement, schema.keyspace, schema.rf, partitioner, udtStmts, null, Collections.emptySet(), true);
             schema.setCassandraVersion(bridge.getVersion());
 
             // write some mutations to CDC CommitLog

cassandra-analytics-cdc/src/test/java/org/apache/cassandra/db/commitlog/BufferingCommitLogReaderTests.java

@@ -89,7 +89,7 @@ public void testReaderSeek(CassandraVersion version)
                            ReplicationFactor.simpleStrategy(1),
                            Partitioner.Murmur3Partitioner,
                            Collections.emptySet(),
-                           null, 0, true);
+                           null, Collections.emptySet(), true);
         int numRows = 1000;
 
         // write some rows to a CommitLog

cassandra-analytics-common/src/main/java/org/apache/cassandra/spark/common/SSTables.java

@@ -22,14 +22,53 @@
 import java.nio.file.Path;
 
 import org.apache.cassandra.bridge.SSTableDescriptor;
+import org.apache.cassandra.spark.data.FileType;
 
 public final class SSTables
 {
+    /**
+     * Suffix identifying the primary SSTable data component, e.g. "-Data.db".
+     * The leading '-' is significant: it excludes SAI per-index components such as
+     * "...+TermsData.db" that also end with "Data.db".
+     */
+    private static final String DATA_COMPONENT_SUFFIX = "-" + FileType.DATA.getFileSuffix();
+
+    /**
+     * Glob matching primary SSTable data components, e.g. "*-Data.db".
+     * Suitable for {@link java.nio.file.Files#newDirectoryStream(Path, String)}.
+     */
+    public static final String DATA_COMPONENT_GLOB = "*" + DATA_COMPONENT_SUFFIX;
+
     private SSTables()
     {
         throw new IllegalStateException(getClass() + " is static utility class and shall not be instantiated");
     }
 
+    /**
+     * Determine whether the given file name is a primary SSTable data component ("<descriptor>-Data.db").
+     * The leading '-' check excludes SAI per-index components such as "...+TermsData.db" which also end with "Data.db".
+     *
+     * @param fileName file name (not a full path)
+     * @return true if the name is a primary data component
+     */
+    public static boolean isDataComponent(String fileName)
+    {
+        return fileName.endsWith(DATA_COMPONENT_SUFFIX);
+    }
+
+    /**
+     * Determine whether the given path is a primary SSTable data component ("<descriptor>-Data.db").
+     *
+     * @param path file path
+     * @return true if the path's file name is a primary data component
+     * @see #isDataComponent(String)
+     */
+    public static boolean isDataComponent(Path path)
+    {
+        Path fileName = path.getFileName();
+        return fileName != null && isDataComponent(fileName.toString());
+    }
+
     /**
      * Get the sstable base name from data file path.
      * For example, the base name of data file '/path/to/table/nb-1-big-Data.db' is 'nb-1-big'

cassandra-analytics-common/src/main/java/org/apache/cassandra/spark/data/CqlTable.java

@@ -57,24 +57,15 @@ public class CqlTable implements Serializable
     private final List<CqlField> staticColumns;
     private final List<CqlField> valueColumns;
     private final transient Map<String, CqlField> columns;
-    private final int indexCount;
-
-    public CqlTable(@NotNull String keyspace,
-                    @NotNull String table,
-                    @NotNull String createStatement,
-                    @NotNull ReplicationFactor replicationFactor,
-                    @NotNull List<CqlField> fields)
-    {
-        this(keyspace, table, createStatement, replicationFactor, fields, Collections.emptySet(), 0);
-    }
+    private final Set<String> indexStatements;
 
     public CqlTable(@NotNull String keyspace,
                     @NotNull String table,
                     @NotNull String createStatement,
                     @NotNull ReplicationFactor replicationFactor,
                     @NotNull List<CqlField> fields,
                     @NotNull Set<CqlField.CqlUdt> udts,
-                    int indexCount)
+                    @NotNull Set<String> indexStatements)
     {
         this.keyspace = keyspace;
         this.table = table;
@@ -87,7 +78,7 @@ public CqlTable(@NotNull String keyspace,
         this.staticColumns = this.fields.stream().filter(CqlField::isStaticColumn).sorted().collect(Collectors.toList());
         this.valueColumns = this.fields.stream().filter(CqlField::isValueColumn).sorted().collect(Collectors.toList());
         this.udts = Collections.unmodifiableSet(udts);
-        this.indexCount = indexCount;
+        this.indexStatements = Collections.unmodifiableSet(indexStatements);
 
         // We use a linked hashmap to guarantee ordering of a 'SELECT * FROM ...'
         this.columns = new LinkedHashMap<>();
@@ -241,9 +232,14 @@ public String createStatement()
         return createStatement;
     }
 
+    public Set<String> indexStatements()
+    {
+        return indexStatements;
+    }
+
     public int indexCount()
     {
-        return indexCount;
+        return indexStatements.size();
     }
 
     /**
@@ -382,8 +378,15 @@ public CqlTable read(Kryo kryo, Input input, Class type)
             {
                 udts.add((CqlField.CqlUdt) CqlField.CqlType.read(input, cassandraTypes));
             }
-            int indexCount = input.readInt();
-            return new CqlTable(keyspace, table, createStatement, replicationFactor, fields, udts, indexCount);
+
+            int numIndexStatements = input.readInt();
+            Set<String> indexStatements = new LinkedHashSet<>(numIndexStatements);
+            for (int idx = 0; idx < numIndexStatements; idx++)
+            {
+                indexStatements.add(input.readString());
+            }
+
+            return new CqlTable(keyspace, table, createStatement, replicationFactor, fields, udts, indexStatements);
         }
 
         @Override
@@ -405,7 +408,13 @@ public void write(Kryo kryo, Output output, CqlTable table)
             {
                 udt.write(output);
             }
-            output.writeInt(table.indexCount());
+
+            Set<String> indexStatements = table.indexStatements();
+            output.writeInt(indexStatements.size());
+            for (String stmt : indexStatements)
+            {
+                output.writeString(stmt);
+            }
         }
     }
 }