Filters for lists (#142)

livebook-dev · Jan 20, 2024 · 4f7e4c1 · 4f7e4c1
1 parent 812faff
commit 4f7e4c1
Show file tree

Hide file tree

Showing 3 changed files with 130 additions and 9 deletions.
diff --git a/lib/assets/data_transform_cell/main.js b/lib/assets/data_transform_cell/main.js
@@ -753,13 +753,9 @@ export async function init(ctx, payload) {
         return data_options ? Object.keys(data_options) : [];
       },
       dataFrameColumnsWithTypes(data_options) {
-        const notAllowedTypes = ["list"];
-        const dataOptions = Object.fromEntries(
-          Object.entries(data_options).filter(
-            ([col, type]) => !notAllowedTypes.includes(type)
-          )
-        );
-        const dataFrameColumns = dataOptions ? Object.entries(dataOptions) : {};
+        const dataFrameColumns = data_options
+          ? Object.entries(data_options)
+          : {};
         const columns = Array.from(dataFrameColumns, ([name, type]) => {
           return { label: `${name} (${type})`, value: name };
         });

diff --git a/lib/kino_explorer/data_transform_cell.ex b/lib/kino_explorer/data_transform_cell.ex
@@ -34,7 +34,7 @@ defmodule KinoExplorer.DataTransformCell do
     "integer" => ["less", "less equal", "equal", "not equal", "greater equal", "greater"],
     "string" => ["equal", "contains", "not contains", "not equal"],
     "time" => ["less", "less equal", "equal", "not equal", "greater equal", "greater"],
-    "list" => []
+    "list" => ["contains", "not contains"]
   }
   @fill_missing_options %{
     "binary" => ["forward", "backward", "max", "min", "scalar"],
@@ -389,10 +389,11 @@ defmodule KinoExplorer.DataTransformCell do
     type = Map.get(data_options, column) || "string"
     default_value = if type == "boolean", do: "true"
     message = if field == "value", do: validation_message(:filters, type, value)
+    filter = if type == "list", do: "contains", else: "equal"
 
     if field == "column" do
       %{
-        "filter" => "equal",
+        "filter" => filter,
         "column" => column,
         "value" => default_value,
         "type" => type,
@@ -615,6 +616,15 @@ defmodule KinoExplorer.DataTransformCell do
   defp build_multiselect([columns]), do: [columns]
   defp build_multiselect(columns), do: [columns]
 
+  defp build_filter([column, filter, value, "list"] = args) do
+    with true <- Enum.all?(args, &(&1 != nil)),
+         {:ok, filter_value} <- cast_typed_value("list", value) do
+      build_filter_for_list(column, filter, filter_value)
+    else
+      _ -> nil
+    end
+  end
+
   defp build_filter([column, filter, "quantile(" <> <<value::bytes-size(1)>> <> ")", type]) do
     build_filter([column, filter, "quantile(0#{value})", type])
   end
@@ -771,6 +781,20 @@ defmodule KinoExplorer.DataTransformCell do
     end
   end
 
+  defp cast_typed_value("list", value) do
+    case Integer.parse(value) do
+      {int, rest} ->
+        if rest == ".0" or rest == "" do
+          {:ok, int}
+        else
+          cast_typed_value("float", value)
+        end
+
+      _ ->
+        nil
+    end
+  end
+
   defp cast_typed_value(type, value)
        when type in ["date", "datetime[ms]", "datetime[μs]", "datetime[ns]"],
        do: to_date(type, value)
@@ -987,4 +1011,12 @@ defmodule KinoExplorer.DataTransformCell do
   defp build_data_options(df) do
     df |> DataFrame.dtypes() |> normalize_dtypes()
   end
+
+  defp build_filter_for_list(column, "contains", value) do
+    {String.to_atom("member?"), [], [quoted_column(column), value]}
+  end
+
+  defp build_filter_for_list(column, "not contains", value) do
+    {String.to_atom("not member?"), [], [quoted_column(column), value]}
+  end
 end
diff --git a/test/kino_explorer/data_transform_cell_test.exs b/test/kino_explorer/data_transform_cell_test.exs
@@ -511,6 +511,99 @@ defmodule KinoExplorer.DataTransformCellTest do
              """
     end
 
+    test "source for a data frame with filter for lists" do
+      attrs =
+        build_attrs(%{
+          filters: [
+            %{
+              "column" => "list_column",
+              "filter" => "contains",
+              "type" => "list",
+              "value" => "3",
+              "active" => true,
+              "operation_type" => "filters"
+            }
+          ]
+        })
+
+      assert DataTransformCell.to_source(attrs) == """
+             people
+             |> Explorer.DataFrame.lazy()
+             |> Explorer.DataFrame.filter(member?(list_column, 3))
+             |> Explorer.DataFrame.collect()\
+             """
+    end
+
+    test "source for a data frame with multiple filter for lists" do
+      attrs =
+        build_attrs(%{
+          filters: [
+            %{
+              "column" => "list_column",
+              "filter" => "contains",
+              "type" => "list",
+              "value" => "3",
+              "active" => true,
+              "operation_type" => "filters"
+            },
+            %{
+              "column" => "list_column",
+              "filter" => "not contains",
+              "type" => "list",
+              "value" => "5",
+              "active" => true,
+              "operation_type" => "filters"
+            }
+          ]
+        })
+
+      assert DataTransformCell.to_source(attrs) == """
+             people
+             |> Explorer.DataFrame.lazy()
+             |> Explorer.DataFrame.filter(member?(list_column, 3) and not member?(list_column, 5))
+             |> Explorer.DataFrame.collect()\
+             """
+    end
+
+    test "do not generate code for invalid filter for lists" do
+      attrs =
+        build_attrs(%{
+          filters: [
+            %{
+              "column" => "list_column",
+              "filter" => "contains",
+              "type" => "list",
+              "value" => "3",
+              "active" => true,
+              "operation_type" => "filters"
+            },
+            %{
+              "column" => "list_column",
+              "filter" => "not contains",
+              "type" => "list",
+              "value" => "5",
+              "active" => true,
+              "operation_type" => "filters"
+            },
+            %{
+              "column" => "list_column",
+              "filter" => "not contains",
+              "type" => "list",
+              "value" => "cat",
+              "active" => true,
+              "operation_type" => "filters"
+            }
+          ]
+        })
+
+      assert DataTransformCell.to_source(attrs) == """
+             people
+             |> Explorer.DataFrame.lazy()
+             |> Explorer.DataFrame.filter(member?(list_column, 3) and not member?(list_column, 5))
+             |> Explorer.DataFrame.collect()\
+             """
+    end
+
     test "source for a data frame with fill_missing" do
       attrs =
         build_attrs(%{