polars-df 0.10.0-x86_64-linux → 0.11.0-x86_64-linux

data/lib/polars/functions/range/date_range.rb

@@ -99,5 +99,97 @@ module Polars
 
       result
     end
+
+    # Create a column of date ranges.
+    #
+    # @param start [Object]
+    #   Lower bound of the date range.
+    # @param stop [Object]
+    #   Upper bound of the date range.
+    # @param interval [Object]
+    #   Interval of the range periods, specified using the Polars duration string language (see "Notes" section below).
+    # @param closed ["both", "left", "right", "none"]
+    #   Define which sides of the range are closed (inclusive).
+    # @param time_unit [nil, "ns", "us", "ms"]
+    #   Time unit of the resulting `Datetime` data type.
+    #   Only takes effect if the output column is of type `Datetime`.
+    # @param time_zone [String]
+    #   Time zone of the resulting `Datetime` data type.
+    #   Only takes effect if the output column is of type `Datetime`.
+    # @param eager [Boolean]
+    #   Evaluate immediately and return a `Series`.
+    #   If set to `false` (default), return an expression instead.
+    #
+    # @return [Object]
+    #
+    # @note
+    #   `interval` is created according to the following string language:
+    #
+    #   - 1ns   (1 nanosecond)
+    #   - 1us   (1 microsecond)
+    #   - 1ms   (1 millisecond)
+    #   - 1s    (1 second)
+    #   - 1m    (1 minute)
+    #   - 1h    (1 hour)
+    #   - 1d    (1 calendar day)
+    #   - 1w    (1 calendar week)
+    #   - 1mo   (1 calendar month)
+    #   - 1q    (1 calendar quarter)
+    #   - 1y    (1 calendar year)
+    #
+    #   Or combine them:
+    #   "3d12h4m25s" # 3 days, 12 hours, 4 minutes, and 25 seconds
+    #
+    #   By "calendar day", we mean the corresponding time on the next day (which may
+    #   not be 24 hours, due to daylight savings). Similarly for "calendar week",
+    #   "calendar month", "calendar quarter", and "calendar year".
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "start" => [Date.new(2022, 1, 1), Date.new(2022, 1, 2)],
+    #       "end" => Date.new(2022, 1, 3)
+    #     }
+    #   )
+    #   df.with_columns(date_range: Polars.date_ranges("start", "end"))
+    #   # =>
+    #   # shape: (2, 3)
+    #   # ┌────────────┬────────────┬─────────────────────────────────┐
+    #   # │ start      ┆ end        ┆ date_range                      │
+    #   # │ ---        ┆ ---        ┆ ---                             │
+    #   # │ date       ┆ date       ┆ list[date]                      │
+    #   # ╞════════════╪════════════╪═════════════════════════════════╡
+    #   # │ 2022-01-01 ┆ 2022-01-03 ┆ [2022-01-01, 2022-01-02, 2022-… │
+    #   # │ 2022-01-02 ┆ 2022-01-03 ┆ [2022-01-02, 2022-01-03]        │
+    #   # └────────────┴────────────┴─────────────────────────────────┘
+    def date_ranges(
+      start,
+      stop,
+      interval = "1d",
+      closed: "both",
+      time_unit: nil,
+      time_zone: nil,
+      eager: false
+    )
+      interval = Utils.parse_interval_argument(interval)
+      if time_unit.nil? && interval.include?("ns")
+        time_unit = "ns"
+      end
+
+      start_rbexpr = Utils.parse_as_expression(start)
+      end_rbexpr = Utils.parse_as_expression(stop)
+
+      result = Utils.wrap_expr(
+        Plr.date_ranges(
+          start_rbexpr, end_rbexpr, interval, closed, time_unit, time_zone
+        )
+      )
+
+      if eager
+        return F.select(result).to_series
+      end
+
+      result
+    end
   end
 end

data/lib/polars/functions/range/datetime_range.rb

