polars-df 0.2.0-x86_64-linux

data/lib/polars/expr_dispatch.rb

@@ -0,0 +1,22 @@
+module Polars
+  # @private
+  module ExprDispatch
+    private
+
+    def self.included(base)
+      base.attr_accessor :_s
+      base.singleton_class.attr_accessor :_accessor
+    end
+
+    def method_missing(method, ...)
+      return super unless self.class.method_defined?(method)
+
+      namespace = self.class._accessor
+
+      s = Utils.wrap_s(_s)
+      expr = Utils.col(s.name)
+      expr = expr.send(namespace) if namespace
+      s.to_frame.select(expr.send(method, ...)).to_series
+    end
+  end
+end

data/lib/polars/functions.rb

@@ -0,0 +1,453 @@
+module Polars
+  module Functions
+    # Convert categorical variables into dummy/indicator variables.
+    #
+    # @param df [DataFrame]
+    #   DataFrame to convert.
+    # @param columns [Array, nil]
+    #   A subset of columns to convert to dummy variables. `nil` means
+    #   "all columns".
+    #
+    # @return [DataFrame]
+    def get_dummies(df, columns: nil)
+      df.to_dummies(columns: columns)
+    end
+
+    # Aggregate multiple Dataframes/Series to a single DataFrame/Series.
+    #
+    # @param items [Object]
+    #   DataFrames/Series/LazyFrames to concatenate.
+    # @param rechunk [Boolean]
+    #   Make sure that all data is in contiguous memory.
+    # @param how ["vertical", "diagonal", "horizontal"]
+    #   Lazy only supports the 'vertical' strategy.
+    #
+    #   - Vertical: applies multiple `vstack` operations.
+    #   - Diagonal: finds a union between the column schemas and fills missing column values with null.
+    #   - Horizontal: stacks Series horizontally and fills with nulls if the lengths don't match.
+    # @param parallel [Boolean]
+    #   Only relevant for LazyFrames. This determines if the concatenated
+    #   lazy computations may be executed in parallel.
+    #
+    # @return [Object]
+    #
+    # @example
+    #   df1 = Polars::DataFrame.new({"a" => [1], "b" => [3]})
+    #   df2 = Polars::DataFrame.new({"a" => [2], "b" => [4]})
+    #   Polars.concat([df1, df2])
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌─────┬─────┐
+    #   # │ a   ┆ b   │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ i64 │
+    #   # ╞═════╪═════╡
+    #   # │ 1   ┆ 3   │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌┤
+    #   # │ 2   ┆ 4   │
+    #   # └─────┴─────┘
+    def concat(items, rechunk: true, how: "vertical", parallel: true)
+      if items.empty?
+        raise ArgumentError, "cannot concat empty list"
+      end
+
+      first = items[0]
+      if first.is_a?(DataFrame)
+        if how == "vertical"
+          out = Utils.wrap_df(_concat_df(items))
+        elsif how == "diagonal"
+          out = Utils.wrap_df(_diag_concat_df(items))
+        elsif how == "horizontal"
+          out = Utils.wrap_df(_hor_concat_df(items))
+        else
+          raise ArgumentError, "how must be one of {{'vertical', 'diagonal', 'horizontal'}}, got #{how}"
+        end
+      elsif first.is_a?(LazyFrame)
+        if how == "vertical"
+          # TODO
+          return Utils.wrap_ldf(_concat_lf(items, rechunk, parallel))
+        else
+          raise ArgumentError, "Lazy only allows 'vertical' concat strategy."
+        end
+      elsif first.is_a?(Series)
+        # TODO
+        out = Utils.wrap_s(_concat_series(items))
+      elsif first.is_a?(Expr)
+        out = first
+        items[1..-1].each do |e|
+          out = out.append(e)
+        end
+      else
+        raise ArgumentError, "did not expect type: #{first.class.name} in 'Polars.concat'."
+      end
+
+      if rechunk
+        out.rechunk
+      else
+        out
+      end
+    end
+
+    # Create a range of type `Datetime` (or `Date`).
+    #
+    # @param low [Object]
+    #   Lower bound of the date range.
+    # @param high [Object]
+    #   Upper bound of the date range.
+    # @param interval [Object]
+    #   Interval periods. It can be a polars duration string, such as `3d12h4m25s`
+    #   representing 3 days, 12 hours, 4 minutes, and 25 seconds.
+    # @param lazy [Boolean]
+    #   Return an expression.
+    # @param closed ["both", "left", "right", "none"]
+    #   Define whether the temporal window interval is closed or not.
+    # @param name [String]
+    #   Name of the output Series.
+    # @param time_unit [nil, "ns", "us", "ms"]
+    #   Set the time unit.
+    # @param time_zone [String]
+    #   Optional timezone
+    #
+    # @return [Object]
+    #
+    # @note
+    #   If both `low` and `high` are passed as date types (not datetime), and the
+    #   interval granularity is no finer than 1d, the returned range is also of
+    #   type date. All other permutations return a datetime Series.
+    #
+    # @example Using polars duration string to specify the interval
+    #   Polars.date_range(Date.new(2022, 1, 1), Date.new(2022, 3, 1), "1mo", name: "drange")
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'drange' [date]
+    #   # [
+    #   #         2022-01-01
+    #   #         2022-02-01
+    #   #         2022-03-01
+    #   # ]
+    #
+    # @example Using `timedelta` object to specify the interval:
+    #   Polars.date_range(
+    #       DateTime.new(1985, 1, 1),
+    #       DateTime.new(1985, 1, 10),
+    #       "1d12h",
+    #       time_unit: "ms"
+    #   )
+    #   # =>
+    #   # shape: (7,)
+    #   # Series: '' [datetime[ms]]
+    #   # [
+    #   #         1985-01-01 00:00:00
+    #   #         1985-01-02 12:00:00
+    #   #         1985-01-04 00:00:00
+    #   #         1985-01-05 12:00:00
+    #   #         1985-01-07 00:00:00
+    #   #         1985-01-08 12:00:00
+    #   #         1985-01-10 00:00:00
+    #   # ]
+    def date_range(
+      low,
+      high,
+      interval,
+      lazy: false,
+      closed: "both",
+      name: nil,
+      time_unit: nil,
+      time_zone: nil
+    )
+      if defined?(ActiveSupport::Duration) && interval.is_a?(ActiveSupport::Duration)
+        raise Todo
+      else
+        interval = interval.to_s
+        if interval.include?(" ")
+          interval = interval.gsub(" ", "")
+        end
+      end
+
+      if low.is_a?(Expr) || high.is_a?(Expr) || lazy
+        low = Utils.expr_to_lit_or_expr(low, str_to_lit: true)
+        high = Utils.expr_to_lit_or_expr(high, str_to_lit: true)
+        return Utils.wrap_expr(
+          _rb_date_range_lazy(low, high, interval, closed, name, time_zone)
+        )
+      end
+
+      low, low_is_date = _ensure_datetime(low)
+      high, high_is_date = _ensure_datetime(high)
+
+      if !time_unit.nil?
+        tu = time_unit
+      elsif interval.include?("ns")
+        tu = "ns"
+      else
+        tu = "us"
+      end
+
+      start = Utils._datetime_to_pl_timestamp(low, tu)
+      stop = Utils._datetime_to_pl_timestamp(high, tu)
+      if name.nil?
+        name = ""
+      end
+
+      dt_range = Utils.wrap_s(
+        _rb_date_range(start, stop, interval, closed, name, tu, time_zone)
+      )
+      if low_is_date && high_is_date && !["h", "m", "s"].any? { |v| _interval_granularity(interval).end_with?(v) }
+        dt_range = dt_range.cast(Date)
+      end
+
+      dt_range
+    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
+
+    #   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]
+    #   Number of elements in the `Series`
+    # @param dtype [Symbol]
+    #   DataType of the elements, defaults to `:f64`
+    #
+    # @return [Series]
+    #
+    # @note
+    #   In the lazy API you should probably not use this, but use `lit(1)`
+    #   instead.
+    def ones(n, dtype: nil)
+      s = Series.new([1.0])
+      if dtype
+        s = s.cast(dtype)
+      end
+      s.new_from_index(0, n)
+    end
+
+    # Return a new Series of given length and type, filled with zeros.
+    #
+    # @param n [Integer]
+    #   Number of elements in the `Series`
+    # @param dtype [Symbol]
+    #   DataType of the elements, defaults to `:f64`
+    #
+    # @return [Series]
+    #
+    # @note
+    #   In the lazy API you should probably not use this, but use `lit(0)`
+    #   instead.
+    def zeros(n, dtype: nil)
+      s = Series.new([0.0])
+      if dtype
+        s = s.cast(dtype)
+      end
+      s.new_from_index(0, n)
+    end
+
+    private
+
+    def _ensure_datetime(value)
+      is_date_type = false
+      if !value.is_a?(::DateTime)
+        value = ::DateTime.new(value.year, value.month, value.day)
+        is_date_type = true
+      end
+      [value, is_date_type]
+    end
+
+    # TODO
+    def _interval_granularity(interval)
+      interval
+    end
+  end
+end