@marimo-team/frontend 0.14.18-dev4 → 0.14.18-dev40

package/src/plugins/impl/data-explorer/functions/types.ts

@@ -0,0 +1,32 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+// Subset of aggregate operations that we support
+type AggregateOp =
+  | "average"
+  | "count"
+  | "distinct"
+  | "max"
+  | "mean"
+  | "median"
+  | "min"
+  | "q1"
+  | "q3"
+  | "stderr"
+  | "stdev"
+  | "sum";
+
+// Subset of time units that we support
+export type TimeUnitOp =
+  | "year"
+  | "month"
+  | "date"
+  | "day"
+  | "hours"
+  | "minutes"
+  | "seconds"
+  | "milliseconds"
+  | "yearmonth"
+  | "yearmonthdate"
+  | "monthdate";
+
+export type FieldFunction = AggregateOp | "bin" | TimeUnitOp;

package/src/plugins/impl/data-explorer/marks.ts

@@ -0,0 +1,21 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import type { SpecQuery } from "compassql/build/src/query/spec";
+import { SHORT_WILDCARD } from "compassql/build/src/wildcard";
+
+export type SpecMark = SpecQuery["mark"] | SHORT_WILDCARD;
+
+export const MARKS = [
+  SHORT_WILDCARD,
+  "area",
+  "bar",
+  "circle",
+  "geoshape",
+  "line",
+  "point",
+  "rect",
+  "rule",
+  "square",
+  "text",
+  "tick",
+  "trail",
+] as const;

package/src/plugins/impl/data-explorer/queries/field-suggestion.ts

@@ -0,0 +1,65 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import type { FieldQuery } from "compassql/build/src/query/encoding";
+import type { Query } from "compassql/build/src/query/query";
+import { SHORT_WILDCARD } from "compassql/build/src/wildcard";
+import type { QueryCreator, ResultingCharts } from "./types";
+
+// This code is adapted and simplified from https://github.com/vega/voyager
+
+function makeFieldSuggestionQueryCreator(params: {
+  type: keyof ResultingCharts;
+  limit: number;
+  additionalFieldQuery: FieldQuery;
+}): QueryCreator {
+  const { type, limit, additionalFieldQuery } = params;
+  return {
+    type,
+    limit: limit,
+    createQuery(query: Query): Query {
+      return {
+        spec: {
+          ...query.spec,
+          encodings: [...query.spec.encodings, additionalFieldQuery],
+        },
+        groupBy: "field",
+        orderBy: ["fieldOrder", "aggregationQuality", "effectiveness"],
+        chooseBy: ["aggregationQuality", "effectiveness"],
+        config: { autoAddCount: true },
+      };
+    },
+  };
+}
+
+export const addCategoricalField = makeFieldSuggestionQueryCreator({
+  type: "addCategoricalField",
+  limit: 4,
+  additionalFieldQuery: {
+    channel: SHORT_WILDCARD,
+    field: SHORT_WILDCARD,
+    type: "nominal",
+  },
+});
+
+export const addQuantitativeField = makeFieldSuggestionQueryCreator({
+  type: "addQuantitativeField",
+  limit: 4,
+  additionalFieldQuery: {
+    channel: SHORT_WILDCARD,
+    bin: SHORT_WILDCARD,
+    aggregate: SHORT_WILDCARD,
+    field: SHORT_WILDCARD,
+    type: "quantitative",
+  },
+});
+
+export const addTemporalField = makeFieldSuggestionQueryCreator({
+  type: "addTemporalField",
+  limit: 2,
+  additionalFieldQuery: {
+    channel: SHORT_WILDCARD,
+    hasFn: true, // Do not show raw time in the summary
+    timeUnit: SHORT_WILDCARD,
+    field: SHORT_WILDCARD,
+    type: "temporal",
+  },
+});

package/src/plugins/impl/data-explorer/queries/histograms.ts