@@ -0,0 +1,149 @@
+module Polars
+  module Functions
+    # Generate a datetime range.
+    #
+    # @param start [Object]
+    #   Lower bound of the datetime range.
+    # @param stop [Object]
+    #   Upper bound of the datetime range.
+    # @param interval [String]
+    #   Interval of the range periods, specified using the Polars duration string language.
+    # @param closed ['both', 'left', 'right', 'none']
+    #   Define which sides of the range are closed (inclusive).
+    # @param time_unit [nil, 'ns', 'us', 'ms']
+    #   Time unit of the resulting `Datetime` data type.
+    # @param time_zone [String]
+    #   Time zone of the resulting `Datetime` data type.
+    # @param eager [Boolean]
+    #   Evaluate immediately and return a `Series`.
+    #   If set to `false` (default), return an expression instead.
+    #
+    # @return [Object]
+    #
+    # @example Using Polars duration string to specify the interval:
+    #   Polars.datetime_range(
+    #     DateTime.new(2022, 1, 1), DateTime.new(2022, 3, 1), "1mo", eager: true
+    #   ).alias("datetime")
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'datetime' [datetime[ns]]
+    #   # [
+    #   #         2022-01-01 00:00:00
+    #   #         2022-02-01 00:00:00
+    #   #         2022-03-01 00:00:00
+    #   # ]
+    #
+    # @example Specifying a time zone:
+    #   Polars.datetime_range(
+    #     DateTime.new(2022, 1, 1),
+    #     DateTime.new(2022, 3, 1),
+    #     "1mo",
+    #     time_zone: "America/New_York",
+    #     eager: true
+    #   ).alias("datetime")
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'datetime' [datetime[ns, America/New_York]]
+    #   # [
+    #   #         2022-01-01 00:00:00 EST
+    #   #         2022-02-01 00:00:00 EST
+    #   #         2022-03-01 00:00:00 EST
+    #   # ]
+    def datetime_range(
+      start,
+      stop,
+      interval = "1d",
+      closed: "both",
+      time_unit: nil,
+      time_zone: nil,
+      eager: false
+    )
+      interval = Utils.parse_interval_argument(interval)
+      if time_unit.nil? && interval.include?("ns")
+        time_unit = "ns"
+      end
+
+      start_rbexpr = Utils.parse_as_expression(start)
+      end_rbexpr = Utils.parse_as_expression(stop)
+      result = Utils.wrap_expr(
+        Plr.datetime_range(
+          start_rbexpr, end_rbexpr, interval, closed, time_unit, time_zone
+        )
+      )
+
+      if eager
+        return Polars.select(result).to_series
+      end
+
+      result
+    end
+
+    # Create a column of datetime ranges.
+    #
+    # @param start [Object]
+    #   Lower bound of the datetime range.
+    # @param stop [Object]
+    #   Upper bound of the datetime range.
+    # @param interval [String]
+    #   Interval of the range periods, specified using the Polars duration string language.
+    # @param closed ['both', 'left', 'right', 'none']
+    #   Define which sides of the range are closed (inclusive).
+    # @param time_unit [nil, 'ns', 'us', 'ms']
+    #   Time unit of the resulting `Datetime` data type.
+    # @param time_zone [String]
+    #   Time zone of the resulting `Datetime` data type.
+    # @param eager [Boolean]
+    #   Evaluate immediately and return a `Series`.
+    #   If set to `false` (default), return an expression instead.
+    #
+    # @return [Object]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "start" => [DateTime.new(2022, 1, 1), DateTime.new(2022, 1, 2)],
+    #       "end" => DateTime.new(2022, 1, 3),
+    #     }
+    #   )
+    #   df.select(datetime_range: Polars.datetime_ranges("start", "end"))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────────────────────────────────┐
+    #   # │ datetime_range                  │
+    #   # │ ---                             │
+    #   # │ list[datetime[ns]]              │
+    #   # ╞═════════════════════════════════╡
+    #   # │ [2022-01-01 00:00:00, 2022-01-… │
+    #   # │ [2022-01-02 00:00:00, 2022-01-… │
+    #   # └─────────────────────────────────┘
+    def datetime_ranges(
+      start,
+      stop,
+      interval: "1d",
+      closed: "both",
+      time_unit: nil,
+      time_zone: nil,
+      eager: false
+    )
+      interval = Utils.parse_interval_argument(interval)
+      if time_unit.nil? && interval.include?("ns")
+        time_unit = "ns"
+      end
+
+      start_rbexpr = Utils.parse_as_expression(start)
+      end_rbexpr = Utils.parse_as_expression(stop)
+
+      result = Utils.wrap_expr(
+        Plr.datetime_ranges(
+          start_rbexpr, end_rbexpr, interval, closed, time_unit, time_zone
+        )
+      )
+
+      if eager
+        return Polars.select(result).to_series
+      end
+
+      result
+    end
+  end
+end

