polars-df 0.1.0

data/lib/polars/lazy_frame.rb

@@ -0,0 +1,139 @@
+module Polars
+  class LazyFrame
+    attr_accessor :_ldf
+
+    def self._from_rbldf(rb_ldf)
+      ldf = LazyFrame.allocate
+      ldf._ldf = rb_ldf
+      ldf
+    end
+
+    def collect(
+      type_coercion: true,
+      predicate_pushdown: true,
+      projection_pushdown: true,
+      simplify_expression: true,
+      string_cache: false,
+      no_optimization: false,
+      slice_pushdown: true,
+      common_subplan_elimination: true,
+      allow_streaming: false
+    )
+      if no_optimization
+        predicate_pushdown = false
+        projection_pushdown = false
+        slice_pushdown = false
+        common_subplan_elimination = false
+      end
+
+      if allow_streaming
+        common_subplan_elimination = false
+      end
+
+      ldf = _ldf.optimization_toggle(
+        type_coercion,
+        predicate_pushdown,
+        projection_pushdown,
+        simplify_expression,
+        slice_pushdown,
+        common_subplan_elimination,
+        allow_streaming
+      )
+      Utils.wrap_df(ldf.collect)
+    end
+
+    def filter(predicate)
+      self.class._from_rbldf(
+        _ldf.filter(
+          Utils.expr_to_lit_or_expr(predicate, str_to_lit: false)._rbexpr
+        )
+      )
+    end
+
+    def select(exprs)
+      exprs = Utils.selection_to_rbexpr_list(exprs)
+      self.class._from_rbldf(_ldf.select(exprs))
+    end
+
+    def groupby(by, maintain_order: false)
+      rbexprs_by = Utils.selection_to_rbexpr_list(by)
+      lgb = _ldf.groupby(rbexprs_by, maintain_order)
+      LazyGroupBy.new(lgb, self.class)
+    end
+
+    def join(
+      other,
+      left_on: nil,
+      right_on: nil,
+      on: nil,
+      how: "inner",
+      suffix: "_right",
+      allow_parallel: true,
+      force_parallel: false
+    )
+      if !other.is_a?(LazyFrame)
+        raise ArgumentError, "Expected a `LazyFrame` as join table, got #{other.class.name}"
+      end
+
+      if how == "cross"
+        return self.class._from_rbldf(
+          _ldf.join(
+            other._ldf, [], [], allow_parallel, force_parallel, how, suffix
+          )
+        )
+      end
+
+      if !on.nil?
+        rbexprs = Utils.selection_to_rbexpr_list(on)
+        rbexprs_left = rbexprs
+        rbexprs_right = rbexprs
+      elsif !left_on.nil? && !right_on.nil?
+        rbexprs_left = Utils.selection_to_rbexpr_list(left_on)
+        rbexprs_right = Utils.selection_to_rbexpr_list(right_on)
+      else
+        raise ArgumentError, "must specify `on` OR `left_on` and `right_on`"
+      end
+
+      self.class._from_rbldf(
+        self._ldf.join(
+          other._ldf,
+          rbexprs_left,
+          rbexprs_right,
+          allow_parallel,
+          force_parallel,
+          how,
+          suffix,
+        )
+      )
+    end
+
+    def with_columns(exprs)
+      exprs =
+        if exprs.nil?
+          []
+        elsif exprs.is_a?(Expr)
+          [exprs]
+        else
+          exprs.to_a
+        end
+
+      rbexprs = []
+      exprs.each do |e|
+        case e
+        when Expr
+          rbexprs << e._rbexpr
+        when Series
+          rbexprs = Internal.lit(e)._rbexpr
+        else
+          raise ArgumentError, "Expected an expression, got #{e}"
+        end
+      end
+
+      self.class._from_rbldf(_ldf.with_columns(rbexprs))
+    end
+
+    def with_column(column)
+      with_columns([column])
+    end
+  end
+end

data/lib/polars/lazy_functions.rb