@@ -0,0 +1,41 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import type { Query } from "compassql/build/src/query/query";
+import { SHORT_WILDCARD } from "compassql/build/src/wildcard";
+import type { QueryCreator } from "./types";
+
+// This code is adapted and simplified from https://github.com/vega/voyager
+
+export const histograms: QueryCreator = {
+  type: "histograms",
+  limit: 12,
+  createQuery(query: Query): Query {
+    return {
+      spec: {
+        data: query.spec.data,
+        mark: SHORT_WILDCARD,
+        transform: query.spec.transform,
+        encodings: [
+          {
+            channel: SHORT_WILDCARD,
+            bin: SHORT_WILDCARD,
+            timeUnit: SHORT_WILDCARD,
+            field: SHORT_WILDCARD,
+            type: SHORT_WILDCARD,
+          },
+          {
+            channel: SHORT_WILDCARD,
+            aggregate: "count",
+            field: "*",
+            type: "quantitative",
+          },
+        ],
+      },
+      groupBy: "fieldTransform",
+      orderBy: ["fieldOrder", "aggregationQuality", "effectiveness"],
+      chooseBy: ["aggregationQuality", "effectiveness"],
+      config: {
+        autoAddCount: false,
+      },
+    };
+  },
+};

package/src/plugins/impl/data-explorer/queries/queries.ts