data/lib/polars/functions/range/time_range.rb

@@ -0,0 +1,141 @@
+module Polars
+  module Functions
+    # Generate a time range.
+    #
+    # @param start [Object]
+    #   Lower bound of the time range.
+    # @param stop [Object]
+    #   Upper bound of the time range.
+    # @param interval [String]
+    #   Interval of the range periods, specified using the Polars duration string language.
+    # @param closed ['both', 'left', 'right', 'none']
+    #   Define which sides of the range are closed (inclusive).
+    # @param eager [Boolean]
+    #   Evaluate immediately and return a `Series`.
+    #   If set to `False` (default), return an expression instead.
+    #
+    # @return [Object]
+    #
+    # @example
+    #   Polars.time_range(
+    #     time(14, 0),
+    #     nil,
+    #     "3h15m",
+    #     eager: true
+    #   ).alias("time")
+    #   # =>
+    #   # shape: (4,)
+    #   # Series: 'time' [time]
+    #   # [
+    #   #         14:00:00
+    #   #         17:15:00
+    #   #         20:30:00
+    #   #         23:45:00
+    #   # ]
+    def time_range(
+      start = nil,
+      stop = nil,
+      interval = "1h",
+      closed: "both",
+      eager: false
+    )
+      interval = Utils.parse_interval_argument(interval)
+      ["y", "mo", "w", "d"].each do |unit|
+        if interval.include?(unit)
+          msg = "invalid interval unit for time_range: found #{unit.inspect}"
+          raise ArgumentError, msg
+        end
+      end
+
+      if start.nil?
+        # start = time(0, 0, 0)
+        raise Todo
+      end
+      if stop.nil?
+        # stop = time(23, 59, 59, 999999)
+        raise Todo
+      end
+
+      start_rbexpr = Utils.parse_as_expression(start)
+      end_rbexpr = Utils.parse_as_expression(stop)
+
+      result = Utils.wrap_expr(Plr.time_range(start_rbexpr, end_rbexpr, interval, closed))
+
+      if eager
+        return Polars.select(result).to_series
+      end
+
+      result
+    end
+
+    # Create a column of time ranges.
+    #
+    # @param start [Object]
+    #   Lower bound of the time range.
+    # @param stop [Object]
+    #   Upper bound of the time range.
+    # @param interval [Integer]
+    #   Interval of the range periods, specified using the Polars duration string language.
+    # @param closed ['both', 'left', 'right', 'none']
+    #   Define which sides of the range are closed (inclusive).
+    # @param eager [Boolean]
+    #   Evaluate immediately and return a `Series`.
+    #   If set to `false` (default), return an expression instead.
+    #
+    # @return [Object]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "start" => [time(9, 0), time(10, 0)],
+    #       "end" => time(11, 0)
+    #     }
+    #   )
+    #   df.with_columns(time_range: Polars.time_ranges("start", "end"))
+    #   # =>
+    #   # shape: (2, 3)
+    #   # ┌──────────┬──────────┬────────────────────────────────┐
+    #   # │ start    ┆ end      ┆ time_range                     │
+    #   # │ ---      ┆ ---      ┆ ---                            │
+    #   # │ time     ┆ time     ┆ list[time]                     │
+    #   # ╞══════════╪══════════╪════════════════════════════════╡
+    #   # │ 09:00:00 ┆ 11:00:00 ┆ [09:00:00, 10:00:00, 11:00:00] │
+    #   # │ 10:00:00 ┆ 11:00:00 ┆ [10:00:00, 11:00:00]           │
+    #   # └──────────┴──────────┴────────────────────────────────┘
+    def time_ranges(
+      start = nil,
+      stop = nil,
+      interval = "1h",
+      closed: "both",
+      eager: false
+    )
+      interval = Utils.parse_interval_argument(interval)
+      ["y", "mo", "w", "d"].each do |unit|
+        if interval.include?(unit)
+          msg = "invalid interval unit for time_range: found #{unit.inspect}"
+          raise ArgumentError, msg
+        end
+      end
+
+      if start.nil?
+        # start = time(0, 0, 0)
+        raise Todo
+      end
+      if stop.nil?
+        # stop = time(23, 59, 59, 999999)
+        raise Todo
+      end
+
+      start_rbexpr = Utils.parse_as_expression(start)
+      end_rbexpr = Utils.parse_as_expression(stop)
+
+      result = Utils.wrap_expr(Plr.time_ranges(start_rbexpr, end_rbexpr, interval, closed))
+
+      if eager
+        return Polars.select(result).to_series
+      end
+
+      result
+    end
+  end
+end

