polars-df 0.2.5 → 0.3.1

data/lib/polars/group_by.rb

@@ -12,6 +12,66 @@ module Polars
       self.maintain_order = maintain_order
     end
 
+    # Allows iteration over the groups of the groupby operation.
+    #
+    # @return [Object]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"foo" => ["a", "a", "b"], "bar" => [1, 2, 3]})
+    #   df.groupby("foo", maintain_order: true).each.to_h
+    #   # =>
+    #   # {"a"=>shape: (2, 2)
+    #   # ┌─────┬─────┐
+    #   # │ foo ┆ bar │
+    #   # │ --- ┆ --- │
+    #   # │ str ┆ i64 │
+    #   # ╞═════╪═════╡
+    #   # │ a   ┆ 1   │
+    #   # │ a   ┆ 2   │
+    #   # └─────┴─────┘, "b"=>shape: (1, 2)
+    #   # ┌─────┬─────┐
+    #   # │ foo ┆ bar │
+    #   # │ --- ┆ --- │
+    #   # │ str ┆ i64 │
+    #   # ╞═════╪═════╡
+    #   # │ b   ┆ 3   │
+    #   # └─────┴─────┘}
+    def each
+      return to_enum(:each) unless block_given?
+
+      temp_col = "__POLARS_GB_GROUP_INDICES"
+      groups_df =
+        Utils.wrap_df(_df)
+          .lazy
+          .with_row_count(name: temp_col)
+          .groupby(by, maintain_order: maintain_order)
+          .agg(Polars.col(temp_col))
+          .collect(no_optimization: true)
+
+      group_names = groups_df.select(Polars.all.exclude(temp_col))
+
+      # When grouping by a single column, group name is a single value
+      # When grouping by multiple columns, group name is a tuple of values
+      if by.is_a?(String) || by.is_a?(Expr)
+        _group_names = group_names.to_series.each
+      else
+        _group_names = group_names.iter_rows
+      end
+
+      _group_indices = groups_df.select(temp_col).to_series
+      _current_index = 0
+
+      while _current_index < _group_indices.length
+        df = _dataframe_class._from_rbdf(_df)
+
+        group_name = _group_names.next
+        group_data = df[_group_indices[_current_index]]
+        _current_index += 1
+
+        yield group_name, group_data
+      end
+    end
+
     # Apply a custom/user-defined function (UDF) over the groups as a sub-DataFrame.
     #
     # Implementing logic using a Ruby function is almost always _significantly_
@@ -45,11 +105,8 @@ module Polars
     #   # │ i64 ┆ str   ┆ str      │
     #   # ╞═════╪═══════╪══════════╡
     #   # │ 1   ┆ green ┆ triangle │
-    #   # ├╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
     #   # │ 2   ┆ green ┆ square   │
-    #   # ├╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
     #   # │ 4   ┆ red   ┆ square   │
-    #   # ├╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
     #   # │ 3   ┆ red   ┆ triangle │
     #   # └─────┴───────┴──────────┘
     # def apply(&f)
@@ -83,7 +140,6 @@ module Polars
     #   # │ str ┆ i64     ┆ i64          │
     #   # ╞═════╪═════════╪══════════════╡
     #   # │ one ┆ 9       ┆ 9            │
-    #   # ├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
     #   # │ two ┆ 6       ┆ 5            │
     #   # └─────┴─────────┴──────────────┘
     def agg(aggs)
@@ -117,15 +173,10 @@ module Polars
     #   # │ str     ┆ i64 │
     #   # ╞═════════╪═════╡
     #   # │ c       ┆ 1   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ c       ┆ 2   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ a       ┆ 3   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ c       ┆ 4   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ a       ┆ 5   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ b       ┆ 6   │
     #   # └─────────┴─────┘
     #
@@ -139,13 +190,9 @@ module Polars
     #   # │ str     ┆ i64 │
     #   # ╞═════════╪═════╡
     #   # │ a       ┆ 3   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ a       ┆ 5   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ b       ┆ 6   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ c       ┆ 1   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ c       ┆ 2   │
     #   # └─────────┴─────┘
     def head(n = 5)
@@ -181,15 +228,10 @@ module Polars
     #   # │ str     ┆ i64 │
     #   # ╞═════════╪═════╡
     #   # │ c       ┆ 1   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ c       ┆ 2   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ a       ┆ 3   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ c       ┆ 4   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ a       ┆ 5   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ b       ┆ 6   │
     #   # └─────────┴─────┘
     #
@@ -203,13 +245,9 @@ module Polars
     #   # │ str     ┆ i64 │
     #   # ╞═════════╪═════╡
     #   # │ a       ┆ 3   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ a       ┆ 5   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ b       ┆ 6   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ c       ┆ 2   │
-    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ c       ┆ 4   │
     #   # └─────────┴─────┘
     def tail(n = 5)
@@ -223,8 +261,6 @@ module Polars
       _dataframe_class._from_rbdf(df._df)
     end
 
-    # pivot is deprecated
-
     # Aggregate the first values in the group.
     #
     # @return [DataFrame]
@@ -247,9 +283,7 @@ module Polars
     #   # │ str    ┆ i64 ┆ f64  ┆ bool  │
     #   # ╞════════╪═════╪══════╪═══════╡
     #   # │ Apple  ┆ 1   ┆ 0.5  ┆ true  │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
     #   # │ Orange ┆ 2   ┆ 0.5  ┆ true  │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
     #   # │ Banana ┆ 4   ┆ 13.0 ┆ false │
     #   # └────────┴─────┴──────┴───────┘
     def first
