polars-df 0.8.0 → 0.9.0

data/lib/polars/testing.rb

@@ -0,0 +1,507 @@
+module Polars
+  module Testing
+    # Assert that the left and right frame are equal.
+    #
+    # Raises a detailed `AssertionError` if the frames differ.
+    # This function is intended for use in unit tests.
+    #
+    # @param left [Object]
+    #   The first DataFrame or LazyFrame to compare.
+    # @param right [Object]
+    #   The second DataFrame or LazyFrame to compare.
+    # @param check_row_order [Boolean]
+    #   Require row order to match.
+    # @param check_column_order [Boolean]
+    #   Require column order to match.
+    # @param check_dtype [Boolean]
+    #   Require data types to match.
+    # @param check_exact [Boolean]
+    #   Require float values to match exactly. If set to `false`, values are considered
+    #   equal when within tolerance of each other (see `rtol` and `atol`).
+    #   Only affects columns with a Float data type.
+    # @param rtol [Float]
+    #   Relative tolerance for inexact checking. Fraction of values in `right`.
+    # @param atol [Float]
+    #   Absolute tolerance for inexact checking.
+    # @param categorical_as_str [Boolean]
+    #   Cast categorical columns to string before comparing. Enabling this helps
+    #   compare columns that do not share the same string cache.
+    #
+    # @return [nil]
+    def assert_frame_equal(
+      left,
+      right,
+      check_row_order: true,
+      check_column_order: true,
+      check_dtype: true,
+      check_exact: false,
+      rtol: 1e-5,
+      atol: 1e-8,
+      categorical_as_str: false
+    )
+      lazy = _assert_correct_input_type(left, right)
+      objects = lazy ? "LazyFrames" : "DataFrames"
+
+      _assert_frame_schema_equal(
+        left,
+        right,
+        check_column_order: check_column_order,
+        check_dtype: check_dtype,
+        objects: objects,
+      )
+
+      if lazy
+        left, right = left.collect, right.collect
+      end
+
+      if left.height != right.height
+        raise_assertion_error(
+          objects, "number of rows does not match", left.height, right.height
+        )
+      end
+
+      if !check_row_order
+        left, right = _sort_dataframes(left, right)
+      end
+
+      left.columns.each do |c|
+        s_left, s_right = left.get_column(c), right.get_column(c)
+        begin
+          _assert_series_values_equal(
+            s_left,
+            s_right,
+            check_exact: check_exact,
+            rtol: rtol,
+            atol: atol,
+            categorical_as_str: categorical_as_str
+          )
+        rescue AssertionError
+          raise_assertion_error(
+            objects,
+            "value mismatch for column #{c.inspect}",
+            s_left.to_a,
+            s_right.to_a
+          )
+        end
+      end
+    end
+
+    # Assert that the left and right frame are **not** equal.
+    #
+    # This function is intended for use in unit tests.
+    #
+    # @param left [Object]
+    #   The first DataFrame or LazyFrame to compare.
+    # @param right [Object]
+    #   The second DataFrame or LazyFrame to compare.
+    # @param check_row_order [Boolean]
+    #   Require row order to match.
+    # @param check_column_order [Boolean]
+    #   Require column order to match.
+    # @param check_dtype [Boolean]
+    #   Require data types to match.
+    # @param check_exact [Boolean]
+    #   Require float values to match exactly. If set to `false`, values are considered
+    #   equal when within tolerance of each other (see `rtol` and `atol`).
+    #   Only affects columns with a Float data type.
+    # @param rtol [Float]
+    #   Relative tolerance for inexact checking. Fraction of values in `right`.
+    # @param atol [Float]
+    #   Absolute tolerance for inexact checking.
+    # @param categorical_as_str [Boolean]
+    #   Cast categorical columns to string before comparing. Enabling this helps
+    #   compare columns that do not share the same string cache.
+    #
+    # @return [nil]
+    def assert_frame_not_equal(
+      left,
+      right,
+      check_row_order: true,
+      check_column_order: true,
+      check_dtype: true,
+      check_exact: false,
+      rtol: 1e-5,
+      atol: 1e-8,
+      categorical_as_str: false
+    )
+      begin
+        assert_frame_equal(
+          left,
+          right,
+          check_column_order: check_column_order,
+          check_row_order: check_row_order,
+          check_dtype: check_dtype,
+          check_exact: check_exact,
+          rtol: rtol,
+          atol: atol,
+          categorical_as_str: categorical_as_str
+        )
+      rescue AssertionError
+        return
+      end
+
+      msg = "frames are equal"
+      raise AssertionError, msg
+    end
+
+    # Assert that the left and right Series are equal.
+    #
+    # Raises a detailed `AssertionError` if the Series differ.
+    # This function is intended for use in unit tests.
+    #
+    # @param left [Object]
+    #   The first Series to compare.
+    # @param right [Object]
+    #   The second Series to compare.
+    # @param check_dtype [Boolean]
+    #   Require data types to match.
+    # @param check_names [Boolean]
+    #   Require names to match.
+    # @param check_exact [Boolean]
+    #   Require float values to match exactly. If set to `false`, values are considered
+    #   equal when within tolerance of each other (see `rtol` and `atol`).
+    #   Only affects columns with a Float data type.
+    # @param rtol [Float]
+    #   Relative tolerance for inexact checking, given as a fraction of the values in
+    #   `right`.
+    # @param atol [Float]
+    #   Absolute tolerance for inexact checking.
+    # @param categorical_as_str [Boolean]
+    #   Cast categorical columns to string before comparing. Enabling this helps
+    #   compare columns that do not share the same string cache.
+    #
+    # @return [nil]
+    def assert_series_equal(
+      left,
+      right,
+      check_dtype: true,
+      check_names: true,
+      check_exact: false,
+      rtol: 1e-5,
+      atol: 1e-8,
+      categorical_as_str: false
+    )
+      if !(left.is_a?(Series) && right.is_a?(Series))
+        raise_assertion_error(
+          "inputs",
+          "unexpected input types",
+          left.class.name,
+          right.class.name
+        )
+      end
+
+      if left.len != right.len
+        raise_assertion_error("Series", "length mismatch", left.len, right.len)
+      end
+
+      if check_names && left.name != right.name
+        raise_assertion_error("Series", "name mismatch", left.name, right.name)
+      end
+
+      if check_dtype && left.dtype != right.dtype
+        raise_assertion_error("Series", "dtype mismatch", left.dtype, right.dtype)
+      end
+
+      _assert_series_values_equal(
+        left,
+        right,
+        check_exact: check_exact,
+        rtol: rtol,
+        atol: atol,
+        categorical_as_str: categorical_as_str
+      )
+    end
+
+    # Assert that the left and right Series are **not** equal.
+    #
+    # This function is intended for use in unit tests.
+    #
+    # @param left [Object]
+    #   The first Series to compare.
+    # @param right [Object]
+    #   The second Series to compare.
+    # @param check_dtype [Boolean]
+    #   Require data types to match.
+    # @param check_names [Boolean]
+    #   Require names to match.
+    # @param check_exact [Boolean]
+    #   Require float values to match exactly. If set to `false`, values are considered
+    #   equal when within tolerance of each other (see `rtol` and `atol`).
+    #   Only affects columns with a Float data type.
+    # @param rtol [Float]
+    #   Relative tolerance for inexact checking, given as a fraction of the values in
+    #   `right`.
+    # @param atol [Float]
+    #   Absolute tolerance for inexact checking.
+    # @param categorical_as_str [Boolean]
+    #   Cast categorical columns to string before comparing. Enabling this helps
+    #   compare columns that do not share the same string cache.
+    #
+    # @return [nil]
+    def assert_series_not_equal(
+      left,
+      right,
+      check_dtype: true,
+      check_names: true,
+      check_exact: false,
+      rtol: 1e-5,
+      atol: 1e-8,
+      categorical_as_str: false
+    )
+      begin
+        assert_series_equal(
+          left,
+          right,
+          check_dtype: check_dtype,
+          check_names: check_names,
+          check_exact: check_exact,
+          rtol: rtol,
+          atol: atol,
+          categorical_as_str: categorical_as_str
+        )
+      rescue AssertionError
+        return
+      end
+
+      msg = "Series are equal"
+      raise AssertionError, msg
+    end
+
+    private
+
+    def _assert_correct_input_type(left, right)
+      if left.is_a?(DataFrame) && right.is_a?(DataFrame)
+        return false
+      elsif left.is_a?(LazyFrame) && right.is_a?(DataFrame)
+        return true
+      else
+        raise_assertion_error(
+          "inputs",
+          "unexpected input types",
+          left.class.name,
+          right.class.name
+        )
+      end
+    end
+
+    def _assert_frame_schema_equal(
+      left,
+      right,
+      check_dtype:,
+      check_column_order:,
+      objects:
+    )
+      left_schema, right_schema = left.schema, right.schema
+
+      # Fast path for equal frames
+      if left_schema == right_schema
+        return
+      end
+
+      # Special error message for when column names do not match
+      if left_schema.keys != right_schema.keys
+        if (left_not_right = right_schema.keys - left_schema.keys).any?
+          msg = "columns #{left_not_right.inspect} in left #{objects[..-1]}, but not in right"
+          raise AssertionError, msg
+        else
+          right_not_left = right_schema.keys - left_schema.keys
+          msg = "columns #{right_not_left.inspect} in right #{objects[..-1]}, but not in left"
+          raise AssertionError, msg
+        end
+      end
+
+      if check_column_order
+        left_columns, right_columns = left_schema.keys, right_schema.keys
+        if left_columns != right_columns
+          detail = "columns are not in the same order"
+          raise_assertion_error(objects, detail, left_columns, right_columns)
+        end
+      end
+
+      if check_dtype
+        left_schema_dict, right_schema_dict = left_schema.to_h, right_schema.to_h
+        if check_column_order || left_schema_dict != right_schema_dict
+          detail = "dtypes do not match"
+          raise_assertion_error(objects, detail, left_schema_dict, right_schema_dict)
+        end
+      end
+    end
+
+    def _sort_dataframes(left, right)
+      by = left.columns
+      begin
+        left = left.sort(by)
+        right = right.sort(by)
+      rescue
+        msg = "cannot set `check_row_order: false` on frame with unsortable columns"
+        raise InvalidAssert, msg
+      end
+      [left, right]
+    end
+
+    def _assert_series_values_equal(
+      left,
+      right,
+      check_exact:,
+      rtol:,
+      atol:,
+      categorical_as_str:
+    )
+      if categorical_as_str
+        if left.dtype == Categorical
+          left = left.cast(String)
+        end
+        if right.dtype == Categorical
+          right = right.cast(String)
+        end
+      end
+
+      # Determine unequal elements
+      begin
+        unequal = left.ne_missing(right)
+      rescue
+        raise_assertion_error(
+          "Series",
+          "incompatible data types",
+          left.dtype,
+          right.dtype
+        )
+      end
+
+      # Check nested dtypes in separate function
+      if _comparing_nested_floats(left.dtype, right.dtype)
+        begin
+          _assert_series_nested_values_equal(
+            left: left.filter(unequal),
+            right: right.filter(unequal),
+            check_exact: check_exact,
+            rtol: rtol,
+            atol: atol,
+            categorical_as_str: categorical_as_str
+          )
+        rescue AssertionError
+          raise_assertion_error(
+            "Series",
+            "nested value mismatch",
+            left.to_a,
+            right.to_a
+          )
+        else
+          return
+        end
+      end
+
+      # If no differences found during exact checking, we're done
+      if !unequal.any
+        return
+      end
+
+      # Only do inexact checking for float types
+      if check_exact || !left.dtype.float? || !right.dtype.float?
+        raise_assertion_error(
+          "Series", "exact value mismatch", left.to_a, right.to_a
+        )
+      end
+
+      _assert_series_null_values_match(left, right)
+      _assert_series_nan_values_match(left, right)
+      _assert_series_values_within_tolerance(
+        left,
+        right,
+        unequal,
+        rtol: rtol,
+        atol: atol
+      )
+    end
+
+    def _assert_series_nested_values_equal(
+      left,
+      right,
+      check_exact:,
+      rtol:,
+      atol:,
+      categorical_as_str:
+    )
+      # compare nested lists element-wise
+      if _comparing_lists(left.dtype, right.dtype)
+        left.zip(right) do |s1, s2|
+          if s1.nil? || s2.nil?
+            raise_assertion_error("Series", "nested value mismatch", s1, s2)
+          end
+
+          _assert_series_values_equal(
+            s1,
+            s2,
+            check_exact: check_exact,
+            rtol: rtol,
+            atol: atol,
+            categorical_as_str: categorical_as_str
+          )
+        end
+
+      # unnest structs as series and compare
+      else
+        ls, rs = left.struct.unnest, right.struct.unnest
+        ls.zip(rs) do |s1, s2|
+          _assert_series_values_equal(
+            s1,
+            s2,
+            check_exact: check_exact,
+            rtol: rtol,
+            atol: atol,
+            categorical_as_str: categorical_as_str
+          )
+        end
+      end
+    end
+
+    def _assert_series_null_values_match(left, right)
+      null_value_mismatch = left.is_null != right.is_null
+      if null_value_mismatch.any
+        raise_assertion_error(
+          "Series", "null value mismatch", left.to_a, right.to_a
+        )
+      end
+    end
+
+    def _assert_series_nan_values_match(left, right)
+      if !_comparing_floats(left.dtype, right.dtype)
+        return
+      end
+      nan_value_mismatch = left.is_nan != right.is_nan
+      if nan_value_mismatch.any
+        raise_assertion_error(
+          "Series",
+          "nan value mismatch",
+          left.to_a,
+          right.to_a
+        )
+      end
+    end
+
+    def _comparing_floats(left, right)
+      left.is_float && right.is_float
+    end
+
+    def _comparing_lists(left, right)
+      [List, Array].include?(left) && [List, Array].include?(right)
+    end
+
+    def _comparing_structs(left, right)
+      left == Struct && right == Struct
+    end
+
+    def _comparing_nested_floats(left, right)
+      if !(_comparing_lists(left, right) || _comparing_structs(left, right))
+        return false
+      end
+
+      left.float? && right.float?
+    end
+
+    def raise_assertion_error(objects, detail, left, right)
+      msg = "#{objects} are different (#{detail})\n[left]:  #{left}\n[right]: #{right}"
+      raise AssertionError, msg
+    end
+  end
+end