@@ -0,0 +1,204 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { recommend } from "compassql/build/src";
+import type {
+  SpecQueryModel,
+  SpecQueryModelGroup,
+} from "compassql/build/src/model";
+import type { Query } from "compassql/build/src/query";
+import { isFieldQuery } from "compassql/build/src/query/encoding";
+import type { SpecQuery } from "compassql/build/src/query/spec";
+import { getTopResultTreeItem, isResultTree } from "compassql/build/src/result";
+import type { Schema } from "compassql/build/src/schema";
+import { contains } from "vega-lite";
+import { NONPOSITION_SCALE_CHANNELS } from "vega-lite/build/src/channel";
+import type { NamedData } from "vega-lite/build/src/data";
+import { type EncodingChannel, fromFieldQuery } from "../encoding";
+import { toSpecQuery } from "../spec";
+import type { ChartSpec } from "../state/types";
+import {
+  addCategoricalField,
+  addQuantitativeField,
+  addTemporalField,
+} from "./field-suggestion";
+import type {
+  PlotFieldInfo,
+  QueryCreator,
+  Result,
+  ResultingCharts,
+  ResultPlot,
+  TopLevelFacetedUnitSpec,
+} from "./types";
+import { hasWildcards, isQueryEmpty } from "./utils";
+
+// This code is adapted and simplified from https://github.com/vega/voyager
+
+const NAMED_DATA: NamedData = { name: "source" };
+
+function getFeaturesForRelatedViewRules(spec: SpecQuery) {
+  let hasOpenPosition = false;
+  let hasStyleChannel = false;
+  let hasOpenFacet = false;
+
+  spec.encodings.forEach((encQ) => {
+    if (encQ.channel === "x" || encQ.channel === "y") {
+      hasOpenPosition = true;
+    } else if (encQ.channel === "row" || encQ.channel === "column") {
+      hasOpenFacet = true;
+    } else if (
+      typeof encQ.channel === "string" &&
+      contains<string>(NONPOSITION_SCALE_CHANNELS, encQ.channel)
+    ) {
+      hasStyleChannel = true;
+    }
+  });
+
+  return {
+    hasOpenPosition,
+    hasStyleChannel,
+    hasOpenFacet,
+  };
+}
+
+export function allRelatedViewResults(
+  query: Query,
+  schema: Schema,
+): Partial<ResultingCharts> {
+  const charts: Partial<ResultingCharts> = {};
+
+  const { hasOpenPosition, hasStyleChannel, hasOpenFacet } =
+    getFeaturesForRelatedViewRules(query.spec);
+
+  if (hasOpenPosition || hasStyleChannel) {
+    charts.addQuantitativeField = relatedViewResult(
+      addQuantitativeField,
+      query,
+      schema,
+    );
+  }
+
+  if (hasOpenPosition || hasStyleChannel || hasOpenFacet) {
+    charts.addCategoricalField = relatedViewResult(
+      addCategoricalField,
+      query,
+      schema,
+    );
+  }
+
+  if (hasOpenPosition) {
+    charts.addTemporalField = relatedViewResult(
+      addTemporalField,
+      query,
+      schema,
+    );
+  }
+
+  // Hidden for now
+  // charts.alternativeEncodings = relatedViewResult(
+  //   alternativeEncodings,
+  //   query,
+  //   schema
+  // );
+  return charts;
+}
+
+export function mainViewResult(
+  mainQuery: Query,
+  schema: Schema,
+): Result | undefined {
+  if (isQueryEmpty(mainQuery.spec)) {
+    return {
+      plots: [],
+      query: mainQuery,
+      limit: 1,
+    };
+  }
+
+  const modelGroup = recommend(mainQuery, schema).result;
+  const plots = fromSpecQueryModelGroup(modelGroup, NAMED_DATA);
+
+  return {
+    plots: [plots[0]],
+    query: mainQuery,
+    limit: 1,
+  };
+}
+
+export function relatedViewResult(
+  queryCreator: QueryCreator,
+  mainQuery: Query,
+  schema: Schema,
+): Result {
+  const query = queryCreator.createQuery(mainQuery);
+
+  const modelGroup = recommend(query, schema).result;
+  const plots = fromSpecQueryModelGroup(modelGroup, NAMED_DATA);
+
+  return {
+    plots,
+    query,
+    limit: queryCreator.limit,
+  };
+}
+
+function fromSpecQueryModelGroup(
+  modelGroup: SpecQueryModelGroup,
+  data: NamedData,
+): ResultPlot[] {
+  return modelGroup.items.map((item) => {
+    if (isResultTree<SpecQueryModel>(item)) {
+      return toPlot(data, getTopResultTreeItem(item));
+    }
+    return toPlot(data, item);
+  });
+}
+
+function toPlot(data: NamedData, specQ: SpecQueryModel): ResultPlot {
+  const fieldInfos = specQ
+    .getEncodings()
+    .filter(isFieldQuery)
+    .map((fieldQ): PlotFieldInfo => {
+      return {
+        fieldDef: fromFieldQuery(fieldQ),
+        channel: fieldQ.channel as EncodingChannel,
+      };
+    });
+
+  return {
+    fieldInfos,
+    spec: specQ.toSpec(data) as TopLevelFacetedUnitSpec,
+  };
+}
+
+export function toQuery(params: {
+  spec: ChartSpec;
+  autoAddCount: boolean;
+}): Query {
+  const { spec, autoAddCount } = params;
+  const specQ = toSpecQuery(spec);
+  const { hasAnyWildcard, hasWildcardFn, hasWildcardField } =
+    hasWildcards(specQ);
+
+  const groupBy = getDefaultGroupBy({ hasWildcardFn, hasWildcardField });
+
+  return {
+    spec: specQ,
+    groupBy,
+    orderBy: ["fieldOrder", "aggregationQuality", "effectiveness"],
+    chooseBy: ["aggregationQuality", "effectiveness"],
+    config: hasAnyWildcard ? { autoAddCount } : undefined,
+  };
+}
+
+function getDefaultGroupBy(args: {
+  hasWildcardField: boolean;
+  hasWildcardFn: boolean;
+}) {
+  const { hasWildcardFn, hasWildcardField } = args;
+
+  return hasWildcardFn
+    ? "fieldTransform"
+    : hasWildcardField
+      ? "field"
+      : "encoding";
+}

package/src/plugins/impl/data-explorer/queries/removeUndefined.ts

@@ -0,0 +1,11 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+export function removeUndefined<T>(obj: T): T {
+  const ret: T = {} as T;
+  for (const key in obj) {
+    if (obj[key] !== undefined) {
+      ret[key] = obj[key];
+    }
+  }
+  return ret;
+}

package/src/plugins/impl/data-explorer/queries/types.ts

