polars-df 0.1.4 → 0.2.0

data/lib/polars/functions.rb

@@ -199,12 +199,201 @@ module Polars
       dt_range
     end
 
-    # def cut
-    # end
+    # Bin values into discrete values.
+    #
+    # @param s [Series]
+    #   Series to bin.
+    # @param bins [Array]
+    #   Bins to create.
+    # @param labels [Array]
+    #   Labels to assign to the bins. If given the length of labels must be
+    #   len(bins) + 1.
+    # @param break_point_label [String]
+    #   Name given to the breakpoint column.
+    # @param category_label [String]
+    #   Name given to the category column.
+    #
+    # @return [DataFrame]
+    #
+    # @note
+    #   This functionality is experimental and may change without it being considered a
+    #   breaking change.
+    #
+    # @example
+    #   a = Polars::Series.new("a", 13.times.map { |i| (-30 + i * 5) / 10.0 })
+    #   Polars.cut(a, [-1, 1])
+    #   # =>
+    #   # shape: (12, 3)
+    #   # ┌──────┬─────────────┬──────────────┐
+    #   # │ a    ┆ break_point ┆ category     │
+    #   # │ ---  ┆ ---         ┆ ---          │
+    #   # │ f64  ┆ f64         ┆ cat          │
+    #   # ╞══════╪═════════════╪══════════════╡
+    #   # │ -3.0 ┆ -1.0        ┆ (-inf, -1.0] │
+    #   # ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ -2.5 ┆ -1.0        ┆ (-inf, -1.0] │
+    #   # ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ -2.0 ┆ -1.0        ┆ (-inf, -1.0] │
+    #   # ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ -1.5 ┆ -1.0        ┆ (-inf, -1.0] │
+    #   # ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ ...  ┆ ...         ┆ ...          │
+    #   # ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ 1.0  ┆ 1.0         ┆ (-1.0, 1.0]  │
+    #   # ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ 1.5  ┆ inf         ┆ (1.0, inf]   │
+    #   # ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ 2.0  ┆ inf         ┆ (1.0, inf]   │
+    #   # ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ 2.5  ┆ inf         ┆ (1.0, inf]   │
+    #   # └──────┴─────────────┴──────────────┘
+    # def cut(
+    #   s,
+    #   bins,
+    #   labels: nil,
+    #   break_point_label: "break_point",
+    #   category_label: "category"
+    # )
+    #   var_nm = s.name
 
-    # def align_frames
+    #   cuts_df = DataFrame.new(
+    #     [
+    #       Series.new(
+    #         break_point_label, bins, dtype: :f64
+    #       ).extend_constant(Float::INFINITY, 1)
+    #     ]
+    #   )
+
+    #   if labels
+    #     if labels.length != bins.length + 1
+    #       raise ArgumentError, "expected more labels"
+    #     end
+    #     cuts_df = cuts_df.with_column(Series.new(category_label, labels))
+    #   else
+    #     cuts_df = cuts_df.with_column(
+    #       Polars.format(
+    #         "({}, {}]",
+    #         Polars.col(break_point_label).shift_and_fill(1, -Float::INFINITY),
+    #         Polars.col(break_point_label)
+    #       ).alias(category_label)
+    #     )
+    #   end
+
+    #   cuts_df = cuts_df.with_column(Polars.col(category_label).cast(:cat))
+
+    #   s.cast(:f64)
+    #     .sort
+    #     .to_frame
+    #     .join_asof(
+    #       cuts_df,
+    #       left_on: var_nm,
+    #       right_on: break_point_label,
+    #       strategy: "forward"
+    #     )
     # end
 