data/lib/polars/utils.rb

@@ -187,7 +187,7 @@ module Polars
       begin
         map_rb_type_to_dtype(data_type)
       rescue TypeError
-        raise ArgumentError, "Conversion of Ruby data type #{data_type} to Polars data type not implemented."
+        raise ArgumentError, "Conversion of Ruby data type #{data_type.inspect} to Polars data type not implemented."
       end
     end
 
@@ -286,20 +286,51 @@ module Polars
       Dir.glob(file).any?
     end
 
-    def self.parse_as_expression(input, str_as_lit: false, structify: false)
+    def self.parse_as_list_of_expressions(*inputs, __structify: false, **named_inputs)
+      exprs = _parse_positional_inputs(inputs, structify: __structify)
+      if named_inputs.any?
+        named_exprs = _parse_named_inputs(named_inputs, structify: __structify)
+        exprs.concat(named_exprs)
+      end
+
+      exprs
+    end
+
+    def self._parse_positional_inputs(inputs, structify: false)
+      inputs_iter = _parse_inputs_as_iterable(inputs)
+      inputs_iter.map { |e| parse_as_expression(e, structify: structify) }
+    end
+
+    def self._parse_inputs_as_iterable(inputs)
+      if inputs.empty?
+        return []
+      end
+
+      if inputs.length == 1 && inputs[0].is_a?(::Array)
+        return inputs[0]
+      end
+
+      inputs
+    end
+
+    def self._parse_named_inputs(named_inputs, structify: false)
+      named_inputs.map do |name, input|
+        parse_as_expression(input, structify: structify)._alias(name.to_s)
+      end
+    end
+
+    def self.parse_as_expression(input, str_as_lit: false, list_as_lit: true, structify: false, dtype: nil)
       if input.is_a?(Expr)
         expr = input
       elsif input.is_a?(::String) && !str_as_lit
         expr = Polars.col(input)
         structify = false
