polars-df 0.15.0 → 0.26.0

data/lib/polars/struct_expr.rb

@@ -26,37 +26,85 @@ module Polars
     #
     # @param name [String]
     #   Name of the field
+    # @param more_names [Array]
+    #   Additional struct field names.
     #
     # @return [Expr]
     #
     # @example
-    #   df = (
-    #     Polars::DataFrame.new(
-    #       {
-    #         "int" => [1, 2],
-    #         "str" => ["a", "b"],
-    #         "bool" => [true, nil],
-    #         "list" => [[1, 2], [3]]
-    #       }
-    #     )
-    #     .to_struct("my_struct")
-    #     .to_frame
-    #   )
-    #   df.select(Polars.col("my_struct").struct.field("str"))
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "aaa" => [1, 2],
+    #       "bbb" => ["ab", "cd"],
+    #       "ccc" => [true, nil],
+    #       "ddd" => [[1, 2], [3]]
+    #     }
+    #   ).select(Polars.struct("aaa", "bbb", "ccc", "ddd").alias("struct_col"))
+    #   df.select(Polars.col("struct_col").struct.field("bbb"))
     #   # =>
     #   # shape: (2, 1)
     #   # ┌─────┐
-    #   # │ str │
+    #   # │ bbb │
     #   # │ --- │
     #   # │ str │
     #   # ╞═════╡
-    #   # │ a   │
-    #   # │ b   │
+    #   # │ ab  │
+    #   # │ cd  │
     #   # └─────┘
-    def field(name)
+    #
+    # @example
+    #   df.select(Polars.col("struct_col").struct.field("aaa", "bbb"))
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌─────┬─────┐
+    #   # │ aaa ┆ bbb │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ str │
+    #   # ╞═════╪═════╡
+    #   # │ 1   ┆ ab  │
+    #   # │ 2   ┆ cd  │
+    #   # └─────┴─────┘
+    def field(name, *more_names)
+      if more_names.any?
+        name = (name.is_a?(::String) ? [name] : name) + more_names
+      end
+      if name.is_a?(::Array)
+        return Utils.wrap_expr(_rbexpr.struct_multiple_fields(name))
+      end
+
       Utils.wrap_expr(_rbexpr.struct_field_by_name(name))
     end
 
+    # Expand the struct into its individual fields.
+    #
+    # Alias for `Expr.struct.field("*")`.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "aaa" => [1, 2],
+    #       "bbb" => ["ab", "cd"],
+    #       "ccc" => [true, nil],
+    #       "ddd" => [[1, 2], [3]]
+    #     }
+    #   ).select(Polars.struct("aaa", "bbb", "ccc", "ddd").alias("struct_col"))
+    #   df.select(Polars.col("struct_col").struct.unnest)
+    #   # =>
+    #   # shape: (2, 4)
+    #   # ┌─────┬─────┬──────┬───────────┐
+    #   # │ aaa ┆ bbb ┆ ccc  ┆ ddd       │
+    #   # │ --- ┆ --- ┆ ---  ┆ ---       │
+    #   # │ i64 ┆ str ┆ bool ┆ list[i64] │
+    #   # ╞═════╪═════╪══════╪═══════════╡
+    #   # │ 1   ┆ ab  ┆ true ┆ [1, 2]    │
+    #   # │ 2   ┆ cd  ┆ null ┆ [3]       │
+    #   # └─────┴─────┴──────┴───────────┘
+    def unnest
+      field("*")
+    end
+
     # Rename the fields of the struct.
     #
     # @param names [Array]
@@ -77,7 +125,7 @@ module Polars
     #     .to_struct("my_struct")
     #     .to_frame
     #   )
-    #   df = df.with_column(
+    #   df = df.with_columns(
     #     Polars.col("my_struct").struct.rename_fields(["INT", "STR", "BOOL", "LIST"])
     #   )
     #   df.select(Polars.col("my_struct").struct.field("INT"))
@@ -94,5 +142,78 @@ module Polars
     def rename_fields(names)
       Utils.wrap_expr(_rbexpr.struct_rename_fields(names))
     end
