Add support for Enum type

.github/workflows/release.yml

@@ -36,18 +36,7 @@ jobs:
           # https://docs.github.com/en/actions/how-tos/write-workflows/choose-where-workflows-run/...
           # choose-the-runner-for-a-job#standard-github-hosted-runners-for-public-repositories
           - { target: aarch64-apple-darwin, os: macos-15 }
-          - { target: aarch64-unknown-linux-gnu, os: ubuntu-22.04, use-cross: true }
-          - { target: aarch64-unknown-linux-musl, os: ubuntu-22.04, use-cross: true }
-          - { target: x86_64-apple-darwin, os: macos-15 }
-          - { target: x86_64-pc-windows-gnu, os: windows-2025, rustflags: "-C target-feature=+fxsr,+sse,+sse2,+sse3,+ssse3,+sse4.1,+sse4.2,+popcnt,+avx,+fma" }
-          - { target: x86_64-pc-windows-gnu, os: windows-2025, variant: "legacy_cpu" }
-          - { target: x86_64-pc-windows-msvc, os: windows-2025, rustflags: "-C target-feature=+fxsr,+sse,+sse2,+sse3,+ssse3,+sse4.1,+sse4.2,+popcnt,+avx,+fma" }
-          - { target: x86_64-pc-windows-msvc, os: windows-2025, variant: "legacy_cpu" }
-          - { target: x86_64-unknown-freebsd, os: ubuntu-22.04, use-cross: true, rustflags: "-C target-feature=+fxsr,+sse,+sse2,+sse3,+ssse3,+sse4.1,+sse4.2,+popcnt,+avx,+fma" }
-          - { target: x86_64-unknown-freebsd, os: ubuntu-22.04, use-cross: true, variant: "legacy_cpu" }
           - { target: x86_64-unknown-linux-gnu, os: ubuntu-22.04, rustflags: "-C target-feature=+fxsr,+sse,+sse2,+sse3,+ssse3,+sse4.1,+sse4.2,+popcnt,+avx,+fma" }
-          - { target: x86_64-unknown-linux-gnu, os: ubuntu-22.04, variant: "legacy_cpu" }
-          - { target: x86_64-unknown-linux-musl, os: ubuntu-22.04, use-cross: true }
 
     steps:
       - name: Checkout source code

README.md

@@ -264,37 +264,17 @@ mix test --only cloud_integration
 
 ## Precompilation
 
-Explorer ships with the NIF code precompiled for the most popular architectures out there.
+Explorer ships with the NIF code precompiled for the supported production and development targets.
 We support the following:
 
 - `aarch64-apple-darwin` - MacOS running on ARM 64 bits CPUs.
-- `aarch64-unknown-linux-gnu` - Linux running on ARM 64 bits CPUs, compiled with GCC.
-- `aarch64-unknown-linux-musl` - Linux running on ARM 64 bits CPUs, compiled with Musl.
-- `x86_64-apple-darwin` - MacOS running on Intel/AMD 64 bits CPUs.
-- `x86_64-pc-windows-msvc` - Windows running on Intel/AMD 64 bits CPUs, compiled with Visual C++.
-- `x86_64-pc-windows-gnu` - Windows running on Intel/AMD 64 bits CPUs, compiled with GCC.
 - `x86_64-unknown-linux-gnu` - Linux running on Intel/AMD 64 bits CPUs, compiled with GCC.
