Send WebAssembly binary over Jupyter WebSocket

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

@@ -1,22 +1,20 @@
 import * as React from "react";
 import { useState, useEffect } from "react";
 import { createRender, useModelState, useModel } from "@anywidget/react";
-import type { Initialize, Render } from "@anywidget/types";
+import type { Initialize, InitializeProps, Render } from "@anywidget/types";
 import Map from "react-map-gl/maplibre";
 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 { initParquetWasmFromBinary } 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,
@@ -190,7 +188,13 @@ function App() {
   );
 }
 
+async function initialize({ model }: InitializeProps): Promise<void> {
+  const parquetWasmBinary: DataView = model.get("_parquet_wasm_content");
+  await initParquetWasmFromBinary(parquetWasmBinary);
+}
+
 const module: { render: Render; initialize?: Initialize } = {
+  initialize,
   render: createRender(App),
 };
 

src/parquet.ts

@@ -1,19 +1,31 @@
-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() {
+// 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): Promise<void> {
   if (WASM_READY) {
     return;
   }
 
-  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;
 }