@@ -0,0 +1,121 @@
+module Polars
+  module LazyFunctions
+    def col(name)
+      name = name.to_s if name.is_a?(Symbol)
+      Utils.wrap_expr(RbExpr.col(name))
+    end
+
+    def std(column, ddof: 1)
+      if column.is_a?(Series)
+        column.std(ddof: ddof)
+      else
+        col(column).std(ddof: ddof)
+      end
+    end
+
+    def var(column, ddof: 1)
+      if column.is_a?(Series)
+        column.var(ddof: ddof)
+      else
+        col(column).var(ddof: ddof)
+      end
+    end
+
+    def max(column)
+      if column.is_a?(Series)
+        column.max
+      elsif column.is_a?(String) || column.is_a?(Symbol)
+        col(column).max
+      else
+        exprs = Utils.selection_to_rbexpr_list(column)
+        # TODO
+        Utils.wrap_expr(_max_exprs(exprs))
+      end
+    end
+
+    def min(column)
+      if column.is_a?(Series)
+        column.min
+      elsif column.is_a?(String) || column.is_a?(Symbol)
+        col(column).min
+      else
+        exprs = Utils.selection_to_rbexpr_list(column)
+        # TODO
+        Utils.wrap_expr(_min_exprs(exprs))
+      end
+    end
+
+    def sum(column)
+      if column.is_a?(Series)
+        column.sum
+      elsif column.is_a?(String) || column.is_a?(Symbol)
+        col(column.to_s).sum
+      elsif column.is_a?(Array)
+        exprs = Utils.selection_to_rbexpr_list(column)
+        # TODO
+        Utils.wrap_expr(_sum_exprs(exprs))
+      else
+        raise "todo"
+      end
+    end
+
+    def mean(column)
+      if column.is_a?(Series)
+        column.mean
+      else
+        col(column).mean
+      end
+    end
+
+    def avg(column)
+      mean(column)
+    end
+
+    def median(column)
+      if column.is_a?(Series)
+        column.median
+      else
+        col(column).median
+      end
+    end
+
+    def lit(value)
+      Utils.wrap_expr(RbExpr.lit(value))
+    end
+
+    def arange(low, high, step: 1, eager: false, dtype: nil)
+      low = Utils.expr_to_lit_or_expr(low, str_to_lit: false)
+      high = Utils.expr_to_lit_or_expr(high, str_to_lit: false)
+      range_expr = Utils.wrap_expr(RbExpr.arange(low._rbexpr, high._rbexpr, step))
+
+      if !dtype.nil? && dtype != "i64"
+        range_expr = range_expr.cast(dtype)
+      end
+
+      if !eager
+        range_expr
+      else
+        DataFrame.new
+          .select(range_expr)
+          .to_series
+          .rename("arange", in_place: true)
+      end
+    end
+
+    def all(name = nil)
+      if name.nil?
+        col("*")
+      elsif name.is_a?(String) || name.is_a?(Symbol)
+        col(name).all
+      else
+        raise "todo"
+      end
+    end
+
+    def when(expr)
+      expr = Utils.expr_to_lit_or_expr(expr)
+      pw = RbExpr.when(expr._rbexpr)
+      When.new(pw)
+    end
+  end
+end

data/lib/polars/lazy_group_by.rb

@@ -0,0 +1,13 @@
+module Polars
+  class LazyGroupBy
+    def initialize(lgb, lazyframe_class)
+      @lgb = lgb
+      @lazyframe_class = lazyframe_class
+    end
+
+    def agg(aggs)
+      rbexprs = Utils.selection_to_rbexpr_list(aggs)
+      @lazyframe_class._from_rbldf(@lgb.agg(rbexprs))
+    end
+  end
+end

data/lib/polars/series.rb

