compute table offsets (#91)

* compute table offsets * flesh out table offsets
geoarrow · Jan 11, 2024 · eb6f6cc · eb6f6cc
1 parent f1c31ef
commit eb6f6cc
Show file tree

Hide file tree

Showing 10 changed files with 133 additions and 26 deletions.
diff --git a/src/arc-layer.ts b/src/arc-layer.ts
@@ -13,7 +13,11 @@ import * as arrow from "apache-arrow";
 import * as ga from "@geoarrow/geoarrow-js";
 import { assignAccessor, extractAccessorsFromProps } from "./utils.js";
 import { child } from "@geoarrow/geoarrow-js";
-import { getPickingInfo } from "./picking.js";
+import {
+  GeoArrowExtraPickingProps,
+  computeChunkOffsets,
+  getPickingInfo,
+} from "./picking.js";
 import { ColorAccessor, FloatAccessor, GeoArrowPickingInfo } from "./types.js";
 import { validateAccessors } from "./validate.js";
 
@@ -108,7 +112,11 @@ export class GeoArrowArcLayer<
   static defaultProps = defaultProps;
   static layerName = "GeoArrowArcLayer";
 
-  getPickingInfo(params: GetPickingInfoParams): GeoArrowPickingInfo {
+  getPickingInfo(
+    params: GetPickingInfoParams & {
+      sourceLayer: { props: GeoArrowExtraPickingProps };
+    },
+  ): GeoArrowPickingInfo {
     return getPickingInfo(params, this.props.data);
   }
 
@@ -137,6 +145,7 @@ export class GeoArrowArcLayer<
       "getSourcePosition",
       "getTargetPosition",
     ]);