@@ -278,9 +312,7 @@ module Polars
     #   # │ str    ┆ i64 ┆ f64  ┆ bool  │
     #   # ╞════════╪═════╪══════╪═══════╡
     #   # │ Apple  ┆ 3   ┆ 10.0 ┆ false │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
     #   # │ Orange ┆ 2   ┆ 0.5  ┆ true  │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
     #   # │ Banana ┆ 5   ┆ 14.0 ┆ true  │
     #   # └────────┴─────┴──────┴───────┘
     def last
@@ -309,9 +341,7 @@ module Polars
     #   # │ str    ┆ i64 ┆ f64  ┆ u32 │
     #   # ╞════════╪═════╪══════╪═════╡
     #   # │ Apple  ┆ 6   ┆ 14.5 ┆ 2   │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ Orange ┆ 2   ┆ 0.5  ┆ 1   │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ Banana ┆ 9   ┆ 27.0 ┆ 1   │
     #   # └────────┴─────┴──────┴─────┘
     def sum
@@ -340,9 +370,7 @@ module Polars
     #   # │ str    ┆ i64 ┆ f64  ┆ bool  │
     #   # ╞════════╪═════╪══════╪═══════╡
     #   # │ Apple  ┆ 1   ┆ 0.5  ┆ false │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
     #   # │ Orange ┆ 2   ┆ 0.5  ┆ true  │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
     #   # │ Banana ┆ 4   ┆ 13.0 ┆ false │
     #   # └────────┴─────┴──────┴───────┘
     def min
@@ -371,9 +399,7 @@ module Polars
     #   # │ str    ┆ i64 ┆ f64  ┆ bool │
     #   # ╞════════╪═════╪══════╪══════╡
     #   # │ Apple  ┆ 3   ┆ 10.0 ┆ true │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┤
     #   # │ Orange ┆ 2   ┆ 0.5  ┆ true │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┤
     #   # │ Banana ┆ 5   ┆ 14.0 ┆ true │
     #   # └────────┴─────┴──────┴──────┘
     def max
@@ -402,9 +428,7 @@ module Polars
     #   # │ str    ┆ u32   │
     #   # ╞════════╪═══════╡
     #   # │ Apple  ┆ 3     │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
     #   # │ Orange ┆ 1     │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
     #   # │ Banana ┆ 2     │
     #   # └────────┴───────┘
     def count
@@ -433,9 +457,7 @@ module Polars
     #   # │ str    ┆ f64 ┆ f64      ┆ f64      │
     #   # ╞════════╪═════╪══════════╪══════════╡
     #   # │ Apple  ┆ 2.0 ┆ 4.833333 ┆ 0.666667 │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
     #   # │ Orange ┆ 2.0 ┆ 0.5      ┆ 1.0      │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
     #   # │ Banana ┆ 4.5 ┆ 13.5     ┆ 0.5      │
     #   # └────────┴─────┴──────────┴──────────┘
     def mean
@@ -463,7 +485,6 @@ module Polars
     #   # │ str    ┆ u32 ┆ u32 │
     #   # ╞════════╪═════╪═════╡
     #   # │ Apple  ┆ 2   ┆ 2   │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
     #   # │ Banana ┆ 3   ┆ 3   │
     #   # └────────┴─────┴─────┘
     def n_unique
@@ -496,9 +517,7 @@ module Polars
     #   # │ str    ┆ f64 ┆ f64  │
     #   # ╞════════╪═════╪══════╡
     #   # │ Apple  ┆ 3.0 ┆ 10.0 │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┤
     #   # │ Orange ┆ 2.0 ┆ 0.5  │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┤
     #   # │ Banana ┆ 5.0 ┆ 14.0 │
     #   # └────────┴─────┴──────┘
     def quantile(quantile, interpolation: "nearest")
@@ -526,7 +545,6 @@ module Polars
     #   # │ str    ┆ f64 ┆ f64  │
     #   # ╞════════╪═════╪══════╡
     #   # │ Apple  ┆ 2.0 ┆ 4.0  │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┤
     #   # │ Banana ┆ 4.0 ┆ 13.0 │
     #   # └────────┴─────┴──────┘
     def median
@@ -542,17 +560,27 @@ module Polars
     #   df.groupby("a", maintain_order: true).agg_list
     #   # =>
     #   # shape: (2, 2)
-    #   # ┌─────┬───────────┐
-    #   # │ a   ┆ b         │
-    #   # │ --- ┆ ---       │
-    #   # │ str ┆ list[i64] │
-    #   # ╞═════╪═══════════╡
-    #   # │ one ┆ [1, 3]    │
-    #   # ├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
-    #   # │ two ┆ [2, 4]    │
-    #   # └─────┴───────────┘
+    #   # ┌─────┬─────────────────┐
+    #   # │ a   ┆ b               │
+    #   # │ --- ┆ ---             │
+    #   # │ str ┆ list[list[i64]] │
+    #   # ╞═════╪═════════════════╡
+    #   # │ one ┆ [[1, 3]]        │
+    #   # │ two ┆ [[2, 4]]        │
+    #   # └─────┴─────────────────┘
     def agg_list
       agg(Polars.all.list)
     end
+
+    # Plot data.
+    #
+    # @return [Vega::LiteChart]
+    def plot(*args, **options)
+      raise ArgumentError, "Multiple groups not supported" if by.is_a?(Array) && by.size > 1
+      # same message as Ruby
+      raise ArgumentError, "unknown keyword: :group" if options.key?(:group)
+
+      Utils.wrap_df(_df).plot(*args, **options, group: by)
+    end
   end
 end

data/lib/polars/io.rb

@@ -817,7 +817,7 @@ module Polars
       if defined?(URI) && file.is_a?(URI)
         require "open-uri"
 
-        file = URI.open(file)
+        file = file.open
       end
 
       yield file