From 4bebdc3886facb4c55eaba1467f4329a5b22758d Mon Sep 17 00:00:00 2001
From: Madhavan Sridharan <madhavan.sridharan@datastax.com>
Date: Wed, 20 Sep 2023 14:19:46 -0400
Subject: [PATCH] Initil draft to support vectors

---
 README.md                                     |  1 +
 pom.xml                                       | 44 ++++++++++++++++++-
 .../java/com/datastax/cdm/data/CqlData.java   |  7 +++
 3 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 632867e5..863f704f 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,7 @@
 [![License Apache2](https://img.shields.io/hexpm/l/plug.svg)](http://www.apache.org/licenses/LICENSE-2.0)
 [![Star on Github](https://img.shields.io/github/stars/datastax/cassandra-data-migrator.svg?style=social)](https://github.com/datastax/cassandra-data-migrator/stargazers)
 ![GitHub release (with filter)](https://img.shields.io/github/v/release/datastax/cassandra-data-migrator?label=latest%20release&color=green&link=!%5BGitHub%20release%20(with%20filter)%5D(https%3A%2F%2Fimg.shields.io%2Fgithub%2Fv%2Frelease%2Fdatastax%2Fcassandra-data-migrator%3Flabel%3Dlatest%2520release%26color%3Dgreen))
+![Docker Pulls](https://img.shields.io/docker/pulls/datastax/cassandra-data-migrator)
 
 # cassandra-data-migrator
 
diff --git a/pom.xml b/pom.xml
index 4707e99b..fd27ed53 100644
--- a/pom.xml
+++ b/pom.xml
@@ -13,9 +13,10 @@
     <spark.version>3.3.1</spark.version>
     <scalatest.version>3.2.12</scalatest.version>
     <connector.version>3.2.0</connector.version>
-    <cassandra.version>3.11.13</cassandra.version>
+    <cassandra.version>3.11.16</cassandra.version>
     <junit.version>5.9.1</junit.version>
     <mockito.version>4.11.0</mockito.version>
+	<java-driver.version>4.17.0</java-driver.version>
   </properties>
 
   <distributionManagement>
@@ -74,7 +75,48 @@
       <groupId>com.datastax.spark</groupId>
       <artifactId>spark-cassandra-connector_${scala.main.version}</artifactId>
       <version>${connector.version}</version>
+	  <exclusions>
+		<exclusion>
+		  <groupId>com.datastax.oss</groupId>
+		  <artifactId>java-driver-core-shaded</artifactId>
+		</exclusion>
+		<exclusion>
+		  <groupId>com.datastax.oss</groupId>
+		  <artifactId>java-driver-mapper-runtime</artifactId>
+		</exclusion>
+		<exclusion>
+		  <groupId>com.datastax.oss</groupId>
+		  <artifactId>java-driver-query-builder</artifactId>
+		</exclusion>
+		<exclusion>
+		  <groupId>com.datastax.oss</groupId>
+		  <artifactId>java-driver-shaded-guava</artifactId>
+		</exclusion>
+	  </exclusions>
     </dependency>
+
+	<!-- Adding Java Driver latest version explicitly -->
+	<dependency>
+	  <groupId>com.datastax.oss</groupId>
+	  <artifactId>java-driver-core-shaded</artifactId>
+	  <version>${java-driver.version}</version>
+	</dependency>
+	<dependency>
+	  <groupId>com.datastax.oss</groupId>
+	  <artifactId>java-driver-mapper-runtime</artifactId>
+	  <version>${java-driver.version}</version>
+	</dependency>
+	<dependency>
+	  <groupId>com.datastax.oss</groupId>
+	  <artifactId>java-driver-query-builder</artifactId>
+	  <version>${java-driver.version}</version>
+	</dependency>
+	<!-- <dependency>
+	  <groupId>com.datastax.oss</groupId>
+	  <artifactId>java-driver-shaded-guava</artifactId>
+	  <version>${java-driver.version}</version>
+	</dependency> -->
+	
     <dependency>
       <groupId>com.github.jnr</groupId>
       <artifactId>jnr-posix</artifactId>
diff --git a/src/main/java/com/datastax/cdm/data/CqlData.java b/src/main/java/com/datastax/cdm/data/CqlData.java
index 5e6db45c..101764fc 100644
--- a/src/main/java/com/datastax/cdm/data/CqlData.java
+++ b/src/main/java/com/datastax/cdm/data/CqlData.java
@@ -13,6 +13,7 @@ public enum Type {
         SET,
         MAP,
         TUPLE,
+        VECTOR,
         UNKNOWN
     }
 
@@ -47,6 +48,7 @@ public static Type toType(DataType dataType) {
         if (dataType instanceof MapType) return Type.MAP;
         if (dataType instanceof TupleType) return Type.TUPLE;
         if (dataType instanceof UserDefinedType) return Type.UDT;
+        if (dataType instanceof VectorType) return Type.VECTOR;
         throw new RuntimeException("Unsupported data type: " + dataType);
     }
 
@@ -60,6 +62,7 @@ public static boolean isCollection(DataType dataType) {
         if (dataType instanceof SetType) return true;
         if (dataType instanceof MapType) return true;
         if (dataType instanceof TupleType) return true;
+        if (dataType instanceof VectorType) return true;
         return false;
     }
 
@@ -70,6 +73,7 @@ public static boolean isFrozen(DataType dataType) {
         if (dataType instanceof SetType) return ((SetType) dataType).isFrozen();
         if (dataType instanceof MapType) return ((MapType) dataType).isFrozen();
         if (dataType instanceof TupleType) return dataType.asCql(true, false).toLowerCase().contains("frozen<");
+        // vector CQL data type doesn't support frozen
         return false;
     }
 
@@ -81,6 +85,7 @@ public static Class getBindClass(DataType dataType) {
         if (dataType instanceof MapType) return java.util.Map.class;
         if (dataType instanceof UserDefinedType) return com.datastax.oss.driver.api.core.data.UdtValue.class;
         if (dataType instanceof TupleType) return com.datastax.oss.driver.api.core.data.TupleValue.class;
+        if (dataType instanceof VectorType) return com.datastax.oss.driver.api.core.data.CqlVector.class;
 
         throw new IllegalArgumentException("Unsupported data type: " + dataType);
     }
@@ -98,6 +103,8 @@ public static List<DataType> extractDataTypesFromCollection(DataType dataType) {
                 return Arrays.asList(((MapType) dataType).getKeyType(), ((MapType) dataType).getValueType());
             case TUPLE:
                 return ((TupleType) dataType).getComponentTypes();
+            case VECTOR:
+            	return Collections.singletonList(((VectorType) dataType).getElementType());
             default:
                 return null;
         }