feat(nodejs): array support (#50)

Author: glasstiger

.github/workflows/build.yml

@@ -5,6 +5,8 @@ on:
     branches:
       - main
   pull_request:
+  schedule:
+    - cron: '15 2,10,18 * * *'
 
 jobs:
   test:

src/buffer/base.ts

@@ -26,8 +26,8 @@ const DEFAULT_MAX_NAME_LENGTH = 127;
 abstract class SenderBufferBase implements SenderBuffer {
   private bufferSize: number;
   private readonly maxBufferSize: number;
-  private buffer: Buffer<ArrayBuffer>;
-  private position: number;
+  protected buffer: Buffer<ArrayBuffer>;
+  protected position: number;
   private endOfLastRow: number;
 
   private hasTable: boolean;
@@ -232,6 +232,15 @@ abstract class SenderBufferBase implements SenderBuffer {
    */
   abstract floatColumn(name: string, value: number): SenderBuffer;
 
+  /**
+   * Writes an array column with its values into the buffer.
+   *
+   * @param {string} name - Column name.
+   * @param {unknown[]} value - Column value, accepts only arrays.
+   * @return {Sender} Returns with a reference to this sender.
+   */
+  abstract arrayColumn(name: string, value: unknown[]): SenderBuffer;
+
   /**
    * Writes a 64-bit signed integer into the buffer. <br>
    * Use it to insert into LONG, INT, SHORT and BYTE columns.
@@ -386,6 +395,10 @@ abstract class SenderBufferBase implements SenderBuffer {
     this.position = this.buffer.writeInt8(data, this.position);
   }
 
+  protected writeInt(data: number) {
+    this.position = this.buffer.writeInt32LE(data, this.position);
+  }
+
   protected writeDouble(data: number) {
     this.position = this.buffer.writeDoubleLE(data, this.position);
   }

src/buffer/bufferv1.ts

@@ -33,6 +33,10 @@ class SenderBufferV1 extends SenderBufferBase {
     );
     return this;
   }
+
+  arrayColumn(): SenderBuffer {
+    throw new Error("Arrays are not supported in protocol v1");
+  }
 }
 
 export { SenderBufferV1 };

src/buffer/bufferv2.ts

@@ -2,8 +2,14 @@
 import { SenderOptions } from "../options";
 import { SenderBuffer } from "./index";
 import { SenderBufferBase } from "./base";
+import { ArrayPrimitive, getDimensions, validateArray } from "../utils";
 
+const COLUMN_TYPE_DOUBLE: number = 10;
+const COLUMN_TYPE_NULL: number = 33;
+
+const ENTITY_TYPE_ARRAY: number = 14;
 const ENTITY_TYPE_DOUBLE: number = 16;
+
 const EQUALS_SIGN: number = "=".charCodeAt(0);
 
 /**
@@ -37,6 +43,91 @@ class SenderBufferV2 extends SenderBufferBase {
     );
     return this;
   }
+
+  arrayColumn(name: string, value: unknown[]): SenderBuffer {
+    const dimensions = getDimensions(value);
+    const type = validateArray(value, dimensions);
+    // only number arrays and NULL supported for now
+    if (type !== "number" && type !== null) {
+      throw new Error(`Unsupported array type [type=${type}]`);
+    }
+
+    this.writeColumn(name, value, () => {
+      this.checkCapacity([], 3);
+      this.writeByte(EQUALS_SIGN);
+      this.writeByte(ENTITY_TYPE_ARRAY);
+
+      if (!value) {
+        this.writeByte(COLUMN_TYPE_NULL);
+      } else {
+        this.writeByte(COLUMN_TYPE_DOUBLE);
+        this.writeArray(value, dimensions, type);
+      }
+    });
+    return this;
+  }
+
+  private writeArray(
+    arr: unknown[],
+    dimensions: number[],
+    type: ArrayPrimitive,
+  ) {
+    this.checkCapacity([], 1 + dimensions.length * 4);
+    this.writeByte(dimensions.length);
+    for (let i = 0; i < dimensions.length; i++) {
+      this.writeInt(dimensions[i]);
+    }
+
+    this.checkCapacity([], SenderBufferV2.arraySize(dimensions, type));
+    this.writeArrayValues(arr, dimensions);
+  }
+
+  private writeArrayValues(arr: unknown[], dimensions: number[]) {
+    if (Array.isArray(arr[0])) {
+      for (let i = 0; i < arr.length; i++) {
+        this.writeArrayValues(arr[i] as unknown[], dimensions);
+      }
+    } else {
+      const type = arr[0] !== undefined ? typeof arr[0] : null;
+      switch (type) {
+        case "number":
+          for (let i = 0; i < arr.length; i++) {
+            this.position = this.buffer.writeDoubleLE(
+              arr[i] as number,
+              this.position,
+            );
+          }
+          break;
+        case null:
+          // empty array
+          break;
+        default:
+          throw new Error(`Unsupported array type [type=${type}]`);
+      }
+    }
+  }
+
+  private static arraySize(dimensions: number[], type: ArrayPrimitive): number {
+    let numOfElements = 1;
+    for (let i = 0; i < dimensions.length; i++) {
+      numOfElements *= dimensions[i];
+    }
+
+    switch (type) {
+      case "number":
+        return numOfElements * 8;
+      case "boolean":
+        return numOfElements;
+      case "string":
+        // in case of string[] capacity check is done separately for each array element
+        return 0;
+      case null:
+        // empty array
+        return 0;
+      default:
+        throw new Error(`Unsupported array type [type=${type}]`);
+    }
+  }
 }
 
 export { SenderBufferV2 };

src/buffer/index.ts

@@ -104,6 +104,8 @@ interface SenderBuffer {
    */
   floatColumn(name: string, value: number): SenderBuffer;
 
+  arrayColumn(name: string, value: unknown[]): SenderBuffer;
+
   /**
    * Writes an integer column with its value into the buffer.
    *

src/sender.ts

@@ -241,6 +241,11 @@ class Sender {
     return this;
   }
 
+  arrayColumn(name: string, value: unknown[]): Sender {
+    this.buffer.arrayColumn(name, value);
+    return this;
+  }
+
   /**
    * Writes an integer column with its value into the buffer of the sender.
    *

src/utils.ts

@@ -1,5 +1,7 @@
 import { Agent } from "undici";
 
+type ArrayPrimitive = "number" | "boolean" | "string" | null;
+
 type TimestampUnit = "ns" | "us" | "ms";
 
 function isBoolean(value: unknown): value is boolean {
@@ -38,6 +40,72 @@ function timestampToNanos(timestamp: bigint, unit: TimestampUnit) {
   }
 }
 
+function getDimensions(data: unknown) {
+  const dimensions: number[] = [];
+  while (Array.isArray(data)) {
+    dimensions.push(data.length);
+    data = data[0];
+  }
+  return dimensions;
+}
+
+function validateArray(data: unknown[], dimensions: number[]): ArrayPrimitive {
+  if (data === null || data === undefined) {
+    return null;
+  }
+  if (!Array.isArray(data)) {
+    throw new Error(
+      `The value must be an array [value=${JSON.stringify(data)}, type=${typeof data}]`,
+    );
+  }
+
+  let expectedType: ArrayPrimitive = null;
+
+  function checkArray(
+    array: unknown[],
+    depth: number = 0,
+    path: string = "",
+  ): void {
+    const expectedLength = dimensions[depth];
+    if (array.length !== expectedLength) {
+      throw new Error(
+        `Lengths of sub-arrays do not match [expected=${expectedLength}, actual=${array.length}, dimensions=[${dimensions}], path=${path}]`,
+      );
+    }
+
+    if (depth < dimensions.length - 1) {
+      // intermediate level, expecting arrays
+      for (let i = 0; i < array.length; i++) {
+        if (!Array.isArray(array[i])) {
+          throw new Error(
+            `Mixed types found [expected=array, current=${typeof array[i]}, path=${path}[${i}]]`,
+          );
+        }
+        checkArray(array[i] as unknown[], depth + 1, `${path}[${i}]`);
+      }
+    } else {
+      // leaf level, expecting primitives
+      if (expectedType === null && array[0] !== undefined) {
+        expectedType = typeof array[0] as ArrayPrimitive;
+      }
+
+      for (let i = 0; i < array.length; i++) {
+        const currentType = typeof array[i] as ArrayPrimitive;
+        if (currentType !== expectedType) {
+          throw new Error(
+            expectedType !== null
+              ? `Mixed types found [expected=${expectedType}, current=${currentType}, path=${path}[${i}]]`
+              : `Unsupported array type [type=${currentType}]`,
+          );
+        }
+      }
+    }
+  }
+
+  checkArray(data);
+  return expectedType;
+}
+
 /**
  * Fetches JSON data from a URL.
  * @template T - The expected type of the JSON response
@@ -83,4 +151,7 @@ export {
   timestampToNanos,
   TimestampUnit,
   fetchJson,
+  getDimensions,
+  validateArray,
+  ArrayPrimitive,
 };