data/lib/polars/group_by.rb

@@ -106,39 +106,104 @@ module Polars
     #   _dataframe_class._from_rbdf(_df.group_by_apply(by, f))
     # end
 
-    # Use multiple aggregations on columns.
+    # Compute aggregations for each group of a group by operation.
     #
-    # This can be combined with complete lazy API and is considered idiomatic polars.
-    #
-    # @param aggs [Object]
-    #   Single / multiple aggregation expression(s).
+    # @param aggs [Array]
+    #   Aggregations to compute for each group of the group by operation,
+    #   specified as positional arguments.
+    #   Accepts expression input. Strings are parsed as column names.
+    # @param named_aggs [Hash]
+    #   Additional aggregations, specified as keyword arguments.
+    #   The resulting columns will be renamed to the keyword used.
     #
     # @return [DataFrame]
     #
-    # @example
+    # @example Compute the aggregation of the columns for each group.
     #   df = Polars::DataFrame.new(
-    #     {"foo" => ["one", "two", "two", "one", "two"], "bar" => [5, 3, 2, 4, 1]}
+    #     {
+    #       "a" => ["a", "b", "a", "b", "c"],
+    #       "b" => [1, 2, 1, 3, 3],
+    #       "c" => [5, 4, 3, 2, 1]
+    #     }
     #   )
-    #   df.group_by("foo", maintain_order: true).agg(
-    #     [
-    #       Polars.sum("bar").suffix("_sum"),
-    #       Polars.col("bar").sort.tail(2).sum.suffix("_tail_sum")
-    #     ]
+    #   df.group_by("a").agg(Polars.col("b"), Polars.col("c"))
+    #   # =>
+    #   # shape: (3, 3)
+    #   # ┌─────┬───────────┬───────────┐
+    #   # │ a   ┆ b         ┆ c         │
+    #   # │ --- ┆ ---       ┆ ---       │
+    #   # │ str ┆ list[i64] ┆ list[i64] │
+    #   # ╞═════╪═══════════╪═══════════╡
+    #   # │ a   ┆ [1, 1]    ┆ [5, 3]    │
+    #   # │ b   ┆ [2, 3]    ┆ [4, 2]    │
+    #   # │ c   ┆ [3]       ┆ [1]       │
+    #   # └─────┴───────────┴───────────┘
+    #
+    # @example Compute the sum of a column for each group.
+    #   df.group_by("a").agg(Polars.col("b").sum)
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────┬─────┐
+    #   # │ a   ┆ b   │
+    #   # │ --- ┆ --- │
+    #   # │ str ┆ i64 │
+    #   # ╞═════╪═════╡
+    #   # │ a   ┆ 2   │
+    #   # │ b   ┆ 5   │
+    #   # │ c   ┆ 3   │
+    #   # └─────┴─────┘
+    #
+    # @example Compute multiple aggregates at once by passing a list of expressions.
+    #   df.group_by("a").agg([Polars.sum("b"), Polars.mean("c")])
+    #   # =>
+    #   # shape: (3, 3)
+    #   # ┌─────┬─────┬─────┐
+    #   # │ a   ┆ b   ┆ c   │
+    #   # │ --- ┆ --- ┆ --- │
+    #   # │ str ┆ i64 ┆ f64 │
+    #   # ╞═════╪═════╪═════╡
+    #   # │ c   ┆ 3   ┆ 1.0 │
+    #   # │ a   ┆ 2   ┆ 4.0 │
+    #   # │ b   ┆ 5   ┆ 3.0 │
+    #   # └─────┴─────┴─────┘
+    #
+    # @example Or use positional arguments to compute multiple aggregations in the same way.
+    #   df.group_by("a").agg(
+    #     Polars.sum("b").name.suffix("_sum"),
+    #     (Polars.col("c") ** 2).mean.name.suffix("_mean_squared")
     #   )
     #   # =>
