Merge branch 'cassandra-5.0' into trunk

* cassandra-5.0:
  CASSANDRA-20609 - CQLSSTableWriter should support setting the format (5.0)

CHANGES.txt

@@ -186,6 +186,7 @@
  * Add the ability to disable bulk loading of SSTables (CASSANDRA-18781)
  * Clean up obsolete functions and simplify cql_version handling in cqlsh (CASSANDRA-18787)
 Merged from 5.0:
+ * CQLSSTableWriter supports setting the format (BTI or Big) (CASSANDRA-20609)
  * Don't allocate in ThreadLocalReadAheadBuffer#close() (CASSANDRA-20551)
  * Ensure RowFilter#isMutableIntersection() properly evaluates numeric ranges on a single column (CASSANDRA-20566)
  * SAI marks an index as non-empty when a partial partition/row modifications is flushed due to repair (CASSANDRA-20567)

src/java/org/apache/cassandra/io/sstable/CQLSSTableWriter.java

@@ -653,6 +653,18 @@ public class CQLSSTableWriter implements Closeable
             return this;
         }
 
+        /**
+         * Specifies the SSTable format this CQLSSTableWriter instance should use for writing.
+         *
+         * @param format The format to use
+         * @return this builder
+         */
+        public Builder withFormat(SSTableFormat<?, ?> format)
+        {
+            this.format = format;
+            return this;
+        }
+
         public CQLSSTableWriter build()
         {
             if (directory == null)

test/unit/org/apache/cassandra/io/sstable/CQLSSTableWriterTest.java

@@ -20,6 +20,8 @@ package org.apache.cassandra.io.sstable;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -35,13 +37,11 @@ import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiPredicate;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-
-import org.apache.cassandra.cql3.constraints.ConstraintViolationException;
-
 import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Rule;
@@ -49,8 +49,10 @@ import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
 
 import com.datastax.driver.core.utils.UUIDs;
+import org.apache.cassandra.Util;
 import org.apache.cassandra.cql3.QueryProcessor;
 import org.apache.cassandra.cql3.UntypedResultSet;
+import org.apache.cassandra.cql3.constraints.ConstraintViolationException;
 import org.apache.cassandra.cql3.functions.types.DataType;
 import org.apache.cassandra.cql3.functions.types.LocalDate;
 import org.apache.cassandra.cql3.functions.types.TypeCodec;
@@ -58,14 +60,17 @@ import org.apache.cassandra.cql3.functions.types.UDTValue;
 import org.apache.cassandra.cql3.functions.types.UserType;
 import org.apache.cassandra.db.marshal.FloatType;
 import org.apache.cassandra.db.marshal.UTF8Type;
+import org.apache.cassandra.dht.ByteOrderedPartitioner;
 import org.apache.cassandra.dht.Murmur3Partitioner;
 import org.apache.cassandra.dht.Range;
 import org.apache.cassandra.dht.Token;
 import org.apache.cassandra.exceptions.InvalidRequestException;
 import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
 import org.apache.cassandra.index.sai.utils.IndexIdentifier;
+import org.apache.cassandra.io.sstable.format.SSTableFormat;
 import org.apache.cassandra.io.sstable.format.SSTableReader;
 import org.apache.cassandra.io.sstable.format.big.BigFormat;
+import org.apache.cassandra.io.sstable.format.bti.BtiFormat;
 import org.apache.cassandra.io.util.File;
 import org.apache.cassandra.io.util.PathUtils;
 import org.apache.cassandra.locator.RangesAtEndpoint;
@@ -82,6 +87,7 @@ import org.apache.cassandra.utils.OutputHandler;
 import org.assertj.core.api.Assertions;
 
 import static org.apache.cassandra.utils.Clock.Global.currentTimeMillis;
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -114,56 +120,85 @@ public abstract class CQLSSTableWriterTest
     }
 
     @Test
-    public void testUnsortedWriter() throws Exception
+    public void testUnsortedWriterBig() throws Exception
     {
-        String schema = "CREATE TABLE " + qualifiedTable + " ("
+        BigFormat format = BigFormat.getInstance();
-                        + "  k int PRIMARY KEY,"
+        testWritingSstableWithFormat(format);
-                        + "  v1 text,"
+    }
-                        + "  v2 int"
-                        + ")";
-        String insert = "INSERT INTO " + qualifiedTable + " (k, v1, v2) VALUES (?, ?, ?)";
-        CQLSSTableWriter writer = CQLSSTableWriter.builder()
-                                                  .inDirectory(dataDir)
-                                                  .forTable(schema)
-                                                  .using(insert).build();
 
