Merge branch 'cassandra-5.0' into trunk

CHANGES.txt

@@ -8,6 +8,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:
+ * Add support of vector type to cqlsh COPY command (CASSANDRA-19118)
  * Make CQLSSTableWriter to support building of SAI indexes (CASSANDRA-18714)
  * Append additional JVM options when using JDK17+ (CASSANDRA-19001)
  * Upgrade Python driver to 3.29.0 (CASSANDRA-19245)

pylib/cqlshlib/copyutil.py

@@ -46,7 +46,7 @@ from queue import Queue
 
 from cassandra import OperationTimedOut
 from cassandra.cluster import Cluster, DefaultConnection
-from cassandra.cqltypes import ReversedType, UserType, VarcharType
+from cassandra.cqltypes import ReversedType, UserType, VarcharType, VectorType
 from cassandra.metadata import protect_name, protect_names, protect_value
 from cassandra.policies import RetryPolicy, WhiteListRoundRobinPolicy, DCAwareRoundRobinPolicy, FallthroughRetryPolicy
 from cassandra.query import BatchStatement, BatchType, SimpleStatement, tuple_factory
@@ -2074,6 +2074,12 @@ class ImportConversion(object):
             return ImmutableDict(frozenset((convert_mandatory(ct.subtypes[0], v[0]), convert(ct.subtypes[1], v[1]))
                                  for v in [split(split_format_str % vv, sep=sep) for vv in split(val)]))
 
