polars-df 0.23.0 → 0.24.0

data/ext/polars/src/lib.rs

@@ -1,5 +1,7 @@
-mod allocator;
+#![allow(clippy::too_many_arguments)]
+
 mod batched_csv;
+mod c_api;
 mod catalog;
 mod conversion;
 mod dataframe;
@@ -19,6 +21,8 @@ mod prelude;
 pub(crate) mod rb_modules;
 mod series;
 mod sql;
+mod testing;
+mod timeout;
 mod utils;
 
 use batched_csv::RbBatchedCsv;
@@ -29,12 +33,11 @@ use error::RbPolarsErr;
 use exceptions::{RbTypeError, RbValueError};
 use expr::RbExpr;
 use expr::datatype::RbDataTypeExpr;
-use expr::rb_exprs_to_exprs;
 use expr::selector::RbSelector;
 use functions::string_cache::RbStringCacheHolder;
 use functions::whenthen::{RbChainedThen, RbChainedWhen, RbThen, RbWhen};
-use interop::arrow::to_ruby::RbArrowArrayStream;
-use lazyframe::{RbLazyFrame, RbOptFlags};
+use interop::arrow::to_rb::{RbArrowArrayStream, RbArrowSchema};
+use lazyframe::{RbInProcessQuery, RbLazyFrame, RbOptFlags};
 use lazygroupby::RbLazyGroupBy;
 use magnus::{Ruby, function, method, prelude::*};
 use series::RbSeries;
@@ -48,9 +51,9 @@ fn re_escape(pattern: String) -> String {
     regex::escape(&pattern)
 }
 