-      elsif [Integer, Float, ::String, Series, ::Date, ::Time, ::DateTime].any? { |cls| input.is_a?(cls) } || input.nil?
-        expr = Polars.lit(input)
-        structify = false
-      elsif input.is_a?(Array)
-        expr = Polars.lit(Polars::Series.new("", [input]))
+      elsif input.is_a?(::Array) && !list_as_lit
+        expr = Polars.lit(Series.new(input), dtype: dtype)
         structify = false
       else
-        raise TypeError, "did not expect value #{input} of type #{input.class.name}, maybe disambiguate with pl.lit or pl.col"
+        expr = Polars.lit(input, dtype: dtype)
+        structify = false
       end
 
       if structify
@@ -320,5 +351,18 @@ module Polars
       end
       ambiguous
     end
+
+    def self._check_arg_is_1byte(arg_name, arg, can_be_empty = false)
+      if arg.is_a?(::String)
+        arg_byte_length = arg.bytesize
+        if can_be_empty
+          if arg_byte_length > 1
+            raise ArgumentError, "#{arg_name} should be a single byte character or empty, but is #{arg_byte_length} bytes long."
+          end
+        elsif arg_byte_length != 1
+          raise ArgumentError, "#{arg_name} should be a single byte character, but is #{arg_byte_length} bytes long."
+        end
+      end
+    end
   end
 end

