Cache JSON when using pydeck local data

frontend/lib/src/components/elements/DeckGlJsonChart/DeckGlJsonChart.test.tsx

@@ -21,22 +21,29 @@ import { shallow } from "@streamlit/lib/src/test_util"
 import { DeckGlJsonChart as DeckGlJsonChartProto } from "@streamlit/lib/src/proto"
 import { NavigationControl } from "react-map-gl"
 import { mockTheme } from "@streamlit/lib/src/mocks/mockTheme"
-import { DeckGlJsonChart, PropsWithHeight } from "./DeckGlJsonChart"
+import { DeckGlJsonChart, PropsWithHeight, State } from "./DeckGlJsonChart"
+
+const mockInitialViewState = {
+  bearing: -27.36,
+  latitude: 52.2323,
+  longitude: -1.415,
+  maxZoom: 15,
+  minZoom: 5,
+  pitch: 40.5,
+  zoom: 6,
+}
+
+const mockElementId = "testId"
+jest.mock("@streamlit/lib/src/theme", () => ({
+  hasLightBackgroundColor: jest.fn(() => false),
+}))
 
 const getProps = (
   elementProps: Partial<DeckGlJsonChartProto> = {},
   initialViewStateProps: Record<string, unknown> = {}
 ): PropsWithHeight => {
   const json = {
-    initialViewState: {
-      bearing: -27.36,
-      latitude: 52.2323,
-      longitude: -1.415,
-      maxZoom: 15,
-      minZoom: 5,
-      pitch: 40.5,
-      zoom: 6,
-    },
+    initialViewState: mockInitialViewState,
     layers: [
       {
         "@@type": "HexagonLayer",
@@ -62,6 +69,7 @@ const getProps = (
 
   return {
     element: DeckGlJsonChartProto.create({
+      elementId: mockElementId,
       json: JSON.stringify(json),
       ...elementProps,
     }),
@@ -94,8 +102,10 @@ describe("DeckGlJsonChart element", () => {
       viewState: { pitch: 5, zoom: 5 },
     })
 
-    // @ts-expect-error
-    wrapper.setProps(getProps({}, { pitch: 40.5, zoom: 10 }))
+    wrapper.setProps(
+      // @ts-expect-error
+      getProps({ elementId: "newTestId" }, { pitch: 40.5, zoom: 10 })
+    )
 
     expect(wrapper.state("viewState")).toStrictEqual({
       pitch: 5,
@@ -142,4 +152,83 @@ describe("DeckGlJsonChart element", () => {
 
     expect(createdTooltip).toBe(false)
   })
+
+  describe("getDeckObject", () => {
+    const mockJsonParse = JSON.parse
+
+    const elementId = "newTestId"
+    const newJson = '{"initialViewState": {"height": "100", "width": "100"}}'
+    const isFullScreen = false
+    const isLightTheme = false
+    const initialized = false
+
+    const getNewState = (overrides?: Partial<State>): State => {
+      const defaultState: State = {
+        pydeckJson: JSON.parse(newJson),
+        isFullScreen,
+        viewState: {},
+        initialized,
+        initialViewState: mockInitialViewState,
+        elementId,
+        isLightTheme,
+      }
+
+      return { ...defaultState, ...overrides }
+    }
+
+    const originalState: State = {
+      pydeckJson: JSON.parse(newJson),
+      isFullScreen,
+      viewState: {},
+      initialized,
+      initialViewState: mockInitialViewState,
+      elementId: mockElementId,
+      isLightTheme,
+    }
+
+    beforeEach(() => {
+      JSON.parse = jest.fn(mockJsonParse)
+    })
+
+    afterEach(() => {
+      JSON.parse = mockJsonParse
+    })
+
+    it("should not call JSON.parse when the element id is the same", () => {
+      DeckGlJsonChart.getDeckObject(getProps(), originalState)
+
+      expect(JSON.parse).not.toHaveBeenCalled()
+
+      // state has different elementId from getProps
+      DeckGlJsonChart.getDeckObject(getProps(), getNewState())
+
+      expect(JSON.parse).toHaveBeenCalled()
+    })
+
+    it("should not call JSON.parse when FullScreen state changes", () => {
+      DeckGlJsonChart.getDeckObject(getProps(), originalState)
+
+      expect(JSON.parse).not.toHaveBeenCalled()
+
+      DeckGlJsonChart.getDeckObject(
+        getProps(),
+        getNewState({ isFullScreen: true })
+      )
+
+      expect(JSON.parse).toHaveBeenCalled()
+    })
+
+    it("should not call JSON.parse when theme state changes", () => {
+      DeckGlJsonChart.getDeckObject(getProps(), originalState)
+
+      expect(JSON.parse).not.toHaveBeenCalled()
+
+      DeckGlJsonChart.getDeckObject(
+        getProps(),
+        getNewState({ isLightTheme: true })
+      )
+
+      expect(JSON.parse).toHaveBeenCalled()
+    })
+  })
 })

frontend/lib/src/components/elements/DeckGlJsonChart/DeckGlJsonChart.tsx

@@ -39,8 +39,6 @@ import { registerLoaders } from "@loaders.gl/core"
 import withFullScreenWrapper from "@streamlit/lib/src/hocs/withFullScreenWrapper"
 import withMapboxToken from "@streamlit/lib/src/hocs/withMapboxToken"
 
-import { notNullOrUndefined } from "@streamlit/lib/src/util/utils"
-
 import { DeckGlJsonChart as DeckGlJsonChartProto } from "@streamlit/lib/src/proto"
 import {
   StyledDeckGlChart,
@@ -83,10 +81,14 @@ export interface PropsWithHeight extends Props {
   height?: number
 }
 
-interface State {
+export interface State {
   viewState: Record<string, unknown>
   initialized: boolean
   initialViewState: Record<string, unknown>
+  elementId: string | undefined
+  pydeckJson: any
+  isFullScreen: boolean
+  isLightTheme: boolean
 }
 
 export const DEFAULT_DECK_GL_HEIGHT = 500
@@ -100,6 +102,10 @@ export class DeckGlJsonChart extends PureComponent<PropsWithHeight, State> {
     },
     initialized: false,
     initialViewState: {},
+    elementId: undefined,
+    pydeckJson: undefined,
+    isFullScreen: false,
+    isLightTheme: hasLightBackgroundColor(this.props.theme),
   }
 
   componentDidMount = (): void => {
@@ -115,7 +121,7 @@ export class DeckGlJsonChart extends PureComponent<PropsWithHeight, State> {
     props: Readonly<PropsWithHeight>,
     state: Partial<State>
   ): Partial<State> | null {
-    const deck = DeckGlJsonChart.getDeckObject(props)
+    const deck = DeckGlJsonChart.getDeckObject(props, state)
 
     // If the ViewState on the server has changed, apply the diff to the current state
     if (!isEqual(deck.initialViewState, state.initialViewState)) {
@@ -144,35 +150,51 @@ export class DeckGlJsonChart extends PureComponent<PropsWithHeight, State> {
     return null
   }
 
-  static getDeckObject = (props: PropsWithHeight): DeckObject => {
+  static getDeckObject = (
+    props: PropsWithHeight,
+    state: Partial<State>
+  ): DeckObject => {
     const { element, width, height, theme } = props
-    const json = JSON.parse(element.json)
+    const { elementId } = state
+
+    const isFullScreen = Boolean(height)
+
+    // Only parse JSON when not transitioning to/from fullscreen, the element id changes, or theme changes
+    if (
+      element.elementId !== elementId ||
+      state.isFullScreen !== isFullScreen ||
+      state.isLightTheme !== hasLightBackgroundColor(theme)
+    ) {
+      state.pydeckJson = JSON.parse(element.json)
+      state.elementId = element.elementId
+    }
 
     // If unset, use either the Mapbox light or dark style based on Streamlit's theme
     // For Mapbox styles, see https://docs.mapbox.com/api/maps/styles/#mapbox-styles
-    if (!notNullOrUndefined(json.mapStyle)) {
-      const mapTheme = hasLightBackgroundColor(theme) ? "light" : "dark"
-      json.mapStyle = `mapbox://styles/mapbox/${mapTheme}-v9`
+    if (!state.pydeckJson?.mapStyle) {
+      state.pydeckJson.mapStyle = `mapbox://styles/mapbox/${
+        hasLightBackgroundColor(theme) ? "light" : "dark"
+      }-v9`
     }
 
-    // The graph dimensions could be set from props ( like withFullscreen ) or
-    // from the generated element object
-    if (height) {
-      json.initialViewState.height = height
-      json.initialViewState.width = width
+    // Set width and height based on the fullscreen state
+    if (isFullScreen) {
+      Object.assign(state.pydeckJson?.initialViewState, { width, height })
     } else {
-      if (!json.initialViewState.height) {
-        json.initialViewState.height = DEFAULT_DECK_GL_HEIGHT
+      if (!state.pydeckJson?.initialViewState?.height) {
+        state.pydeckJson.initialViewState.height = DEFAULT_DECK_GL_HEIGHT
       }
-
       if (element.useContainerWidth) {
-        json.initialViewState.width = width
+        state.pydeckJson.initialViewState.width = width
       }
     }
 
-    delete json.views // We are not using views. This avoids a console warning.
+    state.isFullScreen = isFullScreen
+    state.isLightTheme = hasLightBackgroundColor(theme)
+
+    delete state.pydeckJson?.views // We are not using views. This avoids a console warning.
 
-    return jsonConverter.convert(json)
+    return jsonConverter.convert(state.pydeckJson)
   }
 
   createTooltip = (info: PickingInfo): Record<string, unknown> | boolean => {
@@ -213,9 +235,8 @@ export class DeckGlJsonChart extends PureComponent<PropsWithHeight, State> {
   }
 
   render(): ReactNode {
-    const deck = DeckGlJsonChart.getDeckObject(this.props)
+    const deck = DeckGlJsonChart.getDeckObject(this.props, this.state)
     const { viewState } = this.state
-
     return (
       <StyledDeckGlChart
         className="stDeckGlJsonChart"

lib/streamlit/elements/deck_gl_json_chart.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import hashlib
 import json
 from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, cast
 
@@ -156,12 +157,18 @@ def marshall(
 ) -> None:
     if pydeck_obj is None:
         spec = json.dumps(EMPTY_MAP)
+        element_id = ""
     else:
         spec = pydeck_obj.to_json()
+        json_string = json.dumps(pydeck_obj.to_json())
+        json_bytes = json_string.encode("utf-8")
+        element_id = hashlib.md5(json_bytes).hexdigest()
 
     pydeck_proto.json = spec
     pydeck_proto.use_container_width = use_container_width
 
+    pydeck_proto.element_id = element_id
+
     tooltip = _get_pydeck_tooltip(pydeck_obj)
     if tooltip:
         pydeck_proto.tooltip = json.dumps(tooltip)

proto/streamlit/proto/DeckGlJsonChart.proto

@@ -17,11 +17,14 @@
 syntax = "proto3";
 
 message DeckGlJsonChart {
-  // The dataframe that will be used as the chart's main data source.
+  // The json of the pydeck object (https://deckgl.readthedocs.io/en/latest/deck.html)
   string json = 1;
 
   string tooltip = 2;
 
   // If True, will overwrite the chart width spec to fit to container.
   bool use_container_width = 4;
+
+  // the hash of the json so the the frontend doesn't always have to parse the pydeck json object
+  string element_id = 5;
 }