polars-df 0.3.1-x86_64-linux → 0.5.0-x86_64-linux

data/lib/polars/data_frame.rb

@@ -18,7 +18,10 @@ module Polars
     #   Whether to interpret two-dimensional data as columns or as rows. If `nil`,
     #   the orientation is inferred by matching the columns and data dimensions. If
     #   this does not yield conclusive results, column orientation is used.
-    def initialize(data = nil, columns: nil, orient: nil)
+    def initialize(data = nil, schema: nil, columns: nil, schema_overrides: nil, orient: nil, infer_schema_length: 100, nan_to_null: false)
+      schema ||= columns
+      raise Todo if schema_overrides
+
       # TODO deprecate in favor of read_sql
       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)
@@ -29,14 +32,14 @@ module Polars
       end
 
       if data.nil?
-        self._df = self.class.hash_to_rbdf({}, columns: columns)
+        self._df = self.class.hash_to_rbdf({}, schema: schema, schema_overrides: schema_overrides)
       elsif data.is_a?(Hash)
         data = data.transform_keys { |v| v.is_a?(Symbol) ? v.to_s : v }
-        self._df = self.class.hash_to_rbdf(data, columns: columns)
+        self._df = self.class.hash_to_rbdf(data, schema: schema, schema_overrides: schema_overrides, nan_to_null: nan_to_null)
       elsif data.is_a?(Array)
-        self._df = self.class.sequence_to_rbdf(data, columns: columns, orient: orient)
+        self._df = self.class.sequence_to_rbdf(data, schema: schema, schema_overrides: schema_overrides, orient: orient, infer_schema_length: infer_schema_length)
       elsif data.is_a?(Series)
-        self._df = self.class.series_to_rbdf(data, columns: columns)
+        self._df = self.class.series_to_rbdf(data, schema: schema, schema_overrides: schema_overrides)
       else
         raise ArgumentError, "DataFrame constructor called with unsupported type; got #{data.class.name}"
       end
@@ -56,8 +59,8 @@ module Polars
     end
 
     # @private
-    def self._from_hash(data, columns: nil)
-      _from_rbdf(hash_to_rbdf(data, columns: columns))
+    def self._from_hash(data, schema: nil, schema_overrides: nil)
+      _from_rbdf(hash_to_rbdf(data, schema: schema, schema_overrides: schema_overrides))
     end
 
     # def self._from_records
@@ -97,7 +100,7 @@ module Polars
       eol_char: "\n"
     )
       if Utils.pathlike?(file)
-        path = Utils.format_path(file)
+        path = Utils.normalise_filepath(file)
       else
         path = nil
         # if defined?(StringIO) && file.is_a?(StringIO)
@@ -196,32 +199,56 @@ module Polars
 
     # @private
     def self._read_parquet(
-      file,
+      source,
       columns: nil,
       n_rows: nil,
       parallel: "auto",
       row_count_name: nil,
       row_count_offset: 0,
-      low_memory: false
+      low_memory: false,
+      use_statistics: true,
+      rechunk: true
     )
-      if Utils.pathlike?(file)
-        file = Utils.format_path(file)
+      if Utils.pathlike?(source)
+        source = Utils.normalise_filepath(source)
+      end
+      if columns.is_a?(String)
+        columns = [columns]
       end
 
-      if file.is_a?(String) && file.include?("*")
-        raise Todo
+      if source.is_a?(String) && source.include?("*") && Utils.local_file?(source)
+        scan =
+          Polars.scan_parquet(
+            source,
+            n_rows: n_rows,
+            rechunk: true,
+            parallel: parallel,
+            row_count_name: row_count_name,
+            row_count_offset: row_count_offset,
+            low_memory: low_memory
+          )
+
+        if columns.nil?
+          return self._from_rbdf(scan.collect._df)
+        elsif Utils.is_str_sequence(columns, allow_str: false)
+          return self._from_rbdf(scan.select(columns).collect._df)
+        else
+          raise ArgumentError, "cannot use glob patterns and integer based projection as `columns` argument; Use columns: Array[String]"
+        end
       end
 
       projection, columns = Utils.handle_projection_columns(columns)
       _from_rbdf(
         RbDataFrame.read_parquet(
-          file,
+          source,
           columns,
           projection,
           n_rows,
           parallel,
           Utils._prepare_row_count_args(row_count_name, row_count_offset),
-          low_memory
+          low_memory,
+          use_statistics,
+          rechunk
         )
       )
     end