+    # Align a sequence of frames using the uique values from one or more columns as a key.
+    #
+    # Frames that do not contain the given key values have rows injected (with nulls
+    # filling the non-key columns), and each resulting frame is sorted by the key.
+    #
+    # The original column order of input frames is not changed unless ``select`` is
+    # specified (in which case the final column order is determined from that).
+    #
+    # Note that this does not result in a joined frame - you receive the same number
+    # of frames back that you passed in, but each is now aligned by key and has
+    # the same number of rows.
+    #
+    # @param frames [Array]
+    #   Sequence of DataFrames or LazyFrames.
+    # @param on [Object]
+    #   One or more columns whose unique values will be used to align the frames.
+    # @param select [Object]
+    #   Optional post-alignment column select to constrain and/or order
+    #   the columns returned from the newly aligned frames.
+    # @param reverse [Object]
+    #   Sort the alignment column values in descending order; can be a single
+    #   boolean or a list of booleans associated with each column in `on`.
+    #
+    # @return [Object]
+    #
+    # @example
+    #   df1 = Polars::DataFrame.new(
+    #     {
+    #       "dt" => [Date.new(2022, 9, 1), Date.new(2022, 9, 2), Date.new(2022, 9, 3)],
+    #       "x" => [3.5, 4.0, 1.0],
+    #       "y" => [10.0, 2.5, 1.5]
+    #     }
+    #   )
+    #   df2 = Polars::DataFrame.new(
+    #     {
+    #       "dt" => [Date.new(2022, 9, 2), Date.new(2022, 9, 3), Date.new(2022, 9, 1)],
+    #       "x" => [8.0, 1.0, 3.5],
+    #       "y" => [1.5, 12.0, 5.0]
+    #     }
+    #   )
+    #   df3 = Polars::DataFrame.new(
+    #     {
+    #       "dt" => [Date.new(2022, 9, 3), Date.new(2022, 9, 2)],
+    #       "x" => [2.0, 5.0],
+    #       "y" => [2.5, 2.0]
+    #     }
+    #   )
+    #   af1, af2, af3 = Polars.align_frames(
+    #     df1, df2, df3, on: "dt", select: ["x", "y"]
+    #   )
+    #   (af1 * af2 * af3).fill_null(0).select(Polars.sum(Polars.col("*")).alias("dot"))
+    #   # =>
+    #   # shape: (3, 1)
+    #   # ┌───────┐
+    #   # │ dot   │
+    #   # │ ---   │
+    #   # │ f64   │
+    #   # ╞═══════╡
+    #   # │ 0.0   │
+    #   # ├╌╌╌╌╌╌╌┤
+    #   # │ 167.5 │
+    #   # ├╌╌╌╌╌╌╌┤
+    #   # │ 47.0  │
+    #   # └───────┘
+    def align_frames(
+      *frames,
+      on:,
+      select: nil,
+      reverse: false
+    )
+      if frames.empty?
+        return []
+      elsif frames.map(&:class).uniq.length != 1
+        raise TypeError, "Input frames must be of a consistent type (all LazyFrame or all DataFrame)"
+      end
+
+      # establish the superset of all "on" column values, sort, and cache
+      eager = frames[0].is_a?(DataFrame)
+      alignment_frame = (
+        concat(frames.map { |df| df.lazy.select(on) })
+          .unique(maintain_order: false)
+          .sort(on, reverse: reverse)
+      )
+      alignment_frame = (
+        eager ? alignment_frame.collect.lazy : alignment_frame.cache
+      )
+      # finally, align all frames
+      aligned_frames =
+        frames.map do |df|
+          alignment_frame.join(
+            df.lazy,
+            on: alignment_frame.columns,
+            how: "left"
+          ).select(df.columns)
+        end
+      if !select.nil?
+        aligned_frames = aligned_frames.map { |df| df.select(select) }
+      end
+
+      eager ? aligned_frames.map(&:collect) : aligned_frames
+    end
+
     # Return a new Series of given length and type, filled with ones.
     #
     # @param n [Integer]
@@ -249,8 +438,8 @@ module Polars
 
     def _ensure_datetime(value)
       is_date_type = false
-      if !value.is_a?(DateTime)
-        value = DateTime.new(value.year, value.month, value.day)
+      if !value.is_a?(::DateTime)
+        value = ::DateTime.new(value.year, value.month, value.day)
         is_date_type = true
       end
       [value, is_date_type]

