polars-df 0.7.0-x86_64-linux → 0.9.0-x86_64-linux

data/lib/polars/functions/as_datatype.rb

@@ -0,0 +1,248 @@
+module Polars
+  module Functions
+    # Create polars `Duration` from distinct time components.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "datetime" => [DateTime.new(2022, 1, 1), DateTime.new(2022, 1, 2)],
+    #       "add" => [1, 2]
+    #     }
+    #   )
+    #   df.select(
+    #     [
+    #       (Polars.col("datetime") + Polars.duration(weeks: "add")).alias("add_weeks"),
+    #       (Polars.col("datetime") + Polars.duration(days: "add")).alias("add_days"),
+    #       (Polars.col("datetime") + Polars.duration(seconds: "add")).alias("add_seconds"),
+    #       (Polars.col("datetime") + Polars.duration(milliseconds: "add")).alias(
+    #         "add_milliseconds"
+    #       ),
+    #       (Polars.col("datetime") + Polars.duration(hours: "add")).alias("add_hours")
+    #     ]
+    #   )
+    #   # =>
+    #   # shape: (2, 5)
+    #   # ┌─────────────────────┬─────────────────────┬─────────────────────┬─────────────────────────┬─────────────────────┐
+    #   # │ add_weeks           ┆ add_days            ┆ add_seconds         ┆ add_milliseconds        ┆ add_hours           │
+    #   # │ ---                 ┆ ---                 ┆ ---                 ┆ ---                     ┆ ---                 │
+    #   # │ datetime[ns]        ┆ datetime[ns]        ┆ datetime[ns]        ┆ datetime[ns]            ┆ datetime[ns]        │
+    #   # ╞═════════════════════╪═════════════════════╪═════════════════════╪═════════════════════════╪═════════════════════╡
+    #   # │ 2022-01-08 00:00:00 ┆ 2022-01-02 00:00:00 ┆ 2022-01-01 00:00:01 ┆ 2022-01-01 00:00:00.001 ┆ 2022-01-01 01:00:00 │
+    #   # │ 2022-01-16 00:00:00 ┆ 2022-01-04 00:00:00 ┆ 2022-01-02 00:00:02 ┆ 2022-01-02 00:00:00.002 ┆ 2022-01-02 02:00:00 │
+    #   # └─────────────────────┴─────────────────────┴─────────────────────┴─────────────────────────┴─────────────────────┘
+    def duration(
+      weeks: nil,
+      days: nil,
+      hours: nil,
+      minutes: nil,
+      seconds: nil,
+      milliseconds: nil,
+      microseconds: nil,
+      nanoseconds: nil,
+      time_unit: "us"
+    )
+      if !weeks.nil?
+        weeks = Utils.expr_to_lit_or_expr(weeks, str_to_lit: false)._rbexpr
+      end
+      if !days.nil?
+        days = Utils.expr_to_lit_or_expr(days, str_to_lit: false)._rbexpr
+      end
+      if !hours.nil?
+        hours = Utils.expr_to_lit_or_expr(hours, str_to_lit: false)._rbexpr
+      end
+      if !minutes.nil?
+        minutes = Utils.expr_to_lit_or_expr(minutes, str_to_lit: false)._rbexpr
+      end
+      if !seconds.nil?
+        seconds = Utils.expr_to_lit_or_expr(seconds, str_to_lit: false)._rbexpr
+      end
+      if !milliseconds.nil?
+        milliseconds = Utils.expr_to_lit_or_expr(milliseconds, str_to_lit: false)._rbexpr
+      end
+      if !microseconds.nil?
+        microseconds = Utils.expr_to_lit_or_expr(microseconds, str_to_lit: false)._rbexpr
+      end
+      if !nanoseconds.nil?
+        nanoseconds = Utils.expr_to_lit_or_expr(nanoseconds, str_to_lit: false)._rbexpr
+      end
+
+      Utils.wrap_expr(
+        Plr.duration(
+          weeks,
+          days,
+          hours,
+          minutes,
+          seconds,
+          milliseconds,
+          microseconds,
+          nanoseconds,
+          time_unit
+        )
+      )
+    end
+
+    # Concat the arrays in a Series dtype List in linear time.
+    #
+    # @return [Expr]
+    def concat_list(exprs)
+      exprs = Utils.selection_to_rbexpr_list(exprs)
+      Utils.wrap_expr(Plr.concat_list(exprs))
+    end
+
+    # Collect several columns into a Series of dtype Struct.
+    #
+    # @param exprs [Object]
+    #   Columns/Expressions to collect into a Struct
+    # @param eager [Boolean]
+    #   Evaluate immediately
+    #
+    # @return [Object]
+    #
+    # @example
+    #   Polars::DataFrame.new(
+    #     {
+    #       "int" => [1, 2],
+    #       "str" => ["a", "b"],
+    #       "bool" => [true, nil],
+    #       "list" => [[1, 2], [3]],
+    #     }
+    #   ).select([Polars.struct(Polars.all).alias("my_struct")])
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────────────────────┐
+    #   # │ my_struct           │
+    #   # │ ---                 │
+    #   # │ struct[4]           │
+    #   # ╞═════════════════════╡
+    #   # │ {1,"a",true,[1, 2]} │
+    #   # │ {2,"b",null,[3]}    │
+    #   # └─────────────────────┘
+    #
+    # @example Only collect specific columns as a struct:
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [1, 2, 3, 4], "b" => ["one", "two", "three", "four"], "c" => [9, 8, 7, 6]}
+    #   )
+    #   df.with_column(Polars.struct(Polars.col(["a", "b"])).alias("a_and_b"))
+    #   # =>
+    #   # shape: (4, 4)
+    #   # ┌─────┬───────┬─────┬─────────────┐
+    #   # │ a   ┆ b     ┆ c   ┆ a_and_b     │
+    #   # │ --- ┆ ---   ┆ --- ┆ ---         │
+    #   # │ i64 ┆ str   ┆ i64 ┆ struct[2]   │
+    #   # ╞═════╪═══════╪═════╪═════════════╡
+    #   # │ 1   ┆ one   ┆ 9   ┆ {1,"one"}   │
+    #   # │ 2   ┆ two   ┆ 8   ┆ {2,"two"}   │
+    #   # │ 3   ┆ three ┆ 7   ┆ {3,"three"} │
+    #   # │ 4   ┆ four  ┆ 6   ┆ {4,"four"}  │
+    #   # └─────┴───────┴─────┴─────────────┘
+    def struct(exprs, eager: false)
+      if eager
+        Polars.select(struct(exprs, eager: false)).to_series
+      end
+      exprs = Utils.selection_to_rbexpr_list(exprs)
+      Utils.wrap_expr(Plr.as_struct(exprs))
+    end
+
+    # Horizontally concat Utf8 Series in linear time. Non-Utf8 columns are cast to Utf8.
+    #
+    # @param exprs [Object]
+    #   Columns to concat into a Utf8 Series.
+    # @param sep [String]
+    #   String value that will be used to separate the values.
+    # @param ignore_nulls [Boolean]
+    #   Ignore null values (default).
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "a" => [1, 2, 3],
+    #       "b" => ["dogs", "cats", nil],
+    #       "c" => ["play", "swim", "walk"]
+    #     }
+    #   )
+    #   df.with_columns(
+    #     [
+    #       Polars.concat_str(
+    #         [
+    #           Polars.col("a") * 2,
+    #           Polars.col("b"),
+    #           Polars.col("c")
+    #         ],
+    #         sep: " "
+    #       ).alias("full_sentence")
+    #     ]
+    #   )
+    #   # =>
+    #   # shape: (3, 4)
+    #   # ┌─────┬──────┬──────┬───────────────┐
+    #   # │ a   ┆ b    ┆ c    ┆ full_sentence │
+    #   # │ --- ┆ ---  ┆ ---  ┆ ---           │
+    #   # │ i64 ┆ str  ┆ str  ┆ str           │
+    #   # ╞═════╪══════╪══════╪═══════════════╡
+    #   # │ 1   ┆ dogs ┆ play ┆ 2 dogs play   │
+    #   # │ 2   ┆ cats ┆ swim ┆ 4 cats swim   │
+    #   # │ 3   ┆ null ┆ walk ┆ null          │
+    #   # └─────┴──────┴──────┴───────────────┘
+    def concat_str(exprs, sep: "", ignore_nulls: false)
+      exprs = Utils.selection_to_rbexpr_list(exprs)
+      Utils.wrap_expr(Plr.concat_str(exprs, sep, ignore_nulls))
+    end
+
+    # Format expressions as a string.
+    #
+    # @param fstring [String]
+    #   A string that with placeholders.
+    #   For example: "hello_{}" or "{}_world
+    # @param args [Object]
+    #   Expression(s) that fill the placeholders
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "a": ["a", "b", "c"],
+    #       "b": [1, 2, 3]
+    #     }
+    #   )
+    #   df.select(
+    #     [
+    #       Polars.format("foo_{}_bar_{}", Polars.col("a"), "b").alias("fmt")
+    #     ]
+    #   )
+    #   # =>
+    #   # shape: (3, 1)
+    #   # ┌─────────────┐
+    #   # │ fmt         │
+    #   # │ ---         │
+    #   # │ str         │
+    #   # ╞═════════════╡
+    #   # │ foo_a_bar_1 │
+    #   # │ foo_b_bar_2 │
+    #   # │ foo_c_bar_3 │
+    #   # └─────────────┘
+    def format(fstring, *args)
+      if fstring.scan("{}").length != args.length
+        raise ArgumentError, "number of placeholders should equal the number of arguments"
+      end
+
+      exprs = []
+
+      arguments = args.each
+      fstring.split(/(\{\})/).each do |s|
+        if s == "{}"
+          e = Utils.expr_to_lit_or_expr(arguments.next, str_to_lit: false)
+          exprs << e
+        elsif s.length > 0
+          exprs << lit(s)
+        end
+      end
+
+      concat_str(exprs, sep: "")
+    end
+  end
+end

