CASSANALYTICS-26: Support vector data type

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

@@ -37,6 +37,7 @@
 public abstract class CassandraTypes
 {
     public static final Pattern COLLECTION_PATTERN = Pattern.compile("^(set|list|map|tuple)<(.+)>$", Pattern.CASE_INSENSITIVE);
+    public static final Pattern VECTOR_PATTERN = Pattern.compile("^(vector)<(.+),(.*)>$", Pattern.CASE_INSENSITIVE);
     public static final Pattern FROZEN_PATTERN = Pattern.compile("^frozen<(.*)>$", Pattern.CASE_INSENSITIVE);
 
     private final UDTs udts = new UDTs();
@@ -133,6 +134,8 @@ public List<CqlField.NativeType> supportedTypes()
 
     public abstract CqlField.CqlList list(CqlField.CqlType type);
 
+    public abstract CqlField.CqlVector vector(CqlField.CqlType type, int dimentions);
+
     public abstract CqlField.CqlSet set(CqlField.CqlType type);
 
     public abstract CqlField.CqlMap map(CqlField.CqlType keyType, CqlField.CqlType valueType);
@@ -189,6 +192,14 @@ public CqlField.CqlType parseType(String type, Map<String, CqlField.CqlUdt> udts
                                                                 .map(collectionType -> parseType(collectionType, udts))
                                                                 .toArray(CqlField.CqlType[]::new));
         }