@@ -0,0 +1,261 @@
+module Polars
+  class Series
+    attr_accessor :_s
+
+    def initialize(name = nil, values = nil, dtype: nil, strict: true, nan_to_null: false, dtype_if_empty: nil)
+      # Handle case where values are passed as the first argument
+      if !name.nil? && !name.is_a?(String)
+        if values.nil?
+          values = name
+          name = nil
+        else
+          raise ArgumentError, "Series name must be a string."
+        end
+      end
+
+      name = "" if name.nil?
+
+      if values.nil?
+        self._s = sequence_to_rbseries(name, [], dtype: dtype, dtype_if_empty: dtype_if_empty)
+      elsif values.is_a?(Range)
+        self._s =
+          Polars.arange(
+            values.first,
+            values.last + (values.exclude_end? ? 0 : 1),
+            step: 1,
+            eager: true,
+            dtype: dtype
+          )
+          .rename(name, in_place: true)
+          ._s
+      elsif values.is_a?(Array)
+        self._s = sequence_to_rbseries(name, values, dtype: dtype, strict: strict, dtype_if_empty: dtype_if_empty)
+      else
+        raise ArgumentError, "Series constructor called with unsupported type; got #{values.class.name}"
+      end
+    end
+
+    def self._from_rbseries(s)
+      series = Series.allocate
+      series._s = s
+      series
+    end
+
+    def dtype
+      _s.dtype.to_sym
+    end
+
+    def name
+      _s.name
+    end
+
+    def shape
+      [_s.len]
+    end
+
+    def to_s
+      _s.to_s
+    end
+    alias_method :inspect, :to_s
+
+    def &(other)
+      Utils.wrap_s(_s.bitand(other._s))
+    end
+
+    def |(other)
+      Utils.wrap_s(_s.bitor(other._s))
+    end
+
+    def ^(other)
+      Utils.wrap_s(_s.bitxor(other._s))
+    end
+
+    def +(other)
+     Utils. wrap_s(_s.add(other._s))
+    end
+
+    def -(other)
+      Utils.wrap_s(_s.sub(other._s))
+    end
+
+    def *(other)
+      Utils.wrap_s(_s.mul(other._s))
+    end
+
+    def /(other)
+      Utils.wrap_s(_s.div(other._s))
+    end
+
+    def to_frame
+      Utils.wrap_df(RbDataFrame.new([_s]))
+    end
+
+    def sum
+      _s.sum
+    end
+
+    def mean
+      _s.mean
+    end
+
+    def min
+      _s.min
+    end
+
+    def max
+      _s.max
+    end
+
+    def alias(name)
+      s = dup
+      s._s.rename(name)
+      s
+    end
+
+    def rename(name, in_place: false)
+      if in_place
+        _s.rename(name)
+        self
+      else
+        self.alias(name)
+      end
+    end
+
+    def chunk_lengths
+      _s.chunk_lengths
+    end
+
+    def n_chunks
+      _s.n_chunks
+    end
+
+    def cumsum(reverse: false)
+      Utils.wrap_s(_s.cumsum(reverse))
+    end
+
+    def cummin(reverse: false)
+      Utils.wrap_s(_s.cummin(reverse))
+    end
+
+    def cummax(reverse: false)
+      Utils.wrap_s(_s.cummax(reverse))
+    end
+
+    def limit(n = 10)
+      to_frame().select(Utils.col(name).limit(n)).to_series
+    end
+
+    def slice(offset, length = nil)
+      length = len if length.nil?
+      Utils.wrap_s(_s.slice(offset, length))
+    end
+
+    def append(other)
+      _s.append(other._s)
+      self
+    end
+
+    def filter(predicate)
+      Utils.wrap_s(_s.filter(predicate._s))
+    end
+
+    def head(n = 10)
+      to_frame.select(Utils.col(name).head(n)).to_series
+    end
+
+    def tail(n = 10)
+      to_frame.select(Utils.col(name).tail(n)).to_series
+    end
+
+    def sort(reverse: false, in_place: false)
+      if in_place
+        self._s = _s.sort(reverse)
+        self
+      else
+        Utils.wrap_s(_s.sort(reverse))
+      end
+    end
+
+    def to_a
+      _s.to_a
+    end
+
+    def len
+      _s.len
+    end
+
+    def rechunk(in_place: false)
+      opt_s = _s.rechunk(in_place)
+      in_place ? self : Utils.wrap_s(opt_s)
+    end
+
+    private
+
+    def sequence_to_rbseries(name, values, dtype: nil, strict: true, dtype_if_empty: nil)
+      ruby_dtype = nil
+
+      if (values.nil? || values.empty?) && dtype.nil?
+        if dtype_if_empty
+          # if dtype for empty sequence could be guessed
+          # (e.g comparisons between self and other)
+          dtype = dtype_if_empty
+        else
+          # default to Float32 type
+          dtype = "f32"
+        end
+      end
+
+      # _get_first_non_none
+      value = values.find { |v| !v.nil? }
+
+      if !dtype.nil? && is_polars_dtype(dtype) && ruby_dtype.nil?
+        constructor = polars_type_to_constructor(dtype)
+        rbseries = constructor.call(name, values, strict)
+        return rbseries
+      end
+
+      constructor = rb_type_to_constructor(value.class)
+      constructor.call(name, values, strict)
+    end
+
+    POLARS_TYPE_TO_CONSTRUCTOR = {
+      f32: RbSeries.method(:new_opt_f32),
+      f64: RbSeries.method(:new_opt_f64),
+      i8: RbSeries.method(:new_opt_i8),
+      i16: RbSeries.method(:new_opt_i16),
+      i32: RbSeries.method(:new_opt_i32),
+      i64: RbSeries.method(:new_opt_i64),
+      u8: RbSeries.method(:new_opt_u8),
+      u16: RbSeries.method(:new_opt_u16),
+      u32: RbSeries.method(:new_opt_u32),
+      u64: RbSeries.method(:new_opt_u64),
+      bool: RbSeries.method(:new_opt_bool),
+      str: RbSeries.method(:new_str)
+    }
+
+    def polars_type_to_constructor(dtype)
+      POLARS_TYPE_TO_CONSTRUCTOR.fetch(dtype.to_sym)
+    rescue KeyError
+      raise ArgumentError, "Cannot construct RbSeries for type #{dtype}."
+    end
+
+    RB_TYPE_TO_CONSTRUCTOR = {
+      Float => RbSeries.method(:new_opt_f64),
+      Integer => RbSeries.method(:new_opt_i64),
+      String => RbSeries.method(:new_str),
+      TrueClass => RbSeries.method(:new_opt_bool),
+      FalseClass => RbSeries.method(:new_opt_bool)
+    }
+
+    def rb_type_to_constructor(dtype)
+      RB_TYPE_TO_CONSTRUCTOR.fetch(dtype)
+    rescue KeyError
+      # RbSeries.method(:new_object)
+      raise ArgumentError, "Cannot determine type"
+    end
+
+    def is_polars_dtype(data_type)
+      true
+    end
+  end
+end

