polars-df 0.2.5 → 0.3.1

data/lib/polars/list_expr.rb

@@ -24,7 +24,6 @@ module Polars
     #   # │ u32 │
     #   # ╞═════╡
     #   # │ 2   │
-    #   # ├╌╌╌╌╌┤
     #   # │ 1   │
     #   # └─────┘
     def lengths
@@ -46,7 +45,6 @@ module Polars
     #   # │ i64    │
     #   # ╞════════╡
     #   # │ 1      │
-    #   # ├╌╌╌╌╌╌╌╌┤
     #   # │ 5      │
     #   # └────────┘
     def sum
@@ -68,7 +66,6 @@ module Polars
     #   # │ i64    │
     #   # ╞════════╡
     #   # │ 1      │
-    #   # ├╌╌╌╌╌╌╌╌┤
     #   # │ 3      │
     #   # └────────┘
     def max
@@ -90,7 +87,6 @@ module Polars
     #   # │ i64    │
     #   # ╞════════╡
     #   # │ 1      │
-    #   # ├╌╌╌╌╌╌╌╌┤
     #   # │ 2      │
     #   # └────────┘
     def min
@@ -112,7 +108,6 @@ module Polars
     #   # │ f64    │
     #   # ╞════════╡
     #   # │ 1.0    │
-    #   # ├╌╌╌╌╌╌╌╌┤
     #   # │ 2.5    │
     #   # └────────┘
     def mean
@@ -138,7 +133,6 @@ module Polars
     #   # │ list[i64] │
     #   # ╞═══════════╡
     #   # │ [1, 2, 3] │
-    #   # ├╌╌╌╌╌╌╌╌╌╌╌┤
     #   # │ [1, 2, 9] │
     #   # └───────────┘
     def sort(reverse: false)
@@ -164,7 +158,6 @@ module Polars
     #   # │ list[i64] │
     #   # ╞═══════════╡
     #   # │ [1, 2, 3] │
-    #   # ├╌╌╌╌╌╌╌╌╌╌╌┤
     #   # │ [2, 1, 9] │
     #   # └───────────┘
     def reverse
@@ -218,7 +211,6 @@ module Polars
     #   # │ list[str]       │
     #   # ╞═════════════════╡
     #   # │ ["a", "b", "c"] │
-    #   # ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
     #   # │ ["x", "y", "z"] │
     #   # └─────────────────┘
     def concat(other)
@@ -258,9 +250,7 @@ module Polars
     #   # │ i64  │
     #   # ╞══════╡
     #   # │ 3    │
-    #   # ├╌╌╌╌╌╌┤
     #   # │ null │
-    #   # ├╌╌╌╌╌╌┤
     #   # │ 1    │
     #   # └──────┘
     def get(index)
@@ -290,9 +280,7 @@ module Polars
     #   # │ i64  │
     #   # ╞══════╡
     #   # │ 3    │
-    #   # ├╌╌╌╌╌╌┤
     #   # │ null │
-    #   # ├╌╌╌╌╌╌┤
     #   # │ 1    │
     #   # └──────┘
     def first
@@ -314,9 +302,7 @@ module Polars
     #   # │ i64  │
     #   # ╞══════╡
     #   # │ 1    │
-    #   # ├╌╌╌╌╌╌┤
     #   # │ null │
-    #   # ├╌╌╌╌╌╌┤
     #   # │ 2    │
     #   # └──────┘
     def last
@@ -341,9 +327,7 @@ module Polars
     #   # │ bool  │
     #   # ╞═══════╡
     #   # │ true  │
-    #   # ├╌╌╌╌╌╌╌┤
     #   # │ false │
-    #   # ├╌╌╌╌╌╌╌┤
     #   # │ true  │
     #   # └───────┘
     def contains(item)
@@ -370,7 +354,6 @@ module Polars
     #   # │ str   │
     #   # ╞═══════╡
     #   # │ a b c │
-    #   # ├╌╌╌╌╌╌╌┤
     #   # │ x y   │
     #   # └───────┘
     def join(separator)
@@ -396,7 +379,6 @@ module Polars
     #   # │ u32 │
     #   # ╞═════╡
     #   # │ 0   │
-    #   # ├╌╌╌╌╌┤
     #   # │ 1   │
     #   # └─────┘
     def arg_min
@@ -422,7 +404,6 @@ module Polars
     #   # │ u32 │
     #   # ╞═════╡
     #   # │ 1   │
-    #   # ├╌╌╌╌╌┤
     #   # │ 0   │
     #   # └─────┘
     def arg_max
