share types and function between table and dataframe for download (#6497)

## 📝 Summary

<!--
Provide a concise summary of what this pull request is addressing.

If this PR fixes any issues, list them here by number (e.g., Fixes
#123).
-->

This just helps with future maintainability by reusing types and
functions.

## 🔍 Description of Changes

<!--
Detail the specific changes made in this pull request. Explain the
problem addressed and how it was resolved. If applicable, provide before
and after comparisons, screenshots, or any relevant details to help
reviewers understand the changes easily.
-->

## 📋 Checklist

- [x] I have read the [contributor
guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md).
- [ ] For large changes, or changes that affect the public API: this
change was discussed or approved through an issue, on
[Discord](https://marimo.io/discord?ref=pr), or the community
[discussions](https://github.com/marimo-team/marimo/discussions) (Please
provide a link if applicable).
- [ ] I have added tests for the changes made.
- [x] I have run the code and verified that it works as expected.

Author: Light2Dark

frontend/src/components/data-table/schemas.ts

@@ -0,0 +1,16 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import z from "zod";
+import { rpc } from "@/plugins/core/rpc";
+
+export type DownloadAsArgs = (req: {
+  format: "csv" | "json" | "parquet";
+}) => Promise<string>;
+
+export const DownloadAsSchema = rpc
+  .input(
+    z.object({
+      format: z.enum(["csv", "json", "parquet"]),
+    }),
+  )
+  .output(z.string());

frontend/src/plugins/impl/DataTablePlugin.tsx

@@ -36,6 +36,10 @@ import {
 import { usePanelOwnership } from "@/components/data-table/hooks/use-panel-ownership";
 import { LoadingTable } from "@/components/data-table/loading-table";
 import { RowViewerPanel } from "@/components/data-table/row-viewer-panel/row-viewer";
+import {
+  type DownloadAsArgs,
+  DownloadAsSchema,
+} from "@/components/data-table/schemas";
 import {
   type BinValues,
   type ColumnHeaderStats,
@@ -191,7 +195,7 @@ interface Data<T> {
 
 // eslint-disable-next-line @typescript-eslint/consistent-type-definitions
 type DataTableFunctions = {
-  download_as: (req: { format: "csv" | "json" | "parquet" }) => Promise<string>;
+  download_as: DownloadAsArgs;
   get_column_summaries: <T>(
     opts: ColumnSummariesArgs,
   ) => Promise<ColumnSummaries<T>>;
@@ -265,9 +269,7 @@ export const DataTablePlugin = createPlugin<S>("marimo-table")
     }),
   )
   .withFunctions<DataTableFunctions>({
-    download_as: rpc
-      .input(z.object({ format: z.enum(["csv", "json", "parquet"]) }))
-      .output(z.string()),
+    download_as: DownloadAsSchema,
     get_column_summaries: rpc
       .input(z.object({ precompute: z.boolean() }))
       .output(

frontend/src/plugins/impl/data-frames/DataFramePlugin.tsx

@@ -4,6 +4,10 @@ import { isEqual } from "lodash-es";
 import { Code2Icon, DatabaseIcon, FunctionSquareIcon } from "lucide-react";
 import { type JSX, memo, useEffect, useRef, useState } from "react";
 import { z } from "zod";
+import {
+  type DownloadAsArgs,
+  DownloadAsSchema,
+} from "@/components/data-table/schemas";
 import type { FieldTypesWithExternalType } from "@/components/data-table/types";
 import { ReadonlyCode } from "@/components/editor/code/readonly-python-code";
 import { Spinner } from "@/components/icons/spinner";
@@ -68,7 +72,7 @@ type PluginFunctions = {
     data: TableData<T>;
     total_rows: number;
   }>;
-  download_as: (req: { format: "csv" | "json" | "parquet" }) => Promise<string>;
+  download_as: DownloadAsArgs;
 };
 
 // Value is selection, but it is not currently exposed to the user
@@ -127,13 +131,7 @@ export const DataFramePlugin = createPlugin<S>("marimo-dataframe")
           total_rows: z.number(),
         }),
       ),
-    download_as: rpc
-      .input(
-        z.object({
-          format: z.enum(["csv", "json", "parquet"]),
-        }),
-      )
-      .output(z.string()),
+    download_as: DownloadAsSchema,
   })
   .renderer((props) => (
     <TableProviders>
@@ -152,7 +150,7 @@ interface DataTableProps extends Data, PluginFunctions {
   setValue: (value: S) => void;
   host: HTMLElement;
   showDownload: boolean;
-  download_as: (req: { format: "csv" | "json" | "parquet" }) => Promise<string>;
+  download_as: DownloadAsArgs;
 }
 
 const EMPTY: Transformations = {

marimo/_plugins/ui/_impl/dataframes/dataframe.py

@@ -13,7 +13,6 @@
     Union,
 )
 
-import marimo._output.data.data as mo_data
 from marimo._output.rich_help import mddoc
 from marimo._plugins.ui._core.ui_element import UIElement
 from marimo._plugins.ui._impl.dataframes.transforms.apply import (
@@ -38,6 +37,7 @@
 from marimo._plugins.ui._impl.tables.utils import (
     get_table_manager,
 )
+from marimo._plugins.ui._impl.utils.dataframe import download_as
 from marimo._plugins.validators import (
     validate_no_integer_columns,
     validate_page_size,
@@ -298,18 +298,7 @@ def _download_as(self, args: DownloadAsArgs) -> str:
 
         # Get the table manager for the transformed data
         manager = self._get_cached_table_manager(df, self._limit)
-
-        ext = args.format
-        if ext == "csv":
-            return mo_data.csv(manager.to_csv()).url
-        elif ext == "json":
-            return mo_data.json(manager.to_json()).url
-        elif ext == "parquet":
-            return mo_data.parquet(manager.to_parquet()).url
-        else:
-            raise ValueError(
-                "format must be one of 'csv', 'json', or 'parquet'."
-            )
+        return download_as(manager, args.format)
 
     def _apply_filters_query_sort(
         self,

marimo/_plugins/ui/_impl/table.py

@@ -48,7 +48,11 @@
     TableManager,
 )
 from marimo._plugins.ui._impl.tables.utils import get_table_manager
-from marimo._plugins.ui._impl.utils.dataframe import ListOrTuple, TableData
+from marimo._plugins.ui._impl.utils.dataframe import (
+    ListOrTuple,
+    TableData,
+    download_as,
+)
 from marimo._plugins.validators import (
     validate_no_integer_columns,
     validate_page_size,
@@ -806,17 +810,9 @@ def _download_as(self, args: DownloadAsArgs) -> str:
 
         # Remove the selection column before downloading
         if isinstance(manager_candidate, TableManager):
-            manager = manager_candidate.drop_columns([INDEX_COLUMN_NAME])
-
-            ext = args.format
-            if ext == "csv":
-                return mo_data.csv(manager.to_csv()).url
-            elif ext == "json":
-                return mo_data.json(manager.to_json()).url
-            elif ext == "parquet":
-                return mo_data.parquet(manager.to_parquet()).url
-            else:
-                raise ValueError("format must be one of 'csv' or 'json'.")
+            return download_as(
+                manager_candidate, args.format, drop_marimo_index=True
+            )
         else:
             raise NotImplementedError(
                 "Download is not supported for this table format."

marimo/_plugins/ui/_impl/utils/dataframe.py

@@ -1,16 +1,16 @@
 # Copyright 2024 Marimo. All rights reserved.
 from __future__ import annotations
 
-from typing import (
-    TypeVar,
-    Union,
-)
+from typing import Any, TypeVar, Union
 
 from narwhals.typing import IntoDataFrame
 
 from marimo import _loggers
+from marimo._output.data import data as mo_data
 from marimo._output.mime import MIME
 from marimo._plugins.core.web_component import JSONType
+from marimo._plugins.ui._impl.tables.selection import INDEX_COLUMN_NAME
+from marimo._plugins.ui._impl.tables.table_manager import TableManager
 
 LOGGER = _loggers.marimo_logger()
 
@@ -26,3 +26,35 @@
     dict[str, ListOrTuple[JSONType]],
     IntoDataFrame,
 ]
+
+
+def download_as(
+    manager: TableManager[Any], ext: str, drop_marimo_index: bool = False
+) -> str:
+    """Download the table data in the specified format.
+
+    Args:
+        manager (TableManager[Any]): The table manager to download.
+        ext (str): The format to download the table data in.
+        drop_marimo_index (bool, optional): Whether to drop the marimo selection column.
+            Defaults to False.
+
+    Raises:
+        ValueError: If unrecognized format.
+        NotImplementedError: If the table format is not supported.
+
+    Returns:
+        str: The URL to download the table data.
+    """
+    if drop_marimo_index:
+        # Remove the selection column if exists
+        manager = manager.drop_columns([INDEX_COLUMN_NAME])
+
+    if ext == "csv":
+        return mo_data.csv(manager.to_csv()).url
+    elif ext == "json":
+        return mo_data.json(manager.to_json()).url
+    elif ext == "parquet":
+        return mo_data.parquet(manager.to_parquet()).url
+    else:
+        raise ValueError("format must be one of 'csv', 'json', or 'parquet'.")