polars-df 0.1.0

data/lib/polars/data_frame.rb

@@ -0,0 +1,315 @@
+module Polars
+  class DataFrame
+    attr_accessor :_df
+
+    def initialize(data = nil)
+      if defined?(ActiveRecord) && (data.is_a?(ActiveRecord::Relation) || data.is_a?(ActiveRecord::Result))
+        result = data.is_a?(ActiveRecord::Result) ? data : data.connection.select_all(data.to_sql)
+        data = {}
+        result.columns.each_with_index do |k, i|
+          data[k] = result.rows.map { |r| r[i] }
+        end
+      end
+
+      if data.nil?
+        self._df = hash_to_rbdf({})
+      elsif data.is_a?(Hash)
+        data = data.transform_keys { |v| v.is_a?(Symbol) ? v.to_s : v }
+        self._df = hash_to_rbdf(data)
+      elsif data.is_a?(Array)
+        self._df = sequence_to_rbdf(data)
+      elsif data.is_a?(Series)
+        self._df = series_to_rbdf(data)
+      else
+        raise ArgumentError, "DataFrame constructor called with unsupported type; got #{data.class.name}"
+      end
+    end
+
+    def self._from_rbdf(rb_df)
+      df = DataFrame.allocate
+      df._df = rb_df
+      df
+    end
+
+    def self._read_csv(file, has_header: true)
+      if file.is_a?(String) || (defined?(Pathname) && file.is_a?(Pathname))
+        file = Utils.format_path(file)
+      end
+
+      _from_rbdf(RbDataFrame.read_csv(file, has_header))
+    end
+
+    def self._read_parquet(file)
+      if file.is_a?(String) || (defined?(Pathname) && file.is_a?(Pathname))
+        file = Utils.format_path(file)
+      end
+
+      _from_rbdf(RbDataFrame.read_parquet(file))
+    end
+
+    def self._read_json(file)
+      if file.is_a?(String) || (defined?(Pathname) && file.is_a?(Pathname))
+        file = Utils.format_path(file)
+      end
+
+      _from_rbdf(RbDataFrame.read_json(file))
+    end
+
+    def self._read_ndjson(file)
+      if file.is_a?(String) || (defined?(Pathname) && file.is_a?(Pathname))
+        file = Utils.format_path(file)
+      end
+
+      _from_rbdf(RbDataFrame.read_ndjson(file))
+    end
+
+    def shape
+      _df.shape
+    end
+
+    def height
+      _df.height
+    end
+
+    def width
+      _df.width
+    end
+
+    def columns
+      _df.columns
+    end
+
+    def dtypes
+      _df.dtypes.map(&:to_sym)
+    end
+
+    def to_s
+      _df.to_s
+    end
+    alias_method :inspect, :to_s
+
+    def include?(name)
+      columns.include?(name)
+    end
+
+    def [](name)
+      Utils.wrap_s(_df.column(name))
+    end
+
+    def to_series(index = 0)
+      if index < 0
+        index = columns.length + index
+      end
+      Utils.wrap_s(_df.select_at_idx(index))
+    end
+
+    def write_json(
+      file,
+      pretty: false,
+      row_oriented: false
+    )
+      if file.is_a?(String) || (defined?(Pathname) && file.is_a?(Pathname))
+        file = Utils.format_path(file)
+      end
+
+      _df.write_json(file, pretty, row_oriented)
+      nil
+    end
+
+    def write_ndjson(file)
+      if file.is_a?(String) || (defined?(Pathname) && file.is_a?(Pathname))
+        file = Utils.format_path(file)
+      end
+
+      _df.write_ndjson(file)
+      nil
+    end
+
+    def write_csv(
+      file = nil,
+      has_header: true,
+      sep: ",",
+      quote: '"',
+      batch_size: 1024,
+      datetime_format: nil,
+      date_format: nil,
+      time_format: nil,
+      float_precision: nil,
+      null_value: nil
+    )
+      if sep.length > 1
+        raise ArgumentError, "only single byte separator is allowed"
+      elsif quote.length > 1
+        raise ArgumentError, "only single byte quote char is allowed"
+      elsif null_value == ""
+        null_value = nil
+      end
+
+      if file.nil?
+        buffer = StringIO.new
+        buffer.set_encoding(Encoding::BINARY)
+        _df.write_csv(
+          buffer,
+          has_header,
+          sep.ord,
+          quote.ord,
+          batch_size,
+          datetime_format,
+          date_format,
+          time_format,
+          float_precision,
+          null_value
+        )
+        buffer.rewind
+        return buffer.read.force_encoding(Encoding::UTF_8)
+      end
+
+      if file.is_a?(String) || (defined?(Pathname) && file.is_a?(Pathname))
+        file = Utils.format_path(file)
+      end
+
+      _df.write_csv(
+        file,
+        has_header,
+        sep.ord,
+        quote.ord,
+        batch_size,
+        datetime_format,
+        date_format,
+        time_format,
+        float_precision,
+        null_value,
+      )
+      nil
+    end
+
+    def write_parquet(
+      file,
+      compression: "zstd",
+      compression_level: nil,
+      statistics: false,
+      row_group_size: nil
+    )
+      if compression.nil?
+        compression = "uncompressed"
+      end
+      if file.is_a?(String) || (defined?(Pathname) && file.is_a?(Pathname))
+        file = Utils.format_path(file)
+      end
+
+      _df.write_parquet(
+        file, compression, compression_level, statistics, row_group_size
+      )
+    end
+
+    def filter(predicate)
+      lazy.filter(predicate).collect
+    end
+
+    def sort(by, reverse: false, nulls_last: false)
+      _from_rbdf(_df.sort(by, reverse, nulls_last))
+    end
+
+    def frame_equal(other, null_equal: true)
+      _df.frame_equal(other._df, null_equal)
+    end
+
+    def limit(n = 5)
+      head(n)
+    end
+
+    def head(n = 5)
+      _from_rbdf(_df.head(n))
+    end
+
+    def tail(n = 5)
+      _from_rbdf(_df.tail(n))
+    end
+
+    def groupby(by, maintain_order: false)
+      lazy.groupby(by, maintain_order: maintain_order)
+    end
+
+    def join(other, left_on: nil, right_on: nil, on: nil, how: "inner", suffix: "_right")
+      lazy
+        .join(
+          other.lazy,
+          left_on: left_on,
+          right_on: right_on,
+          on: on,
+          how: how,
+          suffix: suffix,
+        )
+        .collect(no_optimization: true)
+    end
+
+    def with_column(column)
+      lazy
+        .with_column(column)
+        .collect(no_optimization: true, string_cache: false)
+    end
+
+    def lazy
+      wrap_ldf(_df.lazy)
+    end
+
+    def select(exprs)
+      _from_rbdf(
+        lazy
+          .select(exprs)
+          .collect(no_optimization: true, string_cache: false)
+          ._df
+      )
+    end
+
+    def mean(axis: 0, null_strategy: "ignore")
+      case axis
+      when 0
+        _from_rbdf(_df.mean)
+      when 1
+        Utils.wrap_s(_df.hmean(null_strategy))
+      else
+        raise ArgumentError, "Axis should be 0 or 1."
+      end
+    end
+
+    def with_columns(exprs)
+      if !exprs.nil? && !exprs.is_a?(Array)
+        exprs = [exprs]
+      end
+      lazy
+        .with_columns(exprs)
+        .collect(no_optimization: true, string_cache: false)
+    end
+
+    def rechunk
+      _from_rbdf(_df.rechunk)
+    end
+
+    def null_count
+      _from_rbdf(_df.null_count)
+    end
+
+    private
+
+    def hash_to_rbdf(data)
+      RbDataFrame.read_hash(data)
+    end
+
+    def sequence_to_rbdf(data)
+      RbDataFrame.new(data.map(&:_s))
+    end
+
+    def series_to_rbdf(data)
+      RbDataFrame.new([data._s])
+    end
+
+    def wrap_ldf(ldf)
+      LazyFrame._from_rbldf(ldf)
+    end
+
+    def _from_rbdf(rb_df)
+      self.class._from_rbdf(rb_df)
+    end
+  end
+end