data/lib/polars/group_by.rb

@@ -12,7 +12,48 @@ module Polars
       self.maintain_order = maintain_order
     end
 
-    # def apply
+    # Apply a custom/user-defined function (UDF) over the groups as a sub-DataFrame.
+    #
+    # Implementing logic using a Ruby function is almost always _significantly_
+    # slower and more memory intensive than implementing the same logic using
+    # the native expression API because:
+
+    # - The native expression engine runs in Rust; UDFs run in Ruby.
+    # - Use of Ruby UDFs forces the DataFrame to be materialized in memory.
+    # - Polars-native expressions can be parallelised (UDFs cannot).
+    # - Polars-native expressions can be logically optimised (UDFs cannot).
+    #
+    # Wherever possible you should strongly prefer the native expression API
+    # to achieve the best performance.
+    #
+    # @return [DataFrame]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "id" => [0, 1, 2, 3, 4],
+    #       "color" => ["red", "green", "green", "red", "red"],
+    #       "shape" => ["square", "triangle", "square", "triangle", "square"]
+    #     }
+    #   )
+    #   df.groupby("color").apply { |group_df| group_df.sample(2) }
+    #   # =>
+    #   # shape: (4, 3)
+    #   # ┌─────┬───────┬──────────┐
+    #   # │ id  ┆ color ┆ shape    │
+    #   # │ --- ┆ ---   ┆ ---      │
+    #   # │ i64 ┆ str   ┆ str      │
+    #   # ╞═════╪═══════╪══════════╡
+    #   # │ 1   ┆ green ┆ triangle │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ 2   ┆ green ┆ square   │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ 4   ┆ red   ┆ square   │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ 3   ┆ red   ┆ triangle │
+    #   # └─────┴───────┴──────────┘
+    # def apply(&f)
+    #   _dataframe_class._from_rbdf(_df.groupby_apply(by, f))
     # end
 
     # Use multiple aggregations on columns.
@@ -182,8 +223,7 @@ module Polars
       _dataframe_class._from_rbdf(df._df)
     end
 
-    # def pivot
-    # end
+    # pivot is deprecated
 
     # Aggregate the first values in the group.
     #
@@ -294,17 +334,17 @@ module Polars
     #   df.groupby("d", maintain_order: true).min
     #   # =>
     #   # shape: (3, 4)
-    #   # ┌────────┬─────┬──────┬─────┐
-    #   # │ d      ┆ a   ┆ b    ┆ c   │
-    #   # │ ---    ┆ --- ┆ ---  ┆ --- │
-    #   # │ str    ┆ i64 ┆ f64  ┆ u32 │
-    #   # ╞════════╪═════╪══════╪═════╡
-    #   # │ Apple  ┆ 1   ┆ 0.5  ┆ 0   │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌┤
-    #   # │ Orange ┆ 2   ┆ 0.5  ┆ 1   │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌┤
-    #   # │ Banana ┆ 4   ┆ 13.0 ┆ 0   │
-    #   # └────────┴─────┴──────┴─────┘
+    #   # ┌────────┬─────┬──────┬───────┐
+    #   # │ d      ┆ a   ┆ b    ┆ c     │
+    #   # │ ---    ┆ --- ┆ ---  ┆ ---   │
+    #   # │ str    ┆ i64 ┆ f64  ┆ bool  │
+    #   # ╞════════╪═════╪══════╪═══════╡
+    #   # │ Apple  ┆ 1   ┆ 0.5  ┆ false │
+    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
+    #   # │ Orange ┆ 2   ┆ 0.5  ┆ true  │
+    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
+    #   # │ Banana ┆ 4   ┆ 13.0 ┆ false │
+    #   # └────────┴─────┴──────┴───────┘
     def min
       agg(Polars.all.min)
     end
@@ -325,17 +365,17 @@ module Polars
     #   df.groupby("d", maintain_order: true).max
     #   # =>
     #   # shape: (3, 4)
