@marimo-team/frontend 0.14.18-dev5 → 0.14.18-dev50

package/src/plugins/impl/data-frames/panel.tsx

@@ -0,0 +1,308 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { zodResolver } from "@hookform/resolvers/zod";
+import {
+  ArrowUpDownIcon,
+  BracketsIcon,
+  ColumnsIcon,
+  CombineIcon,
+  CopySlashIcon,
+  FileJsonIcon,
+  FilterIcon,
+  FunctionSquareIcon,
+  GroupIcon,
+  MousePointerSquareDashedIcon,
+  PencilIcon,
+  PlusIcon,
+  ShuffleIcon,
+  SquareMousePointerIcon,
+  Trash2Icon,
+} from "lucide-react";
+import React, { type PropsWithChildren, useEffect, useMemo } from "react";
+import { useFieldArray, useForm } from "react-hook-form";
+import useEvent from "react-use-event-hook";
+import type { z } from "zod";
+import {
+  ColumnFetchValuesContext,
+  ColumnInfoContext,
+} from "@/plugins/impl/data-frames/forms/context";
+import { Strings } from "@/utils/strings";
+import { ZodForm } from "../../../components/forms/form";
+import {
+  getDefaults,
+  getUnionLiteral,
+} from "../../../components/forms/form-utils";
+import { Button } from "../../../components/ui/button";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuGroup,
+  DropdownMenuItem,
+  DropdownMenuLabel,
+  DropdownMenuSeparator,
+  DropdownMenuTrigger,
+} from "../../../components/ui/dropdown-menu";
+import { cn } from "../../../utils/cn";
+import { DATAFRAME_FORM_RENDERERS } from "./forms/renderers";
+import {
+  type Transformations,
+  TransformationsSchema,
+  type TransformType,
+  TransformTypeSchema,
+} from "./schema";
+import type { ColumnDataTypes } from "./types";
+import { getUpdatedColumnTypes } from "./utils/getUpdatedColumnTypes";
+
+interface Props {
+  columns: ColumnDataTypes;
+  initialValue: Transformations;
+  onChange: (value: Transformations) => void;
+  onInvalidChange: (value: Transformations) => void;
+  getColumnValues: (req: { column: string }) => Promise<{
+    values: unknown[];
+    too_many_values: boolean;
+  }>;
+}
+
+export const TransformPanel: React.FC<Props> = ({
+  initialValue,
+  columns,
+  onChange,
+  onInvalidChange,
+  getColumnValues,
+}) => {
+  const form = useForm<z.infer<typeof TransformationsSchema>>({
+    resolver: zodResolver(TransformationsSchema),
+    defaultValues: initialValue,
+    mode: "onChange",
+    reValidateMode: "onChange",
+  });
+  const { handleSubmit, watch, control } = form;
+
+  const onSubmit = useEvent((values: z.infer<typeof TransformationsSchema>) => {
+    onChange(values);
+  });
+
+  const onInvalidSubmit = useEvent(
+    (values: z.infer<typeof TransformationsSchema>) => {
+      onInvalidChange(values);
+    },
+  );
+
+  useEffect(() => {
+    const subscription = watch(() => {
+      handleSubmit(onSubmit, () => {
+        onInvalidSubmit(form.getValues());
+      })();
+    });
+    return () => subscription.unsubscribe();
+  }, [handleSubmit, watch, onInvalidSubmit, onSubmit, form]);
+
+  const [selectedTransform, setSelectedTransform] = React.useState<
+    number | undefined
+  >(initialValue.transforms.length > 0 ? 0 : undefined);
+
+  // TODO: This crashes in latest version of react-hook-form
+  const transformsField = useFieldArray({
+    control: control,
+    name: "transforms",
+  });
+
+  const transforms = form.watch("transforms");
+  const selectedTransformType =
+    selectedTransform === undefined
+      ? undefined
+      : transforms[selectedTransform]?.type;
+  const selectedTransformSchema = TransformTypeSchema._def.options.find(
+    (option) => {
+      return getUnionLiteral(option)._def.value === selectedTransformType;
+    },
+  );
+
+  const effectiveColumns = useMemo(() => {
+    const transformsBeforeSelected = transforms.slice(0, selectedTransform);
+    return getUpdatedColumnTypes(transformsBeforeSelected, columns);
+  }, [columns, transforms, selectedTransform]);
+
+  const handleAddTransform = (transform: z.ZodType) => {
+    const next: TransformType = getDefaults(transform);
+    const nextIdx = transformsField.fields.length;
+    transformsField.append(next);
+    setSelectedTransform(nextIdx);
+  };
+
+  return (
+    <ColumnInfoContext value={effectiveColumns}>
+      <ColumnFetchValuesContext value={getColumnValues}>
+        <form
+          onSubmit={(e) => e.preventDefault()}
+          className="flex flex-row max-h-[400px] overflow-hidden bg-background"
+        >
+          <Sidebar
+            items={form.watch("transforms")}
+            selected={selectedTransform}
+            onSelect={(index) => {
+              setSelectedTransform(index);
+            }}
+            onDelete={(index) => {
+              transformsField.remove(index);
+              const indexBefore = index - 1;
+              setSelectedTransform(Math.max(indexBefore, 0));
+            }}
+            onAdd={handleAddTransform}
+          />
+          <div className="flex flex-col flex-1 p-4 overflow-auto min-h-[200px] border-l">
+            {selectedTransform !== undefined && selectedTransformSchema && (
+              <ZodForm
+                key={`transforms.${selectedTransform}`}
+                form={form}
+                schema={selectedTransformSchema}
+                path={`transforms.${selectedTransform}`}
+                renderers={DATAFRAME_FORM_RENDERERS}
+              />
+            )}
+            {(selectedTransform === undefined || !selectedTransformSchema) && (
+              <div className="flex flex-col items-center justify-center grow gap-3">
+                <MousePointerSquareDashedIcon className="w-8 h-8  text-muted-foreground" />
+                <AddTransformDropdown onAdd={handleAddTransform}>
+                  <Button
+                    data-testid="marimo-plugin-data-frames-add-transform"
+                    variant="text"
+                    size="xs"
+                  >
+                    <div className="text-sm">Select a transform to begin</div>
+                  </Button>
+                </AddTransformDropdown>
+              </div>
+            )}
+          </div>
+        </form>
+      </ColumnFetchValuesContext>
+    </ColumnInfoContext>
+  );
+};
+
+interface SidebarProps {
+  items: TransformType[];
+  selected: number | undefined;
+  onSelect: (index: number) => void;
+  onDelete: (index: number) => void;
+  onAdd: (transform: z.ZodType) => void;
+}
+
+export const Sidebar: React.FC<SidebarProps> = ({
+  items,
+  selected,
+  onAdd,
+  onSelect,
+  onDelete,
+}) => {
+  return (
+    <div className="flex flex-col overflow-y-hidden w-[180px] shadow-xs h-full">
+      <div className="flex flex-col overflow-y-auto grow">
+        {items.map((item, idx) => {
+          return (
+            <div
+              key={`${JSON.stringify(item)}-${idx}`}
+              onClick={() => {
+                onSelect(idx);
+              }}
+              className={cn(
+                "flex flex-row min-h-[40px] items-center px-2 cursor-pointer hover:bg-accent/50 text-sm overflow-hidden hover-actions-parent border border-muted border-l-2 border-l-transparent",
+                {
+                  "border-l-primary bg-accent text-accent-foreground":
+                    selected === idx,
+                },
+              )}
+            >
+              <div className="grow text-ellipsis">
+                {Strings.startCase(item.type)}
+              </div>
+              <Trash2Icon
+                className="w-3 h-3 hover-action text-muted-foreground hover:text-destructive"
+                onClick={(e) => {
+                  onDelete(idx);
+                  e.stopPropagation();
+                }}
+              />
+            </div>
+          );
+        })}
+      </div>
+      <div className="flex flex-row shrink-0">
+        <AddTransformDropdown onAdd={onAdd}>
+          <Button
+            data-testid="marimo-plugin-data-frames-add-transform"
+            variant="text"
+            className="w-full rounded-none m-0 hover:text-accent-foreground"
+            size="xs"
+          >
+            <PlusIcon className="w-3 h-3 mr-1" />
+            Add
+          </Button>
+        </AddTransformDropdown>
+      </div>
+    </div>
+  );
+};
+
+const AddTransformDropdown: React.FC<
+  PropsWithChildren<{ onAdd: (transform: z.ZodType) => void }>
+> = ({ onAdd, children }) => {
+  return (
+    <DropdownMenu modal={false}>
+      <DropdownMenuTrigger asChild={true}>{children}</DropdownMenuTrigger>
+      <DropdownMenuContent className="w-56">
+        <DropdownMenuLabel>Add Transform</DropdownMenuLabel>
+        <DropdownMenuSeparator />
+        <DropdownMenuGroup>
+          {Object.values(TransformTypeSchema._def.options).map((type) => {
+            const literal = getUnionLiteral(type);
+            const Icon = ICONS[literal._def.value as TransformType["type"]];
+            return (
+              <DropdownMenuItem
+                key={literal._def.value}
+                onSelect={(evt) => {
+                  evt.stopPropagation();
+                  onAdd(type);
+                }}
+              >
+                <Icon className="w-3.5 h-3.5 mr-2" />
+                <span>{Strings.startCase(literal._def.value)}</span>
+              </DropdownMenuItem>
+            );
+          })}
+          <DropdownMenuItem
+            key="_request_"
+            onSelect={(evt) => {
+              evt.stopPropagation();
+              window.open(
+                "https://github.com/marimo-team/marimo/issues/new?title=New%20dataframe%20transform:&labels=enhancement&template=feature_request.yaml",
+                "_blank",
+              );
+            }}
+          >
+            <span className="underline text-primary text-xs cursor-pointer">
+              Request a transform
+            </span>
+          </DropdownMenuItem>
+        </DropdownMenuGroup>
+      </DropdownMenuContent>
+    </DropdownMenu>
+  );
+};
+
+const ICONS: Record<TransformType["type"], React.FC<{ className?: string }>> = {
+  aggregate: FunctionSquareIcon,
+  column_conversion: ColumnsIcon,
+  filter_rows: FilterIcon,
+  group_by: GroupIcon,
+  rename_column: PencilIcon,
+  select_columns: SquareMousePointerIcon,
+  sort_column: ArrowUpDownIcon,
+  shuffle_rows: ShuffleIcon,
+  sample_rows: CombineIcon,
+  explode_columns: BracketsIcon,
+  expand_dict: FileJsonIcon,
+  unique: CopySlashIcon,
+};

