polars-df 0.23.0 → 0.24.0

data/lib/polars/data_frame_plot.rb

@@ -0,0 +1,173 @@
+module Polars
+  # DataFrame.plot namespace.
+  class DataFramePlot
+    # @private
+    def initialize(df)
+      require "vega"
+
+      @df = df
+    end
+
+    # Draw line plot.
+    #
+    # @param x [String]
+    #   Column with x-coordinates of lines.
+    # @param y [String]
+    #   Column with y-coordinates of lines.
+    # @param color [String]
+    #   Column to color lines by.
+    #
+    # @return [Vega::LiteChart]
+    def line(x, y, color: nil, _type: "line")
+      data = @df[[x, y, color].compact.map(&:to_s).uniq].rows(named: true)
+
+      x_type =
+        if @df[x].dtype.numeric?
+          "quantitative"
+        elsif @df[x].dtype.temporal?
+          "temporal"
+        else
+          "nominal"
+        end
+
+      scale = x_type == "temporal" ? {type: "utc"} : {}
+      encoding = {
+        x: {field: x, type: x_type, scale: scale},
+        y: {field: y, type: "quantitative"}
+      }
+      encoding[:color] = {field: color} if color
+
+      Vega.lite
+        .data(data)
+        .mark(type: _type, tooltip: true, interpolate: "cardinal", point: {size: 60})
+        .encoding(encoding)
+        .config(axis: {labelFontSize: 12})
+    end
+
+    # Draw area plot.
+    #
+    # @param x [String]
+    #   Column with x-coordinates of lines.
+    # @param y [String]
+    #   Column with y-coordinates of lines.
+    # @param color [String]
+    #   Column to color lines by.
+    #
+    # @return [Vega::LiteChart]
+    def area(x, y, color: nil)
+      line(x, y, color: color, _type: "area")
+    end
+
+    # Draw pie chart.
+    #
+    # @param x [String]
+    #   Column with label of slice.
+    # @param y [String]
+    #   Column with size of slice.
+    #
+    # @return [Vega::LiteChart]
+    def pie(x, y)
+      data = @df[[x, y].map(&:to_s).uniq].rows(named: true)
+
+      Vega.lite
+        .data(data)
+        .mark(type: "arc", tooltip: true)
+        .encoding(
+          color: {field: x, type: "nominal", sort: "none", axis: {title: nil}, legend: {labelFontSize: 12}},
+          theta: {field: y, type: "quantitative"}
+        )
+        .view(stroke: nil)
+    end
+
+    # Draw column plot.
+    #
+    # @param x [String]
+    #   Column with x-coordinates of columns.
+    # @param y [String]
+    #   Column with y-coordinates of columns.
+    # @param color [String]
+    #   Column to color columns by.
+    # @param stacked [Boolean]
+    #   Stack columns.
+    #
+    # @return [Vega::LiteChart]
+    def column(x, y, color: nil, stacked: nil)
+      data = @df[[x, y, color].compact.map(&:to_s).uniq].rows(named: true)
+
+      encoding = {
+        x: {field: x, type: "nominal", sort: "none", axis: {labelAngle: 0}},
+        y: {field: y, type: "quantitative"}
+      }
+      if color
+        encoding[:color] = {field: color}
+        encoding[:xOffset] = {field: color} unless stacked
+      end
+
+      Vega.lite
+        .data(data)
+        .mark(type: "bar", tooltip: true)
+        .encoding(encoding)
+        .config(axis: {labelFontSize: 12})
+    end
+
+    # Draw bar plot.
+    #
+    # @param x [String]
+    #   Column with x-coordinates of bars.
+    # @param y [String]
+    #   Column with y-coordinates of bars.
+    # @param color [String]
+    #   Column to color bars by.
+    # @param stacked [Boolean]
+    #   Stack bars.
+    #
+    # @return [Vega::LiteChart]
+    def bar(x, y, color: nil, stacked: nil)
+      data = @df[[x, y, color].compact.map(&:to_s).uniq].rows(named: true)
+
+      encoding = {
+        # TODO determine label angle
+        y: {field: x, type: "nominal", sort: "none", axis: {labelAngle: 0}},
+        x: {field: y, type: "quantitative"}
+      }
+      if color
+        encoding[:color] = {field: color}
+        encoding[:yOffset] = {field: color} unless stacked
+      end
+
+      Vega.lite
+        .data(data)
+        .mark(type: "bar", tooltip: true)
+        .encoding(encoding)
+        .config(axis: {labelFontSize: 12})
+    end
+
+    # Draw scatter plot.
+    #
+    # @param x [String]
+    #   Column with x-coordinates of points.
+    # @param y [String]
+    #   Column with y-coordinates of points.
+    # @param color [String]
+    #   Column to color points by.
+    #
+    # @return [Vega::LiteChart]
+    def scatter(x, y, color: nil)
+      data = @df[[x, y, color].compact.map(&:to_s).uniq].rows(named: true)
+
+      encoding = {
+        x: {field: x, type: "quantitative", scale: {zero: false}},
+        y: {field: y, type: "quantitative", scale: {zero: false}},
+        size: {value: 60}
+      }
+      encoding[:color] = {field: color} if color
+
+      Vega.lite
+        .data(data)
+        .mark(type: "circle", tooltip: true)
+        .encoding(encoding)
+        .config(axis: {labelFontSize: 12})
+    end
+    alias_method :point, :scatter
+  end
+end