@@ -443,7 +424,7 @@ module Polars
     #   s.arr.diff
     #   # =>
     #   # shape: (2,)
-    #   # Series: 'a' [list]
+    #   # Series: 'a' [list[i64]]
     #   # [
     #   #         [null, 1, ... 1]
     #   #         [null, -8, -1]
@@ -464,7 +445,7 @@ module Polars
     #   s.arr.shift
     #   # =>
     #   # shape: (2,)
-    #   # Series: 'a' [list]
+    #   # Series: 'a' [list[i64]]
     #   # [
     #   #         [null, 1, ... 3]
     #   #         [null, 10, 2]
@@ -488,7 +469,7 @@ module Polars
     #   s.arr.slice(1, 2)
     #   # =>
     #   # shape: (2,)
-    #   # Series: 'a' [list]
+    #   # Series: 'a' [list[i64]]
     #   # [
     #   #         [2, 3]
     #   #         [2, 1]
@@ -511,7 +492,7 @@ module Polars
     #   s.arr.head(2)
     #   # =>
     #   # shape: (2,)
-    #   # Series: 'a' [list]
+    #   # Series: 'a' [list[i64]]
     #   # [
     #   #         [1, 2]
     #   #         [10, 2]
@@ -532,7 +513,7 @@ module Polars
     #   s.arr.tail(2)
     #   # =>
     #   # shape: (2,)
-    #   # Series: 'a' [list]
+    #   # Series: 'a' [list[i64]]
     #   # [
     #   #         [3, 4]
     #   #         [2, 1]
@@ -563,7 +544,6 @@ module Polars
     #   # │ struct[3]  │
     #   # ╞════════════╡
     #   # │ {1,2,3}    │
-    #   # ├╌╌╌╌╌╌╌╌╌╌╌╌┤
     #   # │ {1,2,null} │
     #   # └────────────┘
     def to_struct(n_field_strategy: "first_non_null", name_generator: nil)
@@ -598,9 +578,7 @@ module Polars
     #   # │ i64 ┆ i64 ┆ list[f32]  │
     #   # ╞═════╪═════╪════════════╡
     #   # │ 1   ┆ 4   ┆ [1.0, 2.0] │
-    #   # ├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
     #   # │ 8   ┆ 5   ┆ [2.0, 1.0] │
-    #   # ├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
     #   # │ 3   ┆ 2   ┆ [2.0, 1.0] │
     #   # └─────┴─────┴────────────┘
     def eval(expr, parallel: false)

data/lib/polars/list_name_space.rb

@@ -183,7 +183,7 @@ module Polars
     #   s.arr.diff
     #   # =>
     #   # shape: (2,)
-    #   # Series: 'a' [list]
+    #   # Series: 'a' [list[i64]]
     #   # [
     #   #         [null, 1, ... 1]
     #   #         [null, -8, -1]
@@ -204,7 +204,7 @@ module Polars
     #   s.arr.shift
     #   # =>
     #   # shape: (2,)
-    #   # Series: 'a' [list]
+    #   # Series: 'a' [list[i64]]
     #   # [
     #   #         [null, 1, ... 3]
     #   #         [null, 10, 2]
@@ -228,7 +228,7 @@ module Polars
     #   s.arr.slice(1, 2)
     #   # =>
     #   # shape: (2,)
-    #   # Series: 'a' [list]
+    #   # Series: 'a' [list[i64]]
     #   # [
     #   #         [2, 3]
     #   #         [2, 1]
@@ -249,7 +249,7 @@ module Polars
     #   s.arr.head(2)
     #   # =>
     #   # shape: (2,)
-    #   # Series: 'a' [list]
+    #   # Series: 'a' [list[i64]]
     #   # [
     #   #         [1, 2]
     #   #         [10, 2]
@@ -270,7 +270,7 @@ module Polars
     #   s.arr.tail(2)
     #   # =>
     #   # shape: (2,)
-    #   # Series: 'a' [list]
+    #   # Series: 'a' [list[i64]]
     #   # [
     #   #         [3, 4]
     #   #         [2, 1]
@@ -300,7 +300,6 @@ module Polars
     #   # │ struct[3]  │
     #   # ╞════════════╡
     #   # │ {1,2,3}    │
-    #   # ├╌╌╌╌╌╌╌╌╌╌╌╌┤
     #   # │ {1,2,null} │
     #   # └────────────┘
     def to_struct(n_field_strategy: "first_non_null", name_generator: nil)
@@ -334,9 +333,7 @@ module Polars
     #   # │ i64 ┆ i64 ┆ list[f32]  │
     #   # ╞═════╪═════╪════════════╡
     #   # │ 1   ┆ 4   ┆ [1.0, 2.0] │