package/src/plugins/impl/data-frames/schema.ts

@@ -0,0 +1,211 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { z } from "zod";
+import { FieldOptions, randomNumber } from "@/components/forms/options";
+import { DATA_TYPES } from "@/core/kernel/messages";
+import {
+  AGGREGATION_FNS,
+  type ColumnId,
+  NUMPY_DTYPES,
+} from "@/plugins/impl/data-frames/types";
+import {
+  ALL_OPERATORS,
+  isConditionValueValid,
+  type OperatorType,
+} from "./utils/operators";
+
+export const columnToFieldTypesSchema = z.array(
+  z.tuple([z.coerce.string(), z.tuple([z.enum(DATA_TYPES), z.string()])]),
+);
+
+export const column_id = z
+  .string()
+  .min(1, "Required")
+  .or(z.number())
+  .transform((v) => v as ColumnId)
+  .describe(FieldOptions.of({ label: "Column", special: "column_id" }));
+
+export const column_id_array = z
+  .array(column_id.describe(FieldOptions.of({ special: "column_id" })))
+  .min(1, "At least one column is required")
+  .default([])
+  .describe(FieldOptions.of({ label: "Columns" }));
+
+const ColumnConversionTransformSchema = z
+  .object({
+    type: z.literal("column_conversion"),
+    column_id: column_id,
+    data_type: z
+      .enum(NUMPY_DTYPES)
+      .describe(FieldOptions.of({ label: "Data type (numpy)" }))
+      .default("bool"),
+    errors: z
+      .enum(["ignore", "raise"])
+      .default("ignore")
+      .describe(
+        FieldOptions.of({ label: "Handle errors", special: "radio_group" }),
+      ),
+  })
+  .describe(FieldOptions.of({}));
+
+const RenameColumnTransformSchema = z.object({
+  type: z.literal("rename_column"),
+  column_id: column_id,
+  new_column_id: z
+    .string()
+    .min(1, "Required")
+    .transform((v) => v as ColumnId)
+    .describe(FieldOptions.of({ label: "New column name" })),
+});
+
+const SortColumnTransformSchema = z.object({
+  type: z.literal("sort_column"),
+  column_id: column_id,
+  ascending: z
+    .boolean()
+    .describe(FieldOptions.of({ label: "Ascending" }))
+    .default(true),
+  na_position: z
+    .enum(["first", "last"])
+    .describe(
+      FieldOptions.of({ label: "N/A position", special: "radio_group" }),
+    )
+    .default("last"),
+});
+
+export const ConditionSchema = z
+  .object({
+    column_id: column_id,
+    operator: z
+      .enum(Object.keys(ALL_OPERATORS) as [OperatorType, ...OperatorType[]])
+      .describe(FieldOptions.of({ label: " " })),
+    value: z.any().describe(FieldOptions.of({ label: "Value" })),
+  })
+  .describe(FieldOptions.of({ direction: "row", special: "column_filter" }));
+export type ConditionType = z.infer<typeof ConditionSchema>;
+
+const FilterRowsTransformSchema = z.object({
+  type: z.literal("filter_rows"),
+  operation: z
+    .enum(["keep_rows", "remove_rows"])
+    .default("keep_rows")
+    .describe(FieldOptions.of({ special: "radio_group" })),
+  where: z
+    .array(ConditionSchema)
+    .min(1)
+    .transform((value) => {
+      return value.filter((condition) => {
+        return isConditionValueValid(condition.operator, condition.value);
+      });
+    })
+    .describe(FieldOptions.of({ label: "Value" }))
+    .default([{ column_id: "", operator: "==", value: "" }]),
+});
+
+const GroupByTransformSchema = z
+  .object({
+    type: z.literal("group_by"),
+    column_ids: column_id_array,
+    aggregation: z
+      .enum(AGGREGATION_FNS)
+      .default("count")
+      .describe(FieldOptions.of({ label: "Aggregation" })),
+    drop_na: z
+      .boolean()
+      .default(false)
+      .describe(FieldOptions.of({ label: "Drop N/A" })),
+  })
+  .describe(FieldOptions.of({}));
+
+const AggregateTransformSchema = z
+  .object({
+    type: z.literal("aggregate"),
+    column_ids: column_id_array,
+    aggregations: z
+      .array(z.enum(AGGREGATION_FNS))
+      .min(1, "At least one aggregation is required")
+      .default(["count"])
+      .describe(FieldOptions.of({ label: "Aggregations" })),
+  })
+  .describe(FieldOptions.of({ direction: "row" }));
+
+const SelectColumnsTransformSchema = z.object({
+  type: z.literal("select_columns"),
+  column_ids: column_id_array,
+});
+
+const SampleRowsTransformSchema = z.object({
+  type: z.literal("sample_rows"),
+  n: z
+    .number()
+    .positive()
+    .describe(FieldOptions.of({ label: "Number of rows" })),
+  seed: z
+    .number()
+    .default(() => randomNumber())
+    .describe(
+      FieldOptions.of({ label: "Re-sample", special: "random_number_button" }),
+    ),
+  replace: z
+    .boolean()
+    .default(false)
+    .describe(
+      FieldOptions.of({
+        label: "Sample with replacement",
+      }),
+    ),
+});
+
+const ShuffleRowsTransformSchema = z.object({
+  type: z.literal("shuffle_rows"),
+  seed: z
+    .number()
+    .default(() => randomNumber())
+    .describe(
+      FieldOptions.of({ label: "Re-shuffle", special: "random_number_button" }),
+    ),
+});
+
+const ExplodeColumnsTransformSchema = z.object({
+  type: z.literal("explode_columns"),
+  column_ids: column_id_array,
+});
+
+const ExpandDictTransformSchema = z.object({
+  type: z.literal("expand_dict"),
+  column_id: column_id,
+});
+
+const UniqueTransformSchema = z
+  .object({
+    type: z.literal("unique"),
+    column_ids: column_id_array,
+    keep: z
+      .enum(["first", "last", "none", "any"])
+      .default("first")
+      .describe(FieldOptions.of({ label: "Keep" })),
+  })
+  .describe(FieldOptions.of({ direction: "row" }));
+
+export const TransformTypeSchema = z.union([
+  FilterRowsTransformSchema,
+  SelectColumnsTransformSchema,
+  RenameColumnTransformSchema,
+  ColumnConversionTransformSchema,
+  SortColumnTransformSchema,
+  GroupByTransformSchema,
+  AggregateTransformSchema,
+  SampleRowsTransformSchema,
+  ShuffleRowsTransformSchema,
+  ExplodeColumnsTransformSchema,
+  ExpandDictTransformSchema,
+  UniqueTransformSchema,
+]);
+
+export type TransformType = z.infer<typeof TransformTypeSchema>;
+
+export const TransformationsSchema = z.object({
+  transforms: z.array(TransformTypeSchema),
+});
+
+export type Transformations = z.infer<typeof TransformationsSchema>;