+    const tableOffsets = computeChunkOffsets(table.data);
 
     const layers: ArcLayer[] = [];
     for (
@@ -157,6 +166,7 @@ export class GeoArrowArcLayer<
 
         // @ts-expect-error used for picking purposes
         recordBatchIdx,
+        tableOffsets,
 
         id: `${this.props.id}-geoarrow-arc-${recordBatchIdx}`,
         data: {

diff --git a/src/column-layer.ts b/src/column-layer.ts
@@ -18,7 +18,11 @@ import {
 import * as ga from "@geoarrow/geoarrow-js";
 import { ColorAccessor, FloatAccessor, GeoArrowPickingInfo } from "./types.js";
 import { EXTENSION_NAME } from "./constants.js";
-import { getPickingInfo } from "./picking.js";
+import {
+  GeoArrowExtraPickingProps,
+  computeChunkOffsets,
+  getPickingInfo,
+} from "./picking.js";
 import { validateAccessors } from "./validate.js";
 
 /** All properties supported by GeoArrowColumnLayer */
@@ -105,7 +109,11 @@ export class GeoArrowColumnLayer<
   static defaultProps = defaultProps;
   static layerName = "GeoArrowColumnLayer";
 
-  getPickingInfo(params: GetPickingInfoParams): GeoArrowPickingInfo {
+  getPickingInfo(
+    params: GetPickingInfoParams & {
+      sourceLayer: { props: GeoArrowExtraPickingProps };
+    },
+  ): GeoArrowPickingInfo {
     return getPickingInfo(params, this.props.data);
   }
 
@@ -139,6 +147,7 @@ export class GeoArrowColumnLayer<
     const [accessors, otherProps] = extractAccessorsFromProps(this.props, [
       "getPosition",
     ]);
+    const tableOffsets = computeChunkOffsets(table.data);
 
     const layers: ColumnLayer[] = [];
     for (
@@ -158,6 +167,7 @@ export class GeoArrowColumnLayer<
 
         // @ts-expect-error used for picking purposes
         recordBatchIdx,
+        tableOffsets,
 
         id: `${this.props.id}-geoarrow-column-${recordBatchIdx}`,
         data: {

diff --git a/src/h3-hexagon-layer.ts b/src/h3-hexagon-layer.ts
@@ -11,7 +11,11 @@ import type { H3HexagonLayerProps } from "@deck.gl/geo-layers/typed";
 import * as arrow from "apache-arrow";
 import { assignAccessor, extractAccessorsFromProps } from "./utils.js";
 import { GeoArrowPickingInfo } from "./types.js";
-import { getPickingInfo } from "./picking.js";
+import {
+  GeoArrowExtraPickingProps,
+  computeChunkOffsets,
+  getPickingInfo,
+} from "./picking.js";
 import { validateAccessors } from "./validate.js";
 
 /** All properties supported by GeoArrowH3HexagonLayer */
@@ -61,7 +65,11 @@ export class GeoArrowH3HexagonLayer<
   static defaultProps = defaultProps;
   static layerName = "GeoArrowH3HexagonLayer";
 
-  getPickingInfo(params: GetPickingInfoParams): GeoArrowPickingInfo {
+  getPickingInfo(
+    params: GetPickingInfoParams & {
+      sourceLayer: { props: GeoArrowExtraPickingProps };
+    },
+  ): GeoArrowPickingInfo {
     return getPickingInfo(params, this.props.data);
   }
 
@@ -80,6 +88,7 @@ export class GeoArrowH3HexagonLayer<
     const [accessors, otherProps] = extractAccessorsFromProps(this.props, [
       "getHexagon",
     ]);
+    const tableOffsets = computeChunkOffsets(table.data);
 
     const layers: H3HexagonLayer[] = [];
     for (
@@ -98,6 +107,7 @@ export class GeoArrowH3HexagonLayer<
 
         // @ts-expect-error used for picking purposes
         recordBatchIdx,
+        tableOffsets,
 
         id: `${this.props.id}-geoarrow-arc-${recordBatchIdx}`,
 

diff --git a/src/heatmap-layer.ts b/src/heatmap-layer.ts
@@ -18,6 +18,7 @@ import {
 import { FloatAccessor } from "./types.js";
 import { EXTENSION_NAME } from "./constants.js";
 import { validateAccessors } from "./validate.js";
+import { computeChunkOffsets } from "./picking.js";
 
 /** All properties supported by GeoArrowHeatmapLayer */
 export type GeoArrowHeatmapLayerProps = Omit<
@@ -106,6 +107,7 @@ export class GeoArrowHeatmapLayer<
     const [accessors, otherProps] = extractAccessorsFromProps(this.props, [
       "getPosition",
     ]);
+    const tableOffsets = computeChunkOffsets(table.data);
 
     const layers: HeatmapLayer[] = [];
     for (
@@ -125,6 +127,7 @@ export class GeoArrowHeatmapLayer<
 
         // @ts-expect-error used for picking purposes
         recordBatchIdx,
+        tableOffsets,
 
         id: `${this.props.id}-geoarrow-heatmap-${recordBatchIdx}`,
         data: {

diff --git a/src/path-layer.ts b/src/path-layer.ts
@@ -19,7 +19,11 @@ import {
   getMultiLineStringResolvedOffsets,
   invertOffsets,
 } from "./utils.js";
-import { getPickingInfo } from "./picking.js";
+import {
+  GeoArrowExtraPickingProps,
+  computeChunkOffsets,
+  getPickingInfo,
+} from "./picking.js";
 import { ColorAccessor, FloatAccessor, GeoArrowPickingInfo } from "./types.js";
 import { EXTENSION_NAME } from "./constants.js";
 import { validateAccessors } from "./validate.js";
@@ -90,7 +94,11 @@ export class GeoArrowPathLayer<
   static defaultProps = defaultProps;
   static layerName = "GeoArrowPathLayer";
 
-  getPickingInfo(params: GetPickingInfoParams): GeoArrowPickingInfo {
+  getPickingInfo(
+    params: GetPickingInfoParams & {
+      sourceLayer: { props: GeoArrowExtraPickingProps };
+    },
+  ): GeoArrowPickingInfo {
     return getPickingInfo(params, this.props.data);
   }
 
@@ -141,6 +149,7 @@ export class GeoArrowPathLayer<
     const [accessors, otherProps] = extractAccessorsFromProps(this.props, [
       "getPath",
     ]);
+    const tableOffsets = computeChunkOffsets(table.data);
 
     const layers: PathLayer[] = [];
     for (
@@ -163,6 +172,7 @@ export class GeoArrowPathLayer<
 
         // used for picking purposes
         recordBatchIdx,
+        tableOffsets,
 
         id: `${this.props.id}-geoarrow-path-${recordBatchIdx}`,
         data: {
@@ -208,6 +218,7 @@ export class GeoArrowPathLayer<
     const [accessors, otherProps] = extractAccessorsFromProps(this.props, [
       "getPath",
     ]);
+    const tableOffsets = computeChunkOffsets(table.data);
 
     const layers: PathLayer[] = [];
     for (
@@ -237,6 +248,7 @@ export class GeoArrowPathLayer<
 
         // used for picking purposes
         recordBatchIdx,
+        tableOffsets,
         invertedGeomOffsets: invertOffsets(geomOffsets),
 
         id: `${this.props.id}-geoarrow-path-${recordBatchIdx}`,

diff --git a/src/picking.ts b/src/picking.ts
@@ -2,33 +2,40 @@ import * as arrow from "apache-arrow";
 import { GetPickingInfoParams } from "@deck.gl/core/typed";
 import { GeoArrowPickingInfo } from "./types";
 
+export interface GeoArrowExtraPickingProps {
+  recordBatchIdx: number;
+  tableOffsets: Uint32Array;
+  invertedGeomOffsets?: Uint8Array | Uint16Array | Uint32Array;
+}
+
 export function getPickingInfo(
-  { info, sourceLayer }: GetPickingInfoParams,
+  {
+    info,
+    sourceLayer,
+  }: GetPickingInfoParams & {
+    sourceLayer: { props: GeoArrowExtraPickingProps };
+  },
   table: arrow.Table,
 ): GeoArrowPickingInfo {
   // Geometry index as rendered
   let index = info.index;
 
-  // if a MultiPoint dataset, map from the rendered index back to the feature
-  // index
-  // @ts-expect-error `invertedGeomOffsets` is manually set on layer props
+  // if a Multi- geometry dataset, map from the rendered index back to the
+  // feature index
   if (sourceLayer.props.invertedGeomOffsets) {
-    // @ts-expect-error `invertedGeomOffsets` is manually set on layer props
     index = sourceLayer.props.invertedGeomOffsets[index];
   }
 
-  // @ts-expect-error `recordBatchIdx` is manually set on layer props
-  const recordBatchIdx: number = sourceLayer.props.recordBatchIdx;
+  const recordBatchIdx = sourceLayer.props.recordBatchIdx;
+  const tableOffsets = sourceLayer.props.tableOffsets;
+
   const batch = table.batches[recordBatchIdx];
   const row = batch.get(index);
   if (row === null) {
     return info;
   }
 
-  // @ts-expect-error hack: using private method to avoid recomputing via
-  // batch lengths on each iteration
-  const offsets: number[] = table._offsets;
-  const currentBatchOffset = offsets[recordBatchIdx];
+  const currentBatchOffset = tableOffsets[recordBatchIdx];
 
   // Update index to be _global_ index, not within the specific record batch
   index += currentBatchOffset;
@@ -38,3 +45,16 @@ export function getPickingInfo(
     object: row,
   };
 }
+
+// This is vendored from Arrow JS because it's a private API
+export function computeChunkOffsets<T extends arrow.DataType>(
+  chunks: ReadonlyArray<arrow.Data<T>>,
+) {
+  return chunks.reduce(
+    (offsets, chunk, index) => {
+      offsets[index + 1] = offsets[index] + chunk.length;
+      return offsets;
+    },
+    new Uint32Array(chunks.length + 1),
+  );
+}
diff --git a/src/scatterplot-layer.ts b/src/scatterplot-layer.ts
@@ -17,7 +17,11 @@ import {
   getGeometryVector,
   invertOffsets,
 } from "./utils.js";
-import { getPickingInfo } from "./picking.js";
+import {
+  GeoArrowExtraPickingProps,
+  computeChunkOffsets,
+  getPickingInfo,
+} from "./picking.js";
 import { ColorAccessor, FloatAccessor, GeoArrowPickingInfo } from "./types.js";
 import { EXTENSION_NAME } from "./constants.js";
 import { validateAccessors } from "./validate.js";
@@ -91,7 +95,11 @@ export class GeoArrowScatterplotLayer<
   static defaultProps = defaultProps;
   static layerName = "GeoArrowScatterplotLayer";
 
-  getPickingInfo(params: GetPickingInfoParams): GeoArrowPickingInfo {
+  getPickingInfo(
+    params: GetPickingInfoParams & {
+      sourceLayer: { props: GeoArrowExtraPickingProps };
+    },
+  ): GeoArrowPickingInfo {
     return getPickingInfo(params, this.props.data);
   }
 
@@ -137,6 +145,7 @@ export class GeoArrowScatterplotLayer<
     const [accessors, otherProps] = extractAccessorsFromProps(this.props, [
       "getPosition",
     ]);
+    const tableOffsets = computeChunkOffsets(table.data);
 
     const layers: ScatterplotLayer[] = [];
     for (
@@ -156,6 +165,7 @@ export class GeoArrowScatterplotLayer<
 
         // @ts-expect-error used for picking purposes
         recordBatchIdx,
+        tableOffsets,
 
         id: `${this.props.id}-geoarrow-scatterplot-${recordBatchIdx}`,
         data: {
@@ -201,6 +211,7 @@ export class GeoArrowScatterplotLayer<
     const [accessors, otherProps] = extractAccessorsFromProps(this.props, [
       "getPosition",
     ]);
+    const tableOffsets = computeChunkOffsets(table.data);
 
     const layers: ScatterplotLayer[] = [];
     for (
@@ -222,6 +233,7 @@ export class GeoArrowScatterplotLayer<
 
         // @ts-expect-error used for picking purposes
         recordBatchIdx,
+        tableOffsets,
         invertedGeomOffsets: invertOffsets(geomOffsets),
 
         id: `${this.props.id}-geoarrow-scatterplot-${recordBatchIdx}`,