data/lib/polars/data_type_group.rb

@@ -1,4 +1,5 @@
 module Polars
+  # @private
   class DataTypeGroup < Set
   end
 

data/lib/polars/data_types.rb

@@ -278,7 +278,6 @@ module Polars
   # Calendar date and time type.
   class Datetime < TemporalType
     attr_reader :time_unit, :time_zone
-    alias_method :tu, :time_unit
 
     def initialize(time_unit = "us", time_zone = nil)
       @time_unit = time_unit || "us"
@@ -303,7 +302,6 @@ module Polars
   # Time duration/delta type.
   class Duration < TemporalType
     attr_reader :time_unit
-    alias_method :tu, :time_unit
 
     def initialize(time_unit = "us")
       @time_unit = time_unit
@@ -443,7 +441,7 @@ module Polars
     attr_reader :inner
 
     def initialize(inner)
-      @inner = Utils.rb_type_to_dtype(inner)
+      @inner = Utils.parse_into_dtype(inner)
     end
 
     def ==(other)
@@ -463,21 +461,42 @@ module Polars
 
   # Nested list/array type.
   class Array < NestedType
-    attr_reader :inner, :width
+    attr_reader :inner, :size, :shape
 
-    def initialize(inner, width)
-      if width.is_a?(DataType) || (width.is_a?(Class) && width < DataType)
-        inner, width = width, inner
+    def initialize(inner, shape)
+      if shape.nil?
+        msg = "Array constructor is missing the required argument `shape`"
+        raise TypeError, msg
+      end
+
+      inner_parsed = Utils.parse_into_dtype(inner)
+      inner_shape = inner_parsed.is_a?(Array) ? inner_parsed.shape : []
+
+      if shape.is_a?(Integer)
+        @inner = inner_parsed
+        @size = shape
+        @shape = [shape] + inner_shape
+
+      elsif shape.is_a?(::Array) && shape[0].is_a?(Integer)
+        if shape.length > 1
+          inner_parsed = Array.new(inner_parsed, shape[1..])
+        end
+
+        @inner = inner_parsed
+        @size = shape[0]
+        @shape = shape + inner_shape
+
+      else
+        msg = "invalid input for shape: #{shape.inspect}"
+        raise TypeError, msg
       end
-      @inner = Utils.rb_type_to_dtype(inner) if inner
-      @width = width
     end
 
     def ==(other)
       if other.eql?(Array)
         true
       elsif other.is_a?(Array)
-        if @width != other.width
+        if @shape != other.shape
           false
         elsif @inner.nil? || other.inner.nil?
           true
@@ -490,7 +509,7 @@ module Polars
     end
 
     def to_s
-      "#{self.class.name}(#{inner}, width: #{width.inspect})"
+      "#{self.class.name}(#{inner}, shape: #{shape.inspect})"
     end
   end
 
@@ -500,7 +519,7 @@ module Polars
 
     def initialize(name, dtype)
       @name = name
-      @dtype = Utils.rb_type_to_dtype(dtype)
+      @dtype = Utils.parse_into_dtype(dtype)
     end
 
     def ==(other)

data/lib/polars/date_time_expr.rb

@@ -353,7 +353,7 @@ module Polars
     #
     # @return [Expr]
     def combine(time, time_unit: "us")
-      unless time.is_a?(Time) || time.is_a?(Expr)
+      unless time.is_a?(::Time) || time.is_a?(Expr)
         raise TypeError, "expected 'time' to be a Ruby time or Polars expression, found #{time}"
       end
       time = Utils.parse_into_expression(time)
@@ -398,13 +398,21 @@ module Polars
     #   # │ 2020-04-01 00:00:00 ┆ 2020/04/01 00:00:00 │
     #   # │ 2020-05-01 00:00:00 ┆ 2020/05/01 00:00:00 │
     #   # └─────────────────────┴─────────────────────┘
-    def to_string(format)
+    def to_string(format = nil)
+      if format.nil?
+        format = "iso"
+      end
       Utils.wrap_expr(_rbexpr.dt_to_string(format))
     end
 
-    # Format Date/datetime with a formatting rule.
+    # Convert a Date/Time/Datetime column into a String column with the given format.
+    #
+    # Similar to `cast(Polars::String)`, but this method allows you to customize the
+    # formatting of the resulting string.
     #