data/lib/polars/expr.rb

@@ -0,0 +1,233 @@
+module Polars
+  class Expr
+    attr_accessor :_rbexpr
+
+    def self._from_rbexpr(rbexpr)
+      expr = Expr.allocate
+      expr._rbexpr = rbexpr
+      expr
+    end
+
+    def to_s
+      _rbexpr.to_str
+    end
+    alias_method :inspect, :to_s
+
+    def ^(other)
+      wrap_expr(_rbexpr._xor(_to_rbexpr(other)))
+    end
+
+    def &(other)
+      wrap_expr(_rbexpr._and(_to_rbexpr(other)))
+    end
+
+    def |(other)
+      wrap_expr(_rbexpr._or(_to_rbexpr(other)))
+    end
+
+    def *(other)
+      wrap_expr(_rbexpr * _to_rbexpr(other))
+    end
+
+    def >=(other)
+      wrap_expr(_rbexpr.gt_eq(_to_expr(other)._rbexpr))
+    end
+
+    def <=(other)
+      wrap_expr(_rbexpr.lt_eq(_to_expr(other)._rbexpr))
+    end
+
+    def ==(other)
+      wrap_expr(_rbexpr.eq(_to_expr(other)._rbexpr))
+    end
+
+    def !=(other)
+      wrap_expr(_rbexpr.neq(_to_expr(other)._rbexpr))
+    end
+
+    def <(other)
+      wrap_expr(_rbexpr.lt(_to_expr(other)._rbexpr))
+    end
+
+    def >(other)
+      wrap_expr(_rbexpr.gt(_to_expr(other)._rbexpr))
+    end
+
+    def alias(name)
+      wrap_expr(_rbexpr._alias(name))
+    end
+
+    def suffix(suffix)
+      wrap_expr(_rbexpr.suffix(suffix))
+    end
+
+    def is_not
+      wrap_expr(_rbexpr.is_not)
+    end
+
+    def is_null
+      wrap_expr(_rbexpr.is_null)
+    end
+
+    def is_not_null
+      wrap_expr(_rbexpr.is_not_null)
+    end
+
+    def count
+      wrap_expr(_rbexpr.count)
+    end
+
+    def len
+      count
+    end
+
+    def sort(reverse: false, nulls_last: false)
+      wrap_expr(_rbexpr.sort_with(reverse, nulls_last))
+    end
+
+    def sort_by(by, reverse: false)
+      if !by.is_a?(Array)
+        by = [by]
+      end
+      if !reverse.is_a?(Array)
+        reverse = [reverse]
+      end
+      by = Utils.selection_to_rbexpr_list(by)
+
+      wrap_expr(_rbexpr.sort_by(by, reverse))
+    end
+
+    def fill_null(value = nil, strategy: nil, limit: nil)
+      if !value.nil? && !strategy.nil?
+        raise ArgumentError, "cannot specify both 'value' and 'strategy'."
+      elsif value.nil? && strategy.nil?
+        raise ArgumentError, "must specify either a fill 'value' or 'strategy'"
+      elsif ["forward", "backward"].include?(strategy) && !limit.nil?
+        raise ArgumentError, "can only specify 'limit' when strategy is set to 'backward' or 'forward'"
+      end
+
+      if !value.nil?
+        value = Utils.expr_to_lit_or_expr(value, str_to_lit: true)
+        wrap_expr(_rbexpr.fill_null(value._rbexpr))
+      else
+        wrap_expr(_rbexpr.fill_null_with_strategy(strategy, limit))
+      end
+    end
+
+    def fill_nan(fill_value)
+      fill_value = Utils.expr_to_lit_or_expr(fill_value, str_to_lit: true)
+      wrap_expr(_rbexpr.fill_nan(fill_value._rbexpr))
+    end
+
+    def reverse
+      wrap_expr(_rbexpr.reverse)
+    end
+
+    def std(ddof: 1)
+      wrap_expr(_rbexpr.std(ddof))
+    end
+
+    def var(ddof: 1)
+      wrap_expr(_rbexpr.var(ddof))
+    end
+
+    def max
+      wrap_expr(_rbexpr.max)
+    end
+
+    def min
+      wrap_expr(_rbexpr.min)
+    end
+
+    def nan_max
+      wrap_expr(_rbexpr.nan_max)
+    end
+
+    def nan_min
+      wrap_expr(_rbexpr.nan_min)
+    end
+
+    def sum
+      wrap_expr(_rbexpr.sum)
+    end
+
+    def mean
+      wrap_expr(_rbexpr.mean)
+    end
+
+    def median
+      wrap_expr(_rbexpr.median)
+    end
+
+    def product
+      wrap_expr(_rbexpr.product)
+    end
+
+    def n_unique
+      wrap_expr(_rbexpr.n_unique)
+    end
+
+    def unique(maintain_order: false)
+      if maintain_order
+        wrap_expr(_rbexpr.unique_stable)
+      else
+        wrap_expr(_rbexpr.unique)
+      end
+    end
+
+    def first
+      wrap_expr(_rbexpr.first)
+    end
+
+    def last
+      wrap_expr(_rbexpr.last)
+    end
+
+    def over(expr)
+      rbexprs = Utils.selection_to_rbexpr_list(expr)
+      wrap_expr(_rbexpr.over(rbexprs))
+    end
+
+    def filter(predicate)
+      wrap_expr(_rbexpr.filter(predicate._rbexpr))
+    end
+
+    def head(n = 10)
+      wrap_expr(_rbexpr.head(n))
+    end
+
+    def tail(n = 10)
+      wrap_expr(_rbexpr.tail(n))
+    end
+
+    def limit(n = 10)
+      head(n)
+    end
+
+    def interpolate
+      wrap_expr(_rbexpr.interpolate)
+    end
+
+    def list
+      wrap_expr(_rbexpr.list)
+    end
+
+    def str
+      StringExpr.new(self)
+    end
+
+    private
+
+    def wrap_expr(expr)
+      Utils.wrap_expr(expr)
+    end
+
+    def _to_rbexpr(other)
+      _to_expr(other)._rbexpr
+    end
+
+    def _to_expr(other)
+      other.is_a?(Expr) ? other : Utils.lit(other)
+    end
+  end
+end