+        def convert_vector(val, ct=cql_type):
+            string_coordinates = split(val)
+            if len(string_coordinates) != ct.vector_size:
+                raise ParseError("The length of given vector value '%d' is not equal to the vector size from the type definition '%d'" % (len(string_coordinates), ct.vector_size))
+            return [convert_mandatory(ct.subtype, v) for v in string_coordinates]
+
         def convert_user_type(val, ct=cql_type):
             """
             A user type is a dictionary except that we must convert each key into
@@ -2130,6 +2136,7 @@ class ImportConversion(object):
             'map': convert_map,
             'tuple': convert_tuple,
             'frozen': convert_single_subtype,
+            VectorType.typename: convert_vector,
         }
 
         return converters.get(cql_type.typename, convert_unknown)

test/unit/org/apache/cassandra/tools/cqlsh/CqlshTest.java

@@ -18,16 +18,24 @@
 
 package org.apache.cassandra.tools.cqlsh;
 
+import java.io.IOException;
+import java.io.Writer;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+
 import org.junit.BeforeClass;
 import org.junit.Test;
 
 import org.apache.cassandra.cql3.CQLTester;
+import org.apache.cassandra.cql3.UntypedResultSet;
 import org.apache.cassandra.tools.ToolRunner;
 import org.apache.cassandra.tools.ToolRunner.ToolResult;
-import org.hamcrest.CoreMatchers;
 
+import static java.lang.String.format;
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
 
 public class CqlshTest extends CQLTester
 {
@@ -41,7 +49,119 @@ public class CqlshTest extends CQLTester
     public void testKeyspaceRequired()
     {
         ToolResult tool = ToolRunner.invokeCqlsh("SELECT * FROM test");
-        assertThat(tool.getCleanedStderr(), CoreMatchers.containsStringIgnoringCase("No keyspace has been specified"));
+        tool.asserts().errorContains("No keyspace has been specified");
         assertEquals(2, tool.getExitCode());
     }
+
+    @Test
+    public void testCopyFloatVector() throws IOException
+    {
+        assertCopyOfVectorTypeSucceeds("float", 6, new Object[][] {
+            row(1, vector(0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f)),
+            row(2, vector(-0.1f, -0.2f, -0.3f, -0.4f, -0.5f, -0.6f)),
+            row(3, vector(0.9f, 0.8f, 0.7f, 0.6f, 0.5f, 0.4f))
+        });
+
+        assertCopyOfVectorTypeSucceeds("float", 3, new Object[][] {
+            row(1, vector(0.1f, 0.2f, 0.3f)),
+            row(2, vector(-0.4f, -0.5f, -0.6f)),
+            row(3, vector(0.7f, 0.8f, 0.9f))
+        });
+    }
+
+    @Test
+    public void testCopyIntVector() throws IOException
+    {
+        assertCopyOfVectorTypeSucceeds("int", 6, new Object[][] {
+            row(1, vector(1, 2, 3, 4, 5, 6)),
+            row(2, vector(-1, -2, -3, -4, -5, -6)),
+            row(3, vector(9, 8, 7, 6, 5, 4))
+        });
+
+        assertCopyOfVectorTypeSucceeds("int", 3, new Object[][] {
+            row(1, vector(1, 2, 3)),
+            row(2, vector(-4, -5, -6)),
+            row(3, vector(7, 8, 9))
+        });
+    }
+
+    private void assertCopyOfVectorTypeSucceeds(String vectorType, int vectorSize, Object[][] rows) throws IOException
+    {
+        // given a table with a vector column
+        createTable(KEYSPACE, format("CREATE TABLE %%s (id int PRIMARY KEY, embedding_vector vector<%s, %d>)", vectorType, vectorSize));
+        assertTrue("table should be initially empty", execute("SELECT * FROM %s").isEmpty());
+
+        // write the rows into the table
+        for (Object[] row : rows)
+            execute("INSERT INTO %s (id, embedding_vector) VALUES (?, ?)", row);
+
+        // when running COPY TO CSV via cqlsh
+        Path csv = createTempFile("test_copy_to_vector");
+        ToolRunner.ToolResult copyToResult = ToolRunner.invokeCqlsh(format("COPY %s.%s TO '%s'", KEYSPACE, currentTable(), csv.toAbsolutePath()));
+
+        // then all rows should be exported
+        copyToResult.asserts().success();
+        // verify that the exported CSV contains the expected rows
+        assertThat(csv).hasSameTextualContentAs(prepareCSVFile(rows));
+
+        // truncate the table
+        execute("TRUNCATE %s");
+        assertTrue("table should be empty", execute("SELECT * FROM %s").isEmpty());
+
+        // when running COPY FROM via cqlsh
+        ToolRunner.ToolResult copyFromResult = ToolRunner.invokeCqlsh(format("COPY %s.%s FROM '%s'", KEYSPACE, currentTable(), csv.toAbsolutePath()));
+
+        // then all rows should be imported
+        copyFromResult.asserts().success();
+        UntypedResultSet importedRows = execute("SELECT * FROM %s");
+        assertRowsIgnoringOrder(importedRows, rows);
+    }
+
+    @Test
+    public void testCopyOnlyThoseRowsThatMatchVectorTypeSize() throws IOException
+    {
+        // given a table with a vector column and a file containing vector literals
+        createTable(KEYSPACE, "CREATE TABLE %s (id int PRIMARY KEY, embedding_vector vector<int, 6>)");
+        assertTrue("table should be initially empty", execute("SELECT * FROM %s").isEmpty());
+
+        Object[][] rows = {
+            row(1, vector(1, 2, 3, 4, 5, 6)),
+            row(2, vector(1, 2, 3, 4, 5)),
+            row(3, vector(1, 2, 3, 4, 6, 7))
+        };
+
+        Path csv = prepareCSVFile(rows);
+
+        // when running COPY via cqlsh
+        ToolRunner.ToolResult result = ToolRunner.invokeCqlsh(format("COPY %s.%s FROM '%s'", KEYSPACE, currentTable(), csv.toAbsolutePath()));
+
+        // then only rows that match type size should be imported
+        result.asserts().failure();
+        result.asserts().errorContains("The length of given vector value '5' is not equal to the vector size from the type definition '6'");
+        UntypedResultSet importedRows = execute("SELECT * FROM %s");
+        assertRowsIgnoringOrder(importedRows, row(1, vector(1, 2, 3, 4, 5, 6)),
+                                row(3, vector(1, 2, 3, 4, 6, 7)));
+    }
+
+    private static Path prepareCSVFile(Object[][] rows) throws IOException
+    {
+        Path csv = createTempFile("test_copy_from_vector");
+
+        try (Writer out = Files.newBufferedWriter(csv, StandardCharsets.UTF_8))
+        {
+            for (Object[] row : rows)
+            {
+                out.write(String.format("%s,\"%s\"\n", row[0], row[1]));
+            }
+        }
+
+        return csv;
+    }
+
+    private static Path createTempFile(String prefix) throws IOException
+    {
+        Path csv = Files.createTempFile(prefix, ".csv");
+        csv.toFile().deleteOnExit();
+        return csv;
+    }
 }