-    # See [chrono strftime/strptime](https://docs.rs/chrono/latest/chrono/format/strftime/index.html).
+    # @param format [String]
+    #   Format to use, refer to the [chrono strftime documentation](https://docs.rs/chrono/latest/chrono/format/strftime/index.html)
+    #   for specification. Example: `"%y-%m-%d"`.
     #
     # @return [Expr]
     #
@@ -451,8 +459,8 @@ module Polars
     #   # │ 2020-04-01 00:00:00 ┆ Wednesday ┆ April      │
     #   # │ 2020-05-01 00:00:00 ┆ Friday    ┆ May        │
     #   # └─────────────────────┴───────────┴────────────┘
-    def strftime(fmt)
-      Utils.wrap_expr(_rbexpr.strftime(fmt))
+    def strftime(format)
+      Utils.wrap_expr(_rbexpr.strftime(format))
     end
 
     # Extract the millennium from underlying representation.
@@ -717,6 +725,35 @@ module Polars
       Utils.wrap_expr(_rbexpr.dt_month)
     end
 
+    # Extract the number of days in the month from the underlying Date representation.
+    #
+    # Applies to Date and Datetime columns.
+    #
+    # Returns the number of days in the month.
+    # The return value ranges from 28 to 31.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"date" => [Date.new(2001, 1, 1), Date.new(2001, 2, 1), Date.new(2000, 2, 1)]}
+    #   )
+    #   df.with_columns(Polars.col("date").dt.days_in_month.alias("days_in_month"))
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌────────────┬───────────────┐
+    #   # │ date       ┆ days_in_month │
+    #   # │ ---        ┆ ---           │
+    #   # │ date       ┆ i8            │
+    #   # ╞════════════╪═══════════════╡
+    #   # │ 2001-01-01 ┆ 31            │
+    #   # │ 2001-02-01 ┆ 28            │
+    #   # │ 2000-02-01 ┆ 29            │
+    #   # └────────────┴───────────────┘
+    def days_in_month
+      Utils.wrap_expr(_rbexpr.dt_days_in_month)
+    end
+
     # Extract the week from the underlying Date representation.
     #
     # Applies to Date and Datetime columns.
@@ -896,13 +933,6 @@ module Polars
       Utils.wrap_expr(_rbexpr.dt_date)
     end
 
-    # Datetime
-    #
-    # @return [Expr]
-    def datetime
-      Utils.wrap_expr(_rbexpr.dt_datetime)
-    end
-
     # Extract hour from underlying DateTime representation.
     #
     # Applies to Datetime columns.
@@ -1190,7 +1220,7 @@ module Polars
       elsif time_unit == "s"
         timestamp("ms").floordiv(F.lit(1000, dtype: Int64))
       elsif time_unit == "d"
-        Utils.wrap_expr(_rbexpr).cast(:date).cast(:i32)
+        Utils.wrap_expr(_rbexpr).cast(Date).cast(Int32)
       else
         raise ArgumentError, "time_unit must be one of {'ns', 'us', 'ms', 's', 'd'}, got #{time_unit.inspect}"
       end
@@ -1228,47 +1258,6 @@ module Polars
       Utils.wrap_expr(_rbexpr.dt_timestamp(time_unit))
     end
 
-    # Set time unit of a Series of dtype Datetime or Duration.
-    #
-    # This does not modify underlying data, and should be used to fix an incorrect
-    # time unit.
-    #
-    # @param time_unit ["ns", "us", "ms"]
-    #   Time unit for the `Datetime` Series.
-    #
-    # @return [Expr]
-    #
-    # @example
-    #   df = Polars::DataFrame.new(
-    #     {
-    #       "date" => Polars.datetime_range(
-    #         Time.utc(2001, 1, 1),
-    #         Time.utc(2001, 1, 3),
-    #         "1d",
-    #         time_unit: "ns",
-    #         eager: true
-    #       )
-    #     }
-    #   )
-    #   df.select(
-    #     Polars.col("date"),
-    #     Polars.col("date").dt.with_time_unit("us").alias("time_unit_us")
-    #   )
-    #   # =>
-    #   # shape: (3, 2)
-    #   # ┌─────────────────────┬───────────────────────┐
-    #   # │ date                ┆ time_unit_us          │
-    #   # │ ---                 ┆ ---                   │
-    #   # │ datetime[ns]        ┆ datetime[μs]          │
-    #   # ╞═════════════════════╪═══════════════════════╡
-    #   # │ 2001-01-01 00:00:00 ┆ +32971-04-28 00:00:00 │
-    #   # │ 2001-01-02 00:00:00 ┆ +32974-01-22 00:00:00 │
-    #   # │ 2001-01-03 00:00:00 ┆ +32976-10-18 00:00:00 │
-    #   # └─────────────────────┴───────────────────────┘
-    def with_time_unit(time_unit)
-      Utils.wrap_expr(_rbexpr.dt_with_time_unit(time_unit))
-    end
-
     # Cast the underlying data to another time unit. This may lose precision.
     #
     # @param time_unit ["ns", "us", "ms"]