-        writer.addRow(0, "test1", 24);
+    @Test
-        writer.addRow(1, "test2", 44);
+    public void testUnsortedWriterBti() throws Exception
-        writer.addRow(2, "test3", 42);
+    {
-        writer.addRow(ImmutableMap.<String, Object>of("k", 3, "v2", 12));
+        SSTableFormat<?, ?> btiFormat = new BtiFormat.BtiFormatFactory().getInstance(Collections.emptyMap());
+        testWritingSstableWithFormat(btiFormat);
+    }
 
-        writer.close();
+    private void testWritingSstableWithFormat(SSTableFormat<?, ?> format) throws Exception
-
+    {
-        loadSSTables(dataDir, keyspace, table);
+        try (AutoCloseable ignored = Util.switchPartitioner(ByteOrderedPartitioner.instance))
-
-        if (verifyDataAfterLoading)
         {
-            UntypedResultSet rs = QueryProcessor.executeInternal("SELECT * FROM " + qualifiedTable);
+            String schema = "CREATE TABLE " + qualifiedTable + " ("
-            assertEquals(4, rs.size());
+                            + "  k int PRIMARY KEY,"
+                            + "  v1 text,"
+                            + "  v2 int"
+                            + ")";
+            String insert = "INSERT INTO " + qualifiedTable + " (k, v1, v2) VALUES (?, ?, ?)";
+            CQLSSTableWriter writer = CQLSSTableWriter.builder()
+                                                      .inDirectory(dataDir)
+                                                      .forTable(schema)
+                                                      .withFormat(format)
+                                                      .using(insert).build();
 
-            Iterator<UntypedResultSet.Row> iter = rs.iterator();
+            writer.addRow(0, "test1", 24);
-            UntypedResultSet.Row row;
+            writer.addRow(1, "test2", 44);
+            writer.addRow(2, "test3", 42);
+            writer.addRow(ImmutableMap.<String, Object>of("k", 3, "v2", 12));
 
-            row = iter.next();
+            writer.close();
-            assertEquals(0, row.getInt("k"));
-            assertEquals("test1", row.getString("v1"));
-            assertEquals(24, row.getInt("v2"));
 
-            row = iter.next();
+            validateFilesAreInFormat(format);
-            assertEquals(1, row.getInt("k"));
+            loadSSTables(dataDir, keyspace, table);
-            assertEquals("test2", row.getString("v1"));
-            //assertFalse(row.has("v2"));
-            assertEquals(44, row.getInt("v2"));
 
-            row = iter.next();
+            if (verifyDataAfterLoading)
-            assertEquals(2, row.getInt("k"));
+            {
-            assertEquals("test3", row.getString("v1"));
+                UntypedResultSet rs = QueryProcessor.executeInternal("SELECT * FROM " + qualifiedTable);
-            assertEquals(42, row.getInt("v2"));
+                assertEquals(4, rs.size());
 
-            row = iter.next();
+                Iterator<UntypedResultSet.Row> iter = rs.iterator();
-            assertEquals(3, row.getInt("k"));
+                UntypedResultSet.Row row;
-            assertEquals(null, row.getBytes("v1")); // Using getBytes because we know it won't NPE
+
-            assertEquals(12, row.getInt("v2"));
+                row = iter.next();
+                assertEquals(0, row.getInt("k"));
+                assertEquals("test1", row.getString("v1"));
+                assertEquals(24, row.getInt("v2"));
+
+                row = iter.next();
+                assertEquals(1, row.getInt("k"));
+                assertEquals("test2", row.getString("v1"));
+                assertEquals(44, row.getInt("v2"));
+
+                row = iter.next();
+                assertEquals(2, row.getInt("k"));
+                assertEquals("test3", row.getString("v1"));
+                assertEquals(42, row.getInt("v2"));
+
+                row = iter.next();
+                assertEquals(3, row.getInt("k"));
+                assertFalse(row.has("v1"));
+                assertEquals(12, row.getInt("v2"));
+            }
+        }
+    }
+
+    private void validateFilesAreInFormat(SSTableFormat<?, ?> format) throws IOException
+    {
+        try (Stream<Path> dataFilePaths = Files.list(dataDir.toPath()).filter(p -> p.toString().endsWith("Data.db")))
+        {
+            dataFilePaths.forEach(dataFilePath -> {
+                File dataFile = new File(dataFilePath.toFile());
+                Descriptor descriptor = Descriptor.fromFile(dataFile);
+                assertEquals(format, descriptor.version.format);
+            });
         }
     }
 
@@ -1611,9 +1646,9 @@ public abstract class CQLSSTableWriterTest
             {
                 assertEquals(cnt, row.getInt("k"));
                 List<Float> vector = row.getVector("v1", FloatType.instance, 5);
-                Assertions.assertThat(vector).hasSize(5);
+                assertThat(vector).hasSize(5);
                 final float floatCount = (float)cnt;
-                Assertions.assertThat(vector).allMatch(val -> val == floatCount);
+                assertThat(vector).allMatch(val -> val == floatCount);
                 cnt++;
             }
         }