-    #   # ├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
     #   # │ 8   ┆ 5   ┆ [2.0, 1.0] │
-    #   # ├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
     #   # │ 3   ┆ 2   ┆ [2.0, 1.0] │
     #   # └─────┴─────┴────────────┘
     def eval(expr, parallel: false)

data/lib/polars/plot.rb

@@ -0,0 +1,109 @@
+module Polars
+  module Plot
+    # Plot data.
+    #
+    # @return [Vega::LiteChart]
+    def plot(x = nil, y = nil, type: nil, group: nil, stacked: nil)
+      require "vega"
+
+      raise ArgumentError, "Must specify columns" if columns.size != 2 && (!x || !y)
+      x ||= columns[0]
+      y ||= columns[1]
+      type ||= begin
+        if self[x].numeric? && self[y].numeric?
+          "scatter"
+        elsif self[x].utf8? && self[y].numeric?
+          "column"
+        elsif (self[x].dtype == Date || self[x].dtype.is_a?(Datetime)) && self[y].numeric?
+          "line"
+        else
+          raise "Cannot determine type. Use the type option."
+        end
+      end
+      df = self[(group.nil? ? [x, y] : [x, y, group]).map(&:to_s).uniq]
+      data = df.rows(named: true)
+
+      case type
+      when "line", "area"
+        x_type =
+          if df[x].numeric?
+            "quantitative"
+          elsif df[x].datelike?
+            "temporal"
+          else
+            "nominal"
+          end
+
+        scale = x_type == "temporal" ? {type: "utc"} : {}
+        encoding = {
+          x: {field: x, type: x_type, scale: scale},
+          y: {field: y, type: "quantitative"}
+        }
+        encoding[:color] = {field: group} if group
+
+        Vega.lite
+          .data(data)
+          .mark(type: type, tooltip: true, interpolate: "cardinal", point: {size: 60})
+          .encoding(encoding)
+          .config(axis: {labelFontSize: 12})
+      when "pie"
+        raise ArgumentError, "Cannot use group option with pie chart" unless group.nil?
+
+        Vega.lite
+          .data(data)
+          .mark(type: "arc", tooltip: true)
+          .encoding(
+            color: {field: x, type: "nominal", sort: "none", axis: {title: nil}, legend: {labelFontSize: 12}},
+            theta: {field: y, type: "quantitative"}
+          )
+          .view(stroke: nil)
+      when "column"
+        encoding = {
+          x: {field: x, type: "nominal", sort: "none", axis: {labelAngle: 0}},
+          y: {field: y, type: "quantitative"}
+        }
+        if group
+          encoding[:color] = {field: group}
+          encoding[:xOffset] = {field: group} unless stacked
+        end
+
+        Vega.lite
+          .data(data)
+          .mark(type: "bar", tooltip: true)
+          .encoding(encoding)
+          .config(axis: {labelFontSize: 12})
+      when "bar"
+        encoding = {
+          # TODO determine label angle
+          y: {field: x, type: "nominal", sort: "none", axis: {labelAngle: 0}},
+          x: {field: y, type: "quantitative"}
+        }
+        if group
+          encoding[:color] = {field: group}
+          encoding[:yOffset] = {field: group} unless stacked
+        end
+
+        Vega.lite
+          .data(data)
+          .mark(type: "bar", tooltip: true)
+          .encoding(encoding)
+          .config(axis: {labelFontSize: 12})
+      when "scatter"
+        encoding = {
+          x: {field: x, type: "quantitative", scale: {zero: false}},
+          y: {field: y, type: "quantitative", scale: {zero: false}},
+          size: {value: 60}
+        }
+        encoding[:color] = {field: group} if group
+
+        Vega.lite
+          .data(data)
+          .mark(type: "circle", tooltip: true)
+          .encoding(encoding)
+          .config(axis: {labelFontSize: 12})
+      else
+        raise ArgumentError, "Invalid type: #{type}"
+      end
+    end
+  end
+end

data/lib/polars/series.rb

@@ -259,6 +259,17 @@ module Polars
       0 - self
     end
 
+    # Returns an enumerator.
+    #
+    # @return [Object]
+    def each
+      return to_enum(:each) unless block_given?
+
+      length.times do |i|
+        yield self[i]
+      end
+    end
+
     # Returns elements of the Series.
     #
     # @return [Object]
@@ -432,15 +443,10 @@ module Polars
     #   # │ str        ┆ f64      │
     #   # ╞════════════╪══════════╡
     #   # │ min        ┆ 1.0      │
