polars-df 0.1.0 → 0.1.2

data/lib/polars/lazy_frame.rb

@@ -1,13 +1,203 @@
 module Polars
+  # Representation of a Lazy computation graph/query againat a DataFrame.
   class LazyFrame
+    # @private
     attr_accessor :_ldf
 
+    # @private
     def self._from_rbldf(rb_ldf)
       ldf = LazyFrame.allocate
       ldf._ldf = rb_ldf
       ldf
     end
 
+    # @private
+    def self._scan_csv(
+      file,
+      has_header: true,
+      sep: ",",
+      comment_char: nil,
+      quote_char: '"',
+      skip_rows: 0,
+      dtypes: nil,
+      null_values: nil,
+      ignore_errors: false,
+      cache: true,
+      with_column_names: nil,
+      infer_schema_length: 100,
+      n_rows: nil,
+      encoding: "utf8",
+      low_memory: false,
+      rechunk: true,
+      skip_rows_after_header: 0,
+      row_count_name: nil,
+      row_count_offset: 0,
+      parse_dates: false,
+      eol_char: "\n"
+    )
+      dtype_list = nil
+      if !dtypes.nil?
+        dtype_list = []
+        dtypes.each do |k, v|
+          dtype_list << [k, Utils.rb_type_to_dtype(v)]
+        end
+      end
+      processed_null_values = Utils._process_null_values(null_values)
+
+      _from_rbldf(
+        RbLazyFrame.new_from_csv(
+          file,
+          sep,
+          has_header,
+          ignore_errors,
+          skip_rows,
+          n_rows,
+          cache,
+          dtype_list,
+          low_memory,
+          comment_char,
+          quote_char,
+          processed_null_values,
+          infer_schema_length,
+          with_column_names,
+          rechunk,
+          skip_rows_after_header,
+          encoding,
+          Utils._prepare_row_count_args(row_count_name, row_count_offset),
+          parse_dates,
+          eol_char
+        )
+      )
+    end
+
+    # @private
+    def self._scan_parquet(
+      file,
+      n_rows: nil,
+      cache: true,
+      parallel: "auto",
+      rechunk: true,
+      row_count_name: nil,
+      row_count_offset: 0,
+      storage_options: nil,
+      low_memory: false
+    )
+      _from_rbldf(
+        RbLazyFrame.new_from_parquet(
+          file,
+          n_rows,
+          cache,
+          parallel,
+          rechunk,
+          Utils._prepare_row_count_args(row_count_name, row_count_offset),
+          low_memory
+        )
+      )
+    end
+
+    # @private
+    def self._scan_ipc(
+      file,
+      n_rows: nil,
+      cache: true,
+      rechunk: true,
+      row_count_name: nil,
+      row_count_offset: 0,
+      storage_options: nil,
+      memory_map: true
+    )
+      if file.is_a?(String) || (defined?(Pathname) && file.is_a?(Pathname))
+        file = Utils.format_path(file)
+      end
+
+      _from_rbldf(
+        RbLazyFrame.new_from_ipc(
+          file,
+          n_rows,
+          cache,
+          rechunk,
+          Utils._prepare_row_count_args(row_count_name, row_count_offset),
+          memory_map
+        )
+      )
+    end
+
+    # @private
+    def self._scan_ndjson(
+      file,
+      infer_schema_length: nil,
+      batch_size: nil,
+      n_rows: nil,
+      low_memory: false,
+      rechunk: true,
+      row_count_name: nil,
+      row_count_offset: 0
+    )
+      _from_rbldf(
+        RbLazyFrame.new_from_ndjson(
+          file,
+          infer_schema_length,
+          batch_size,
+          n_rows,
+          low_memory,
+          rechunk,
+          Utils._prepare_row_count_args(row_count_name, row_count_offset)
+        )
+      )
+    end
+
+    # def self.from_json
+    # end
+
+    # def self.read_json
+    # end
+
+    # def columns
+    # end
+
+    # def dtypes
+    # end
+
+    # def schema
+    # end
+
+    # def width
+    # end
+
+    # def include?(key)
+    # end
+
+    # clone handled by initialize_copy
+
+    # def [](item)
+    # end
+
+    # def to_s
+    # end
+    # alias_method :inspect, :to_s
+
+    # def write_json
+    # end
+
+    # def pipe
+    # end
+
+    # def describe_plan
+    # end
+
+    # def describe_optimized_plan
+    # end
+
+    # def show_graph
+    # end
+
+    # def sort
+    # end
+
+    # def profile
+    # end
+
+    #
     def collect(
       type_coercion: true,
       predicate_pushdown: true,
@@ -42,8 +232,23 @@ module Polars
       Utils.wrap_df(ldf.collect)
     end
 
+    # def fetch
+    # end
+
+    #
+    def lazy
+      self
+    end
+
+    # def cache
+    # end
+
+    # def cleared
+    # end
+
+    #
     def filter(predicate)
-      self.class._from_rbldf(
+      _from_rbldf(
         _ldf.filter(
           Utils.expr_to_lit_or_expr(predicate, str_to_lit: false)._rbexpr
         )
@@ -52,7 +257,7 @@ module Polars
 
     def select(exprs)
       exprs = Utils.selection_to_rbexpr_list(exprs)
-      self.class._from_rbldf(_ldf.select(exprs))
+      _from_rbldf(_ldf.select(exprs))
     end
 
     def groupby(by, maintain_order: false)
@@ -61,6 +266,16 @@ module Polars
       LazyGroupBy.new(lgb, self.class)
     end
 
+    # def groupby_rolling
+    # end
+
+    # def groupby_dynamic
+    # end
+
+    # def join_asof
+    # end
+
+    #
     def join(
       other,
       left_on: nil,
@@ -76,7 +291,7 @@ module Polars
       end
 
       if how == "cross"
-        return self.class._from_rbldf(
+        return _from_rbldf(
           _ldf.join(
             other._ldf, [], [], allow_parallel, force_parallel, how, suffix
           )
@@ -94,7 +309,7 @@ module Polars
         raise ArgumentError, "must specify `on` OR `left_on` and `right_on`"
       end
 
-      self.class._from_rbldf(
+      _from_rbldf(
         self._ldf.join(
           other._ldf,
           rbexprs_left,
@@ -123,17 +338,134 @@ module Polars
         when Expr
           rbexprs << e._rbexpr
         when Series
-          rbexprs = Internal.lit(e)._rbexpr
+          rbexprs = Utils.lit(e)._rbexpr
         else
           raise ArgumentError, "Expected an expression, got #{e}"
         end
       end
 
-      self.class._from_rbldf(_ldf.with_columns(rbexprs))
+      _from_rbldf(_ldf.with_columns(rbexprs))
     end
 
+    # def with_context
+    # end
+
+    #
     def with_column(column)
       with_columns([column])
     end
+
+    # def drop
+    # end
+
+    #
+    def rename(mapping)
+      existing = mapping.keys
+      _new = mapping.values
+      _from_rbldf(_ldf.rename(existing, _new))
+    end
+
+    # def reverse
+    # end
+
+    # def shift
+    # end
+
+    # def shift_and_fill
+    # end
+
+    # def slice
+    # end
+
+    # def limit
+    # end
+
+    # def head
+    # end
+
+    # def tail
+    # end
+
+    # def last
+    # end
+
+    # def first
+    # end
+
+    # def with_row_count
+    # end
+
+    # def take_every
+    # end
+
+    # def fill_null
+    # end
+
+    #
+    def fill_nan(fill_value)
+      if !fill_value.is_a?(Expr)
+        fill_value = Utils.lit(fill_value)
+      end
+      _from_rbldf(_ldf.fill_nan(fill_value._rbexpr))
+    end
+
+    # def std
+    # end
+
+    # def var
+    # end
+
+    # def max
+    # end
+
+    # def min
+    # end
+
+    # def sum
+    # end
+
+    # def mean
+    # end
+
+    # def median
+    # end
+
+    # def quantile
+    # end
+
+    #
+    def explode(columns)
+      columns = Utils.selection_to_rbexpr_list(columns)
+      _from_rbldf(_ldf.explode(columns))
+    end
+
+    # def unique
+    # end
+
+    # def drop_nulls
+    # end
+
+    # def melt
+    # end
+
+    # def map
+    # end
+
+    # def interpolate
+    # end
+
+    # def unnest
+    # end
+
+    private
+
+    def initialize_copy(other)
+      super
+      self._ldf = _ldf._clone
+    end
+
+    def _from_rbldf(rb_ldf)
+      self.class._from_rbldf(rb_ldf)
+    end
   end
 end

data/lib/polars/lazy_functions.rb

@@ -1,10 +1,45 @@
 module Polars
   module LazyFunctions
     def col(name)
-      name = name.to_s if name.is_a?(Symbol)
-      Utils.wrap_expr(RbExpr.col(name))
+      if name.is_a?(Series)
+        name = name.to_a
+      end
+
+      if name.is_a?(Array)
+        if name.length == 0 || name[0].is_a?(String) || name[0].is_a?(Symbol)
+          name = name.map { |v| v.is_a?(Symbol) ? v.to_s : v }
+          Utils.wrap_expr(RbExpr.cols(name))
+        elsif Utils.is_polars_dtype(name[0])
+          raise Todo
+          # Utils.wrap_expr(_dtype_cols(name))
+        else
+          raise ArgumentError, "Expected list values to be all `str` or all `DataType`"
+        end
+      else
+        name = name.to_s if name.is_a?(Symbol)
+        Utils.wrap_expr(RbExpr.col(name))
+      end
     end
 
+    def element
+      col("")
+    end
+
+    def count(column = nil)
+      if column.nil?
+        return Utils.wrap_expr(RbExpr.count)
+      end
+
+      if column.is_a?(Series)
+        column.len
+      else
+        col(column).count
+      end
+    end
+
+    # def to_list
+    # end
+
     def std(column, ddof: 1)
       if column.is_a?(Series)
         column.std(ddof: ddof)
@@ -55,7 +90,7 @@ module Polars
         # TODO
         Utils.wrap_expr(_sum_exprs(exprs))
       else
-        raise "todo"
+        raise Todo
       end
     end
 
@@ -79,10 +114,97 @@ module Polars
       end
     end
 
+    # def n_unique
+    # end
+
+    def first(column = nil)
+      if column.nil?
+        return Utils.wrap_expr(RbExpr.first)
+      end
+
+      if column.is_a?(Series)
+        if column.len > 0
+          column[0]
+        else
+          raise IndexError, "The series is empty, so no first value can be returned."
+        end
+      else
+        col(column).first
+      end
+    end
+
+    # def last
+    # end
+
+    # def head
+    # end
+
+    # def tail
+    # end
+
     def lit(value)
       Utils.wrap_expr(RbExpr.lit(value))
     end
 
+    # def cumsum
+    # end
+
+    # def spearman_rank_corr
+    # end
+
+    # def pearson_corr
+    # end
+
+    # def cov
+    # end
+
+    # def map
+    # end
+
+    # def apply
+    # end
+
+    def fold(acc, f, exprs)
+      acc = Utils.expr_to_lit_or_expr(acc, str_to_lit: true)
+      if exprs.is_a?(Expr)
+        exprs = [exprs]
+      end
+
+      exprs = Utils.selection_to_rbexpr_list(exprs)
+      Utils.wrap_expr(RbExpr.fold(acc._rbexpr, f, exprs))
+    end
+
+    # def reduce
+    # end
+
+    # def cumfold
+    # end
+
+    # def cumreduce
+    # end
+
+    # def any
+    # end
+
+    # def exclude
+    # 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 groups
+    # end
+
+    # def quantile
+    # 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)
@@ -102,16 +224,41 @@ module Polars
       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
+    # def argsort_by
+    # end
+
+    # def duration
+    # end
+
+    # def format
+    # end
+
+    def concat_list(exprs)
+      exprs = Utils.selection_to_rbexpr_list(exprs)
+      Utils.wrap_expr(RbExpr.concat_lst(exprs))
     end
 
+    # def collect_all
+    # end
+
+    # def select
+    # end
+
+    # def struct
+    # end
+
+    # def repeat
+    # end
+
+    # def arg_where
+    # end
+
+    # def coalesce
+    # end
+
+    # def from_epoch
+    # end
+
     def when(expr)
       expr = Utils.expr_to_lit_or_expr(expr)
       pw = RbExpr.when(expr._rbexpr)

data/lib/polars/list_expr.rb

@@ -0,0 +1,108 @@
+module Polars
+  class ListExpr
+    attr_accessor :_rbexpr
+
+    def initialize(expr)
+      self._rbexpr = expr._rbexpr
+    end
+
+    def lengths
+      Utils.wrap_expr(_rbexpr.arr_lengths)
+    end
+
+    def sum
+      Utils.wrap_expr(_rbexpr.lst_sum)
+    end
+
+    def max
+      Utils.wrap_expr(_rbexpr.lst_max)
+    end
+
+    def min
+      Utils.wrap_expr(_rbexpr.lst_min)
+    end
+
+    def mean
+      Utils.wrap_expr(_rbexpr.lst_mean)
+    end
+
+    def sort(reverse: false)
+      Utils.wrap_expr(_rbexpr.lst_sort(reverse))
+    end
+
+    def reverse
+      Utils.wrap_expr(_rbexpr.lst_reverse)
+    end
+
+    def unique
+      Utils.wrap_expr(_rbexpr.lst_unique)
+    end
+
+    # def concat
+    # end
+
+    def get(index)
+      index = Utils.expr_to_lit_or_expr(index, str_to_lit: false)._rbexpr
+      Utils.wrap_expr(_rbexpr.lst_get(index))
+    end
+
+    def [](item)
+      get(item)
+    end
+
+    def first
+      get(0)
+    end
+
+    def last
+      get(-1)
+    end
+
+    def contains(item)
+      Utils.wrap_expr(_rbexpr.arr_contains(Utils.expr_to_lit_or_expr(item)._rbexpr))
+    end
+
+    def join(separator)
+      Utils.wrap_expr(_rbexpr.lst_join(separator))
+    end
+
+    def arg_min
+      Utils.wrap_expr(_rbexpr.lst_arg_min)
+    end
+
+    def arg_max
+      Utils.wrap_expr(_rbexpr.lst_arg_max)
+    end
+
+    def diff(n: 1, null_behavior: "ignore")
+      Utils.wrap_expr(_rbexpr.lst_diff(n, null_behavior))
+    end
+
+    def shift(periods = 1)
+      Utils.wrap_expr(_rbexpr.lst_shift(periods))
+    end
+
+    def slice(offset, length = nil)
+      offset = Utils.expr_to_lit_or_expr(offset, str_to_lit: false)._rbexpr
+      length = Utils.expr_to_lit_or_expr(length, str_to_lit: false)._rbexpr
+      Utils.wrap_expr(_rbexpr.lst_slice(offset, length))
+    end
+
+    def head(n = 5)
+      slice(0, n)
+    end
+
+    def tail(n = 5)
+      offset = -Utils.expr_to_lit_or_expr(n, str_to_lit: false)
+      slice(offset, n)
+    end
+
+    # def to_struct(n_field_strategy: "first_non_null", name_generator: nil)
+    #   Utils.wrap_expr(_rbexpr.lst_to_struct(n_field_strategy, name_generator))
+    # end
+
+    def eval(expr, parallel: false)
+       Utils.wrap_expr(_rbexpr.lst_eval(expr._rbexpr, parallel))
+    end
+  end
+end