Reduce memory footprint of data indexing

bokehjs/src/lib/core/util/spatial.ts

@@ -8,17 +8,34 @@ export type IndexedRect = Rect & {i: number}
 export class SpatialIndex {
   private readonly index: FlatBush | null = null
 
-  constructor(private readonly points: IndexedRect[]) {
-    if (points.length > 0) {
-      this.index = new FlatBush(points.length)
+  constructor(size: number) {
+    if (size > 0) {
+      this.index = new FlatBush(size)
+    }
+  }
 
-      for (const p of points) {
-        const {x0, y0, x1, y1} = p
-        this.index.add(x0, y0, x1, y1)
-      }
+  static from(points: IndexedRect[]): SpatialIndex {
+    const index = new SpatialIndex(points.length)
 
-      this.index.finish()
+    for (const p of points) {
+      const {x0, y0, x1, y1} = p
+      index.add(x0, y0, x1, y1)
     }
+
+    index.finish()
+    return index
+  }
+
+  add(x0: number, y0: number, x1: number, y1: number): void {
+    this.index?.add(x0, y0, x1, y1)
+  }
+
+  add_empty(): void {
+    this.index?.add(Infinity, Infinity, -Infinity, -Infinity)
+  }
+
+  finish(): void {
+    this.index?.finish()
   }
 
   protected _normalize(rect: Rect): Rect {
@@ -39,17 +56,34 @@ export class SpatialIndex {
     }
   }
 
-  search(rect: Rect): IndexedRect[] {
+  indices(rect: Rect): number[] {
     if (this.index == null)
       return []
     else {
       const {x0, y0, x1, y1} = this._normalize(rect)
-      const indices = this.index.search(x0, y0, x1, y1)
-      return indices.map((j) => this.points[j])
+      return this.index.search(x0, y0, x1, y1)
     }
   }
 
-  indices(rect: Rect): number[] {
-    return this.search(rect).map(({i}) => i)
+  bounds(rect: Rect): Rect {
+    const bounds = empty()
+
+    for (const i of this.indices(rect)) {
+      const boxes = (this.index as any)._boxes as Float64Array
+      const x1 = boxes[4*i + 0]
+      const y1 = boxes[4*i + 1]
+      const x0 = boxes[4*i + 2]
+      const y0 = boxes[4*i + 3]
+      if (x0 < bounds.x0)
+        bounds.x0 = x0
+      if (x1 > bounds.x1)
+        bounds.x1 = x1
+      if (y0 < bounds.y0)
+        bounds.y0 = y0
+      if (y1 > bounds.y1)
+        bounds.y1 = y1
+    }
+
+    return bounds
   }
 }

bokehjs/src/lib/models/glyphs/bezier.ts

@@ -101,18 +101,18 @@ export class BezierView extends GlyphView {
   model: Bezier
   visuals: Bezier.Visuals
 
-  protected _index_data(): SpatialIndex {
-    const points = []
-    for (let i = 0, end = this._x0.length; i < end; i++) {
-      if (isNaN(this._x0[i] + this._x1[i] + this._y0[i] + this._y1[i] + this._cx0[i] + this._cy0[i] + this._cx1[i] + this._cy1[i]))
-        continue
+  protected _index_data(index: SpatialIndex): void {
+    const {data_size} = this
 
-      const [x0, y0, x1, y1] = _cbb(this._x0[i],  this._y0[i],  this._x1[i],  this._y1[i],
-                                    this._cx0[i], this._cy0[i], this._cx1[i], this._cy1[i])
-      points.push({x0, y0, x1, y1, i})
+    for (let i = 0; i < data_size; i++) {
+      if (isNaN(this._x0[i] + this._x1[i] + this._y0[i] + this._y1[i] + this._cx0[i] + this._cy0[i] + this._cx1[i] + this._cy1[i]))
+        index.add_empty()
+      else {
+        const [x0, y0, x1, y1] = _cbb(this._x0[i],  this._y0[i],  this._x1[i],  this._y1[i],
+                                      this._cx0[i], this._cy0[i], this._cx1[i], this._cy1[i])
+        index.add(x0, y0, x1, y1)
+      }
     }
-
-    return new SpatialIndex(points)
   }
 
   protected _render(ctx: Context2d, indices: number[],

bokehjs/src/lib/models/glyphs/box.ts

@@ -50,23 +50,17 @@ export abstract class BoxView extends GlyphView {
 
   protected abstract _lrtb(i: number): [number, number, number, number]
 
-  protected _index_box(len: number): SpatialIndex {
-    const points = []
+  protected _index_data(index: SpatialIndex): void {
+    const {min, max} = Math
+    const {data_size} = this
 
-    for (let i = 0; i < len; i++) {
+    for (let i = 0; i < data_size; i++) {
       const [l, r, t, b] = this._lrtb(i)
       if (isNaN(l + r + t + b) || !isFinite(l + r + t + b))
-        continue
-      points.push({
-        x0: Math.min(l, r),
-        y0: Math.min(t, b),
-        x1: Math.max(r, l),
-        y1: Math.max(t, b),
-        i,
-      })
+        index.add_empty()
+      else
+        index.add(min(l, r), min(t, b), max(r, l), max(t, b))
     }
-
-    return new SpatialIndex(points)
   }
 
   protected _render(ctx: Context2d, indices: number[],

bokehjs/src/lib/models/glyphs/glyph.ts

@@ -44,10 +44,28 @@ export abstract class GlyphView extends View {
     return this.glglyph != null
   }
 
-  index: SpatialIndex
+  private _index: SpatialIndex | null = null
+
+  private _data_size: number | null = null
 
   protected _nohit_warned: Set<geometry.Geometry["type"]> = new Set()
 
+  get index(): SpatialIndex {
+    const {_index} = this
+    if (_index != null)
+      return _index
+    else
+      throw new Error(`${this}.index_data() wasn't called`)
+  }
+
+  get data_size(): number {
+    const {_data_size} = this
+    if (_data_size != null)
+      return _data_size
+    else
+      throw new Error(`${this}.set_data() wasn't called`)
+  }
+
   initialize(): void {
     super.initialize()
     this.visuals = new visuals.Visuals(this.model)
@@ -115,25 +133,9 @@ export abstract class GlyphView extends View {
   }
 
   log_bounds(): Rect {
-    const bb = bbox.empty()
-
-    const positive_x_bbs = this.index.search(bbox.positive_x())
-    for (const x of positive_x_bbs) {
-      if (x.x0 < bb.x0)
-        bb.x0 = x.x0
-      if (x.x1 > bb.x1)
-        bb.x1 = x.x1
-    }
-
-    const positive_y_bbs = this.index.search(bbox.positive_y())
-    for (const y of positive_y_bbs) {
-      if (y.y0 < bb.y0)
-        bb.y0 = y.y0
-      if (y.y1 > bb.y1)
-        bb.y1 = y.y1
-    }
-
-    return this._bounds(bb)
+    const {x0, x1} = this.index.bounds(bbox.positive_x())
+    const {y0, y1} = this.index.bounds(bbox.positive_y())
+    return this._bounds({x0, y0, x1, y1})
   }
 
   get_anchor_point(anchor: Anchor, i: number, [sx, sy]: [number, number]): {x: number, y: number} | null {
@@ -240,6 +242,16 @@ export abstract class GlyphView extends View {
       data = data_subset
     }
 
+    let data_size = Infinity
+    for (const k in data) {
+      if (k.charAt(0) === '_')
+        data_size = Math.min(data_size, (data[k] as any).length)
+    }
+    if (data_size != Infinity)
+      this._data_size = data_size
+    else
+      this._data_size = 0 // XXX: this only happens in degenerate unit tests
+
     const self = this as any
     extend(self, data)
 
@@ -300,10 +312,13 @@ export abstract class GlyphView extends View {
 
   protected _set_data(_indices: number[] | null): void {}
 
-  protected abstract _index_data(): SpatialIndex
+  protected abstract _index_data(index: SpatialIndex): void
 
   index_data(): void {
-    this.index = this._index_data()
+    const index = new SpatialIndex(this.data_size)
+    this._index_data(index)
+    index.finish()
+    this._index = index
   }
 
   mask_data(indices: number[]): number[] {

bokehjs/src/lib/models/glyphs/harea.ts

@@ -2,7 +2,7 @@ import {PointGeometry} from 'core/geometry'
 import {Arrayable} from "core/types"
 import {Area, AreaView, AreaData} from "./area"
 import {Context2d} from "core/util/canvas"
-import {SpatialIndex, IndexedRect} from "core/util/spatial"
+import {SpatialIndex} from "core/util/spatial"
 import * as hittest from "core/hittest"
 import * as p from "core/properties"
 import {Selection} from "../selections/selection"
@@ -23,21 +23,20 @@ export class HAreaView extends AreaView {
   model: HArea
   visuals: HArea.Visuals
 
-  protected _index_data(): SpatialIndex {
-    const points: IndexedRect[] = []
+  protected _index_data(index: SpatialIndex): void {
+    const {min, max} = Math
+    const {data_size} = this
 
-    for (let i = 0, end = this._x1.length; i < end; i++) {
+    for (let i = 0; i < data_size; i++) {
       const x1 = this._x1[i]
       const x2 = this._x2[i]
       const y = this._y[i]
 
       if (isNaN(x1 + x2 + y) || !isFinite(x1 + x2 + y))
-        continue
-
-      points.push({x0: Math.min(x1, x2), y0: y, x1: Math.max(x1, x2), y1: y, i})
+        index.add_empty()
+      else
+        index.add(min(x1, x2), y, max(x1, x2), y)
     }
-
-    return new SpatialIndex(points)
   }
 
   protected _inner(ctx: Context2d, sx1: Arrayable<number>, sx2: Arrayable<number>, sy: Arrayable<number>, func: (this: Context2d) => void): void {

bokehjs/src/lib/models/glyphs/hbar.ts

@@ -1,7 +1,6 @@
 import {Box, BoxView, BoxData} from "./box"
 import {Arrayable} from "core/types"
 import * as p from "core/properties"
-import {SpatialIndex} from "core/util/spatial"
 
 export interface HBarData extends BoxData {
   _left: Arrayable<number>
@@ -33,10 +32,6 @@ export class HBarView extends BoxView {
     return this.sy[i]
   }
 
-  protected _index_data(): SpatialIndex {
-    return this._index_box(this._y.length)
-  }
-
   protected _lrtb(i: number): [number, number, number, number] {
     const l = Math.min(this._left[i], this._right[i])
     const r = Math.max(this._left[i], this._right[i])

bokehjs/src/lib/models/glyphs/hex_tile.ts

@@ -63,7 +63,7 @@ export class HexTileView extends GlyphView {
     }
   }
 
-  protected _index_data(): SpatialIndex {
+  protected _index_data(index: SpatialIndex): void {
     let ysize = this.model.size
     let xsize = Math.sqrt(3)*ysize/2
 
@@ -73,15 +73,17 @@ export class HexTileView extends GlyphView {
     } else
       xsize /= this.model.aspect_scale
 
-    const points = []
-    for (let i = 0; i < this._x.length; i++) {
+    const {data_size} = this
+
+    for (let i = 0; i < data_size; i++) {
       const x = this._x[i]
       const y = this._y[i]
-      if (isNaN(x+y) || !isFinite(x+y))
-        continue
-      points.push({x0: x-xsize, y0: y-ysize, x1: x+xsize, y1: y+ysize, i})
+
+      if (isNaN(x + y) || !isFinite(x + y))
+        index.add_empty()
+      else
+        index.add(x - xsize, y - ysize, x + xsize, y + ysize)
     }
-    return new SpatialIndex(points)
   }
 
   // overriding map_data instead of _map_data because the default automatic mappings

bokehjs/src/lib/models/glyphs/image_base.ts

@@ -83,16 +83,16 @@ export abstract class ImageBaseView extends XYGlyphView {
     }
   }
 
-  _index_data(): SpatialIndex {
-    const points = []
-    for (let i = 0, end = this._x.length; i < end; i++) {
+  protected _index_data(index: SpatialIndex): void {
+    const {data_size} = this
+
+    for (let i = 0; i < data_size; i++) {
       const [l, r, t, b] = this._lrtb(i)
-      if (isNaN(l + r + t + b) || !isFinite(l + r + t + b)) {
-        continue
-      }
-      points.push({x0: l, y0: b, x1: r, y1: t, i})
+      if (isNaN(l + r + t + b) || !isFinite(l + r + t + b))
+        index.add_empty()
+      else
+        index.add(l, b, r, t)
     }
-    return new SpatialIndex(points)
   }
 
   _lrtb(i: number): [number, number, number, number]{

bokehjs/src/lib/models/glyphs/image_url.ts

@@ -40,8 +40,13 @@ export class ImageURLView extends XYGlyphView {
     this.connect(this.model.properties.global_alpha.change, () => this.renderer.request_render())
   }
 
-  protected _index_data(): SpatialIndex {
-    return new SpatialIndex([])
+  protected _index_data(index: SpatialIndex): void {
+    const {data_size} = this
+
+    for (let i = 0; i < data_size; i++) {
+      // TODO: add a proper implementation (same as ImageBase?)
+      index.add_empty()
+    }
   }
 
   protected _set_data(): void {