data/lib/polars/version.rb

@@ -1,4 +1,4 @@
 module Polars
   # @private
-  VERSION = "0.8.0"
+  VERSION = "0.9.0"
 end

data/lib/polars.rb

@@ -30,10 +30,22 @@ require_relative "polars/dynamic_group_by"
 require_relative "polars/exceptions"
 require_relative "polars/expr"
 require_relative "polars/functions"
+require_relative "polars/functions/as_datatype"
+require_relative "polars/functions/col"
+require_relative "polars/functions/eager"
+require_relative "polars/functions/lazy"
+require_relative "polars/functions/len"
+require_relative "polars/functions/lit"
+require_relative "polars/functions/random"
+require_relative "polars/functions/repeat"
+require_relative "polars/functions/whenthen"
+require_relative "polars/functions/aggregation/horizontal"
+require_relative "polars/functions/aggregation/vertical"
+require_relative "polars/functions/range/date_range"
+require_relative "polars/functions/range/int_range"
 require_relative "polars/group_by"
 require_relative "polars/io"
 require_relative "polars/lazy_frame"
-require_relative "polars/lazy_functions"
 require_relative "polars/lazy_group_by"
 require_relative "polars/list_expr"
 require_relative "polars/list_name_space"
@@ -43,10 +55,12 @@ require_relative "polars/rolling_group_by"
 require_relative "polars/series"
 require_relative "polars/slice"
 require_relative "polars/sql_context"
+require_relative "polars/string_cache"
 require_relative "polars/string_expr"
 require_relative "polars/string_name_space"
 require_relative "polars/struct_expr"
 require_relative "polars/struct_name_space"
+require_relative "polars/testing"
 require_relative "polars/utils"
 require_relative "polars/version"
 require_relative "polars/when"
@@ -56,5 +70,4 @@ module Polars
   extend Convert
   extend Functions
   extend IO
-  extend LazyFunctions
 end