Send WebAssembly binary over Jupyter WebSocket

Author: kylebarron

lonboard/_map.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import sys
+from functools import lru_cache
 from pathlib import Path
 from typing import IO, TYPE_CHECKING, Optional, Sequence, TextIO, Union
 
@@ -46,6 +47,13 @@
 """
 
 
+@lru_cache
+def _load_parquet_wasm_binary() -> bytes:
+    """Load the gzipped parquet-wasm binary blob"""
+    with open(bundler_output_dir / "arrow2_bg.wasm.gz", "rb") as f:
+        return f.read()
+
+
 class Map(BaseAnyWidget):
     """
     The top-level class used to display a map in a Jupyter Widget.
@@ -91,10 +99,14 @@ def __init__(
         if isinstance(layers, BaseLayer):
             layers = [layers]
 
-        super().__init__(layers=layers, **kwargs)
+        _parquet_wasm_content = _load_parquet_wasm_binary()
+        super().__init__(
+            layers=layers, _parquet_wasm_content=_parquet_wasm_content, **kwargs
+        )
 
     _esm = bundler_output_dir / "index.js"
     _css = bundler_output_dir / "index.css"
+    _parquet_wasm_content = traitlets.Bytes(allow_none=False).tag(sync=True)
 
     view_state = ViewStateTrait()
     """

package.json

@@ -31,7 +31,7 @@
     "vitest": "^1.4.0"
   },
   "scripts": {
-    "build": "node ./build.mjs",
+    "build": "node ./build.mjs && gzip -c node_modules/parquet-wasm/esm/arrow2_bg.wasm > lonboard/static/arrow2_bg.wasm.gz",
     "build:watch": "nodemon --watch src/ --exec \"npm run build\" --ext js,json,ts,tsx,css",
     "fmt:check": "prettier './src/**/*.{ts,tsx,css}' --check",
     "fmt": "prettier './src/**/*.{ts,tsx,css}' --write",

pyproject.toml

@@ -6,7 +6,13 @@ authors = ["Kyle Barron <[email protected]>"]
 license = "MIT"
 readme = "README.md"
 packages = [{ include = "lonboard" }]
-include = ["lonboard/static/*.js", "lonboard/static/*.css", "MANIFEST.in"]
+include = [
+    "lonboard/static/*.js",
+    "lonboard/static/*.css",
+    "lonboard/static/*.wasm",
+    "lonboard/static/*.wasm.gz",
+    "MANIFEST.in",
+]
 
 [tool.poetry.dependencies]
 python = "^3.8"

src/index.tsx

@@ -7,16 +7,14 @@ import DeckGL from "@deck.gl/react/typed";
 import { MapViewState, type Layer } from "@deck.gl/core/typed";
 import { BaseLayerModel, initializeLayer } from "./model/index.js";
 import type { WidgetModel } from "@jupyter-widgets/base";
-import { initParquetWasm } from "./parquet.js";
+import { useParquetWasm } from "./parquet.js";
 import { getTooltip } from "./tooltip/index.js";
 import { isDefined, loadChildModels } from "./util.js";
 import { v4 as uuidv4 } from "uuid";
 import { Message } from "./types.js";
 import { flyTo } from "./actions/fly-to.js";
 import { useViewStateDebounced } from "./state";
 
-await initParquetWasm();
-
 const DEFAULT_INITIAL_VIEW_STATE = {
   latitude: 10,
   longitude: 0,
@@ -65,6 +63,10 @@ async function getChildModelState(
 function App() {
   let model = useModel();
 
+  let [parquetWasmBinary] = useModelState<DataView | null>(
+    "_parquet_wasm_content",
+  );
+  let [parquetWasmReady] = useParquetWasm(parquetWasmBinary);
   let [mapStyle] = useModelState<string>("basemap_style");
   let [mapHeight] = useModelState<number>("_height");
   let [showTooltip] = useModelState<boolean>("show_tooltip");
@@ -108,7 +110,15 @@ function App() {
   let [stateCounter, setStateCounter] = useState<Date>(new Date());
 
   useEffect(() => {
+    if (!parquetWasmReady) {
+      return;
+    }
+
     const callback = async () => {
+      if (!parquetWasmReady) {
+        throw new Error("inside callback but parquetWasm not ready!");
+      }
+
       const childModels = await loadChildModels(
         model.widget_manager,
         childLayerIds,
@@ -122,7 +132,7 @@ function App() {
       setSubModelState(newSubModelState);
     };
     callback().catch(console.error);
-  }, [childLayerIds]);
+  }, [parquetWasmReady, childLayerIds]);
 
   const layers: Layer[] = [];
   for (const subModel of Object.values(subModelState)) {

src/parquet.ts

@@ -1,20 +1,36 @@
 import { useEffect, useState } from "react";
-import _initParquetWasm, { readParquet } from "parquet-wasm/esm/arrow2";
+import { initSync, readParquet } from "parquet-wasm/esm/arrow2";
 import * as arrow from "apache-arrow";
 
-// NOTE: this version must be synced exactly with the parquet-wasm version in
-// use.
-const PARQUET_WASM_VERSION = "0.5.0";
-const PARQUET_WASM_CDN_URL = `https://cdn.jsdelivr.net/npm/parquet-wasm@${PARQUET_WASM_VERSION}/esm/arrow2_bg.wasm`;
 let WASM_READY: boolean = false;
 
-export async function initParquetWasm() {
-  if (WASM_READY) {
-    return;
+// https://developer.mozilla.org/en-US/docs/Web/API/Compression_Streams_API
+async function decompressBlob(blob: Blob) {
+  const ds = new DecompressionStream("gzip");
+  const decompressedStream = blob.stream().pipeThrough(ds);
+  return await new Response(decompressedStream).blob();
+}
+
+/**
+ * Initialize parquet-wasm from an existing WASM binary blob.
+ * It is expected that this WASM has been gzipped
+ *
+ * @return Whether initialization succeeded
+ */
+export async function initParquetWasmFromBinary(
+  view: DataView | null,
+): Promise<boolean> {
+  if (!view) {
+    return false;
   }
 
-  await _initParquetWasm(PARQUET_WASM_CDN_URL);
+  let blob = new Blob([view]);
+  const decompressedBlob = await decompressBlob(blob);
+  const decompressedBuffer = await decompressedBlob.arrayBuffer();
+
+  initSync(decompressedBuffer);
   WASM_READY = true;
+  return true;
 }
 
 /**
@@ -58,3 +74,21 @@ export function parseParquetBuffers(dataViews: DataView[]): arrow.Table {
 
   return new arrow.Table(batches);
 }
+
+export function useParquetWasm(view: DataView | null): [boolean] {
+  const [wasmReady, setWasmReady] = useState<boolean>(false);
+
+  // Init parquet wasm
+  useEffect(() => {
+    const callback = async () => {
+      const succeeded = await initParquetWasmFromBinary(view);
+      if (succeeded) {
+        setWasmReady(true);
+      }
+    };
+
+    callback();
+  }, []);
+
+  return [wasmReady];
+}