package/src/plugins/impl/data-frames/types.ts

@@ -0,0 +1,91 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { DataType } from "@/core/kernel/messages";
+import { logNever } from "@/utils/assertNever";
+
+/**
+ * Strongly typed string/number
+ */
+export type ColumnId = (string | number) & { __columnId: "columnId" };
+
+export type NumPyType = string;
+
+/**
+ * List of column Id and their data types
+ *
+ * We cannot use a js map, since maps don't preserve keys as ints (e.g. "1" and 1 are the same key)
+ */
+export type RawColumnDataTypes = Array<[ColumnId, [DataType, NumPyType]]>;
+/**
+ * Map of column Id and their data types
+ * ES6 maps preserve keys as ints (e.g. "1" and 1 are different keys)
+ */
+export type ColumnDataTypes = Map<ColumnId, NumPyType>;
+
+export const NUMPY_DTYPES = [
+  "int8",
+  "int16",
+  "int32",
+  "int64",
+  "uint8",
+  "uint16",
+  "uint32",
+  "uint64",
+  "float16",
+  "float32",
+  "float64",
+  "complex64",
+  "complex128",
+  "bool",
+  "object",
+  "str",
+  "unicode",
+  "datetime64",
+  "timedelta64",
+] as const;
+
+type NumpyType = (typeof NUMPY_DTYPES)[number];
+
+export function numpyTypeToDataType(nptype: NumpyType): DataType {
+  nptype = nptype.toLowerCase() as NumpyType;
+
+  switch (nptype) {
+    case "int8":
+    case "int16":
+    case "int32":
+    case "int64":
+    case "uint8":
+    case "uint16":
+    case "uint32":
+    case "uint64":
+      return "integer";
+    case "float16":
+    case "float32":
+    case "float64":
+      return "number";
+    case "complex64":
+    case "complex128":
+    case "object":
+    case "str":
+    case "unicode":
+      return "string";
+    case "bool":
+      return "boolean";
+    case "datetime64":
+      return "datetime";
+    case "timedelta64":
+      return "date";
+    default:
+      logNever(nptype);
+      return "unknown";
+  }
+}
+
+export const AGGREGATION_FNS = [
+  "count",
+  "sum",
+  "mean",
+  "median",
+  "min",
+  "max",
+] as const;