@@ -0,0 +1,54 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import type { Query } from "compassql/build/src/query";
+import type { VisualizationSpec } from "react-vega";
+import type { NamedData } from "vega-lite/build/src/data";
+import type { FacetedUnitSpec, TopLevel } from "vega-lite/build/src/spec";
+import type { EncodingChannel, FieldDefinition } from "../encoding";
+
+export interface PlotFieldInfo {
+  fieldDef: FieldDefinition;
+  channel: EncodingChannel;
+}
+
+export interface ResultPlot {
+  fieldInfos: PlotFieldInfo[];
+
+  /**
+   * Spec to be used for rendering.
+   */
+  spec: VisualizationSpec;
+}
+
+export interface Result {
+  plots: ResultPlot[] | null;
+
+  query: Query;
+
+  limit: number;
+}
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export type TopLevelFacetedUnitSpec = TopLevel<FacetedUnitSpec<any, any>> & {
+  data: NamedData;
+};
+
+export interface ResultingCharts {
+  main: Result;
+
+  histograms: Result;
+
+  addCategoricalField: Result;
+  addQuantitativeField: Result;
+  addTemporalField: Result;
+
+  alternativeEncodings: Result;
+  summaries: Result;
+}
+
+export interface QueryCreator {
+  type: keyof ResultingCharts;
+
+  limit: number;
+
+  createQuery(query: Query): Query;
+}

package/src/plugins/impl/data-explorer/queries/utils.ts

@@ -0,0 +1,85 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import {
+  isAutoCountQuery,
+  isValueQuery,
+} from "compassql/build/src/query/encoding";
+import type { SpecQuery } from "compassql/build/src/query/spec";
+import { isWildcard } from "compassql/build/src/wildcard";
+import type { EncodingChannel } from "../encoding";
+
+// This code is adapted and simplified from https://github.com/vega/voyager
+
+export function isQuerySpecific(spec: SpecQuery) {
+  return !hasWildcards(spec).hasAnyWildcard;
+}
+
+/**
+ * A query is full if it has "x", "y", and "color" channels
+ * and a facet
+ */
+export function isQueryFull(spec: SpecQuery) {
+  // Must have all these
+  const encodings: EncodingChannel[] = ["x", "y", "color"];
+  // Must have one of these
+  const facets: EncodingChannel[] = ["row", "column"];
+
+  const listedEncodings = new Set(spec.encodings.map((encQ) => encQ.channel));
+  const hasAllChannels = encodings.every((channel) =>
+    listedEncodings.has(channel),
+  );
+  const hasFacet = facets.some((facet) => listedEncodings.has(facet));
+
+  return hasAllChannels && hasFacet;
+}
+
+export function isQueryEmpty(spec: SpecQuery) {
+  return spec.encodings.length === 0;
+}
+
+interface HasWildcard {
+  hasAnyWildcard: boolean;
+  hasWildcardField: boolean;
+  hasWildcardFn: boolean;
+  hasWildcardChannel: boolean;
+}
+
+export function hasWildcards(spec: SpecQuery): HasWildcard {
+  let hasWildcardField = false;
+  let hasWildcardFn = false;
+  let hasWildcardChannel = false;
+
+  for (const encQ of spec.encodings) {
+    if (isValueQuery(encQ)) {
+      continue;
+    }
+
+    if (isAutoCountQuery(encQ)) {
+      if (isWildcard(encQ.autoCount)) {
+        hasWildcardFn = true;
+      }
+    } else {
+      // encQ is FieldQuery
+      if (isWildcard(encQ.field)) {
+        hasWildcardField = true;
+      }
+
+      if (
+        isWildcard(encQ.aggregate) ||
+        isWildcard(encQ.bin) ||
+        isWildcard(encQ.timeUnit)
+      ) {
+        hasWildcardFn = true;
+      }
+
+      if (isWildcard(encQ.channel)) {
+        hasWildcardChannel = true;
+      }
+    }
+  }
+  return {
+    hasAnyWildcard: hasWildcardChannel || hasWildcardField || hasWildcardFn,
+    hasWildcardField,
+    hasWildcardFn,
+    hasWildcardChannel,
+  };
+}

package/src/plugins/impl/data-explorer/spec.ts