data/lib/polars/functions.rb

@@ -0,0 +1,45 @@
+module Polars
+  module Functions
+    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(_concat_df(items))
+        elsif how == "diagonal"
+          out = Utils.wrap_df(_diag_concat_df(items))
+        elsif how == "horizontal"
+          out = Utils.wrap_df(_hor_concat_df(items))
+        else
+          raise ArgumentError, "how must be one of {{'vertical', 'diagonal', 'horizontal'}}, got #{how}"
+        end
+      elsif first.is_a?(LazyFrame)
+        if how == "vertical"
+          # TODO
+          return Utils.wrap_ldf(_concat_lf(items, rechunk, parallel))
+        else
+          raise ArgumentError, "Lazy only allows 'vertical' concat strategy."
+        end
+      elsif first.is_a?(Series)
+        # TODO
+        out = Utils.wrap_s(_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
+  end
+end

data/lib/polars/io.rb

@@ -0,0 +1,39 @@
+module Polars
+  module IO
+    def read_csv(file, has_header: true)
+      _prepare_file_arg(file) do |data|
+        DataFrame._read_csv(data, has_header: has_header)
+      end
+    end
+
+    def read_parquet(file)
+      _prepare_file_arg(file) do |data|
+        DataFrame._read_parquet(data)
+      end
+    end
+
+    def read_json(file)
+      DataFrame._read_json(file)
+    end
+
+    def read_ndjson(file)
+      DataFrame._read_ndjson(file)
+    end
+
+    private
+
+    def _prepare_file_arg(file)
+      if file.is_a?(String) && file =~ /\Ahttps?:\/\//
+        raise ArgumentError, "use URI(...) for remote files"
+      end
+
+      if defined?(URI) && file.is_a?(URI)
+        require "open-uri"
+
+        file = URI.open(file)
+      end
+
+      yield file
+    end
+  end
+end