polars-df 0.15.0 → 0.26.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_expr.rb

@@ -0,0 +1,52 @@
+module Polars
+  # A lazily instantiated `DataType` that can be used in an `Expr`.
+  class DataTypeExpr
+    # @private
+    attr_accessor :_rbdatatype_expr
+
+    # @private
+    def self._from_rbdatatype_expr(rbdatatype_expr)
+      slf = new
+      slf._rbdatatype_expr = rbdatatype_expr
+      slf
+    end
+
+    # Materialize the `DataTypeExpr` in a specific context.
+    #
+    # This is a useful function when debugging datatype expressions.
+    #
+    # @return [DataType]
+    #
+    # @example
+    #   lf = Polars::LazyFrame.new(
+    #     {
+    #       "a" => [1, 2, 3]
+    #     }
+    #   )
+    #   Polars.dtype_of("a").collect_dtype(lf)
+    #   # => Polars::Int64
+    #
+    # @example
+    #   Polars.dtype_of("a").collect_dtype({"a" => Polars::String})
+    #   # => Polars::String
+    def collect_dtype(
+      context
+    )
+      schema = nil
+      if context.is_a?(Schema)
+         schema = context
+      elsif context.is_a?(Hash)
+        schema = Schema.new(context)
+      elsif context.is_a?(DataFrame)
+        schema = context.schema
+      elsif context.is_a?(LazyFrame)
+        schema = context.collect_schema
+      else
+        msg = "DataTypeExpr.collect_dtype did not expect #{context.inspect}"
+        raise TypeError, msg
+      end
+
+      _rbdatatype_expr.collect_dtype(schema.to_h)
+    end
+  end
+end

data/lib/polars/data_type_group.rb

@@ -1,7 +1,9 @@
 module Polars
+  # @private
   class DataTypeGroup < Set
   end
 
+  # @private
   SIGNED_INTEGER_DTYPES = DataTypeGroup.new(
     [
       Int8,
@@ -10,6 +12,7 @@ module Polars
       Int64
     ]
   )
+  # @private
   UNSIGNED_INTEGER_DTYPES = DataTypeGroup.new(
     [
       UInt8,
@@ -18,10 +21,13 @@ module Polars
       UInt64
     ]
   )
+  # @private
   INTEGER_DTYPES = (
     SIGNED_INTEGER_DTYPES | UNSIGNED_INTEGER_DTYPES
   )
+  # @private
   FLOAT_DTYPES = DataTypeGroup.new([Float32, Float64])
+  # @private
   NUMERIC_DTYPES = DataTypeGroup.new(
     FLOAT_DTYPES + INTEGER_DTYPES | [Decimal]
   )

data/lib/polars/data_types.rb

@@ -99,12 +99,34 @@ module Polars
       self < NestedType
     end
 
+    # Return a `DataTypeExpr` with a static `DataType`.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   Polars::Int16.new.to_dtype_expr.collect_dtype({})
+    #   # => Polars::Int16
+    def self.to_dtype_expr
+      DataTypeExpr._from_rbdatatype_expr(RbDataTypeExpr.from_dtype(self))
+    end
+
     [:numeric?, :decimal?, :integer?, :signed_integer?, :unsigned_integer?, :float?, :temporal?, :nested?].each do |v|
       define_method(v) do
         self.class.public_send(v)
       end
     end
 
+    # Return a `DataTypeExpr` with a static `DataType`.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   Polars::Int16.new.to_dtype_expr.collect_dtype({})
+    #   # => Polars::Int16
+    def to_dtype_expr
+      DataTypeExpr._from_rbdatatype_expr(RbDataTypeExpr.from_dtype(self))
+    end
+
     # Returns a string representing the data type.
     #
     # @return [String]
@@ -167,6 +189,10 @@ module Polars
   class Int64 < SignedIntegerType
   end
 
+  # 128-bit signed integer type.
+  class Int128 < SignedIntegerType
+  end
+
   # 8-bit unsigned integer type.
   class UInt8 < UnsignedIntegerType
   end
@@ -183,6 +209,14 @@ module Polars
   class UInt64 < UnsignedIntegerType
   end
 
+  # 128-bit unsigned integer type.
+  class UInt128 < UnsignedIntegerType
+  end
+
+  # 16-bit floating point type.
+  class Float16 < FloatType
+  end
+
   # 32-bit floating point type.
   class Float32 < FloatType
   end
@@ -197,7 +231,11 @@ module Polars
   class Decimal < NumericType
     attr_reader :precision, :scale
 
-    def initialize(precision, scale)
+    def initialize(precision = nil, scale = 0)
+      if precision.nil?
+        precision = 38
+      end
+
       @precision = precision
       @scale = scale
     end
@@ -244,7 +282,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"
@@ -269,7 +306,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
@@ -290,12 +326,55 @@ module Polars
     end
   end
 
+  # A named collection of categories for `Categorical`.
+  #
+  # Two categories are considered equal (and will use the same physical mapping of
+  # categories to strings) if they have the same name, namespace and physical backing
+  # type, even if they are created in separate calls to `Categories`.
+  #
+  # @note
+  #   This functionality is currently considered **unstable**. It may be
+  #   changed at any point without it being considered a breaking change.
+  class Categories
+    attr_accessor :_categories
+
+    def initialize(name = nil)
+      if name.nil? || name == ""
+        self._categories = RbCategories.global_categories
+        return
+      end
+
+      raise Todo
+    end
+
+    # @private
+    def self._from_rb_categories(rb_categories)
+      slf = new
+      slf._categories = rb_categories
+      slf
+    end
+  end
+
   # A categorical encoding of a set of strings.
   class Categorical < DataType
-    attr_reader :ordering
+    attr_reader :ordering, :categories
 
-    def initialize(ordering = "physical")
-      @ordering = ordering
+    def initialize(ordering = "physical", **kwargs)
+      if ordering.is_a?(Categories)
+        @ordering = "lexical"
+        @categories = ordering
+        # assert kwargs.length == 0
+        return
+      end
+
+      @ordering = "lexical"
+      if kwargs[:categories]
+        # assert kwargs.length == 1
+        @categories = kwargs[:categories]
+      else
+        # assert kwargs.length == 0
+        @categories = Categories.new
+      end
     end
   end
 
@@ -311,7 +390,7 @@ module Polars
       end
 
       if categories.empty?
-        self.categories = Series.new("category", [], dtype: String)
+        @categories = Series.new("category", [], dtype: String)
         return
       end
 
@@ -353,7 +432,7 @@ module Polars
   class Object < DataType
   end
 
-  # Type representing Null / None values.
+  # Type representing Null / nil values.
   class Null < DataType
   end
 
@@ -366,7 +445,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)
@@ -386,21 +465,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
@@ -413,7 +513,7 @@ module Polars
     end
 
     def to_s
-      "#{self.class.name}(#{inner}, width: #{width.inspect})"
+      "#{self.class.name}(#{inner}, shape: #{shape.inspect})"
     end
   end
 
@@ -423,7 +523,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)