@@ -0,0 +1,28 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import type { EncodingQuery } from "compassql/build/src/query/encoding";
+import type { SpecQuery } from "compassql/build/src/query/spec";
+import {
+  type EncodingChannel,
+  type SpecificEncoding,
+  toFieldQuery,
+} from "./encoding";
+import type { ChartSpec } from "./state/types";
+
+export function toSpecQuery(spec: ChartSpec): SpecQuery {
+  return {
+    mark: spec.mark,
+    encodings: adaptEncodings(spec.encoding),
+    config: spec.config,
+  };
+}
+
+function adaptEncodings(
+  encoding: SpecificEncoding | undefined,
+): EncodingQuery[] {
+  if (!encoding) {
+    return [];
+  }
+  return Object.entries(encoding).map(([channel, def]) => {
+    return toFieldQuery(def, channel as EncodingChannel);
+  });
+}

package/src/plugins/impl/data-explorer/state/reducer.ts

@@ -0,0 +1,92 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { Schema } from "compassql/build/src/schema";
+import { SHORT_WILDCARD } from "compassql/build/src/wildcard";
+import { atom } from "jotai";
+import { createReducerAndAtoms } from "@/utils/createReducer";
+import type { SpecificEncoding } from "../encoding";
+import type { SpecMark } from "../marks";
+import { histograms } from "../queries/histograms";
+import {
+  allRelatedViewResults,
+  mainViewResult,
+  relatedViewResult,
+  toQuery,
+} from "../queries/queries";
+import { removeUndefined } from "../queries/removeUndefined";
+import type { ResultingCharts } from "../queries/types";
+import { isQueryEmpty, isQueryFull, isQuerySpecific } from "../queries/utils";
+import type { ChartSpec } from "./types";
+
+function initialState(): ChartSpec {
+  return {
+    mark: SHORT_WILDCARD,
+    encoding: {},
+    config: {},
+    schema: null,
+  };
+}
+
+const { valueAtom: chartSpecAtom, useActions } = createReducerAndAtoms(
+  initialState,
+  {
+    setSchema: (state, schema: Schema) => {
+      return { ...state, schema };
+    },
+    setMark: (state, mark: SpecMark) => {
+      return { ...state, mark };
+    },
+    setEncoding: (state, encoding: SpecificEncoding) => {
+      // Merge and remove undefined values
+      const nextEncoding = removeUndefined({
+        ...state.encoding,
+        ...encoding,
+      });
+
+      return { ...state, encoding: nextEncoding };
+    },
+    set: (state, next: ChartSpec) => {
+      // remove schema
+      const { schema, ...rest } = next;
+      return { ...state, ...rest };
+    },
+  },
+);
+
+export { chartSpecAtom };
+
+export function useChartSpecActions() {
+  return useActions();
+}
+
+export const relatedChartSpecsAtom = atom<Partial<ResultingCharts>>((get) => {
+  const spec = get(chartSpecAtom);
+
+  // No schema
+  if (!spec.schema) {
+    return {};
+  }
+
+  const query = toQuery({ spec, autoAddCount: true });
+
+  // If the query is empty, we can show the main view and histograms
+  if (isQueryEmpty(query.spec)) {
+    return {
+      main: mainViewResult(query, spec.schema),
+      histograms: relatedViewResult(histograms, query, spec.schema),
+    };
+  }
+
+  // If the query is specific, we can show the main view and related views
+  if (isQuerySpecific(query.spec) && !isQueryFull(query.spec)) {
+    return {
+      main: mainViewResult(query, spec.schema),
+      ...allRelatedViewResults(query, spec.schema),
+    };
+  }
+
+  // Otherwise, we can only show the main view
+  return {
+    main: mainViewResult(query, spec.schema),
+  };
+});

package/src/plugins/impl/data-explorer/state/types.ts

@@ -0,0 +1,28 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { SpecQuery } from "compassql/build/src/query/spec";
+import type { Schema } from "compassql/build/src/schema";
+import type { SpecificEncoding } from "../encoding";
+import type { SpecMark } from "../marks";
+
+export interface ChartSpec {
+  /**
+   * The type of the mark.
+   */
+  mark: SpecMark;
+
+  /**
+   * Mapping between specific encoding channels and encoding definitions.
+   */
+  encoding: SpecificEncoding;
+
+  /**
+   * Configuration of the chart.
+   */
+  config: SpecQuery["config"];
+
+  /**
+   * The schema for the data.
+   */
+  schema?: Schema | null;
+}