+
+    # Convert this struct to a string column with json values.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   Polars::DataFrame.new(
+    #     {"a" => [{"a" => [1, 2], "b" => [45]}, {"a" => [9, 1, 3], "b" => nil}]}
+    #   ).with_columns(Polars.col("a").struct.json_encode.alias("encoded"))
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌──────────────────┬────────────────────────┐
+    #   # │ a                ┆ encoded                │
+    #   # │ ---              ┆ ---                    │
+    #   # │ struct[2]        ┆ str                    │
+    #   # ╞══════════════════╪════════════════════════╡
+    #   # │ {[1, 2],[45]}    ┆ {"a":[1,2],"b":[45]}   │
+    #   # │ {[9, 1, 3],null} ┆ {"a":[9,1,3],"b":null} │
+    #   # └──────────────────┴────────────────────────┘
+    def json_encode
+      Utils.wrap_expr(_rbexpr.struct_json_encode)
+    end
+
+    # Add or overwrite fields of this struct.
+    #
+    # This is similar to `with_columns` on `DataFrame`.
+    #
+    # @param exprs [Array]
+    #   Field(s) to add, specified as positional arguments.
+    #   Accepts expression input. Strings are parsed as column names, other
+    #   non-expression inputs are parsed as literals.
+    # @param named_exprs [Hash]
+    #   Additional fields to add, specified as keyword arguments.
+    #   The columns will be renamed to the keyword used.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "coords" => [{"x" => 1, "y" => 4}, {"x" => 4, "y" => 9}, {"x" => 9, "y" => 16}],
+    #       "multiply" => [10, 2, 3]
+    #     }
+    #   )
+    #   df.with_columns(
+    #     Polars.col("coords").struct.with_fields(
+    #       Polars.field("x").sqrt,
+    #       y_mul: Polars.field("y") * Polars.col("multiply")
+    #     )
+    #   )
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────────────┬──────────┐
+    #   # │ coords      ┆ multiply │
+    #   # │ ---         ┆ ---      │
+    #   # │ struct[3]   ┆ i64      │
+    #   # ╞═════════════╪══════════╡
+    #   # │ {1.0,4,40}  ┆ 10       │
+    #   # │ {2.0,9,18}  ┆ 2        │
+    #   # │ {3.0,16,48} ┆ 3        │
+    #   # └─────────────┴──────────┘
+    def with_fields(
+      *exprs,
+      **named_exprs
+    )
+      structify = ENV.fetch("POLARS_AUTO_STRUCTIFY", 0).to_i != 0
+
+      rbexprs = Utils.parse_into_list_of_expressions(
+        *exprs, **named_exprs, __structify: structify
+      )
+
+      Utils.wrap_expr(_rbexpr.struct_with_fields(rbexprs))
+    end
   end
 end

data/lib/polars/struct_name_space.rb

@@ -80,7 +80,7 @@ module Polars
       super
     end
 
-    # Get the struct definition as a name/dtype schema dict.
+    # Get the struct definition as a name/dtype schema hash.
     #
     # @return [Object]
     #
@@ -116,5 +116,23 @@ module Polars
     def unnest
       Utils.wrap_df(_s.struct_unnest)
     end
+
+    # Convert this struct to a string column with json values.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [{"a" => [1, 2], "b" => [45]}, {"a" => [9, 1, 3], "b" => nil}])
+    #   s.struct.json_encode
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'a' [str]
+    #   # [
+    #   #         "{"a":[1,2],"b":[45]}"
+    #   #         "{"a":[9,1,3],"b":null}"
+    #   # ]
+    def json_encode
+      super
+    end
   end
 end

data/lib/polars/testing.rb

@@ -40,50 +40,22 @@ module Polars
       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
+      Plr.assert_dataframe_equal_rb(
+        left._df,
+        right._df,
+        check_row_order,
+        check_column_order,
+        check_dtype,
+        check_exact,
+        rtol,
+        atol,
+        categorical_as_str,
+      )
     end
 
     # Assert that the left and right frame are **not** equal.
@@ -157,6 +129,8 @@ module Polars
     #   Require data types to match.
     # @param check_names [Boolean]
     #   Require names to match.
+    # @param check_order [Boolean]
+    #   Requires elements to appear in the same order.
     # @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`).
@@ -176,6 +150,7 @@ module Polars
       right,
       check_dtype: true,
       check_names: true,
+      check_order: true,
       check_exact: false,
       rtol: 1e-5,
       atol: 1e-8,
@@ -190,25 +165,16 @@ module Polars
         )
       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
+      Plr.assert_series_equal_rb(
+        left._s,
+        right._s,
+        check_dtype,
+        check_names,
+        check_order,
+        check_exact,
+        rtol,
+        atol,
+        categorical_as_str
       )
     end
 
@@ -284,221 +250,6 @@ module Polars
       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

data/lib/polars/utils/constants.rb

@@ -1,5 +1,7 @@
 module Polars
   module Utils
+    U32_MAX = 2**32 - 1
+
     SECONDS_PER_DAY = 86_400
     SECONDS_PER_HOUR = 3_600
     NS_PER_SECOND = 1_000_000_000