NODEJS-655: Initial cut at vector support (#407)

Author: absurdfarce

doc/features/datatypes/README.md

@@ -29,6 +29,7 @@ tuple|[Tuple](tuples)
 uuid|[Uuid](uuids)
 varchar|String
 varint|[Integer](numerical)
+vector|[Float32Array](collections)
 
 ## Encoding data 
 

doc/features/datatypes/collections/README.md

@@ -69,4 +69,23 @@ client.execute('SELECT map_val FROM tbl')
   .then(function (result) {
     console.log(result.rows[0]['map_val'] instanceof Map); // true
   });
-```
+```
+
+### Vector
+
+As of version 4.7.0 the driver also includes support for the vector type available in Cassandra 5.0.  Vectors are represented as instances of
+the [Float32Array] class.  For example, to create and write to a vector with three dimensions you might do something like the following:
+
+```javascript
+await c.connect()
+  .then(() => c.execute("drop keyspace if exists test"))
+  .then(() => c.execute("create KEYSPACE test WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}"))
+  .then(() => c.execute("create table test.foo(i varint primary key, j vector<float,3>)"))
+  .then(() => c.execute("create custom index ann_index on test.foo(j) using 'StorageAttachedIndex'"))
+
+  // Base inserts using simple and prepared statements
+  .then(() => c.execute(`insert into test.foo (i, j) values (?, ?)`, [cassandra.types.Integer.fromInt(1), new Float32Array([8, 2.3, 58])]))
+  .then(() => c.execute(`insert into test.foo (i, j) values (?, ?)`, [cassandra.types.Integer.fromInt(5), new Float32Array([23, 18, 3.9])], {prepare: true}));
+```
+
+[Float32Array]: https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Float32Array

lib/encoder.js

@@ -69,7 +69,8 @@ const cqlNames = Object.freeze({
   map: 'map',
   tuple: 'tuple',
   empty: 'empty',
-  duration: 'duration'
+  duration: 'duration',
+  vector: 'vector'
 });
 const singleTypeNames = Object.freeze({
   'org.apache.cassandra.db.marshal.UTF8Type':           dataTypes.varchar,
@@ -93,6 +94,7 @@ const singleTypeNames = Object.freeze({
   'org.apache.cassandra.db.marshal.IntegerType':        dataTypes.varint,
   'org.apache.cassandra.db.marshal.CounterColumnType':  dataTypes.counter
 });
+const singleTypeNamesByDataType = invertObject(singleTypeNames);
 const singleFqTypeNamesLength = Object.keys(singleTypeNames).reduce(function (previous, current) {
   return current.length > previous ? current.length : previous;
 }, 0);
@@ -102,7 +104,8 @@ const customTypeNames = Object.freeze({
   lineString: 'org.apache.cassandra.db.marshal.LineStringType',
   point: 'org.apache.cassandra.db.marshal.PointType',
   polygon: 'org.apache.cassandra.db.marshal.PolygonType',
-  dateRange: 'org.apache.cassandra.db.marshal.DateRangeType'
+  dateRange: 'org.apache.cassandra.db.marshal.DateRangeType',
+  vector: 'org.apache.cassandra.db.marshal.VectorType'
 });
 
 const nullValueBuffer = utils.allocBufferFromArray([255, 255, 255, 255]);
@@ -186,6 +189,16 @@ function defineInstanceMembers() {
     return this.handleBuffer(bytes);
   };
   this.decodeCustom = function (bytes, typeName) {
+
+    // Make sure we actually have something to process in typeName before we go any further
+    if (!typeName || typeName.length === 0) {
+      return this.handleBuffer(bytes);
+    }
+
+    // Special handling for vector custom types (since they have args)
+    if (typeName.startsWith(customTypeNames.vector)) {
+      return this.decodeVector(bytes, this.parseVectorTypeArgs(typeName, customTypeNames.vector, this.parseFqTypeName));
+    }
     const handler = customDecoders[typeName];
     if (handler) {
       return handler.call(this, bytes);
@@ -695,12 +708,17 @@ function defineInstanceMembers() {
     }
     return value;
   };
-  this.encodeCustom = function (value, name) {
-    const handler = customEncoders[name];
+  this.encodeCustom = function (value, customTypeName) {
+
+    // Special handling for vector custom types (since they have args)
+    if (customTypeName.startsWith(customTypeNames.vector)) {
+      return this.encodeVector(value, this.parseVectorTypeArgs(customTypeName, customTypeNames.vector, this.parseFqTypeName));
+    }
+    const handler = customEncoders[customTypeName];
     if (handler) {
       return handler.call(this, value);
     }
-    throw new TypeError('No encoding handler found for type ' + name);
+    throw new TypeError('No encoding handler found for type ' + customTypeName);
   };
   /**
    * @param {Boolean} value
@@ -890,6 +908,92 @@ function defineInstanceMembers() {
     return Buffer.concat(parts, totalLength);
   };
 
+  this.decodeVector = function(buffer, params) {
+    const subtype = params["subtype"];
+    const dimensions = params["dimensions"];
+    const elemLength = 4; // TODO: figure this out based on the subtype
+    const expectedLength = buffer.length / elemLength;
+    if ((elemLength * dimensions) !== buffer.length) {
+      throw new TypeError(`Expected buffer of subtype ${subtype} with dimensions ${dimensions} to be of size ${expectedLength}, observed size ${buffer.length}`);
+    }
+    const rv = [];
+    let offset = 0;
+    for (let i = 0; i < dimensions; i++) {
+      offset = i * elemLength;
+      rv[i] = this.decode(buffer.slice(offset, offset + elemLength), subtype);
+    }
+    return new Float32Array(rv);
+  };
+
+  /**
+   * @param {CqlVector} value
+   * @param {Object} params
+   */
+  this.encodeVector = function(value, params) {
+
+    // Evaluate params to encodeVector(), returning the computed subtype
+    function evalParams() {
+
+      if (!(value instanceof Float32Array)) {
+        throw new TypeError("Driver only supports vectors of 4 byte floating point values");
+      }
+
+      // Perform client-side validation iff we were actually supplied with meaningful type info.  In practice
+      // this will only occur when using prepared statements.
+      if (params.hasOwnProperty("subtype") && params.hasOwnProperty("dimensions")) {
+
+        const subtype = params["subtype"];
+        const dimensions = params["dimensions"];
+        if (value.length !== dimensions) {
+          throw new TypeError(`Expected vector with ${dimensions} dimensions, observed size of ${value.length}`);
+        }
+        if (subtype.code !== dataTypes.float) {
+          throw new TypeError("Driver only supports vectors of 4 byte floating point values");
+        }
+        return subtype;
+      }
+
+      return { code: dataTypes.float };
+    }
+
+    if (!Encoder.isTypedArray(value)) {
+      throw new TypeError('Expected TypedArray subclass, obtained ' + util.inspect(value));
+    }
+    if (value.length === 0) {
+      throw new TypeError("Cannot encode empty array as vector");
+    }
+
+    const subtype = evalParams();
+
+    // TypedArrays are _not_ JS arrays so explicitly convert them here before trying to write them
+    // into a buffer
+    const elems = [];
+    for (const elem of value) {
+      elems.push(this.encode(elem, subtype));
+    }
+    return Buffer.concat(elems);
+  };
+
+  /**
+   * Extract the (typed) arguments from a vector type
+   *
+   * @param {String} typeName
+   * @param {String} stringToExclude Leading string indicating this is a vector type (to be excluded when eval'ing args)
+   * @param {Function} subtypeResolveFn Function used to resolve subtype type; varies depending on type naming convention
+   * @returns {Object}
+   * @internal
+   */
+  this.parseVectorTypeArgs = function(typeName, stringToExclude, subtypeResolveFn) {
+
+    const argsStartIndex = stringToExclude.length + 1;
+    const argsLength = typeName.length - (stringToExclude.length + 2);
+    const params = parseParams(typeName, argsStartIndex, argsLength);
+    if (params.length === 2) {
+      return {subtype: subtypeResolveFn(params[0]), dimensions: parseInt(params[1], 10)};
+    }
+    throw new TypeError('Not a valid type ' + typeName);
+  };
+
   /**
    * If not provided, it uses the array of buffers or the parameters and hints to build the routingKey
    * @param {Array} params
@@ -1139,6 +1243,19 @@ function defineInstanceMembers() {
       return dataType;
     }
 
+    if (typeName.indexOf(cqlNames.vector, startIndex) === startIndex) {
+      // It's a vector, so record the subtype and dimension.
+      dataType.code = dataTypes.custom;
+
+      // parseVectorTypeArgs is not an async function but we are.  To keep things simple let's ask the
+      // function to just return whatever it finds for an arg and we'll eval it after the fact
+      const params = this.parseVectorTypeArgs(typeName, cqlNames.vector, (arg) => arg );
+      params["subtype"] = await this.parseTypeName(keyspace, params["subtype"]);
+      dataType.info = params;
+
+      return dataType;
+    }
+
     const quoted = typeName.indexOf('"', startIndex) === startIndex;
     if (quoted) {
       // Remove quotes
@@ -1214,7 +1331,7 @@ function defineInstanceMembers() {
       }
     };
     startIndex = startIndex || 0;
-    let innerTypes;
+    let params;
     if (!length) {
       length = typeName.length;
     }
@@ -1254,12 +1371,12 @@ function defineInstanceMembers() {
       //move cursor across the name and bypass the parenthesis
       startIndex += complexTypeNames.list.length + 1;
       length -= complexTypeNames.list.length + 2;
-      innerTypes = parseParams(typeName, startIndex, length);
-      if (innerTypes.length !== 1) {
+      params = parseParams(typeName, startIndex, length);
+      if (params.length !== 1) {
         throw new TypeError('Not a valid type ' + typeName);
       }
       dataType.code = dataTypes.list;
-      dataType.info = this.parseFqTypeName(innerTypes[0]);
+      dataType.info = this.parseFqTypeName(params[0]);
       return dataType;
     }
     if (typeName.indexOf(complexTypeNames.set, startIndex) === startIndex) {
@@ -1268,27 +1385,27 @@ function defineInstanceMembers() {
       //move cursor across the name and bypass the parenthesis
       startIndex += complexTypeNames.set.length + 1;
       length -= complexTypeNames.set.length + 2;
-      innerTypes = parseParams(typeName, startIndex, length);
-      if (innerTypes.length !== 1)
+      params = parseParams(typeName, startIndex, length);
+      if (params.length !== 1)
       {
         throw new TypeError('Not a valid type ' + typeName);
       }
       dataType.code = dataTypes.set;
-      dataType.info = this.parseFqTypeName(innerTypes[0]);
+      dataType.info = this.parseFqTypeName(params[0]);
       return dataType;
     }
     if (typeName.indexOf(complexTypeNames.map, startIndex) === startIndex) {
       //org.apache.cassandra.db.marshal.MapType(keyType,valueType)
       //move cursor across the name and bypass the parenthesis
       startIndex += complexTypeNames.map.length + 1;
       length -= complexTypeNames.map.length + 2;
-      innerTypes = parseParams(typeName, startIndex, length);
+      params = parseParams(typeName, startIndex, length);
       //It should contain the key and value types
-      if (innerTypes.length !== 2) {
+      if (params.length !== 2) {
         throw new TypeError('Not a valid type ' + typeName);
       }
       dataType.code = dataTypes.map;
-      dataType.info = [this.parseFqTypeName(innerTypes[0]), this.parseFqTypeName(innerTypes[1])];
+      dataType.info = [this.parseFqTypeName(params[0]), this.parseFqTypeName(params[1])];
       return dataType;
     }
     if (typeName.indexOf(complexTypeNames.udt, startIndex) === startIndex) {
@@ -1301,12 +1418,19 @@ function defineInstanceMembers() {
       //move cursor across the name and bypass the parenthesis
       startIndex += complexTypeNames.tuple.length + 1;
       length -= complexTypeNames.tuple.length + 2;
-      innerTypes = parseParams(typeName, startIndex, length);
-      if (innerTypes.length < 1) {
+      params = parseParams(typeName, startIndex, length);
+      if (params.length < 1) {
         throw new TypeError('Not a valid type ' + typeName);
       }
       dataType.code = dataTypes.tuple;
-      dataType.info = innerTypes.map(x => this.parseFqTypeName(x));
+      dataType.info = params.map(x => this.parseFqTypeName(x));
+      return dataType;
+    }
+
+    if (typeName.indexOf(customTypeNames.vector, startIndex) === startIndex) {
+      // It's a vector, so record the subtype and dimension.
+      dataType.code = dataTypes.custom;
+      dataType.info = this.parseVectorTypeArgs(typeName, customTypeNames.vector, this.parseFqTypeName);
       return dataType;
     }
 
@@ -1628,6 +1752,12 @@ Encoder.guessDataType = function (value) {
     code = dataTypes.custom;
     info = customTypeNames.duration;
   }
+  // Map JS TypedArrays onto vectors
+  else if (Encoder.isTypedArray(value)) {
+    code = dataTypes.custom;
+    // TODO: another area that we have to generalize if we ever need to support vector subtypes other than float
+    info = buildParameterizedCustomType(customTypeNames.vector, [singleTypeNamesByDataType[dataTypes.float], value.length]);
+  }
   else if (Array.isArray(value)) {
     code = dataTypes.list;
   }
@@ -1862,4 +1992,24 @@ function concatRoutingKey(parts, totalLength) {
   return routingKey;
 }
 
+function buildParameterizedCustomType(customTypeName, args) {
+  return `${customTypeName}(${args.join(',')})`;
+}
+
+function invertObject(obj) {
+  const rv = {};
+  for(const k in obj){
+    if (Object.hasOwn(obj,k)) {
+      rv[obj[k]] = k;
+    }
+  }
+  return rv;
+}
+Encoder.isTypedArray = function(arg) {
+  // The TypedArray superclass isn't available directly so to detect an instance of a TypedArray
+  // subclass we have to access the prototype of a concrete instance.  There's nothing magical about
+  // Uint8Array here; we could just as easily use any of the other TypedArray subclasses.
+  return (arg instanceof Object.getPrototypeOf(Uint8Array));
+};
+
 module.exports = Encoder;

lib/types/index.js

@@ -29,7 +29,6 @@ const utils = require('../utils');
  */
 const Long = require('long');
 
-
 /**
  * Consistency levels
  * @type {Object}