@@ -229,7 +256,7 @@ module Polars
     # @private
     def self._read_avro(file, columns: nil, n_rows: nil)
       if Utils.pathlike?(file)
-        file = Utils.format_path(file)
+        file = Utils.normalise_filepath(file)
       end
       projection, columns = Utils.handle_projection_columns(columns)
       _from_rbdf(RbDataFrame.read_avro(file, columns, projection, n_rows))
@@ -246,7 +273,7 @@ module Polars
       memory_map: true
     )
       if Utils.pathlike?(file)
-        file = Utils.format_path(file)
+        file = Utils.normalise_filepath(file)
       end
       if columns.is_a?(String)
         columns = [columns]
@@ -272,7 +299,7 @@ module Polars
     # @private
     def self._read_json(file)
       if Utils.pathlike?(file)
-        file = Utils.format_path(file)
+        file = Utils.normalise_filepath(file)
       end
 
       _from_rbdf(RbDataFrame.read_json(file))
@@ -281,7 +308,7 @@ module Polars
     # @private
     def self._read_ndjson(file)
       if Utils.pathlike?(file)
-        file = Utils.format_path(file)
+        file = Utils.normalise_filepath(file)
       end
 
       _from_rbdf(RbDataFrame.read_ndjson(file))
@@ -312,6 +339,7 @@ module Polars
     end
     alias_method :count, :height
     alias_method :length, :height
+    alias_method :size, :height
 
     # Get the width of the DataFrame.
     #
@@ -522,6 +550,13 @@ module Polars
     end
     alias_method :inspect, :to_s
 
+    # Returns an array representing the DataFrame
+    #
+    # @return [Array]
+    def to_a
+      rows(named: true)
+    end
+
     # Check if DataFrame includes column.
     #
     # @return [Boolean]
@@ -631,7 +666,7 @@ module Polars
       end
 
       # Ruby-specific
-      if item.is_a?(Expr)
+      if item.is_a?(Expr) || item.is_a?(Series)
         return filter(item)
       end
 
@@ -641,15 +676,42 @@ module Polars
     # Set item.
     #
     # @return [Object]
-    #
-    # def []=(key, value)
-    #   if key.is_a?(String)
-    #     raise TypeError, "'DataFrame' object does not support 'Series' assignment by index. Use 'DataFrame.with_columns'"
-    #   end
+    def []=(*key, value)
+      if key.length == 1
+        key = key.first
+      elsif key.length != 2
+        raise ArgumentError, "wrong number of arguments (given #{key.length + 1}, expected 2..3)"
+      end
 
-    #   raise Todo
-    # end
+      if Utils.strlike?(key)
+        if value.is_a?(Array) || (defined?(Numo::NArray) && value.is_a?(Numo::NArray))
+          value = Series.new(value)
+        elsif !value.is_a?(Series)
+          value = Polars.lit(value)
+        end
+        self._df = with_column(value.alias(key.to_s))._df
+      elsif key.is_a?(Array)
+        row_selection, col_selection = key
+
+        if Utils.strlike?(col_selection)
+          s = self[col_selection]
+        elsif col_selection.is_a?(Integer)
+          raise Todo
+        else
+          raise ArgumentError, "column selection not understood: #{col_selection}"
+        end
+
+        s[row_selection] = value
 
+        if col_selection.is_a?(Integer)
+          replace_at_idx(col_selection, s)
+        elsif Utils.strlike?(col_selection)
+          replace(col_selection, s)
+        end
+      else
+        raise Todo
+      end
+    end
 
     # Return the dataframe as a scalar.
     #
@@ -774,7 +836,7 @@ module Polars
       row_oriented: false
     )
       if Utils.pathlike?(file)