data/lib/polars/functions/col.rb

@@ -0,0 +1,47 @@
+module Polars
+  module Functions
+    # Return an expression representing a column in a DataFrame.
+    #
+    # @return [Expr]
+    def col(name, *more_names)
+      if more_names.any?
+        if Utils.strlike?(name)
+          names_str = [name]
+          names_str.concat(more_names)
+          return Utils.wrap_expr(Plr.cols(names_str.map(&:to_s)))
+        elsif Utils.is_polars_dtype(name)
+          dtypes = [name]
+          dtypes.concat(more_names)
+          return Utils.wrap_expr(Plr.dtype_cols(dtypes))
+        else
+          msg = "invalid input for `col`\n\nExpected `str` or `DataType`, got #{name.class.name}."
+          raise TypeError, msg
+        end
+      end
+
+      if Utils.strlike?(name)
+        Utils.wrap_expr(Plr.col(name.to_s))
+      elsif Utils.is_polars_dtype(name)
+        Utils.wrap_expr(Plr.dtype_cols([name]))
+      elsif name.is_a?(::Array)
+        names = Array(name)
+        if names.empty?
+          return Utils.wrap_expr(Plr.cols(names))
+        end
+
+        item = names[0]
+        if Utils.strlike?(item)
+          Utils.wrap_expr(Plr.cols(names.map(&:to_s)))
+        elsif Utils.is_polars_dtype(item)
+          Utils.wrap_expr(Plr.dtype_cols(names))
+        else
+          msg = "invalid input for `col`\n\nExpected iterable of type `str` or `DataType`, got iterable of type #{item.class.name}."
+          raise TypeError, msg
+        end
+      else
+        msg = "invalid input for `col`\n\nExpected `str` or `DataType`, got #{name.class.name}."
+        raise TypeError, msg
+      end
+    end
+  end
+end