-    #   # ┌────────┬─────┬──────┬─────┐
-    #   # │ d      ┆ a   ┆ b    ┆ c   │
-    #   # │ ---    ┆ --- ┆ ---  ┆ --- │
-    #   # │ str    ┆ i64 ┆ f64  ┆ u32 │
-    #   # ╞════════╪═════╪══════╪═════╡
-    #   # │ Apple  ┆ 3   ┆ 10.0 ┆ 1   │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌┤
-    #   # │ Orange ┆ 2   ┆ 0.5  ┆ 1   │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌┤
-    #   # │ Banana ┆ 5   ┆ 14.0 ┆ 1   │
-    #   # └────────┴─────┴──────┴─────┘
+    #   # ┌────────┬─────┬──────┬──────┐
+    #   # │ d      ┆ a   ┆ b    ┆ c    │
+    #   # │ ---    ┆ --- ┆ ---  ┆ ---  │
+    #   # │ str    ┆ i64 ┆ f64  ┆ bool │
+    #   # ╞════════╪═════╪══════╪══════╡
+    #   # │ Apple  ┆ 3   ┆ 10.0 ┆ true │
+    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┤
+    #   # │ Orange ┆ 2   ┆ 0.5  ┆ true │
+    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌┤
+    #   # │ Banana ┆ 5   ┆ 14.0 ┆ true │
+    #   # └────────┴─────┴──────┴──────┘
     def max
       agg(Polars.all.max)
     end
@@ -387,17 +427,17 @@ module Polars
     #   df.groupby("d", maintain_order: true).mean
     #   # =>
     #   # shape: (3, 4)
-    #   # ┌────────┬─────┬──────────┬──────┐
-    #   # │ d      ┆ a   ┆ b        ┆ c    │
-    #   # │ ---    ┆ --- ┆ ---      ┆ ---  │
-    #   # │ str    ┆ f64 ┆ f64      ┆ bool │
-    #   # ╞════════╪═════╪══════════╪══════╡
-    #   # │ Apple  ┆ 2.0 ┆ 4.833333 ┆ null │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┤
-    #   # │ Orange ┆ 2.0 ┆ 0.5      ┆ null │
-    #   # ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┤
-    #   # │ Banana ┆ 4.5 ┆ 13.5     ┆ null │
-    #   # └────────┴─────┴──────────┴──────┘
+    #   # ┌────────┬─────┬──────────┬──────────┐
+    #   # │ d      ┆ a   ┆ b        ┆ c        │
+    #   # │ ---    ┆ --- ┆ ---      ┆ ---      │
+    #   # │ 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
       agg(Polars.all.mean)
     end

data/lib/polars/io.rb

@@ -59,7 +59,7 @@ module Polars
     #   Lossy means that invalid utf8 values are replaced with `�`
     #   characters. When using other encodings than `utf8` or
     #   `utf8-lossy`, the input is first decoded im memory with
-    #   python.
+    #   Ruby.
     # @param low_memory [Boolean]
     #   Reduce memory usage at expense of performance.
     # @param rechunk [Boolean]
@@ -451,8 +451,24 @@ module Polars
       )
     end
 
-    # def read_avro
-    # end
+    # Read into a DataFrame from Apache Avro format.
+    #
+    # @param file [Object]
+    #   Path to a file or a file-like object.
+    # @param columns [Object]
+    #   Columns to select. Accepts a list of column indices (starting at zero) or a list
+    #   of column names.
+    # @param n_rows [Integer]
+    #   Stop reading from Apache Avro file after reading ``n_rows``.
+    #
+    # @return [DataFrame]
+    def read_avro(file, columns: nil, n_rows: nil)
+      if file.is_a?(String) || (defined?(Pathname) && file.is_a?(Pathname))
+        file = Utils.format_path(file)
+      end
+
+      DataFrame._read_avro(file, n_rows: n_rows, columns: columns)
+    end
 
     # Read into a DataFrame from Arrow IPC (Feather v2) file.
     #