+        Matcher vectorMatcher = VECTOR_PATTERN.matcher(type);
+        if (vectorMatcher.find())
+        {
+            // CQL vector
+            String subType = vectorMatcher.group(2);
+            int dimensions = Integer.parseInt(vectorMatcher.group(3).trim());
+            return vector(parseType(subType, udts), dimensions);
+        }
         Matcher frozenMatcher = FROZEN_PATTERN.matcher(type);
         if (frozenMatcher.find())
         {

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

@@ -67,7 +67,7 @@ public interface CqlType extends Serializable
     {
         enum InternalType
         {
-            NativeCql, Set, List, Map, Frozen, Udt, Tuple;
+            NativeCql, Set, List, Map, Frozen, Udt, Tuple, Vector;
 
             public static InternalType fromString(String name)
             {
@@ -77,6 +77,8 @@ public static InternalType fromString(String name)
                         return Set;
                     case "list":
                         return List;
+                    case "vector":
+                        return Vector;
                     case "map":
                         return Map;
                     case "tuple":
@@ -237,6 +239,10 @@ public interface CqlList extends CqlCollection
     {
     }
 
+    public interface CqlVector extends CqlCollection
+    {
+    }
+
     public interface CqlTuple extends CqlCollection
     {
         ByteBuffer serializeTuple(Object[] values);

cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/endtoend/DataTypeTests.java

@@ -33,17 +33,21 @@
 import org.junit.jupiter.params.provider.MethodSource;
 
 import org.apache.cassandra.bridge.CassandraBridge;
+import org.apache.cassandra.bridge.CassandraVersion;
 import org.apache.cassandra.spark.TestUtils;
 import org.apache.cassandra.spark.Tester;
 import org.apache.cassandra.spark.data.CqlField;
 import org.apache.cassandra.spark.utils.RandomUtils;
 import org.apache.cassandra.spark.utils.test.TestSchema;
 import org.apache.spark.sql.Row;
+import org.quicktheories.core.Gen;
 import scala.collection.mutable.AbstractSeq;
 
 import static org.apache.cassandra.spark.utils.ScalaConversionUtils.mutableSeqAsJavaList;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assumptions.assumeThat;
 import static org.quicktheories.QuickTheory.qt;
+import static org.quicktheories.generators.SourceDSL.arbitrary;
 
 @Tag("Sequential")
 public class DataTypeTests
@@ -103,6 +107,104 @@ public void testSet(CassandraBridge bridge)
             );
     }
 
+    @ParameterizedTest
+    @MethodSource("org.apache.cassandra.bridge.VersionRunner#bridges")
+    public void testVector(CassandraBridge bridge)
+    {
+        assumeThat(bridge.getVersion().versionNumber()).isGreaterThanOrEqualTo(CassandraVersion.FIVEZERO.versionNumber());
+        qt().forAll(vectorSupportedTypes(bridge))
+            .checkAssert(type ->
+                         Tester.builder(TestSchema.builder(bridge)
+                                                  .withPartitionKey("pk", bridge.uuid())
+                                                  .withColumn("a", bridge.vector(type, 10)))
+                               .withExpectedRowCountPerSSTable(Tester.DEFAULT_NUM_ROWS)
+                               .run(bridge.getVersion())
+            );
+    }
+
+    @ParameterizedTest
+    @MethodSource("org.apache.cassandra.bridge.VersionRunner#bridges")
+    public void testVectorVector(CassandraBridge bridge)
+    {
+        assumeThat(bridge.getVersion().versionNumber()).isGreaterThanOrEqualTo(CassandraVersion.FIVEZERO.versionNumber());
+        qt().forAll(vectorSupportedTypes(bridge))
+            .checkAssert(type ->
+                         Tester.builder(TestSchema.builder(bridge)
+                                                  .withPartitionKey("pk", bridge.uuid())
+                                                  .withColumn("a", bridge.vector(bridge.vector(type, 2), 5)))
+                               .withExpectedRowCountPerSSTable(Tester.DEFAULT_NUM_ROWS)
+                               .run(bridge.getVersion())
+            );
+    }
+
+    @ParameterizedTest
+    @MethodSource("org.apache.cassandra.bridge.VersionRunner#bridges")
+    public void testVectorList(CassandraBridge bridge)
+    {
+        assumeThat(bridge.getVersion().versionNumber()).isGreaterThanOrEqualTo(CassandraVersion.FIVEZERO.versionNumber());
+        qt().forAll(vectorSupportedTypes(bridge))
+            .checkAssert(type ->
+                         Tester.builder(TestSchema.builder(bridge)
+                                                  .withPartitionKey("pk", bridge.uuid())
+                                                  .withColumn("a", bridge.vector(bridge.list(type), 3)))
+                               .withExpectedRowCountPerSSTable(Tester.DEFAULT_NUM_ROWS)
+                               .run(bridge.getVersion())
+            );
+    }
+
+    @ParameterizedTest
+    @MethodSource("org.apache.cassandra.bridge.VersionRunner#bridges")
+    public void testVectorUDT(CassandraBridge bridge)
+    {
+        // pk -> a vector<frozen<nested_udt<x int, y type, z int>>, 10>
+        // Test vector of UDTs
+        assumeThat(bridge.getVersion().versionNumber()).isGreaterThanOrEqualTo(CassandraVersion.FIVEZERO.versionNumber());
+        qt().withExamples(10)
+            .forAll(vectorSupportedTypes(bridge))
+            .checkAssert(type ->
+                         Tester.builder(TestSchema.builder(bridge)
+                                                  .withPartitionKey("pk", bridge.uuid())
+                                                  .withColumn("a", bridge.vector(
+                                                                                bridge.udt("keyspace", "nested_udt")
+                                                                                      .withField("x", bridge.aInt())
+                                                                                      .withField("y", type)
+                                                                                      .withField("z", bridge.aInt())
+                                                                                      .build().frozen(),
+                                                                                10)))
+                               .run(bridge.getVersion())
+            );
+    }
+
+    @ParameterizedTest
+    @MethodSource("org.apache.cassandra.bridge.VersionRunner#bridges")
+    public void testVectorTuple(CassandraBridge bridge)
+    {
+        // pk -> a vector<frozen<tuple<type, float, text>>, 7>
+        // Test tuple nested within vector
+        assumeThat(bridge.getVersion().versionNumber()).isGreaterThanOrEqualTo(CassandraVersion.FIVEZERO.versionNumber());
+        qt().withExamples(10)
+            .forAll(vectorSupportedTypes(bridge))
+            .checkAssert(type ->
+                         Tester.builder(TestSchema.builder(bridge)
+                                                  .withPartitionKey("pk", bridge.uuid())
+                                                  .withColumn("a", bridge.vector(bridge.tuple(type,
+                                                                                           bridge.aFloat(),
+                                                                                           bridge.text()).frozen(), 7)))
+                               .run(bridge.getVersion())
+            );
+    }
+
+    private static Gen<CqlField.NativeType> vectorSupportedTypes(CassandraBridge bridge)
+    {
+        // TODO: Vector of durations fail, because we cannot replace DurationSerializer with
+        //  AnalyticsDurationSerializer across all serializers used by VectorType.
+        // TODO: Fix for CASSANDRA-20979 required.
+        List<CqlField.NativeType> supportedTypes = bridge.supportedTypes().stream()
+                                                         .filter(t -> !t.equals(bridge.date()) && !t.equals(bridge.time()) && !t.equals(bridge.duration()))
+                                                         .collect(Collectors.toList());
+        return arbitrary().pick(supportedTypes);
+    }
+
     @ParameterizedTest
     @MethodSource("org.apache.cassandra.bridge.VersionRunner#bridges")
     public void testList(CassandraBridge bridge)

cassandra-analytics-spark-four-zero-converter/src/main/java/org/apache/cassandra/spark/data/converter/SparkSqlTypeConverterImplementation.java

@@ -166,6 +166,10 @@ protected static SparkType getOrThrow(CqlField.CqlType cqlType)
             {
                 return new SparkSet(INSTANCE, (CqlField.CqlSet) cqlType);
             }
+            else if (cqlType instanceof CqlField.CqlVector)
+            {
+                return new SparkList(INSTANCE, (CqlField.CqlVector) cqlType);
+            }
             else if (cqlType instanceof CqlField.CqlList)
             {
                 return new SparkList(INSTANCE, (CqlField.CqlList) cqlType);

cassandra-bridge/src/main/java/org/apache/cassandra/bridge/CassandraBridge.java

@@ -316,6 +316,11 @@ public CqlField.CqlList list(CqlField.CqlType type)
         return cassandraTypes().list(type);
     }
 
+    public CqlField.CqlVector vector(CqlField.CqlType type, int dimensions)
+    {
+        return cassandraTypes().vector(type, dimensions);
+    }
+
     public CqlField.CqlSet set(CqlField.CqlType type)
     {
         return cassandraTypes().set(type);

cassandra-five-zero-bridge/src/test/java/org/apache/cassandra/spark/data/converter/types/VectorTypeTests.java

@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.cassandra.spark.data.converter.types;
+
+import java.util.List;
+import java.util.Set;
+
+import org.junit.jupiter.api.Test;
+
+import org.apache.cassandra.bridge.CassandraBridgeImplementation;
+import org.apache.cassandra.spark.data.complex.CqlList;
+import org.apache.cassandra.spark.data.complex.CqlVector;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class VectorTypeTests
+{
+    private static final CassandraBridgeImplementation BRIDGE = new CassandraBridgeImplementation();
+
+    @Test
+    public void testSimpleTypeConversion()
+    {
+        CqlVector cqlVector = new CqlVector(org.apache.cassandra.spark.data.types.Float.INSTANCE, 3);
+        Object cqlWriterObj = cqlVector.convertForCqlWriter(List.of(3.14f, 0.0f, -1f), BRIDGE.getVersion(), false);
+        assertThat(cqlWriterObj).isInstanceOf(List.class);
+        List<Float> cqlWriterList = (List<Float>) cqlWriterObj;
+        assertThat(cqlWriterList).containsExactly(3.14f, 0.0f, -1f);
+    }
+
+    @Test
+    public void testComplexTypeConversion()
+    {
+        CqlVector cqlVector = new CqlVector(CqlList.set(org.apache.cassandra.spark.data.types.Float.INSTANCE), 3);
+        Object cqlWriterObj = cqlVector.convertForCqlWriter(List.of(Set.of(3.14f, 0f), Set.of(1f), Set.of()), BRIDGE.getVersion(), false);
+        assertThat(cqlWriterObj).isInstanceOf(List.class);
+        List<? extends Set<Float>> cqlWriterList = (List<? extends Set<Float>>) cqlWriterObj;
+        assertThat(cqlWriterList).hasSize(3);
+        assertThat(cqlWriterList.get(0)).containsExactlyInAnyOrder(3.14f, 0f);
+        assertThat(cqlWriterList.get(1)).containsExactly(1f);
+        assertThat(cqlWriterList.get(2)).isEmpty();
+    }
+}

cassandra-five-zero-types/src/main/java/org/apache/cassandra/bridge/CassandraTypesImplementation.java

@@ -32,6 +32,8 @@
 import org.apache.cassandra.dht.Murmur3Partitioner;
 import org.apache.cassandra.locator.SimpleSnitch;
 import org.apache.cassandra.security.EncryptionContext;
+import org.apache.cassandra.spark.data.CqlField;
+import org.apache.cassandra.spark.data.complex.CqlVector;
 
 public class CassandraTypesImplementation extends AbstractCassandraTypes
 {
@@ -88,4 +90,10 @@ protected static void setupCommitLogConfigs(Path path)
         DatabaseDescriptor.getRawConfig().commitlog_total_space = new DataStorageSpec.IntMebibytesBound(1024);
         DatabaseDescriptor.setCommitLogSegmentMgrProvider(commitLog -> new CommitLogSegmentManagerStandard(commitLog, commitLogPath.toString()));
     }
+
+    @Override
+    public CqlField.CqlVector vector(CqlField.CqlType type, int dimensions)
+    {
+        return new CqlVector(type, dimensions);
+    }
 }