@marimo-team/frontend 0.25.3-dev13 → 0.25.3-dev134

package/src/components/data-table/charts/chart-spec/spec.ts

@@ -51,13 +51,19 @@ export type ErrorMessage = TypedString<"ErrorMessage">;
 export const X_AXIS_REQUIRED = "X-axis column is required" as ErrorMessage;
 export const Y_AXIS_REQUIRED = "Y-axis column is required" as ErrorMessage;
 
-export function createSpecWithoutData(
-  chartType: ChartType,
-  formValues: ChartSchemaType,
-  theme: ResolvedTheme,
-  width: number | "container",
-  height: number,
-): TopLevelSpec | ErrorMessage {
+export function createSpecWithoutData({
+  chartType,
+  formValues,
+  theme,
+  width,
+  height,
+}: {
+  chartType: ChartType;
+  formValues: ChartSchemaType;
+  theme: ResolvedTheme;
+  width: number | "container";
+  height: number;
+}): TopLevelSpec | ErrorMessage {
   const {
     xColumn,
     yColumn,
@@ -85,27 +91,25 @@ export function createSpecWithoutData(
   const yEncodingKey = "y";
 
   // Create encodings
-  const xEncoding = getAxisEncoding(
-    xColumn,
-    formValues.xAxis?.bin,
-    getFieldLabel(formValues.xAxis?.label),
-    colorByColumn?.field && horizontal ? stacking : undefined,
+  const xEncoding = getAxisEncoding(xColumn, {
+    binValues: formValues.xAxis?.bin,
+    label: getFieldLabel(formValues.xAxis?.label),
+    stack: colorByColumn?.field && horizontal ? stacking : undefined,
     chartType,
-  );
+  });
 
   let defaultYAggregation: ValidAggregationFn = DEFAULT_AGGREGATION;
   if (yColumn?.selectedDataType === "string") {
     defaultYAggregation = "count";
   }
 
-  const yEncoding = getAxisEncoding(
-    yColumn,
-    formValues.yAxis?.bin,
-    getFieldLabel(formValues.yAxis?.label),
-    colorByColumn?.field && !horizontal ? stacking : undefined,
+  const yEncoding = getAxisEncoding(yColumn, {
+    binValues: formValues.yAxis?.bin,
+    label: getFieldLabel(formValues.yAxis?.label),
+    stack: colorByColumn?.field && !horizontal ? stacking : undefined,
     chartType,
-    defaultYAggregation,
-  );
+    defaultAggregate: defaultYAggregation,
+  });
 
   const rowFacet = facet?.row.field
     ? getFacetEncoding(facet.row, chartType)
@@ -115,14 +119,14 @@ export function createSpecWithoutData(
     : undefined;
 
   const colorByEncoding = getColorEncoding(chartType, formValues);
-  const baseSpec = getBaseSpec(
+  const baseSpec = getBaseSpec({
     chartType,
     formValues,
     theme,
     width,
     height,
     title,
-  );
+  });
   const baseEncoding: Encoding<Field> = {
     [xEncodingKey]: horizontal ? yEncoding : xEncoding,
     [yEncodingKey]: horizontal ? xEncoding : yEncoding,
@@ -160,11 +164,19 @@ export function augmentSpecWithData(
 
 export function getAxisEncoding(
   column: NonNullable<z.infer<typeof AxisSchema>>,
-  binValues: z.infer<typeof BinSchema> | undefined,
-  label: string | undefined,
-  stack: boolean | undefined,
-  chartType: ChartType,
-  defaultAggregate?: ValidAggregationFn,
+  {
+    binValues,
+    label,
+    stack,
+    chartType,
+    defaultAggregate,
+  }: {
+    binValues: z.infer<typeof BinSchema> | undefined;
+    label: string | undefined;
+    stack: boolean | undefined;
+    chartType: ChartType;
+    defaultAggregate?: ValidAggregationFn;
+  },
 ): PositionDef<string> {
   const selectedDataType = column.selectedDataType || "string";
 
@@ -236,13 +248,12 @@ function getPieChartSpec(
     return "Size by column is required" as ErrorMessage;
   }
 
-  const thetaEncoding: PolarDef<string> = getAxisEncoding(
-    yColumn,
-    formValues.xAxis?.bin,
-    getFieldLabel(formValues.xAxis?.label),
-    undefined,
-    ChartType.PIE,
-  );
+  const thetaEncoding: PolarDef<string> = getAxisEncoding(yColumn, {
+    binValues: formValues.xAxis?.bin,
+    label: getFieldLabel(formValues.xAxis?.label),
+    stack: undefined,
+    chartType: ChartType.PIE,
+  });
 
   const colorEncoding: ColorDef<string> = {
     field: escapeFieldName(colorByColumn.field),
@@ -252,7 +263,14 @@ function getPieChartSpec(
   };
 
   return {
-    ...getBaseSpec(ChartType.PIE, formValues, theme, width, height, title),
+    ...getBaseSpec({
+      chartType: ChartType.PIE,
+      formValues,
+      theme,
+      width,
+      height,
+      title,
+    }),
     mark: {
       type: convertChartTypeToMark(ChartType.PIE),
       innerRadius: formValues.style?.innerRadius,
@@ -270,14 +288,21 @@ function getPieChartSpec(
   };
 }
 
-function getBaseSpec(
-  chartType: ChartType,
-  formValues: ChartSchemaType,
-  theme: ResolvedTheme,
-  width: number | "container",
-  height: number,
-  title?: string,
-): BaseSpec {
+function getBaseSpec({
+  chartType,
+  formValues,
+  theme,
+  width,
+  height,
+  title,
+}: {
+  chartType: ChartType;
+  formValues: ChartSchemaType;
+  theme: ResolvedTheme;
+  width: number | "container";
+  height: number;
+  title?: string;
+}): BaseSpec {
   let gridLines = formValues.style?.gridLines ?? false;
   // Scatter charts have grid lines by default
   if (chartType === ChartType.SCATTER) {

package/src/components/data-table/charts/charts.tsx

@@ -13,8 +13,8 @@ import {
   XIcon,
 } from "lucide-react";
 import type { JSX } from "react";
-import React, { useMemo, useRef, useState } from "react";
-import { type UseFormReturn, useForm } from "react-hook-form";
+import React, { useMemo, useState } from "react";
+import { type UseFormReturn, useForm, useWatch } from "react-hook-form";
 import useResizeObserver from "use-resize-observer";
 import { PythonIcon } from "@/components/editor/cell/code/icons";
 import { Alert, AlertDescription } from "@/components/ui/alert";
@@ -48,7 +48,12 @@ import { HeatmapForm } from "./forms/heatmap";
 import { PieForm } from "./forms/pie";
 import { LazyChart } from "./lazy-chart";
 import { ChartSchema, type ChartSchemaType, getChartDefaults } from "./schemas";
-import { getChartTabName, type TabName, tabsStorageAtom } from "./storage";
+import {
+  createTabId,
+  getChartTabName,
+  type TabName,
+  tabsStorageAtom,
+} from "./storage";
 import { ChartType } from "./types";
 
 const NEW_CHART_TYPE = "bar" as ChartType;
@@ -103,26 +108,28 @@ export const TablePanel: React.FC<TablePanelProps> = ({
 
   const [selectedTab, setSelectedTab] = useState(DEFAULT_TAB_NAME);
   const [tabCounter, setTabCounter] = useState(tabs.length);
-  const prevDisplayHeader = useRef(displayHeader);
+  const [prevDisplayHeader, setPrevDisplayHeader] = useState(displayHeader);
+
+  if (displayHeader !== prevDisplayHeader) {
+    setPrevDisplayHeader(displayHeader);
+  }
 
   // Auto-create a default chart tab when chart builder opens with no tabs
-  if (
-    displayHeader &&
-    !prevDisplayHeader.current &&
-    tabs.length === 0 &&
-    cellId
-  ) {
-    prevDisplayHeader.current = displayHeader;
+  if (displayHeader && !prevDisplayHeader && tabs.length === 0 && cellId) {
     const tabName = getChartTabName(0, NEW_CHART_TYPE);
     const newTabs = new Map(tabsMap);
     newTabs.set(cellId, [
-      { tabName, chartType: NEW_CHART_TYPE, config: getChartDefaults() },
+      {
+        id: createTabId(),
+        tabName,
+        chartType: NEW_CHART_TYPE,
+        config: getChartDefaults(),
+      },
     ]);
     saveTabsMap(newTabs);
     setTabCounter(1);
     setSelectedTab(tabName);
   }
-  prevDisplayHeader.current = displayHeader;
 
   if (!displayHeader || (tabs.length === 0 && !displayHeader)) {
     return dataTable;
@@ -138,6 +145,7 @@ export const TablePanel: React.FC<TablePanelProps> = ({
     newTabs.set(cellId, [
       ...tabs,
       {
+        id: createTabId(),
         tabName,
         chartType: NEW_CHART_TYPE,
         config: getChartDefaults(),
@@ -236,9 +244,9 @@ export const TablePanel: React.FC<TablePanelProps> = ({
         >
           Table
         </TabsTrigger>
-        {tabs.map((tab, idx) => (
+        {tabs.map((tab) => (
           <TabsTrigger
-            key={idx}
+            key={tab.id}
             className="text-xs"
             value={tab.tabName}
             onClick={() => setSelectedTab(tab.tabName)}
@@ -266,7 +274,7 @@ export const TablePanel: React.FC<TablePanelProps> = ({
       <TabsContent className="mt-1 overflow-hidden" value={DEFAULT_TAB_NAME}>
         {dataTable}
       </TabsContent>
-      {tabs.map((tab, idx) => {
+      {tabs.map((tab) => {
         const saveChart = (formValues: ChartSchemaType) => {
           saveTabChart({
             tabName: tab.tabName,
@@ -278,7 +286,11 @@ export const TablePanel: React.FC<TablePanelProps> = ({
           saveTabChartType(tab.tabName, chartType);
         };
         return (
-          <TabsContent key={idx} value={tab.tabName} className="h-[400px] mt-0">
+          <TabsContent
+            key={tab.id}
+            value={tab.tabName}
+            className="h-[400px] mt-0"
+          >
             <ChartPanel
               tableData={data}
               chartConfig={tab.config}
@@ -320,9 +332,10 @@ export const ChartPanel: React.FC<{
   saveChartType,
   getDataUrl,
   fieldTypes,
-  hasSchema = Boolean(fieldTypes?.length),
+  hasSchema: providedHasSchema,
   isLargeDataset,
 }) => {
+  const hasSchema = providedHasSchema ?? Boolean(fieldTypes?.length);
   const { theme } = useTheme();
   const form = useForm<ChartSchemaType>({
     defaultValues: chartConfig ?? getChartDefaults(),
@@ -370,20 +383,17 @@ export const ChartPanel: React.FC<{
     // Re-run when the data table changes
   }, [tableData, renderLargeCharts, csvFieldTypes]);
 
-  const formValues = form.watch();
+  // useWatch returns a new object only when a value changes. It is typed as
+  // deeply partial, but the form always starts from complete defaults.
+  const formValues = useWatch({ control: form.control }) as ChartSchemaType;
 
-  // This ensures the chart re-renders when the actual values change
-  const memoizedFormValues = useMemo(() => {
-    return structuredClone(formValues);
-  }, [formValues]);
-
-  const specWithoutData = createSpecWithoutData(
-    selectedChartType,
-    memoizedFormValues,
+  const specWithoutData = createSpecWithoutData({
+    chartType: selectedChartType,
+    formValues,
     theme,
-    "container",
-    CHART_HEIGHT,
-  );
+    width: "container",
+    height: CHART_HEIGHT,
+  });
 
   // Prevent unnecessary re-renders of the chart
   const memoizedChart = useMemo(() => {

package/src/components/data-table/charts/storage.ts

@@ -7,13 +7,16 @@ import { Logger } from "@/utils/Logger";
 import { NotebookScopedLocalStorage } from "@/utils/storage/typed";
 import { capitalize } from "@/utils/strings";
 import type { TypedString } from "@/utils/typed";
+import { generateUUID } from "@/utils/uuid";
 import { ChartSchema, type ChartSchemaType } from "./schemas";
 import type { ChartType } from "./types";
 
 export type TabName = TypedString<"TabName">;
+export type TabId = TypedString<"TabId">;
 export const KEY = "marimo:charts:v2";
 
 interface TabStorage {
+  id: TabId;
   tabName: TabName; // unique within cell
   chartType: ChartType;
   config: ChartSchemaType;
@@ -21,11 +24,15 @@ interface TabStorage {
 
 type TabStorageMap = Map<CellId, TabStorage[]>;
 
-const TabStorageSchema = z.object({
-  tabName: z.string().transform((name) => name as TabName),
-  chartType: z.string().transform((type) => type as ChartType),
-  config: ChartSchema,
-});
+const TabStorageSchema = z
+  .object({
+    id: z.string().optional(),
+    tabName: z.string().transform((name) => name as TabName),
+    chartType: z.string().transform((type) => type as ChartType),
+    config: ChartSchema,
+  })
+  // Stored tabs without an id use their name, which is unique within a cell.
+  .transform((tab) => ({ ...tab, id: (tab.id ?? tab.tabName) as TabId }));
 const TabStorageEntriesSchema = z.array(
   z.tuple([
     z.string().transform((name) => name as CellId),
@@ -64,6 +71,10 @@ export const tabsStorageAtom = atomWithStorage<TabStorageMap>(
   mapStorage,
 );
 
+export function createTabId(): TabId {
+  return generateUUID() as TabId;
+}
+
 /**
  * Convenience function to get the tab name for a given tab number and chart type
  */

package/src/components/data-table/column-explorer-panel/column-explorer.tsx

@@ -345,7 +345,7 @@ const ColumnPreview = ({
   const chart = chart_spec && renderChart(chart_spec, theme);
 
   const addDataframeChart = chart_code && (
-    <AddDataframeChart chartCode={chart_code} />
+    <AddDataframeChart chartCode={chart_code} columnName={columnName} />
   );
 
   return (

package/src/components/data-table/column-summary/chart-spec-model.tsx

@@ -119,16 +119,16 @@ export class ColumnChartSpecModel<T> {
   private columnBinValues = new Map<ColumnName, BinValues>();
   private columnValueCounts = new Map<ColumnName, ValueCounts>();
 
-  public static readonly EMPTY = new ColumnChartSpecModel(
-    [],
-    new Map(),
-    {},
-    {},
-    {},
-    {
+  public static readonly EMPTY = new ColumnChartSpecModel({
+    data: [],
+    fieldTypes: new Map(),
+    stats: {},
+    binValues: {},
+    valueCounts: {},
+    opts: {
       includeCharts: false,
     },
-  );
+  });
 
   private dataSpec: TopLevelSpec["data"];
 
@@ -144,16 +144,23 @@ export class ColumnChartSpecModel<T> {
     includeCharts: boolean;
   };
 
-  public constructor(
-    data: T[] | string,
-    fieldTypes: FieldTypes,
-    stats: Record<ColumnName, Partial<ColumnHeaderStats>>,
-    binValues: Record<ColumnName, BinValues>,
-    valueCounts: Record<ColumnName, ValueCounts>,
+  public constructor({
+    data,
+    fieldTypes,
+    stats,
+    binValues,
+    valueCounts,
+    opts,
+  }: {
+    data: T[] | string;
+    fieldTypes: FieldTypes;
+    stats: Record<ColumnName, Partial<ColumnHeaderStats>>;
+    binValues: Record<ColumnName, BinValues>;
+    valueCounts: Record<ColumnName, ValueCounts>;
     opts: {
       includeCharts: boolean;
-    },
-  ) {
+    };
+  }) {
     this.fieldTypes = fieldTypes;
     this.stats = stats;
     this.binValues = binValues;

package/src/components/data-table/column-summary/column-summary.tsx

@@ -16,7 +16,7 @@ export const ColumnChartContext = React.createContext<
   ColumnChartSpecModel<unknown>
 >(ColumnChartSpecModel.EMPTY);
 
-interface Props<TData, TValue> {
+interface Props {
   columnId: string;
 }
 
@@ -24,9 +24,7 @@ interface Props<TData, TValue> {
 // for all calls with the same key.
 const batchedLoader = createBatchedLoader();
 
-export const TableColumnSummary = <TData, TValue>({
-  columnId,
-}: Props<TData, TValue>) => {
+export const TableColumnSummary = ({ columnId }: Props) => {
   const { locale } = useLocale();
 
   const chartSpecModel = React.use(ColumnChartContext);

package/src/components/data-table/column-visibility-dropdown.tsx

@@ -6,7 +6,7 @@
 
 import type { Table } from "@tanstack/react-table";
 import { Columns3Icon, EyeIcon, EyeOffIcon, ScanEyeIcon } from "lucide-react";
-import React, { useMemo } from "react";
+import React from "react";
 import { ColumnName } from "@/components/datasources/components";
 import { Button } from "@/components/ui/button";
 import {
@@ -24,6 +24,7 @@ import {
 } from "@/components/ui/popover";
 import { type BulkAction, useSelectList } from "@/components/ui/select-core";
 import type { DataType } from "@/core/kernel/messages";
+import { Arrays } from "@/utils/arrays";
 import { cn } from "@/utils/cn";
 import { Events } from "@/utils/events";
 import { smartMatchFilter } from "@/utils/smartMatch";
@@ -67,9 +68,8 @@ export const ColumnVisibilityDropdown = <TData,>({
   const hideableIds = userColumns
     .filter((column) => column.getCanHide())
     .map((column) => column.id);
-  const visibleHideableCount = hideableIds.filter(
-    (id) => !hiddenIds.includes(id),
-  ).length;
+  const visibleHideableColumnIds = Arrays.difference(hideableIds, hiddenIds);
+  const visibleHideableCount = visibleHideableColumnIds.length;
 
   const applyHidden = (next: string[] | string | null) => {
     const hidden = new Set(Array.isArray(next) ? next : []);
@@ -106,15 +106,10 @@ export const ColumnVisibilityDropdown = <TData,>({
           .filter((option) => !option.disabled)
           .map((option) => option.value)
       : [];
-  const showOnlyColumnState = useMemo(
-    () => ({
-      visibleHideableColumnIds: hideableIds.filter(
-        (id) => !hiddenIds.includes(id),
-      ),
-      hideableColumnIdSet: new Set(hideableIds),
-    }),
-    [hideableIds, hiddenIds],
-  );
+  const showOnlyColumnState = {
+    visibleHideableColumnIds,
+    hideableColumnIdSet: new Set(hideableIds),
+  };
   const hasSearchBulkActions =
     showOnlyMatchIds.length > 0 || matchingActions.length > 0;