-    #   # shape: (2, 3)
-    #   # ┌─────┬─────────┬──────────────┐
-    #   # │ foo ┆ bar_sum ┆ bar_tail_sum │
-    #   # │ --- ┆ ---     ┆ ---          │
-    #   # │ str ┆ i64     ┆ i64          │
-    #   # ╞═════╪═════════╪══════════════╡
-    #   # │ one ┆ 9       ┆ 9            │
-    #   # │ two ┆ 6       ┆ 5            │
-    #   # └─────┴─────────┴──────────────┘
-    def agg(aggs)
+    #   # shape: (3, 3)
+    #   # ┌─────┬───────┬────────────────┐
+    #   # │ a   ┆ b_sum ┆ c_mean_squared │
+    #   # │ --- ┆ ---   ┆ ---            │
+    #   # │ str ┆ i64   ┆ f64            │
+    #   # ╞═════╪═══════╪════════════════╡
+    #   # │ a   ┆ 2     ┆ 17.0           │
+    #   # │ c   ┆ 3     ┆ 1.0            │
+    #   # │ b   ┆ 5     ┆ 10.0           │
+    #   # └─────┴───────┴────────────────┘
+    #
+    # @example Use keyword arguments to easily name your expression inputs.
+    #   df.group_by("a").agg(
+    #     b_sum: Polars.sum("b"),
+    #     c_mean_squared: (Polars.col("c") ** 2).mean
+    #   )
+    #   # =>
+    #   # shape: (3, 3)
+    #   # ┌─────┬───────┬────────────────┐
+    #   # │ a   ┆ b_sum ┆ c_mean_squared │
+    #   # │ --- ┆ ---   ┆ ---            │
+    #   # │ str ┆ i64   ┆ f64            │
+    #   # ╞═════╪═══════╪════════════════╡
+    #   # │ a   ┆ 2     ┆ 17.0           │
+    #   # │ c   ┆ 3     ┆ 1.0            │
+    #   # │ b   ┆ 5     ┆ 10.0           │
+    #   # └─────┴───────┴────────────────┘
+    def agg(*aggs, **named_aggs)
       @df.lazy
         .group_by(@by, maintain_order: @maintain_order)
-        .agg(aggs)
+        .agg(*aggs, **named_aggs)
         .collect(no_optimization: true)
     end
 

data/lib/polars/io/avro.rb

@@ -0,0 +1,24 @@
+module Polars
+  module IO
+    # Read into a DataFrame from Apache Avro format.
+    #
+    # @param source [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(source, columns: nil, n_rows: nil)
+      if Utils.pathlike?(source)
+        source = Utils.normalize_filepath(source)
+      end
+      projection, column_names = Utils.handle_projection_columns(columns)
+
+      rbdf = RbDataFrame.read_avro(source, column_names, projection, n_rows)
+      Utils.wrap_df(rbdf)
+    end
+  end
+end