data/lib/polars/functions/eager.rb

@@ -0,0 +1,182 @@
+module Polars
+  module Functions
+    # 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", "vertical_relaxed", "diagonal", "horizontal"]
+    #   LazyFrames do not support the `horizontal` 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(Plr.concat_df(items))
+        elsif how == "diagonal"
+          out = Utils.wrap_df(Plr.concat_df_diagonal(items))
+        elsif how == "horizontal"
+          out = Utils.wrap_df(Plr.concat_df_horizontal(items))
+        else
+          raise ArgumentError, "how must be one of {{'vertical', 'diagonal', 'horizontal'}}, got #{how}"
+        end
+      elsif first.is_a?(LazyFrame)
+        if how == "vertical"
+          return Utils.wrap_ldf(Plr.concat_lf(items, rechunk, parallel, false))
+        elsif how == "vertical_relaxed"
+          return Utils.wrap_ldf(Plr.concat_lf(items, rechunk, parallel, true))
+        elsif how == "diagonal"
+          return Utils.wrap_ldf(Plr.concat_lf_diagonal(items, rechunk, parallel, false))
+        else
+          raise ArgumentError, "Lazy only allows 'vertical', 'vertical_relaxed', and 'diagonal' concat strategy."
+        end
+      elsif first.is_a?(Series)
+        # TODO
+        out = Utils.wrap_s(Plr.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
+
+    # 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
+  end
+end