-- `x86_64-unknown-linux-musl` - Linux running on Intel/AMD 64 bits CPUs, compiled with Musl.
-- `x86_64-unknown-freebsd` - FreeBSD running on Intel/AMD 64 bits.
 
 This means that Explorer is going to work without the need to compile it from source.
 
 This currently **only works for Hex releases**. For more information on how it works, please
 check the [RustlerPrecompiled project](https://hexdocs.pm/rustler_precompiled).
 
-### Legacy CPUs
-
-We ship some of the precompiled artifacts with modern CPU features enabled by default. But in
-case your computer is not compatible with them, you can set an application environment that is
-going to be read at compile time, enabling the legacy variants of artifacts.
-
-```elixir
-config :explorer, use_legacy_artifacts: true
-```
-
-If you see the error message "Illegal instruction" after your project compiles, you need to
-enable the legacy artifacts.
-
 ### Features disabled
 
 Some of the features cannot be compiled to some targets, because one of the dependencies

checksum-Elixir.Explorer.PolarsBackend.Native.exs

@@ -0,0 +1,6 @@
+%{
+  "libexplorer-v0.13.0-nif-2.15-aarch64-apple-darwin.so.tar.gz" =>
+    "sha256:3e565666bd519a58dcdd715714e555cc14101883ebdde4e3e1292adfd03723d9",
+  "libexplorer-v0.13.0-nif-2.15-x86_64-unknown-linux-gnu.so.tar.gz" =>
+    "sha256:20c192064ea27db483f7e21affa322f503526e5091d30f4e471153682610bb9d"
+}

lib/explorer/polars_backend/native.ex

@@ -8,48 +8,14 @@ defmodule Explorer.PolarsBackend.Native do
   # We want "debug" in dev and test because it's faster to compile.
   mode = if Mix.env() in [:dev, :test], do: :debug, else: :release
 
-  use_legacy =
-    Application.compile_env(
-      :explorer,
-      :use_legacy_artifacts,
-      System.get_env("EXPLORER_USE_LEGACY_ARTIFACTS") in ["true", "1"]
-    )
-
-  variants_for_linux = [
-    legacy_cpu: fn ->
-      # These are the same from the release workflow.
-      # See the meaning in: https://unix.stackexchange.com/a/43540
-      needed_caps = ~w[fxsr sse sse2 ssse3 sse4_1 sse4_2 popcnt avx fma]
-
-      use_legacy or
-        (is_nil(use_legacy) and
-           not Explorer.ComptimeUtils.cpu_with_all_caps?(needed_caps))
-    end
-  ]
-
-  other_variants = [legacy_cpu: fn -> use_legacy end]
-
   use RustlerPrecompiled,
     otp_app: :explorer,
     version: version,
     base_url: "#{github_url}/releases/download/v#{version}",
     targets: ~w(
       aarch64-apple-darwin
-      aarch64-unknown-linux-gnu
-      aarch64-unknown-linux-musl
-      x86_64-apple-darwin
-      x86_64-pc-windows-msvc
-      x86_64-pc-windows-gnu
       x86_64-unknown-linux-gnu
-      x86_64-unknown-linux-musl
-      x86_64-unknown-freebsd
     ),
-    variants: %{
-      "x86_64-unknown-linux-gnu" => variants_for_linux,
-      "x86_64-pc-windows-msvc" => other_variants,
-      "x86_64-pc-windows-gnu" => other_variants,
-      "x86_64-unknown-freebsd" => other_variants
-    },
     # We don't use any features of newer NIF versions, so 2.15 is enough.
     nif_versions: ["2.15"],
     mode: mode,
@@ -399,6 +365,7 @@ defmodule Explorer.PolarsBackend.Native do
   def s_from_list_str(_name, _val), do: err()
   def s_from_list_binary(_name, _val), do: err()
   def s_from_list_categories(_name, _val), do: err()
+  def s_from_list_enum(_name, _val, _categories), do: err()
   def s_from_list_decimal(_name, _val, _precision, _scale), do: err()
   def s_from_list_of_series(_name, _val, _dtype), do: err()
   def s_from_list_of_series_as_structs(_name, _val, _dtype), do: err()

lib/explorer/polars_backend/series.ex

@@ -14,6 +14,9 @@ defmodule Explorer.PolarsBackend.Series do
   defguardp is_non_finite(n) when n in [:nan, :infinity, :neg_infinity]
   defguardp is_numeric(n) when is_number(n) or is_non_finite(n)
 
+  defguardp is_enum_dtype(dtype)
+            when is_tuple(dtype) and tuple_size(dtype) == 2 and elem(dtype, 0) == :enum
+
   @integer_types Explorer.Shared.integer_types()
 
   # Conversion
@@ -62,8 +65,9 @@ defmodule Explorer.PolarsBackend.Series do
   def size(series), do: Shared.apply_series(series, :s_size)
 
   @impl true
-  def categories(%Series{dtype: :category} = series),
-    do: Shared.apply_series(series, :s_categories)
+  def categories(%Series{dtype: dtype} = series)
+      when dtype == :category or is_enum_dtype(dtype),
+      do: Shared.apply_series(series, :s_categories)
 
   @impl true
   def categorise(%Series{dtype: {integer_type, _}} = series, %Series{dtype: dtype} = categories)

lib/explorer/polars_backend/shared.ex

@@ -182,6 +182,7 @@ defmodule Explorer.PolarsBackend.Shared do
       :boolean -> Native.s_from_list_bool(name, list)
       :string -> Native.s_from_list_str(name, list)
       :category -> Native.s_from_list_categories(name, list)
+      {:enum, categories} -> apply(:s_from_list_enum, [name, list, categories])
       :date -> apply(:s_from_list_date, [name, list])
       :time -> apply(:s_from_list_time, [name, list])
       {:naive_datetime, precision} -> apply(:s_from_list_naive_datetime, [name, list, precision])

lib/explorer/series.ex

@@ -16,6 +16,8 @@ defmodule Explorer.Series do
     * `:binary` - Binaries (sequences of bytes)
     * `:boolean` - Boolean
     * `:category` - Strings but represented internally as integers
+    * `{:enum, categories}` - Strings restricted to the given categories and represented
+      internally as integers
     * `:date` - Date type that unwraps to `Elixir.Date`
     * `{:naive_datetime, precision}` - Naive DateTime type with millisecond/microsecond/nanosecond
       precision that unwraps to `Elixir.NaiveDateTime`
@@ -172,6 +174,7 @@ defmodule Explorer.Series do
           | datetime_dtype
           | decimal_dtype
           | duration_dtype
+          | enum_dtype
           | float_dtype
           | list_dtype
           | naive_datetime_dtype
@@ -181,6 +184,7 @@ defmodule Explorer.Series do
 
   @type time_unit :: :nanosecond | :microsecond | :millisecond
   @type time_zone :: String.t()
+  @type enum_dtype :: {:enum, [String.t()]}
   @type naive_datetime_dtype :: {:naive_datetime, time_unit}
   @type datetime_dtype :: {:datetime, time_unit, time_zone}
   @type duration_dtype :: {:duration, time_unit}
@@ -232,6 +236,11 @@ defmodule Explorer.Series do
                  |> K.and(elem(dtype, 0) == :decimal)
                  |> K.and(elem(dtype, 2) |> K.is_integer())
 
+  defguardp is_enum_dtype(dtype)
+            when is_tuple(dtype)
+                 |> K.and(tuple_size(dtype) == 2)
+                 |> K.and(elem(dtype, 0) == :enum)
+
   defguardp is_numeric_dtype(dtype)
             when K.or(K.in(dtype, @numeric_dtypes), is_decimal_dtype(dtype))
 
@@ -473,6 +482,14 @@ defmodule Explorer.Series do
         category ["EUA", "Brazil", "Poland"]
       >
 
+  You can also create an enum series by providing its allowed categories:
+
+      iex> Explorer.Series.from_list(["EUA", "Brazil"], dtype: {:enum, ["EUA", "Brazil", "Poland"]})
+      #Explorer.Series<
+        Polars[2]
+        enum ["EUA", "Brazil"]
+      >
+
   If you need to create a series of dates, you can pass `Date` structs, but also
   a series of integers representing days since Unix Epoch:
 
@@ -927,7 +944,7 @@ defmodule Explorer.Series do
   @doc type: :conversion
   @spec to_iovec(series :: Series.t()) :: [binary]
   def to_iovec(%Series{dtype: dtype} = series) do
-    if is_io_dtype(dtype) do
+    if K.or(is_io_dtype(dtype), is_enum_dtype(dtype)) do
       apply_series(series, :to_iovec, [], false)
     else
       raise ArgumentError, "cannot convert series of dtype #{inspect(dtype)} into iovec"
@@ -1294,6 +1311,7 @@ defmodule Explorer.Series do
   def iotype(%Series{dtype: dtype}) do
     case dtype do
       :category -> {:u, 32}
+      {:enum, _categories} -> {:u, 32}
       {:decimal, _, _} -> {:s, 128}
       other -> Shared.dtype_to_iotype(other)
     end
@@ -1324,8 +1342,12 @@ defmodule Explorer.Series do
   """
   @doc type: :introspection
   @spec categories(series :: Series.t()) :: Series.t()
-  def categories(%Series{dtype: :category} = series), do: apply_series(series, :categories)
-  def categories(%Series{dtype: dtype}), do: dtype_error("categories/1", dtype, [:category])
+  def categories(%Series{dtype: dtype} = series)
+      when K.or(dtype == :category, is_enum_dtype(dtype)),
+      do: apply_series(series, :categories)
+
+  def categories(%Series{dtype: dtype}),
+    do: dtype_error("categories/1", dtype, [:category, :enum])
 
   @doc """
   Categorise a series of integers or strings according to `categories`.
@@ -4930,15 +4952,15 @@ defmodule Explorer.Series do
         Polars[3 x 3]
         values f64 [1.0, 2.0, 3.0]
         break_point f64 [1.5, 2.5, Inf]
-        category string ["(-inf, 1.5]", "(1.5, 2.5]", "(2.5, inf]"]
+        category enum ["(-inf, 1.5]", "(1.5, 2.5]", "(2.5, inf]"]
       >
 
       iex> s = Explorer.Series.from_list([1.0, 2.0, 3.0])
       iex> Explorer.Series.cut(s, [1.5, 2.5])
       #Explorer.DataFrame<
         Polars[3 x 2]
         values f64 [1.0, 2.0, 3.0]
-        category string ["(-inf, 1.5]", "(1.5, 2.5]", "(2.5, inf]"]
+        category enum ["(-inf, 1.5]", "(1.5, 2.5]", "(2.5, inf]"]
       >
   """
   @doc type: :aggregation
@@ -4998,7 +5020,7 @@ defmodule Explorer.Series do
       #Explorer.DataFrame<
         Polars[5 x 2]
         values f64 [1.0, 2.0, 3.0, 4.0, 5.0]
-        category string ["(-inf, 2]", "(-inf, 2]", "(2, 4]", "(2, 4]", "(4, inf]"]
+        category category ["(-inf, 2]", "(-inf, 2]", "(2, 4]", "(2, 4]", "(4, inf]"]
       >
   """
   @doc type: :aggregation

lib/explorer/shared.ex

@@ -43,7 +43,12 @@ defmodule Explorer.Shared do
   """
   def dtypes do
     @scalar_types ++
-      [{:list, :any}, {:struct, :any}, {:decimal, :pos_integer, :pos_integer}]
+      [
+        {:enum, :any},
+        {:list, :any},
+        {:struct, :any},
+        {:decimal, :pos_integer, :pos_integer}
+      ]
   end
 
   @doc """
@@ -78,6 +83,12 @@ defmodule Explorer.Shared do
   def normalise_dtype({:datetime, p, tz} = dtype) when p in @precisions and is_binary(tz),
     do: dtype
 
+  def normalise_dtype({:enum, categories}) when is_list(categories) do
+    if Enum.all?(categories, &is_binary/1) and Enum.uniq(categories) == categories do
+      {:enum, categories}
+    end
+  end
+
   def normalise_dtype(dtype) when dtype in @scalar_types, do: dtype
   def normalise_dtype(dtype) when dtype in [:float, :f64], do: {:f, 64}
   def normalise_dtype(dtype) when dtype in [:integer, :s64], do: {:s, 64}
@@ -538,6 +549,7 @@ defmodule Explorer.Shared do
   def dtype_to_string({:naive_datetime, p}), do: "naive_datetime[#{precision_string(p)}]"
   def dtype_to_string({:datetime, p, tz}), do: "datetime[#{precision_string(p)}, #{tz}]"
   def dtype_to_string({:duration, p}), do: "duration[#{precision_string(p)}]"
+  def dtype_to_string({:enum, _categories}), do: "enum"
   def dtype_to_string({:list, dtype}), do: "list[" <> dtype_to_string(dtype) <> "]"
   def dtype_to_string({:struct, fields}), do: "struct[#{length(fields)}]"
   def dtype_to_string({:f, size}), do: "f" <> Integer.to_string(size)

mix.exs

@@ -1,8 +1,8 @@
 defmodule Explorer.MixProject do
   use Mix.Project
 
-  @source_url "https://github.com/elixir-nx/explorer"
-  @version "0.12.0-dev"
+  @source_url "https://github.com/liveflow-io/explorer"
+  @version "0.13.0"
   @dev? String.ends_with?(@version, "-dev")
   @force_build? System.get_env("EXPLORER_BUILD") in ["1", "true"]
 

native/explorer/src/datatypes/ex_dtypes.rs

@@ -3,6 +3,7 @@ use polars::datatypes::CategoricalPhysical;
 use polars::datatypes::Categories;
 use polars::datatypes::DataType;
 use polars::datatypes::Field;
+use polars::datatypes::FrozenCategories;
 use polars::datatypes::PlSmallStr;
 use polars::datatypes::TimeUnit;
 use rustler::NifTaggedEnum;
@@ -42,6 +43,7 @@ pub enum ExSeriesDtype {
     Binary,
     Boolean,
     Category,
+    Enum(Vec<String>),
     Date,
     F(u8),
     S(u8),
@@ -66,6 +68,13 @@ impl TryFrom<&DataType> for ExSeriesDtype {
             DataType::Binary => Ok(ExSeriesDtype::Binary),
             DataType::Boolean => Ok(ExSeriesDtype::Boolean),
             DataType::Categorical(_, _) => Ok(ExSeriesDtype::Category),
+            DataType::Enum(categories, _) => Ok(ExSeriesDtype::Enum(
+                categories
+                    .categories()
+                    .values_iter()
+                    .map(|category| category.to_string())
+                    .collect(),
+            )),
             DataType::Date => Ok(ExSeriesDtype::Date),
             DataType::Float64 => Ok(ExSeriesDtype::F(64)),
             DataType::Float32 => Ok(ExSeriesDtype::F(32)),
@@ -128,6 +137,9 @@ impl TryFrom<&ExSeriesDtype> for DataType {
                 );
                 Ok(DataType::from_categories(cats.clone()))
             }
+            ExSeriesDtype::Enum(categories) => Ok(DataType::from_frozen_categories(
+                FrozenCategories::new(categories.iter().map(|category| category.as_str()))?,
+            )),
             ExSeriesDtype::Date => Ok(DataType::Date),
             ExSeriesDtype::F(64) => Ok(DataType::Float64),
             ExSeriesDtype::F(32) => Ok(DataType::Float32),

native/explorer/src/encoding.rs

@@ -709,7 +709,9 @@ pub fn resource_term_from_value<'b>(
             encode_datetime(v, time_unit, time_zone.parse::<Tz>().unwrap(), env)
         }
         AnyValue::Duration(v, time_unit) => encode_duration(v, time_unit, env),
-        AnyValue::Categorical(idx, mapping) => Ok(mapping.cat_to_str(idx).encode(env)),
+        AnyValue::Categorical(idx, mapping) | AnyValue::Enum(idx, mapping) => {
+            Ok(mapping.cat_to_str(idx).encode(env))
+        }
         AnyValue::List(series) => list_from_series(ExSeries::new(series), env),
         AnyValue::Struct(_, _, fields) => v
             ._iter_struct_av()
@@ -773,7 +775,7 @@ pub fn list_from_series(s: ExSeries, env: Env) -> Result<Term, ExplorerError> {
 
         DataType::Binary => generic_binary_series_to_list(&s.resource, &s, env),
         DataType::String => generic_string_series_to_list(&s, env),
-        DataType::Categorical(_, _) => categorical_series_to_list(&s, env),
+        DataType::Categorical(_, _) | DataType::Enum(_, _) => categorical_series_to_list(&s, env),
 
         DataType::List(_inner_dtype) => s
             .list()?
@@ -828,6 +830,29 @@ pub fn iovec_from_series(s: ExSeries, env: Env) -> Result<Term, ExplorerError> {
         DataType::Categorical(_, _) => {
             series_to_iovec!(resource, s.cast(&DataType::UInt32)?.u32()?, env, u32)
         }
+        DataType::Enum(categories, _) => {
+            let enum_categories = categories
+                .categories()
+                .values_iter()
+                .enumerate()
+                .map(|(idx, category)| (category, idx as u32))
+                .collect::<HashMap<_, _>>();
+            let strings = s.cast(&DataType::String)?;
+            let values: Vec<Option<u32>> = strings
+                .str()?
+                .iter()
+                .map(|category| {
+                    category.and_then(|category| enum_categories.get(category).copied())
+                })
+                .collect();
+            let mut bin = OwnedBinary::new(values.len() * mem::size_of::<u32>()).unwrap();
+            for (idx, value) in values.into_iter().enumerate() {
+                let offset = idx * mem::size_of::<u32>();
+                bin.as_mut_slice()[offset..offset + mem::size_of::<u32>()]
+                    .copy_from_slice(&value.unwrap_or_default().to_ne_bytes());
+            }
+            Ok([bin.release(env)].encode(env))
+        }
         dt => panic!("to_iovec/1 not implemented for {dt:?}"),
     }
 }