data/lib/polars/string_expr.rb

@@ -0,0 +1,17 @@
+module Polars
+  class StringExpr
+    attr_accessor :_rbexpr
+
+    def initialize(expr)
+      self._rbexpr = expr._rbexpr
+    end
+
+    def lengths
+      Utils.wrap_expr(_rbexpr.str_lengths)
+    end
+
+    def contains(pattern, literal: false)
+      Utils.wrap_expr(_rbexpr.str_contains(pattern, literal))
+    end
+  end
+end

data/lib/polars/utils.rb

@@ -0,0 +1,47 @@
+module Polars
+  module Utils
+    def self.wrap_s(s)
+      Series._from_rbseries(s)
+    end
+
+    def self.wrap_df(df)
+      DataFrame._from_rbdf(df)
+    end
+
+    def self.wrap_expr(rbexpr)
+      Expr._from_rbexpr(rbexpr)
+    end
+
+    def self.col(name)
+      Polars.col(name)
+    end
+
+    def self.selection_to_rbexpr_list(exprs)
+      if exprs.is_a?(String) || exprs.is_a?(Expr) || exprs.is_a?(Series)
+        exprs = [exprs]
+      end
+
+      exprs.map { |e| expr_to_lit_or_expr(e, str_to_lit: false)._rbexpr }
+    end
+
+    def self.expr_to_lit_or_expr(expr, str_to_lit: true)
+      if expr.is_a?(String) && !str_to_lit
+        col(expr)
+      elsif expr.is_a?(Integer) || expr.is_a?(Float) || expr.is_a?(String) || expr.is_a?(Series) || expr.nil?
+        lit(expr)
+      elsif expr.is_a?(Expr)
+        expr
+      else
+        raise ArgumentError, "did not expect value #{expr} of type #{expr.class.name}, maybe disambiguate with Polars.lit or Polars.col"
+      end
+    end
+
+    def self.lit(value)
+      Polars.lit(value)
+    end
+
+    def self.format_path(path)
+      File.expand_path(path)
+    end
+  end
+end

data/lib/polars/version.rb

@@ -0,0 +1,3 @@
+module Polars
+  VERSION = "0.1.0"
+end

data/lib/polars/when.rb

@@ -0,0 +1,15 @@
+module Polars
+  class When
+    attr_accessor :_rbwhen
+
+    def initialize(rbwhen)
+      self._rbwhen = rbwhen
+    end
+
+    def then(expr)
+      expr = Utils.expr_to_lit_or_expr(expr)
+      rbwhenthen = _rbwhen._then(expr._rbexpr)
+      WhenThen.new(rbwhenthen)
+    end
+  end
+end

data/lib/polars/when_then.rb

@@ -0,0 +1,18 @@
+module Polars
+  class WhenThen
+    attr_accessor :_rbwhenthen
+
+    def initialize(rbwhenthen)
+      self._rbwhenthen = rbwhenthen
+    end
+
+    def when(predicate)
+      WhenThenThen.new(_rbwhenthen.when(predicate._rbexpr))
+    end
+
+    def otherwise(expr)
+      expr = Utils.expr_to_lit_or_expr(expr)
+      Utils.wrap_expr(_rbwhenthen.otherwise(expr._rbexpr))
+    end
+  end
+end

data/lib/polars-df.rb

@@ -0,0 +1 @@
+require "polars"

data/lib/polars.rb

@@ -0,0 +1,25 @@
+# ext
+require "polars/polars"
+
+# modules
+require "polars/data_frame"
+require "polars/expr"
+require "polars/functions"
+require "polars/lazy_frame"
+require "polars/lazy_functions"
+require "polars/lazy_group_by"
+require "polars/io"
+require "polars/series"
+require "polars/string_expr"
+require "polars/utils"
+require "polars/version"
+require "polars/when"
+require "polars/when_then"
+
+module Polars
+  class Error < StandardError; end
+
+  extend Functions
+  extend IO
+  extend LazyFunctions
+end