-    #   # ├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
     #   # │ max        ┆ 5.0      │
-    #   # ├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
     #   # │ null_count ┆ 0.0      │
-    #   # ├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
     #   # │ mean       ┆ 3.0      │
-    #   # ├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
     #   # │ std        ┆ 1.581139 │
-    #   # ├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
     #   # │ count      ┆ 5.0      │
     #   # └────────────┴──────────┘
     #
@@ -455,9 +461,7 @@ module Polars
     #   # │ str        ┆ i64   │
     #   # ╞════════════╪═══════╡
     #   # │ unique     ┆ 4     │
-    #   # ├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
     #   # │ null_count ┆ 1     │
-    #   # ├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
     #   # │ count      ┆ 5     │
     #   # └────────────┴───────┘
     def describe
@@ -660,13 +664,11 @@ module Polars
     #   # │ u8  ┆ u8  ┆ u8  │
     #   # ╞═════╪═════╪═════╡
     #   # │ 1   ┆ 0   ┆ 0   │
-    #   # ├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ 0   ┆ 1   ┆ 0   │
-    #   # ├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ 0   ┆ 0   ┆ 1   │
     #   # └─────┴─────┴─────┘
-    def to_dummies
-      Utils.wrap_df(_s.to_dummies)
+    def to_dummies(separator: "_")
+      Utils.wrap_df(_s.to_dummies(separator))
     end
 
     # Count the unique values in a Series.
@@ -687,9 +689,7 @@ module Polars
     #   # │ i64 ┆ u32    │
     #   # ╞═════╪════════╡
     #   # │ 1   ┆ 1      │
-    #   # ├╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
     #   # │ 2   ┆ 2      │
-    #   # ├╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
     #   # │ 3   ┆ 1      │
     #   # └─────┴────────┘
     def value_counts(sort: false)
@@ -1285,8 +1285,12 @@ module Polars
     #   Expression or scalar value.
     #
     # @return [Integer]
-    def search_sorted(element)
-      Polars.select(Polars.lit(self).search_sorted(element))[0, 0]
+    def search_sorted(element, side: "any")
+      if element.is_a?(Integer) || element.is_a?(Float)
+        return Polars.select(Polars.lit(self).search_sorted(element, side: side)).item
+      end
+      element = Series.new(element)
+      Polars.select(Polars.lit(self).search_sorted(element, side: side)).to_series
     end
 
     # Get unique elements in series.
@@ -1500,7 +1504,7 @@ module Polars
     #   sets = Polars::Series.new("sets", [[1, 2, 3], [1, 2], [9, 10]])
     #   # =>
     #   # shape: (3,)
-    #   # Series: 'sets' [list]
+    #   # Series: 'sets' [list[i64]]
     #   # [
     #   #         [1, 2, 3]
     #   #         [1, 2]
@@ -1772,8 +1776,9 @@ module Polars
     #   s.is_datelike
     #   # => true
     def is_datelike
-      [Date, Datetime, Duration, Time].include?(dtype)
+      [Date, Time].include?(dtype) || dtype.is_a?(Datetime) || dtype.is_a?(Duration)
     end
+    alias_method :datelike?, :is_datelike
 
     # Check if this Series has floating point numbers.
     #
@@ -1819,8 +1824,45 @@ module Polars
     # def view
     # end
 
-    # def to_numo
-    # end
+    # Convert this Series to a Numo array. This operation clones data but is completely safe.
+    #
+    # @return [Numo::NArray]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [1, 2, 3])
+    #   s.to_numo
+    #   # =>
+    #   # Numo::Int64#shape=[3]
+    #   # [1, 2, 3]
+    def to_numo
+      if !has_validity
+        if is_datelike
+          Numo::RObject.cast(to_a)
+        elsif is_numeric
+          # TODO make more efficient
+          {
+            UInt8 => Numo::UInt8,
+            UInt16 => Numo::UInt16,
+            UInt32 => Numo::UInt32,
+            UInt64 => Numo::UInt64,
+            Int8 => Numo::Int8,
+            Int16 => Numo::Int16,
+            Int32 => Numo::Int32,
+            Int64 => Numo::Int64,
+            Float32 => Numo::SFloat,
+            Float64 => Numo::DFloat
+          }.fetch(dtype).cast(to_a)
+        elsif is_boolean
+          Numo::Bit.cast(to_a)
+        else
+          _s.to_numo
+        end
+      elsif is_datelike
+        Numo::RObject.cast(to_a)
+      else
+        _s.to_numo
+      end
+    end
 
     # Set masked values.
     #