@@ -1451,7 +1440,7 @@ module Polars
     #   df.select(
     #     [
     #       Polars.col("date"),
-    #       Polars.col("date").diff.dt.days.alias("days_diff")
+    #       Polars.col("date").diff.dt.total_days.alias("days_diff")
     #     ]
     #   )
     #   # =>
@@ -1468,7 +1457,6 @@ module Polars
     def total_days(fractional: false)
       Utils.wrap_expr(_rbexpr.dt_total_days(fractional))
     end
-    alias_method :days, :total_days
 
     # Extract the hours from a Duration type.
     #
@@ -1488,7 +1476,7 @@ module Polars
     #   df.select(
     #     [
     #       Polars.col("date"),
-    #       Polars.col("date").diff.dt.hours.alias("hours_diff")
+    #       Polars.col("date").diff.dt.total_hours.alias("hours_diff")
     #     ]
     #   )
     #   # =>
@@ -1506,7 +1494,6 @@ module Polars
     def total_hours(fractional: false)
       Utils.wrap_expr(_rbexpr.dt_total_hours(fractional))
     end
-    alias_method :hours, :total_hours
 
     # Extract the minutes from a Duration type.
     #
@@ -1526,7 +1513,7 @@ module Polars
     #   df.select(
     #     [
     #       Polars.col("date"),
-    #       Polars.col("date").diff.dt.minutes.alias("minutes_diff")
+    #       Polars.col("date").diff.dt.total_minutes.alias("minutes_diff")
     #     ]
     #   )
     #   # =>
@@ -1544,7 +1531,6 @@ module Polars
     def total_minutes(fractional: false)
       Utils.wrap_expr(_rbexpr.dt_total_minutes(fractional))
     end
-    alias_method :minutes, :total_minutes
 
     # Extract the seconds from a Duration type.
     #
@@ -1564,7 +1550,7 @@ module Polars
     #   df.select(
     #     [
     #       Polars.col("date"),
-    #       Polars.col("date").diff.dt.seconds.alias("seconds_diff")
+    #       Polars.col("date").diff.dt.total_seconds.alias("seconds_diff")
     #     ]
     #   )
     #   # =>
@@ -1583,7 +1569,6 @@ module Polars
     def total_seconds(fractional: false)
       Utils.wrap_expr(_rbexpr.dt_total_seconds(fractional))
     end
-    alias_method :seconds, :total_seconds
 
     # Extract the milliseconds from a Duration type.
     #
@@ -1603,7 +1588,7 @@ module Polars
     #   df.select(
     #     [
     #       Polars.col("date"),
-    #       Polars.col("date").diff.dt.milliseconds.alias("milliseconds_diff")
+    #       Polars.col("date").diff.dt.total_milliseconds.alias("milliseconds_diff")
     #     ]
     #   )
     #   # =>
@@ -1628,7 +1613,6 @@ module Polars
     def total_milliseconds(fractional: false)
       Utils.wrap_expr(_rbexpr.dt_total_milliseconds(fractional))
     end
-    alias_method :milliseconds, :total_milliseconds
 
     # Extract the microseconds from a Duration type.
     #
@@ -1648,7 +1632,7 @@ module Polars
     #   df.select(
     #     [
     #       Polars.col("date"),
-    #       Polars.col("date").diff.dt.microseconds.alias("microseconds_diff")
+    #       Polars.col("date").diff.dt.total_microseconds.alias("microseconds_diff")
     #     ]
     #   )
     #   # =>
@@ -1673,7 +1657,6 @@ module Polars
     def total_microseconds(fractional: false)
       Utils.wrap_expr(_rbexpr.dt_total_microseconds(fractional))
     end
-    alias_method :microseconds, :total_microseconds
 
     # Extract the nanoseconds from a Duration type.
     #
@@ -1693,7 +1676,7 @@ module Polars
     #   df.select(
     #     [
     #       Polars.col("date"),
-    #       Polars.col("date").diff.dt.nanoseconds.alias("nanoseconds_diff")
+    #       Polars.col("date").diff.dt.total_nanoseconds.alias("nanoseconds_diff")
     #     ]
     #   )
     #   # =>
@@ -1718,7 +1701,6 @@ module Polars
     def total_nanoseconds(fractional: false)
       Utils.wrap_expr(_rbexpr.dt_total_nanoseconds(fractional))
     end
-    alias_method :nanoseconds, :total_nanoseconds
 
     # Offset this date by a relative time offset.
     #