-        file = Utils.format_path(file)
+        file = Utils.normalise_filepath(file)
       end
 
       _df.write_json(file, pretty, row_oriented)
@@ -789,7 +851,7 @@ module Polars
     # @return [nil]
     def write_ndjson(file)
       if Utils.pathlike?(file)
-        file = Utils.format_path(file)
+        file = Utils.normalise_filepath(file)
       end
 
       _df.write_ndjson(file)
@@ -879,7 +941,7 @@ module Polars
       end
 
       if Utils.pathlike?(file)
-        file = Utils.format_path(file)
+        file = Utils.normalise_filepath(file)
       end
 
       _df.write_csv(
@@ -917,7 +979,7 @@ module Polars
         compression = "uncompressed"
       end
       if Utils.pathlike?(file)
-        file = Utils.format_path(file)
+        file = Utils.normalise_filepath(file)
       end
 
       _df.write_avro(file, compression)
@@ -936,7 +998,7 @@ module Polars
         compression = "uncompressed"
       end
       if Utils.pathlike?(file)
-        file = Utils.format_path(file)
+        file = Utils.normalise_filepath(file)
       end
 
       _df.write_ipc(file, compression)
@@ -978,7 +1040,7 @@ module Polars
         compression = "uncompressed"
       end
       if Utils.pathlike?(file)
-        file = Utils.format_path(file)
+        file = Utils.normalise_filepath(file)
       end
 
       _df.write_parquet(
@@ -1438,6 +1500,20 @@ module Polars
       end
     end
 
+    # Sort the DataFrame by column in-place.
+    #
+    # @param by [String]
+    #   By which column to sort.
+    # @param reverse [Boolean]
+    #   Reverse/descending sort.
+    # @param nulls_last [Boolean]
+    #   Place null values last. Can only be used if sorted by a single column.
+    #
+    # @return [DataFrame]
+    def sort!(by, reverse: false, nulls_last: false)
+      self._df = sort(by, reverse: reverse, nulls_last: nulls_last)._df
+    end
+
     # Check if DataFrame is equal to other.
     #
     # @param other [DataFrame]
@@ -1495,7 +1571,7 @@ module Polars
     #   # │ 30  ┆ 6   │
     #   # └─────┴─────┘
     def replace(column, new_col)
-      _df.replace(column, new_col._s)
+      _df.replace(column.to_s, new_col._s)
       self
     end
 
@@ -1836,7 +1912,7 @@ module Polars
     #     "2020-01-08 23:16:43"
     #   ]
     #   df = Polars::DataFrame.new({"dt" => dates, "a" => [3, 7, 5, 9, 2, 1]}).with_column(
-    #     Polars.col("dt").str.strptime(:datetime)
+    #     Polars.col("dt").str.strptime(Polars::Datetime)
     #   )
     #   df.groupby_rolling(index_column: "dt", period: "2d").agg(
     #     [
@@ -2767,6 +2843,16 @@ module Polars
       Utils.wrap_s(_df.drop_in_place(name))
     end
 
+    # Drop in place if exists.
+    #
+    # @param name [Object]
+    #   Column to drop.
+    #
+    # @return [Series]
+    def delete(name)
+      drop_in_place(name) if include?(name)
+    end
+
     # Create an empty copy of the current DataFrame.
     #
     # Returns a DataFrame with identical schema but no data.
@@ -3042,24 +3128,28 @@ module Polars
       if aggregate_fn.is_a?(String)
         case aggregate_fn
         when "first"
-          aggregate_fn = Polars.element.first
+          aggregate_expr = Polars.element.first._rbexpr
         when "sum"
-          aggregate_fn = Polars.element.sum
+          aggregate_expr = Polars.element.sum._rbexpr
         when "max"
-          aggregate_fn = Polars.element.max
+          aggregate_expr = Polars.element.max._rbexpr
         when "min"
-          aggregate_fn = Polars.element.min
+          aggregate_expr = Polars.element.min._rbexpr
         when "mean"
-          aggregate_fn = Polars.element.mean
+          aggregate_expr = Polars.element.mean._rbexpr
         when "median"
-          aggregate_fn = Polars.element.median
+          aggregate_expr = Polars.element.median._rbexpr
         when "last"
-          aggregate_fn = Polars.element.last
+          aggregate_expr = Polars.element.last._rbexpr
         when "count"
-          aggregate_fn = Polars.count
+          aggregate_expr = Polars.count._rbexpr
         else
           raise ArgumentError, "Argument aggregate fn: '#{aggregate_fn}' was not expected."
         end
+      elsif aggregate_fn.nil?
+        aggregate_expr = nil
+      else
+        aggregate_expr = aggregate_function._rbexpr
       end
 
       _from_rbdf(
@@ -3067,9 +3157,9 @@ module Polars
           values,
           index,
           columns,
-          aggregate_fn._rbexpr,
           maintain_order,
           sort_columns,
+          aggregate_expr,
           separator
         )
       )
@@ -3174,7 +3264,7 @@ module Polars
     #   # │ B    ┆ 1    │
     #   # │ C    ┆ 2    │
     #   # │ D    ┆ 3    │
-    #   # │ ...  ┆ ...  │
+    #   # │ E    ┆ 4    │
     #   # │ F    ┆ 5    │
     #   # │ G    ┆ 6    │
     #   # │ H    ┆ 7    │
@@ -4053,15 +4143,12 @@ module Polars
     #   # │ 5   ┆ 3.0 ┆ true  │
     #   # └─────┴─────┴───────┘
     def unique(maintain_order: true, subset: nil, keep: "first")
-      if !subset.nil?
-        if subset.is_a?(String)
-          subset = [subset]
-        elsif !subset.is_a?(Array)
-          subset = subset.to_a
-        end
-      end
-
-      _from_rbdf(_df.unique(maintain_order, subset, keep))
+      self._from_rbdf(
+        lazy
+          .unique(maintain_order: maintain_order, subset: subset, keep: keep)
+          .collect(no_optimization: true)
+          ._df
+      )
     end
 
     # Return the number of unique rows, or the number of unique row-subsets.
@@ -4403,7 +4490,7 @@ module Polars
       end
     end
 
-    # Returns an iterator over the DataFrame of rows of python-native values.
+    # Returns an iterator over the DataFrame of rows of Ruby-native values.
     #
     # @param named [Boolean]
     #   Return hashes instead of arrays. The hashes are a mapping of
@@ -4464,6 +4551,24 @@ module Polars
       end
     end
 
+    # Returns an iterator over the DataFrame of rows of Ruby-native values.
+    #
+    # @param named [Boolean]
+    #   Return hashes instead of arrays. The hashes are a mapping of
+    #   column name to row value. This is more expensive than returning an
+    #   array, but allows for accessing values by column name.
+    # @param buffer_size [Integer]
+    #   Determines the number of rows that are buffered internally while iterating
+    #   over the data; you should only modify this in very specific cases where the
+    #   default value is determined not to be a good fit to your access pattern, as
+    #   the speedup from using the buffer is significant (~2-4x). Setting this
+    #   value to zero disables row buffering.
+    #
+    # @return [Object]
+    def each_row(named: true, buffer_size: 500, &block)
+      iter_rows(named: named, buffer_size: buffer_size, &block)
+    end
+
     # Shrink DataFrame memory usage.
     #
     # Shrinks to fit the exact capacity needed to hold the data.
@@ -4717,20 +4822,63 @@ module Polars
     end
 
     # @private
-    def self.hash_to_rbdf(data, columns: nil)
-      if !columns.nil?
-        columns, dtypes = _unpack_columns(columns, lookup_names: data.keys)
+    def self.expand_hash_scalars(data, schema_overrides: nil, order: nil, nan_to_null: false)
+      updated_data = {}
+      unless data.empty?
+        dtypes = schema_overrides || {}
+        array_len = data.values.map { |val| Utils.arrlen(val) || 0 }.max
+        if array_len > 0
+          data.each do |name, val|
+            dtype = dtypes[name]
+            if val.is_a?(Hash) && dtype != Struct
+              updated_data[name] = DataFrame.new(val).to_struct(name)
+            elsif !Utils.arrlen(val).nil?
+              updated_data[name] = Series.new(String.new(name), val, dtype: dtype)
+            elsif val.nil? || [Integer, Float, TrueClass, FalseClass, String, ::Date, ::DateTime, ::Time].any? { |cls| val.is_a?(cls) }
+              dtype = Polars::Float64 if val.nil? && dtype.nil?
+              updated_data[name] = Series.new(String.new(name), [val], dtype: dtype).extend_constant(val, array_len - 1)
+            else
+              raise Todo
+            end
+          end
+        elsif data.values.all? { |val| Utils.arrlen(val) == 0 }
+          data.each do |name, val|
+            updated_data[name] = Series.new(name, val, dtype: dtypes[name])
+          end
+        elsif data.values.all? { |val| Utils.arrlen(val).nil? }
+          data.each do |name, val|
+            updated_data[name] = Series.new(name, [val], dtype: dtypes[name])
+          end
+        end
+      end
+      updated_data
+    end
 
-        if data.empty? && dtypes
-          data_series = columns.map { |name| Series.new(name, [], dtype: dtypes[name])._s }
-        else
-          data_series = data.map { |name, values| Series.new(name, values, dtype: dtypes[name])._s }
+    # @private
+    def self.hash_to_rbdf(data, schema: nil, schema_overrides: nil, nan_to_null: nil)
+      if schema.is_a?(Hash) && !data.empty?
+        if !data.all? { |col, _| schema[col] }
+          raise ArgumentError, "The given column-schema names do not match the data dictionary"
         end
-        data_series = _handle_columns_arg(data_series, columns: columns)
-        return RbDataFrame.new(data_series)
+
+        data = schema.to_h { |col| [col, data[col]] }
+      end
+
+      column_names, schema_overrides = _unpack_schema(
+        schema, lookup_names: data.keys, schema_overrides: schema_overrides
+      )
+      if column_names.empty?
+        column_names = data.keys
+      end
+
+      if data.empty? && !schema_overrides.empty?
+        data_series = column_names.map { |name| Series.new(name, [], dtype: schema_overrides[name], nan_to_null: nan_to_null)._s }
+      else
+        data_series = expand_hash_scalars(data, schema_overrides: schema_overrides, nan_to_null: nan_to_null).values.map(&:_s)
       end
 
-      RbDataFrame.read_hash(data)
+      data_series = _handle_columns_arg(data_series, columns: column_names, from_hash: true)
+      RbDataFrame.new(data_series)
     end
 
     # @private
@@ -4739,14 +4887,12 @@ module Polars
     end
 
     # @private
-    def self._unpack_columns(columns, schema_overrides: nil, lookup_names: nil, n_expected: nil)
-      raise Todo if schema_overrides
-
-      if columns.is_a?(Hash)
-        columns = columns.to_a
+    def self._unpack_schema(schema, schema_overrides: nil, n_expected: nil, lookup_names: nil, include_overrides_in_columns: false)
+      if schema.is_a?(Hash)
+        schema = schema.to_a
       end
       column_names =
-        (columns || []).map.with_index do |col, i|
+        (schema || []).map.with_index do |col, i|
           if col.is_a?(String)
             col || "column_#{i}"
           else
@@ -4759,21 +4905,38 @@ module Polars
       # TODO zip_longest
       lookup = column_names.zip(lookup_names || []).to_h
 
-      [
-        column_names,
-        (columns || []).select { |col| !col.is_a?(String) && col[1] }.to_h do |col|
+      column_dtypes =
+        (schema || []).select { |col| !col.is_a?(String) && col[1] }.to_h do |col|
           [lookup[col[0]] || col[0], col[1]]
         end
-      ]
+
+      if schema_overrides
+        raise Todo
+      end
+
+      column_dtypes.each do |col, dtype|
+        if !Utils.is_polars_dtype(dtype, include_unknown: true) && !dtype.nil?
+          column_dtypes[col] = Utils.rb_type_to_dtype(dtype)
+        end
+      end
+
+      [column_names, column_dtypes]
     end
 
-    def self._handle_columns_arg(data, columns: nil)
-      if columns.nil?
+    def self._handle_columns_arg(data, columns: nil, from_hash: false)
+      if columns.nil? || columns.empty?
         data
       else
         if data.empty?
           columns.map { |c| Series.new(c, nil)._s }
         elsif data.length == columns.length
+          if from_hash
+            series_map = data.to_h { |s| [s.name, s] }
+            if columns.all? { |col| series_map.key?(col) }
+              return columns.map { |col| series_map[col] }
+            end
+          end
+
           columns.each_with_index do |c, i|
             # not in-place?
             data[i].rename(c)
@@ -4788,7 +4951,7 @@ module Polars
     def self._post_apply_columns(rbdf, columns, structs: nil, schema_overrides: nil)
       rbdf_columns = rbdf.columns
       rbdf_dtypes = rbdf.dtypes
-      columns, dtypes = _unpack_columns(
+      columns, dtypes = _unpack_schema(
         (columns || rbdf_columns), schema_overrides: schema_overrides
       )
       column_subset = []
@@ -4826,20 +4989,23 @@ module Polars
     end
 
     # @private
-    def self.sequence_to_rbdf(data, columns: nil, orient: nil, infer_schema_length: 50)
+    def self.sequence_to_rbdf(data, schema: nil, schema_overrides: nil, orient: nil, infer_schema_length: 50)
+      raise Todo if schema_overrides
+      columns = schema
+
       if data.length == 0
-        return hash_to_rbdf({}, columns: columns)
+        return hash_to_rbdf({}, schema: schema, schema_overrides: schema_overrides)
       end
 
       if data[0].is_a?(Series)
         # series_names = data.map(&:name)
-        # columns, dtypes = _unpack_columns(columns || series_names, n_expected: data.length)
+        # columns, dtypes = _unpack_schema(columns || series_names, n_expected: data.length)
         data_series = []
         data.each do |s|
           data_series << s._s
         end
       elsif data[0].is_a?(Hash)
-        column_names, dtypes = _unpack_columns(columns)
+        column_names, dtypes = _unpack_schema(columns)
         schema_overrides = dtypes ? include_unknowns(dtypes, column_names) : nil
         rbdf = RbDataFrame.read_hashes(data, infer_schema_length, schema_overrides)
         if column_names
@@ -4865,11 +5031,21 @@ module Polars
     end
 
     # @private
-    def self.series_to_rbdf(data, columns: nil)
-      if columns
-        raise Todo
+    def self.series_to_rbdf(data, schema: nil, schema_overrides: nil)
+      data_series = [data._s]
+      series_name = data_series.map(&:name)
+      column_names, schema_overrides = _unpack_schema(
+        schema || series_name, schema_overrides: schema_overrides, n_expected: 1
+      )
+      if schema_overrides.any?
+        new_dtype = schema_overrides.values[0]
+        if new_dtype != data.dtype
+          data_series[0] = data_series[0].cast(new_dtype, true)
+        end
       end
-      RbDataFrame.new([data._s])
+
+      data_series = _handle_columns_arg(data_series, columns: column_names)
+      RbDataFrame.new(data_series)
     end
 
     def wrap_ldf(ldf)

data/lib/polars/data_types.rb

@@ -84,20 +84,22 @@ module Polars
 
   # Calendar date and time type.
   class Datetime < TemporalType
-    attr_reader :tu
+    attr_reader :time_unit, :time_zone
+    alias_method :tu, :time_unit
 
     def initialize(time_unit = "us", time_zone = nil)
-      @tu = time_unit || "us"
+      @time_unit = time_unit || "us"
       @time_zone = time_zone
     end
   end
 
   # Time duration/delta type.
   class Duration < TemporalType
-    attr_reader :tu
+    attr_reader :time_unit
+    alias_method :tu, :time_unit
 
     def initialize(time_unit = "us")
-      @tu = time_unit
+      @time_unit = time_unit
     end
   end