Parse WKB to GeoArrow

package.json

@@ -44,9 +44,13 @@
     "src/"
   ],
   "peerDependencies": {
+    "@loaders.gl/schema": "*",
+    "@loaders.gl/wkt": "*",
     "apache-arrow": ">=15"
   },
   "devDependencies": {
+    "@loaders.gl/schema": "^4.3.3",
+    "@loaders.gl/wkt": "^4.3.3",
     "@rollup/plugin-node-resolve": "^15.2.3",
     "@rollup/plugin-terser": "^0.4.3",
     "@rollup/plugin-typescript": "^11.1.2",

src/algorithm/utils/assert.ts

@@ -4,6 +4,6 @@ export function assert(condition: boolean, message?: string) {
   }
 }
 
-export function assertFalse(): never {
-  throw new Error(`assertion failed`);
+export function assertFalse(message?: string): never {
+  throw new Error(`assertion failed ${message}`);
 }

src/child.ts

@@ -2,9 +2,7 @@
  * Strongly typed accessors for children, since arrow.Data.children[] is untyped
  */
 
-import { Data } from "apache-arrow/data";
-import { Vector } from "apache-arrow/vector";
-import { Float } from "apache-arrow/type";
+import { Data, Vector, type Float } from "apache-arrow";
 import {
   LineStringData,
   MultiLineStringData,

src/data.ts

@@ -1,5 +1,6 @@
-import { Data } from "apache-arrow/data";
+import { Data } from "apache-arrow";
 import {
+  WKB,
   Point,
   LineString,
   Polygon,
@@ -27,6 +28,7 @@ export type GeoArrowData =
   | MultiPointData
   | MultiLineStringData
   | MultiPolygonData;
+export type WKBData = Data<WKB>;
 
 export function isPointData(data: Data): data is PointData {
   return isPoint(data.type);

src/index.ts

@@ -4,3 +4,4 @@ export * as data from "./data.js";
 export * as type from "./type.js";
 export * as vector from "./vector.js";
 export * as worker from "./worker";
+export * as io from "./io";

src/io/index.ts

@@ -0,0 +1 @@
+export { parseWkb, WKBType } from "./wkb.js";

src/io/wkb.ts

@@ -0,0 +1,269 @@
+import { makeData, Field, FixedSizeList, Float64, List } from "apache-arrow";
+import {
+  GeoArrowData,
+  LineStringData,
+  PointData,
+  PolygonData,
+  WKBData,
+} from "../data";
+import { WKBLoader } from "@loaders.gl/wkt";
+import type {
+  BinaryGeometry,
+  BinaryPointGeometry,
+  BinaryLineGeometry,
+  BinaryPolygonGeometry,
+} from "@loaders.gl/schema";
+import { assert, assertFalse } from "../algorithm/utils/assert";
+
+export enum WKBType {
+  Point,
+  LineString,
+  Polygon,
+  MultiPoint,
+  MultiLineString,
+  MultiPolygon,
+}
+
+/**
+ * Parse an Arrow array of WKB
+ *
+ * @return  {[type]}  [return description]
+ */
+export function parseWkb(
+  data: WKBData,
+  type: WKBType,
+  dim: number,
+): GeoArrowData {
+  const parsedGeometries: BinaryGeometry[] = [];
+
+  for (const item of iterBinary(data)) {
+    if (item === null) {
+      throw new Error("Null entries are not currently supported");
+    }
+    const arrayBuffer = copyViewToArrayBuffer(item);
+    const parsed = WKBLoader.parseSync(arrayBuffer, {
+      wkb: { shape: "binary-geometry" },
+    }) as BinaryGeometry;
+    parsedGeometries.push(parsed);
+  }
+
+  switch (type) {
+    case WKBType.Point:
+      return repackPoints(parsedGeometries as BinaryPointGeometry[], dim);
+
+    case WKBType.LineString:
+      return repackLineStrings(parsedGeometries as BinaryLineGeometry[], dim);
+
+    case WKBType.Polygon:
+      return repackPolygons(parsedGeometries as BinaryPolygonGeometry[], dim);
+
+    default:
+      assertFalse("Not yet implemented for this geometry type");
+  }
+}
+
+function* iterBinary(data: WKBData): IterableIterator<Uint8Array | null> {
+  const values = data.values;
+  const valueOffsets = data.valueOffsets;
+  for (let i = 0; i < data.length; i++) {
+    if (!data.getValid(i)) {
+      yield null;
+    } else {
+      const startOffset = valueOffsets[i];
+      const endOffset = valueOffsets[i + 1];
+      yield values.subarray(startOffset, endOffset);
+    }
+  }
+}
+
+// TODO: update loaders.gl parseWKB to accept Uint8Array, not just ArrayBuffer.
+function copyViewToArrayBuffer(view: Uint8Array): ArrayBuffer {
+  const buffer = new ArrayBuffer(view.byteLength);
+  new Uint8Array(buffer).set(view);
+  return buffer;
+}
+
+function repackPoints(geoms: BinaryPointGeometry[], dim: number): PointData {
+  const geomLength = geoms.length;
+  const coords = new Float64Array(geomLength * dim);
+  let coordOffset = 0;
+  for (const geom of geoms) {
+    assert(geom.positions.value instanceof Float64Array);
+    coords.set(geom.positions.value, coordOffset * dim);
+    coordOffset += 1;
+  }
+
+  const coordsData = makeData({
+    type: new Float64(),
+    data: coords,
+  });
+  return makeData({
+    type: new FixedSizeList(
+      dim,
+      new Field(coordFieldName(dim), new Float64(), false),
+    ),
+    child: coordsData,
+  });
+}
+
+type LineStringCapacity = {
+  coordCapacty: number;
+  geomCapacity: number;
+};
+
+function repackLineStrings(
+  geoms: BinaryLineGeometry[],
+  dim: number,
+): LineStringData {
+  const capacity = inferLineStringCapacity(geoms);
+  const coords = new Float64Array(capacity.coordCapacty * dim);
+  const geomOffsets = new Int32Array(capacity.geomCapacity + 1);
+
+  let geomIndex = 0;
+  let coordOffset = 0;
+  for (const geom of geoms) {
+    assert(geom.positions.value instanceof Float64Array);
+    const numCoords = geom.positions.value.length / geom.positions.size;
+    coords.set(geom.positions.value, coordOffset * dim);
+    geomIndex += 1;
+    coordOffset += numCoords;
+
+    // Note this is after we've added one
+    geomOffsets[geomIndex] = coordOffset;
+  }
+
+  const coordsData = makeData({
+    type: new Float64(),
+    data: coords,
+  });
+  const verticesData = makeData({
+    type: new FixedSizeList(
+      dim,
+      new Field(coordFieldName(dim), coordsData.type, false),
+    ),
+    child: coordsData,
+  });
+  return makeData({
+    type: new List(new Field("vertices", verticesData.type, false)),
+    valueOffsets: geomOffsets,
+    child: verticesData,
+  });
+}
+
+function inferLineStringCapacity(
+  geoms: BinaryLineGeometry[],
+): LineStringCapacity {
+  let capacity: LineStringCapacity = {
+    coordCapacty: 0,
+    geomCapacity: 0,
+  };
+
+  // TODO: check geom.pathIndices to validate that we have a LineString and not
+  // a multi line string
+  for (const geom of geoms) {
+    capacity.geomCapacity += 1;
+    capacity.coordCapacty += geom.positions.value.length / geom.positions.size;
+  }
+
+  return capacity;
+}
+
+type PolygonCapacity = {
+  coordCapacty: number;
+  /** This is what loaders.gl calls `primitivePolygonIndices` */
+  ringCapacity: number;
+  geomCapacity: number;
+};
+
+function repackPolygons(
+  geoms: BinaryPolygonGeometry[],
+  dim: number,
+): PolygonData {
+  const capacity = inferPolygonCapacity(geoms);
+  const coords = new Float64Array(capacity.coordCapacty * dim);
+  const ringOffsets = new Int32Array(capacity.ringCapacity + 1);
+  const geomOffsets = new Int32Array(capacity.geomCapacity + 1);
+
+  let geomIndex = 0;
+  let coordOffset = 0;
+  let ringOffset = 0;
+
+  for (const geom of geoms) {
+    assert(geom.positions.value instanceof Float64Array);
+    const numCoords = geom.positions.value.length / geom.positions.size;
+
+    coords.set(geom.positions.value, coordOffset * dim);
+    coordOffset += numCoords;
+
+    for (
+      let ringIdx = 0;
+      ringIdx < geom.primitivePolygonIndices.value.length - 1;
+      ringIdx++
+    ) {
+      ringOffsets[ringOffset + 1] =
+        ringOffsets[ringOffset] +
+        (geom.primitivePolygonIndices.value[ringIdx + 1] -
+          geom.primitivePolygonIndices.value[ringIdx]);
+      ringOffset += 1;
+    }
+
+    geomOffsets[geomIndex + 1] = ringOffset;
+    geomIndex += 1;
+  }
+
+  const coordsData = makeData({
+    type: new Float64(),
+    data: coords,
+  });
+  const verticesData = makeData({
+    type: new FixedSizeList(
+      dim,
+      new Field(coordFieldName(dim), coordsData.type, false),
+    ),
+    child: coordsData,
+  });
+
+  const ringsData = makeData({
+    type: new List(new Field("vertices", verticesData.type, false)),
+    valueOffsets: ringOffsets,
+    child: verticesData,
+  });
+
+  return makeData({
+    type: new List(new Field("rings", ringsData.type, false)),
+    valueOffsets: geomOffsets,
+    child: ringsData,
+  });
+}
+
+function inferPolygonCapacity(geoms: BinaryPolygonGeometry[]): PolygonCapacity {
+  let capacity: PolygonCapacity = {
+    coordCapacty: 0,
+    ringCapacity: 0,
+    geomCapacity: 0,
+  };
+
+  // TODO: check geom.polygonIndices to validate that we have a Polygon and not
+  // a MultiPolygon
+  for (const geom of geoms) {
+    capacity.geomCapacity += 1;
+    assert(
+      geom.primitivePolygonIndices.value.length >= 1,
+      "Expected primitivePolygonIndices to always have length at least 1",
+    );
+    capacity.ringCapacity += geom.primitivePolygonIndices.value.length - 1;
+    capacity.coordCapacty += geom.positions.value.length / geom.positions.size;
+  }
+
+  return capacity;
+}
+
+function coordFieldName(dim: number): "xy" | "xyz" {
+  if (dim === 2) {
+    return "xy";
+  } else if (dim === 3) {
+    return "xyz";
+  } else {
+    assertFalse("expected dimension of 2 or 3");
+  }
+}

src/type.ts

@@ -1,10 +1,11 @@
 import {
-  Struct,
-  Float,
-  List,
-  FixedSizeList,
+  type Binary,
+  type Struct,
+  type Float,
+  type List,
+  type FixedSizeList,
   DataType,
-} from "apache-arrow/type";
+} from "apache-arrow";
 
 // Note: this apparently has to be arrow.Float and not arrow.Float64 to ensure
 // that recreating a data instance with arrow.makeData type checks using the
@@ -29,6 +30,7 @@ export type GeoArrowType =
   | MultiPoint
   | MultiLineString
   | MultiPolygon;
+export type WKB = Binary;
 
 /** Check that the given type is a Point data type */
 export function isPoint(type: DataType): type is Point {

src/vector.ts

@@ -1,4 +1,4 @@
-import { Vector } from "apache-arrow/vector";
+import { Vector } from "apache-arrow";
 import {
   Point,
   LineString,

src/worker/hard-clone.ts

@@ -1,7 +1,4 @@
-import { Data } from "apache-arrow/data";
-import { DataType } from "apache-arrow/type";
-import { Vector } from "apache-arrow/vector";
-import { BufferType } from "apache-arrow/enum";
+import { Data, DataType, Vector, BufferType } from "apache-arrow";
 import type { Buffers } from "apache-arrow/data";
 
 type TypedArray =

src/worker/rehydrate.ts

@@ -1,5 +1,9 @@
 import {
-  DataType,
+  BufferType,
+  Type,
+  Data,
+  Vector,
+  Field,
   Null,
   Int,
   Float,
@@ -18,11 +22,8 @@ import {
   FixedSizeList,
   Map_,
   Duration,
-} from "apache-arrow/type";
-import { BufferType, Type } from "apache-arrow/enum";
-import { Data } from "apache-arrow/data";
-import { Vector } from "apache-arrow/vector";
-import { Field } from "apache-arrow/schema";
+  type DataType,
+} from "apache-arrow";
 import type { Buffers } from "apache-arrow/data";
 import { Polygon, isPolygon } from "../type";
 import { PolygonData } from "../data";