-#[magnus::init]
+#[magnus::init(name = "polars")]
 fn init(ruby: &Ruby) -> RbResult<()> {
-    crate::on_startup::register_startup_deps();
+    crate::on_startup::register_startup_deps(true);
 
     let module = ruby.define_module("Polars")?;
 
@@ -70,14 +73,10 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_singleton_method("from_rows", function!(RbDataFrame::from_rows, 3))?;
     class.define_singleton_method("from_hashes", function!(RbDataFrame::from_hashes, 5))?;
     class.define_singleton_method("read_json", function!(RbDataFrame::read_json, 4))?;
-    class.define_singleton_method("read_ndjson", function!(RbDataFrame::read_ndjson, 4))?;
     class.define_method("estimated_size", method!(RbDataFrame::estimated_size, 0))?;
     class.define_method("dtype_strings", method!(RbDataFrame::dtype_strings, 0))?;
     class.define_method("write_avro", method!(RbDataFrame::write_avro, 3))?;
     class.define_method("write_json", method!(RbDataFrame::write_json, 1))?;
-    class.define_method("write_ndjson", method!(RbDataFrame::write_ndjson, 1))?;
-    class.define_method("write_csv", method!(RbDataFrame::write_csv, 10))?;
-    class.define_method("write_ipc", method!(RbDataFrame::write_ipc, 5))?;
     class.define_method(
         "write_ipc_stream",
         method!(RbDataFrame::write_ipc_stream, 3),
@@ -89,7 +88,6 @@ fn init(ruby: &Ruby) -> RbResult<()> {
         method!(RbDataFrame::__arrow_c_stream__, 0),
     )?;
     class.define_method("to_numo", method!(RbDataFrame::to_numo, 0))?;
-    class.define_method("write_parquet", method!(RbDataFrame::write_parquet, 6))?;
     class.define_method("add", method!(RbDataFrame::add, 1))?;
     class.define_method("sub", method!(RbDataFrame::sub, 1))?;
     class.define_method("div", method!(RbDataFrame::div, 1))?;
@@ -121,17 +119,17 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("vstack_mut", method!(RbDataFrame::vstack_mut, 1))?;
     class.define_method("vstack", method!(RbDataFrame::vstack, 1))?;
     class.define_method("drop_in_place", method!(RbDataFrame::drop_in_place, 1))?;
-    class.define_method("select_at_idx", method!(RbDataFrame::select_at_idx, 1))?;
+    class.define_method("to_series", method!(RbDataFrame::to_series, 1))?;
     class.define_method(
         "get_column_index",
         method!(RbDataFrame::get_column_index, 1),
     )?;
     class.define_method("get_column", method!(RbDataFrame::get_column, 1))?;
     class.define_method("select", method!(RbDataFrame::select, 1))?;
-    class.define_method("take", method!(RbDataFrame::gather, 1))?;
+    class.define_method("gather", method!(RbDataFrame::gather, 1))?;
     class.define_method(
-        "take_with_series",
-        method!(RbDataFrame::take_with_series, 1),
+        "gather_with_series",
+        method!(RbDataFrame::gather_with_series, 1),
     )?;
     class.define_method("replace", method!(RbDataFrame::replace, 2))?;
     class.define_method("replace_column", method!(RbDataFrame::replace_column, 2))?;
@@ -204,6 +202,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("unique_stable", method!(RbExpr::unique_stable, 0))?;
     class.define_method("first", method!(RbExpr::first, 0))?;
     class.define_method("last", method!(RbExpr::last, 0))?;
+    class.define_method("item", method!(RbExpr::item, 1))?;
     class.define_method("implode", method!(RbExpr::implode, 0))?;
     class.define_method("quantile", method!(RbExpr::quantile, 2))?;
     class.define_method("cut", method!(RbExpr::cut, 4))?;
@@ -217,7 +216,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("value_counts", method!(RbExpr::value_counts, 4))?;
     class.define_method("unique_counts", method!(RbExpr::unique_counts, 0))?;
     class.define_method("null_count", method!(RbExpr::null_count, 0))?;
-    class.define_method("cast", method!(RbExpr::cast, 2))?;
+    class.define_method("cast", method!(RbExpr::cast, 3))?;
     class.define_method("sort_with", method!(RbExpr::sort_with, 2))?;
     class.define_method("arg_sort", method!(RbExpr::arg_sort, 2))?;
     class.define_method("top_k", method!(RbExpr::top_k, 1))?;
@@ -282,7 +281,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("radians", method!(RbExpr::radians, 0))?;
     class.define_method("sign", method!(RbExpr::sign, 0))?;
     class.define_method("is_duplicated", method!(RbExpr::is_duplicated, 0))?;
-    class.define_method("over", method!(RbExpr::over, 1))?;
+    class.define_method("over", method!(RbExpr::over, 5))?;
     class.define_method("rolling", method!(RbExpr::rolling, 4))?;
     class.define_method("and_", method!(RbExpr::and_, 1))?;
     class.define_method("or_", method!(RbExpr::or_, 1))?;
@@ -356,6 +355,8 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("arr_tail", method!(RbExpr::arr_tail, 2))?;
     class.define_method("arr_shift", method!(RbExpr::arr_shift, 1))?;
     class.define_method("arr_explode", method!(RbExpr::arr_explode, 0))?;
+    class.define_method("arr_eval", method!(RbExpr::arr_eval, 2))?;
+    class.define_method("arr_agg", method!(RbExpr::arr_agg, 1))?;
     class.define_method("binary_contains", method!(RbExpr::bin_contains, 1))?;
     class.define_method("binary_ends_with", method!(RbExpr::bin_ends_with, 1))?;
     class.define_method("binary_starts_with", method!(RbExpr::bin_starts_with, 1))?;
@@ -418,6 +419,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("dt_iso_year", method!(RbExpr::dt_iso_year, 0))?;
     class.define_method("dt_quarter", method!(RbExpr::dt_quarter, 0))?;
     class.define_method("dt_month", method!(RbExpr::dt_month, 0))?;
+    class.define_method("dt_days_in_month", method!(RbExpr::dt_days_in_month, 0))?;
     class.define_method("dt_week", method!(RbExpr::dt_week, 0))?;
     class.define_method("dt_weekday", method!(RbExpr::dt_weekday, 0))?;
     class.define_method("dt_day", method!(RbExpr::dt_day, 0))?;
@@ -450,7 +452,6 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("dt_timestamp", method!(RbExpr::dt_timestamp, 1))?;
     class.define_method("dt_to_string", method!(RbExpr::dt_to_string, 1))?;
     class.define_method("dt_offset_by", method!(RbExpr::dt_offset_by, 1))?;
-    class.define_method("dt_with_time_unit", method!(RbExpr::dt_with_time_unit, 1))?;
     class.define_method(
         "dt_convert_time_zone",
         method!(RbExpr::dt_convert_time_zone, 1),
@@ -493,6 +494,8 @@ fn init(ruby: &Ruby) -> RbResult<()> {
         "rolling_quantile_by",
         method!(RbExpr::rolling_quantile_by, 6),
     )?;
+    class.define_method("rolling_rank", method!(RbExpr::rolling_rank, 5))?;
+    class.define_method("rolling_rank_by", method!(RbExpr::rolling_rank_by, 6))?;
     class.define_method("rolling_skew", method!(RbExpr::rolling_skew, 4))?;
     class.define_method("rolling_kurtosis", method!(RbExpr::rolling_kurtosis, 5))?;
     class.define_method("lower_bound", method!(RbExpr::lower_bound, 0))?;
@@ -529,6 +532,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("list_shift", method!(RbExpr::list_shift, 1))?;
     class.define_method("list_slice", method!(RbExpr::list_slice, 2))?;
     class.define_method("list_eval", method!(RbExpr::list_eval, 1))?;
+    class.define_method("list_agg", method!(RbExpr::list_agg, 1))?;
     class.define_method("list_filter", method!(RbExpr::list_filter, 1))?;
     class.define_method("cumulative_eval", method!(RbExpr::cumulative_eval, 2))?;
     class.define_method("list_to_struct", method!(RbExpr::list_to_struct, 1))?;
@@ -561,6 +565,10 @@ fn init(ruby: &Ruby) -> RbResult<()> {
         "struct_field_by_name",
         method!(RbExpr::struct_field_by_name, 1),
     )?;
+    class.define_method(
+        "struct_multiple_fields",
+        method!(RbExpr::struct_multiple_fields, 1),
+    )?;
     class.define_method(
         "struct_field_by_index",
         method!(RbExpr::struct_field_by_index, 1),
@@ -637,6 +645,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     )?;
     class.define_method("meta_is_literal", method!(RbExpr::meta_is_literal, 1))?;
     class.define_method("meta_tree_format", method!(RbExpr::meta_tree_format, 1))?;
+    class.define_method("meta_show_graph", method!(RbExpr::meta_show_graph, 1))?;
 
     // name
     class.define_method("name_keep", method!(RbExpr::name_keep, 0))?;
@@ -645,6 +654,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("name_suffix", method!(RbExpr::name_suffix, 1))?;
     class.define_method("name_to_lowercase", method!(RbExpr::name_to_lowercase, 0))?;
     class.define_method("name_to_uppercase", method!(RbExpr::name_to_uppercase, 0))?;
+    class.define_method("name_replace", method!(RbExpr::name_replace, 3))?;
     class.define_method("name_prefix_fields", method!(RbExpr::name_prefix_fields, 1))?;
     class.define_method("name_suffix_fields", method!(RbExpr::name_suffix_fields, 1))?;
 
@@ -654,6 +664,8 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_singleton_method("len", function!(functions::lazy::len, 0))?;
     class.define_singleton_method("fold", function!(functions::lazy::fold, 5))?;
     class.define_singleton_method("cum_fold", function!(functions::lazy::cum_fold, 6))?;
+    class.define_singleton_method("cum_reduce", function!(functions::lazy::cum_reduce, 4))?;
+    class.define_singleton_method("datetime", function!(functions::lazy::datetime, 10))?;
     class.define_singleton_method("lit", function!(functions::lazy::lit, 3))?;
     class.define_singleton_method("int_range", function!(functions::range::int_range, 4))?;
     class.define_singleton_method("int_ranges", function!(functions::range::int_ranges, 4))?;
@@ -666,13 +678,13 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_singleton_method("sql_expr", function!(functions::lazy::sql_expr, 1))?;
     class.define_singleton_method("cov", function!(functions::lazy::cov, 3))?;
     class.define_singleton_method("arctan2", function!(functions::lazy::arctan2, 2))?;
-    class.define_singleton_method("arctan2d", function!(functions::lazy::arctan2d, 2))?;
     class.define_singleton_method("rolling_corr", function!(functions::lazy::rolling_corr, 5))?;
     class.define_singleton_method("rolling_cov", function!(functions::lazy::rolling_cov, 5))?;
     class.define_singleton_method("arg_sort_by", function!(functions::lazy::arg_sort_by, 5))?;
     class.define_singleton_method("when", function!(functions::whenthen::when, 1))?;
     class.define_singleton_method("concat_str", function!(functions::lazy::concat_str, 3))?;
     class.define_singleton_method("concat_list", function!(functions::lazy::concat_list, 1))?;
+    class.define_singleton_method("concat_arr", function!(functions::lazy::concat_arr, 1))?;
     class.define_singleton_method(
         "business_day_count",
         function!(functions::business::business_day_count, 4),
@@ -731,13 +743,13 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_singleton_method("ipc_schema", function!(functions::io::read_ipc_schema, 1))?;
     class.define_singleton_method(
         "read_parquet_metadata",
-        function!(functions::io::read_parquet_metadata, 1),
+        function!(functions::io::read_parquet_metadata, 4),
     )?;
     class.define_singleton_method(
         "parquet_schema",
         function!(functions::io::read_parquet_schema, 1),
     )?;
-    class.define_singleton_method("collect_all", function!(functions::lazy::collect_all, 1))?;
+    class.define_singleton_method("collect_all", function!(functions::lazy::collect_all, 3))?;
     class.define_singleton_method("date_range", function!(functions::range::date_range, 4))?;
     class.define_singleton_method("date_ranges", function!(functions::range::date_ranges, 4))?;
     class.define_singleton_method(
@@ -750,6 +762,10 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     )?;
     class.define_singleton_method("time_range", function!(functions::range::time_range, 4))?;
     class.define_singleton_method("time_ranges", function!(functions::range::time_ranges, 4))?;
+    class.define_singleton_method(
+        "linear_spaces",
+        function!(functions::range::linear_spaces, 5),
+    )?;
     class.define_singleton_method(
         "dtype_str_repr",
         function!(functions::misc::dtype_str_repr, 1),
@@ -819,6 +835,30 @@ fn init(ruby: &Ruby) -> RbResult<()> {
         function!(functions::random::set_random_seed, 1),
     )?;
     class.define_singleton_method("re_escape", function!(re_escape, 1))?;
+    class.define_singleton_method(
+        "escape_regex",
+        function!(functions::strings::escape_regex, 1),
+    )?;
+    class.define_singleton_method(
+        "assert_dataframe_equal_rb",
+        function!(testing::assert_dataframe_equal_rb, 9),
+    )?;
+    class.define_singleton_method(
+        "assert_series_equal_rb",
+        function!(testing::assert_series_equal_rb, 9),
+    )?;
+    class.define_singleton_method(
+        "get_engine_affinity",
+        function!(functions::utils::rb_get_engine_affinity, 0),
+    )?;
+    class.define_singleton_method(
+        "polars_schema_to_rbcapsule",
+        function!(interop::arrow::to_rb::polars_schema_to_rbcapsule, 1),
+    )?;
+    class.define_singleton_method(
+        "init_polars_schema_from_arrow_c_schema",
+        function!(interop::arrow::init_polars_schema_from_arrow_c_schema, 2),
+    )?;
 
     let class = module.define_class("RbLazyFrame", ruby.class_object())?;
     #[cfg(feature = "serialize_binary")]
@@ -826,6 +866,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
         "serialize_binary",
         method!(RbLazyFrame::serialize_binary, 1),
     )?;
+    class.define_method("serialize_json", method!(RbLazyFrame::serialize_json, 1))?;
     #[cfg(feature = "serialize_binary")]
     class.define_singleton_method(
         "deserialize_binary",
@@ -837,7 +878,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     )?;
     class.define_singleton_method(
         "new_from_ndjson",
-        function!(RbLazyFrame::new_from_ndjson, 8),
+        function!(RbLazyFrame::new_from_ndjson, -1),
     )?;
     class.define_singleton_method("new_from_csv", function!(RbLazyFrame::new_from_csv, -1))?;
     class.define_singleton_method(
@@ -845,26 +886,39 @@ fn init(ruby: &Ruby) -> RbResult<()> {
         function!(RbLazyFrame::new_from_parquet, 6),
     )?;
     class.define_singleton_method("new_from_ipc", function!(RbLazyFrame::new_from_ipc, 3))?;
-    class.define_method("write_json", method!(RbLazyFrame::write_json, 1))?;
     class.define_method("describe_plan", method!(RbLazyFrame::describe_plan, 0))?;
     class.define_method(
         "describe_optimized_plan",
         method!(RbLazyFrame::describe_optimized_plan, 0),
     )?;
     class.define_method(
-        "optimization_toggle",
-        method!(RbLazyFrame::optimization_toggle, 9),
+        "describe_plan_tree",
+        method!(RbLazyFrame::describe_plan_tree, 0),
+    )?;
+    class.define_method(
+        "describe_optimized_plan_tree",
+        method!(RbLazyFrame::describe_optimized_plan_tree, 0),
+    )?;
+    class.define_method("to_dot", method!(RbLazyFrame::to_dot, 1))?;
+    class.define_method(
+        "to_dot_streaming_phys",
+        method!(RbLazyFrame::to_dot_streaming_phys, 1),
     )?;
     class.define_method("sort", method!(RbLazyFrame::sort, 5))?;
     class.define_method("sort_by_exprs", method!(RbLazyFrame::sort_by_exprs, 5))?;
     class.define_method("top_k", method!(RbLazyFrame::top_k, 3))?;
     class.define_method("bottom_k", method!(RbLazyFrame::bottom_k, 3))?;
     class.define_method("cache", method!(RbLazyFrame::cache, 0))?;
-    class.define_method("collect", method!(RbLazyFrame::collect, 0))?;
-    class.define_method("sink_parquet", method!(RbLazyFrame::sink_parquet, 9))?;
-    class.define_method("sink_ipc", method!(RbLazyFrame::sink_ipc, 5))?;
+    class.define_method(
+        "with_optimizations",
+        method!(RbLazyFrame::with_optimizations, 1),
+    )?;
+    class.define_method("profile", method!(RbLazyFrame::profile, 0))?;
+    class.define_method("collect", method!(RbLazyFrame::collect, 1))?;
+    class.define_method("sink_parquet", method!(RbLazyFrame::sink_parquet, 12))?;
+    class.define_method("sink_ipc", method!(RbLazyFrame::sink_ipc, 7))?;
     class.define_method("sink_csv", method!(RbLazyFrame::sink_csv, -1))?;
-    class.define_method("sink_json", method!(RbLazyFrame::sink_json, 4))?;
+    class.define_method("sink_json", method!(RbLazyFrame::sink_json, 5))?;
     class.define_method("filter", method!(RbLazyFrame::filter, 1))?;
     class.define_method("remove", method!(RbLazyFrame::remove, 1))?;
     class.define_method("select", method!(RbLazyFrame::select, 1))?;
@@ -875,7 +929,6 @@ fn init(ruby: &Ruby) -> RbResult<()> {
         "group_by_dynamic",
         method!(RbLazyFrame::group_by_dynamic, 9),
     )?;
-    class.define_method("with_context", method!(RbLazyFrame::with_context, 1))?;
     class.define_method("join_asof", method!(RbLazyFrame::join_asof, 14))?;
     class.define_method("join", method!(RbLazyFrame::join, 11))?;
     class.define_method("join_where", method!(RbLazyFrame::join_where, 3))?;
@@ -914,6 +967,19 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("unnest", method!(RbLazyFrame::unnest, 2))?;
     class.define_method("count", method!(RbLazyFrame::count, 0))?;
     class.define_method("merge_sorted", method!(RbLazyFrame::merge_sorted, 2))?;
+    class.define_method("hint_sorted", method!(RbLazyFrame::hint_sorted, 3))?;
+    class.define_method(
+        "collect_concurrently",
+        method!(RbLazyFrame::collect_concurrently, 0),
+    )?;
+
+    let class = module.define_class("RbInProcessQuery", ruby.class_object())?;
+    class.define_method("cancel", method!(RbInProcessQuery::cancel, 0))?;
+    class.define_method("fetch", method!(RbInProcessQuery::fetch, 0))?;
+    class.define_method(
+        "fetch_blocking",
+        method!(RbInProcessQuery::fetch_blocking, 0),
+    )?;
 
     let class = module.define_class("RbLazyGroupBy", ruby.class_object())?;
     class.define_method("agg", method!(RbLazyGroupBy::agg, 1))?;
@@ -949,7 +1015,6 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_singleton_method("new_series_list", function!(RbSeries::new_series_list, 3))?;
     class.define_singleton_method("new_array", function!(RbSeries::new_array, 5))?;
     class.define_singleton_method("new_decimal", function!(RbSeries::new_decimal, 3))?;
-    class.define_singleton_method("repeat", function!(RbSeries::repeat, 4))?;
     class.define_singleton_method(
         "from_arrow_c_stream",
         function!(RbSeries::from_arrow_c_stream, 1),
@@ -975,6 +1040,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("cat_is_local", method!(RbSeries::cat_is_local, 0))?;
     class.define_method("cat_to_local", method!(RbSeries::cat_to_local, 0))?;
     class.define_method("estimated_size", method!(RbSeries::estimated_size, 0))?;
+    class.define_method("reshape", method!(RbSeries::reshape, 1))?;
     class.define_method("get_fmt", method!(RbSeries::get_fmt, 2))?;
     class.define_method("rechunk", method!(RbSeries::rechunk, 1))?;
     class.define_method("get_index", method!(RbSeries::get_index, 1))?;
@@ -986,7 +1052,6 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("name", method!(RbSeries::name, 0))?;
     class.define_method("rename", method!(RbSeries::rename, 1))?;
     class.define_method("dtype", method!(RbSeries::dtype, 0))?;
-    class.define_method("inner_dtype", method!(RbSeries::inner_dtype, 0))?;
     class.define_method("set_sorted", method!(RbSeries::set_sorted_flag, 1))?;
     class.define_method("mean", method!(RbSeries::mean, 0))?;
     class.define_method("max", method!(RbSeries::max, 0))?;
@@ -1015,11 +1080,12 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("all", method!(RbSeries::all, 1))?;
     class.define_method("arg_min", method!(RbSeries::arg_min, 0))?;
     class.define_method("arg_max", method!(RbSeries::arg_max, 0))?;
-    class.define_method("take_with_series", method!(RbSeries::take_with_series, 1))?;
+    class.define_method(
+        "gather_with_series",
+        method!(RbSeries::gather_with_series, 1),
+    )?;
     class.define_method("null_count", method!(RbSeries::null_count, 0))?;
     class.define_method("has_nulls", method!(RbSeries::has_nulls, 0))?;
-    class.define_method("sample_n", method!(RbSeries::sample_n, 4))?;
-    class.define_method("sample_frac", method!(RbSeries::sample_frac, 4))?;
     class.define_method("equals", method!(RbSeries::equals, 4))?;
     class.define_method("eq", method!(RbSeries::eq, 1))?;
     class.define_method("neq", method!(RbSeries::neq, 1))?;
@@ -1034,7 +1100,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
         method!(RbSeries::str_to_decimal_infer, 1),
     )?;
     class.define_method("str_json_decode", method!(RbSeries::str_json_decode, 1))?;
-    class.define_method("to_s", method!(RbSeries::to_s, 0))?;
+    class.define_method("to_s", method!(RbSeries::as_str, 0))?;
     class.define_method("len", method!(RbSeries::len, 0))?;
     class.define_method("to_a", method!(RbSeries::to_a, 0))?;
     class.define_method("median", method!(RbSeries::median, 0))?;
@@ -1049,11 +1115,12 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("dot", method!(RbSeries::dot, 1))?;
     class.define_method("skew", method!(RbSeries::skew, 1))?;
     class.define_method("kurtosis", method!(RbSeries::kurtosis, 2))?;
-    class.define_method("cast", method!(RbSeries::cast, 2))?;
+    class.define_method("cast", method!(RbSeries::cast, 3))?;
     class.define_method("get_chunks", method!(RbSeries::get_chunks, 0))?;
     class.define_method("is_sorted", method!(RbSeries::is_sorted, 2))?;
     class.define_method("clear", method!(RbSeries::clear, 0))?;
-    class.define_method("time_unit", method!(RbSeries::time_unit, 0))?;
+    class.define_method("head", method!(RbSeries::head, 1))?;
+    class.define_method("tail", method!(RbSeries::tail, 1))?;
     class.define_method("scatter", method!(RbSeries::scatter, 2))?;
 
     // set
@@ -1138,6 +1205,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("eq_i16", method!(RbSeries::eq_i16, 1))?;
     class.define_method("eq_i32", method!(RbSeries::eq_i32, 1))?;
     class.define_method("eq_i64", method!(RbSeries::eq_i64, 1))?;
+    class.define_method("eq_i128", method!(RbSeries::eq_i128, 1))?;
     class.define_method("eq_f32", method!(RbSeries::eq_f32, 1))?;
     class.define_method("eq_f64", method!(RbSeries::eq_f64, 1))?;
 
@@ -1150,6 +1218,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("neq_i16", method!(RbSeries::neq_i16, 1))?;
     class.define_method("neq_i32", method!(RbSeries::neq_i32, 1))?;
     class.define_method("neq_i64", method!(RbSeries::neq_i64, 1))?;
+    class.define_method("neq_i128", method!(RbSeries::neq_i128, 1))?;
     class.define_method("neq_f32", method!(RbSeries::neq_f32, 1))?;
     class.define_method("neq_f64", method!(RbSeries::neq_f64, 1))?;
 
@@ -1162,6 +1231,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("gt_i16", method!(RbSeries::gt_i16, 1))?;
     class.define_method("gt_i32", method!(RbSeries::gt_i32, 1))?;
     class.define_method("gt_i64", method!(RbSeries::gt_i64, 1))?;
+    class.define_method("gt_i128", method!(RbSeries::gt_i128, 1))?;
     class.define_method("gt_f32", method!(RbSeries::gt_f32, 1))?;
     class.define_method("gt_f64", method!(RbSeries::gt_f64, 1))?;
 
@@ -1174,6 +1244,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("gt_eq_i16", method!(RbSeries::gt_eq_i16, 1))?;
     class.define_method("gt_eq_i32", method!(RbSeries::gt_eq_i32, 1))?;
     class.define_method("gt_eq_i64", method!(RbSeries::gt_eq_i64, 1))?;
+    class.define_method("gt_eq_i128", method!(RbSeries::gt_eq_i128, 1))?;
     class.define_method("gt_eq_f32", method!(RbSeries::gt_eq_f32, 1))?;
     class.define_method("gt_eq_f64", method!(RbSeries::gt_eq_f64, 1))?;
 
@@ -1186,6 +1257,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("lt_i16", method!(RbSeries::lt_i16, 1))?;
     class.define_method("lt_i32", method!(RbSeries::lt_i32, 1))?;
     class.define_method("lt_i64", method!(RbSeries::lt_i64, 1))?;
+    class.define_method("lt_i128", method!(RbSeries::lt_i128, 1))?;
     class.define_method("lt_f32", method!(RbSeries::lt_f32, 1))?;
     class.define_method("lt_f64", method!(RbSeries::lt_f64, 1))?;
 
@@ -1198,6 +1270,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("lt_eq_i16", method!(RbSeries::lt_eq_i16, 1))?;
     class.define_method("lt_eq_i32", method!(RbSeries::lt_eq_i32, 1))?;
     class.define_method("lt_eq_i64", method!(RbSeries::lt_eq_i64, 1))?;
+    class.define_method("lt_eq_i128", method!(RbSeries::lt_eq_i128, 1))?;
     class.define_method("lt_eq_f32", method!(RbSeries::lt_eq_f32, 1))?;
     class.define_method("lt_eq_f64", method!(RbSeries::lt_eq_f64, 1))?;
 
@@ -1212,9 +1285,6 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     // npy
     class.define_method("to_numo", method!(RbSeries::to_numo, 0))?;
 
-    // extra
-    class.define_method("extend_constant", method!(RbSeries::extend_constant, 2))?;
-
     // when then
     let class = module.define_class("RbWhen", ruby.class_object())?;
     class.define_method("then", method!(RbWhen::then, 1))?;
@@ -1246,6 +1316,10 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     let class = module.define_class("ArrowArrayStream", ruby.class_object())?;
     class.define_method("to_i", method!(RbArrowArrayStream::to_i, 0))?;
 
+    // arrow schema
+    let class = module.define_class("ArrowSchema", ruby.class_object())?;
+    class.define_method("to_i", method!(RbArrowSchema::to_i, 0))?;
+
     // catalog
     let class = module.define_class("RbCatalogClient", ruby.class_object())?;
     class.define_singleton_method("new", function!(RbCatalogClient::new, 2))?;
@@ -1336,92 +1410,83 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_singleton_method("default", function!(RbOptFlags::default, 0))?;
     class.define_method("no_optimizations", method!(RbOptFlags::no_optimizations, 0))?;
     class.define_method("copy", method!(RbOptFlags::copy, 0))?;
+    class.define_method("type_coercion", method!(RbOptFlags::get_type_coercion, 0))?;
+    class.define_method("type_coercion=", method!(RbOptFlags::set_type_coercion, 1))?;
+    class.define_method("type_check", method!(RbOptFlags::get_type_check, 0))?;
+    class.define_method("type_check=", method!(RbOptFlags::set_type_check, 1))?;
     class.define_method(
-        "get_type_coercion",
-        method!(RbOptFlags::get_type_coercion, 0),
-    )?;
-    class.define_method(
-        "set_type_coercion",
-        method!(RbOptFlags::set_type_coercion, 1),
-    )?;
-    class.define_method("get_type_check", method!(RbOptFlags::get_type_check, 0))?;
-    class.define_method("set_type_check", method!(RbOptFlags::set_type_check, 1))?;
-    class.define_method(
-        "get_projection_pushdown",
+        "projection_pushdown",
         method!(RbOptFlags::get_projection_pushdown, 0),
     )?;
     class.define_method(
-        "set_projection_pushdown",
+        "projection_pushdown=",
         method!(RbOptFlags::set_projection_pushdown, 1),
     )?;
     class.define_method(
-        "get_predicate_pushdown",
+        "predicate_pushdown",
         method!(RbOptFlags::get_predicate_pushdown, 0),
     )?;
     class.define_method(
-        "set_predicate_pushdown",
+        "predicate_pushdown=",
         method!(RbOptFlags::set_predicate_pushdown, 1),
     )?;
     class.define_method(
-        "get_cluster_with_columns",
+        "cluster_with_columns",
         method!(RbOptFlags::get_cluster_with_columns, 0),
     )?;
     class.define_method(
-        "set_cluster_with_columns",
+        "cluster_with_columns=",
         method!(RbOptFlags::set_cluster_with_columns, 1),
     )?;
     class.define_method(
-        "get_simplify_expression",
+        "simplify_expression",
         method!(RbOptFlags::get_simplify_expression, 0),
     )?;
     class.define_method(
-        "set_simplify_expression",
+        "simplify_expression=",
         method!(RbOptFlags::set_simplify_expression, 1),
     )?;
+    class.define_method("slice_pushdown", method!(RbOptFlags::get_slice_pushdown, 0))?;
     class.define_method(
-        "get_slice_pushdown",
-        method!(RbOptFlags::get_slice_pushdown, 0),
-    )?;
-    class.define_method(
-        "set_slice_pushdown",
+        "slice_pushdown=",
         method!(RbOptFlags::set_slice_pushdown, 1),
     )?;
     class.define_method(
-        "get_comm_subplan_elim",
+        "comm_subplan_elim",
         method!(RbOptFlags::get_comm_subplan_elim, 0),
     )?;
     class.define_method(
-        "set_comm_subplan_elim",
+        "comm_subplan_elim=",
         method!(RbOptFlags::set_comm_subplan_elim, 1),
     )?;
     class.define_method(
-        "get_comm_subexpr_elim",
+        "comm_subexpr_elim",
         method!(RbOptFlags::get_comm_subexpr_elim, 0),
     )?;
     class.define_method(
-        "set_comm_subexpr_elim",
+        "comm_subexpr_elim=",
         method!(RbOptFlags::set_comm_subexpr_elim, 1),
     )?;
     class.define_method(
-        "get_check_order_observe",
+        "check_order_observe",
         method!(RbOptFlags::get_check_order_observe, 0),
     )?;
     class.define_method(
-        "set_check_order_observe",
+        "check_order_observe=",
         method!(RbOptFlags::set_check_order_observe, 1),
     )?;
     class.define_method(
-        "get_fast_projection",
+        "fast_projection",
         method!(RbOptFlags::get_fast_projection, 0),
     )?;
     class.define_method(
-        "set_fast_projection",
+        "fast_projection=",
         method!(RbOptFlags::set_fast_projection, 1),
     )?;
-    class.define_method("get_eager", method!(RbOptFlags::get_eager, 0))?;
-    class.define_method("set_eager", method!(RbOptFlags::set_eager, 1))?;
-    class.define_method("get_streaming", method!(RbOptFlags::get_streaming, 0))?;
-    class.define_method("set_streaming", method!(RbOptFlags::set_streaming, 1))?;
+    class.define_method("eager", method!(RbOptFlags::get_eager, 0))?;
+    class.define_method("eager=", method!(RbOptFlags::set_eager, 1))?;
+    class.define_method("streaming", method!(RbOptFlags::get_streaming, 0))?;
+    class.define_method("streaming=", method!(RbOptFlags::set_streaming, 1))?;
 
     Ok(())
 }

data/ext/polars/src/map/dataframe.rs

@@ -80,19 +80,19 @@ pub fn apply_lambda_unknown<'a>(
                 .as_value(),
                 false,
             ));
-        // } else if out.is_kind_of(ruby.class_string()) {
-        //     let first_value = String::try_convert(out).ok();
-        //     return Ok((
-        //         RbSeries::new(
-        //             apply_lambda_with_utf8_out_type(df, lambda, null_count, first_value)
-        //                 .into_series(),
-        //         )
-        //         .into(),
-        //         false,
-        //     ));
+        } else if out.is_kind_of(ruby.class_string()) {
+            let first_value = String::try_convert(out).ok();
+            return Ok((
+                ruby.obj_wrap(RbSeries::new(
+                    apply_lambda_with_utf8_out_type(df, lambda, null_count, first_value.as_deref())
+                        .into_series(),
+                ))
+                .as_value(),
+                false,
+            ));
         } else if out.respond_to("_s", true)? {
             let rb_rbseries: Obj<RbSeries> = out.funcall("_s", ()).unwrap();
-            let series = rb_rbseries.series.borrow();
+            let series = rb_rbseries.series.read();
             let dt = series.dtype();
             return Ok((
                 ruby.obj_wrap(RbSeries::new(
@@ -245,7 +245,7 @@ pub fn apply_lambda_with_list_out_type(
                 Ok(val) => match val.funcall::<_, _, Value>("_s", ()) {
                     Ok(val) => Obj::<RbSeries>::try_convert(val)
                         .ok()
-                        .map(|ps| ps.series.borrow().clone()),
+                        .map(|ps| ps.series.read().clone()),
                     Err(_) => {
                         if val.is_nil() {
                             None