polars-df 0.9.0 → 0.10.0

data/ext/polars/src/expr/general.rs

@@ -267,7 +267,7 @@ impl RbExpr {
     pub fn sort_with(&self, descending: bool, nulls_last: bool) -> Self {
         self.clone()
             .inner
-            .sort_with(SortOptions {
+            .sort(SortOptions {
                 descending,
                 nulls_last,
                 multithreaded: true,
@@ -323,9 +323,28 @@ impl RbExpr {
         self.clone().inner.gather(idx.inner.clone()).into()
     }
 
-    pub fn sort_by(&self, by: RArray, reverse: Vec<bool>) -> RbResult<Self> {
+    pub fn sort_by(
+        &self,
+        by: RArray,
+        descending: Vec<bool>,
+        nulls_last: bool,
+        multithreaded: bool,
+        maintain_order: bool,
+    ) -> RbResult<Self> {
         let by = rb_exprs_to_exprs(by)?;
-        Ok(self.clone().inner.sort_by(by, reverse).into())
+        Ok(self
+            .clone()
+            .inner
+            .sort_by(
+                by,
+                SortMultipleOptions {
+                    descending,
+                    nulls_last,
+                    multithreaded,
+                    maintain_order,
+                },
+            )
+            .into())
     }
 
     pub fn backward_fill(&self, limit: FillNullLimit) -> Self {

data/ext/polars/src/expr/list.rs

@@ -51,8 +51,12 @@ impl RbExpr {
             .into()
     }
 
-    pub fn list_get(&self, index: &RbExpr) -> Self {
-        self.inner.clone().list().get(index.inner.clone()).into()
+    pub fn list_get(&self, index: &RbExpr, null_on_oob: bool) -> Self {
+        self.inner
+            .clone()
+            .list()
+            .get(index.inner.clone(), null_on_oob)
+            .into()
     }
 
     pub fn list_join(&self, separator: &RbExpr, ignore_nulls: bool) -> Self {

data/ext/polars/src/expr/string.rs

@@ -244,12 +244,12 @@ impl RbExpr {
             .into()
     }
 
-    pub fn str_to_integer(&self, base: u32, strict: bool) -> Self {
+    pub fn str_to_integer(&self, base: &Self, strict: bool) -> Self {
         self.inner
             .clone()
             .str()
-            .to_integer(base, strict)
-            .with_fmt("str.parse_int")
+            .to_integer(base.inner.clone(), strict)
+            .with_fmt("str.to_integer")
             .into()
     }
 

data/ext/polars/src/file.rs

@@ -1,20 +1,167 @@
-use magnus::{exception, prelude::*, Error, RString, Value};
-use polars::io::mmap::MmapBytesReader;
 use std::fs::File;
-use std::io::Cursor;
+use std::io;
+use std::io::{BufReader, Cursor, Read, Seek, SeekFrom, Write};
 use std::path::PathBuf;
 
+use magnus::{exception, prelude::*, Error, RString, Value};
+use polars::io::mmap::MmapBytesReader;
+
+use crate::error::RbPolarsErr;
+use crate::prelude::resolve_homedir;
 use crate::RbResult;
 
-pub fn get_file_like(f: Value, truncate: bool) -> RbResult<File> {
-    let str_slice = PathBuf::try_convert(f)?;
-    let f = if truncate {
-        File::create(str_slice)
-            .map_err(|e| Error::new(exception::runtime_error(), e.to_string()))?
+#[derive(Clone)]
+pub struct RbFileLikeObject {
+    inner: Value,
+}
+
+/// Wraps a `Value`, and implements read, seek, and write for it.
+impl RbFileLikeObject {
+    /// Creates an instance of a `RbFileLikeObject` from a `Value`.
+    /// To assert the object has the required methods methods,
+    /// instantiate it with `RbFileLikeObject::require`
+    pub fn new(object: Value) -> Self {
+        RbFileLikeObject { inner: object }
+    }
+
+    pub fn as_buffer(&self) -> std::io::Cursor<Vec<u8>> {
+        let data = self.as_file_buffer().into_inner();
+        std::io::Cursor::new(data)
+    }
+
+    pub fn as_file_buffer(&self) -> Cursor<Vec<u8>> {
+        let bytes = self
+            .inner
+            .funcall::<_, _, RString>("read", ())
+            .expect("no read method found");
+
+        let buf = unsafe { bytes.as_slice() }.to_vec();
+
+        Cursor::new(buf)
+    }
+
+    /// Same as `RbFileLikeObject::new`, but validates that the underlying
+    /// ruby object has a `read`, `write`, and `seek` methods in respect to parameters.
+    /// Will return a `TypeError` if object does not have `read`, `seek`, and `write` methods.
+    pub fn with_requirements(object: Value, read: bool, write: bool, seek: bool) -> RbResult<Self> {
+        if read && !object.respond_to("read", false)? {
+            return Err(Error::new(
+                exception::type_error(),
+                "Object does not have a .read() method.",
+            ));
+        }
+
+        if seek && !object.respond_to("seek", false)? {
+            return Err(Error::new(
+                exception::type_error(),
+                "Object does not have a .seek() method.",
+            ));
+        }
+
+        if write && !object.respond_to("write", false)? {
+            return Err(Error::new(
+                exception::type_error(),
+                "Object does not have a .write() method.",
+            ));
+        }
+
+        Ok(RbFileLikeObject::new(object))
+    }
+}
+
+/// Extracts a string repr from, and returns an IO error to send back to rust.
+fn rberr_to_io_err(e: Error) -> io::Error {
+    io::Error::new(io::ErrorKind::Other, e.to_string())
+}
+
+impl Read for RbFileLikeObject {
+    fn read(&mut self, mut buf: &mut [u8]) -> Result<usize, io::Error> {
+        let bytes = self
+            .inner
+            .funcall::<_, _, RString>("read", (buf.len(),))
+            .map_err(rberr_to_io_err)?;
+
+        buf.write_all(unsafe { bytes.as_slice() })?;
+
+        Ok(bytes.len())
+    }
+}
+
+impl Write for RbFileLikeObject {
+    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
+        let rbbytes = RString::from_slice(buf);
+
+        let number_bytes_written = self
+            .inner
+            .funcall::<_, _, usize>("write", (rbbytes,))
+            .map_err(rberr_to_io_err)?;
+
+        Ok(number_bytes_written)
+    }
+
+    fn flush(&mut self) -> Result<(), io::Error> {
+        self.inner
+            .funcall::<_, _, Value>("flush", ())
+            .map_err(rberr_to_io_err)?;
+
+        Ok(())
+    }
+}
+
+impl Seek for RbFileLikeObject {
+    fn seek(&mut self, pos: SeekFrom) -> Result<u64, io::Error> {
+        let (whence, offset) = match pos {
+            SeekFrom::Start(i) => (0, i as i64),
+            SeekFrom::Current(i) => (1, i),
+            SeekFrom::End(i) => (2, i),
+        };
+
+        let new_position = self
+            .inner
+            .funcall("seek", (offset, whence))
+            .map_err(rberr_to_io_err)?;
+
+        Ok(new_position)
+    }
+}
+
+pub trait FileLike: Read + Write + Seek {}
+
+impl FileLike for File {}
+impl FileLike for RbFileLikeObject {}
+
+pub enum EitherRustRubyFile {
+    Rb(RbFileLikeObject),
+    Rust(BufReader<File>),
+}
+
+///
+/// # Arguments
+/// * `truncate` - open or create a new file.
+pub fn get_either_file(rb_f: Value, truncate: bool) -> RbResult<EitherRustRubyFile> {
+    if let Ok(rstring) = RString::try_convert(rb_f) {
+        let s = unsafe { rstring.as_str() }?;
+        let file_path = std::path::Path::new(&s);
+        let file_path = resolve_homedir(file_path);
+        let f = if truncate {
+            File::create(file_path).map_err(RbPolarsErr::io)?
+        } else {
+            polars_utils::open_file(&file_path).map_err(RbPolarsErr::from)?
+        };
+        let reader = BufReader::new(f);
+        Ok(EitherRustRubyFile::Rust(reader))
     } else {
-        File::open(str_slice).map_err(|e| Error::new(exception::runtime_error(), e.to_string()))?
-    };
-    Ok(f)
+        let f = RbFileLikeObject::with_requirements(rb_f, !truncate, truncate, !truncate)?;
+        Ok(EitherRustRubyFile::Rb(f))
+    }
+}
+
+pub fn get_file_like(f: Value, truncate: bool) -> RbResult<Box<dyn FileLike>> {
+    use EitherRustRubyFile::*;
+    match get_either_file(f, truncate)? {
+        Rb(f) => Ok(Box::new(f)),
+        Rust(f) => Ok(Box::new(f.into_inner())),
+    }
 }
 
 pub fn get_mmap_bytes_reader(rb_f: Value) -> RbResult<Box<dyn MmapBytesReader>> {

data/ext/polars/src/functions/lazy.rs

@@ -55,9 +55,24 @@ pub fn rolling_cov(
     .into()
 }
 
-pub fn arg_sort_by(by: RArray, descending: Vec<bool>) -> RbResult<RbExpr> {
+pub fn arg_sort_by(
+    by: RArray,
+    descending: Vec<bool>,
+    nulls_last: bool,
+    multithreaded: bool,
+    maintain_order: bool,
+) -> RbResult<RbExpr> {
     let by = rb_exprs_to_exprs(by)?;
-    Ok(dsl::arg_sort_by(by, &descending).into())
+    Ok(dsl::arg_sort_by(
+        by,
+        SortMultipleOptions {
+            descending,
+            nulls_last,
+            multithreaded,
+            maintain_order,
+        },
+    )
+    .into())
 }
 
 pub fn arg_where(condition: &RbExpr) -> RbExpr {
@@ -324,6 +339,6 @@ pub fn spearman_rank_corr(a: &RbExpr, b: &RbExpr, ddof: u8, propagate_nans: bool
 }
 
 pub fn sql_expr(sql: String) -> RbResult<RbExpr> {
-    let expr = polars::sql::sql_expr(&sql).map_err(RbPolarsErr::from)?;
+    let expr = polars::sql::sql_expr(sql).map_err(RbPolarsErr::from)?;
     Ok(expr.into())
 }

data/ext/polars/src/functions/whenthen.rs

@@ -2,42 +2,72 @@ use polars::lazy::dsl;
 
 use crate::RbExpr;
 
+pub fn when(condition: &RbExpr) -> RbWhen {
+    RbWhen {
+        inner: dsl::when(condition.inner.clone()),
+    }
+}
+
 #[magnus::wrap(class = "Polars::RbWhen")]
 #[derive(Clone)]
 pub struct RbWhen {
     pub inner: dsl::When,
 }
 
-impl From<dsl::When> for RbWhen {
-    fn from(inner: dsl::When) -> Self {
-        RbWhen { inner }
-    }
-}
-
-#[magnus::wrap(class = "Polars::RbWhenThen")]
+#[magnus::wrap(class = "Polars::RbThen")]
 #[derive(Clone)]
 pub struct RbThen {
     pub inner: dsl::Then,
 }
 
-impl From<dsl::Then> for RbThen {
-    fn from(inner: dsl::Then) -> Self {
-        RbThen { inner }
-    }
+#[magnus::wrap(class = "Polars::RbChainedWhen")]
+#[derive(Clone)]
+pub struct RbChainedWhen {
+    pub inner: dsl::ChainedWhen,
+}
+
+#[magnus::wrap(class = "Polars::RbChainedThen")]
+#[derive(Clone)]
+pub struct RbChainedThen {
+    pub inner: dsl::ChainedThen,
 }
 
 impl RbWhen {
-    pub fn then(&self, expr: &RbExpr) -> RbThen {
-        self.inner.clone().then(expr.inner.clone()).into()
+    pub fn then(&self, statement: &RbExpr) -> RbThen {
+        RbThen {
+            inner: self.inner.clone().then(statement.inner.clone()),
+        }
     }
 }
 
 impl RbThen {
-    pub fn overwise(&self, expr: &RbExpr) -> RbExpr {
-        self.inner.clone().otherwise(expr.inner.clone()).into()
+    pub fn when(&self, condition: &RbExpr) -> RbChainedWhen {
+        RbChainedWhen {
+            inner: self.inner.clone().when(condition.inner.clone()),
+        }
+    }
+
+    pub fn otherwise(&self, statement: &RbExpr) -> RbExpr {
+        self.inner.clone().otherwise(statement.inner.clone()).into()
     }
 }
 
-pub fn when(predicate: &RbExpr) -> RbWhen {
-    dsl::when(predicate.inner.clone()).into()
+impl RbChainedWhen {
+    pub fn then(&self, statement: &RbExpr) -> RbChainedThen {
+        RbChainedThen {
+            inner: self.inner.clone().then(statement.inner.clone()),
+        }
+    }
+}
+
+impl RbChainedThen {
+    pub fn when(&self, condition: &RbExpr) -> RbChainedWhen {
+        RbChainedWhen {
+            inner: self.inner.clone().when(condition.inner.clone()),
+        }
+    }
+
+    pub fn otherwise(&self, statement: &RbExpr) -> RbExpr {
+        self.inner.clone().otherwise(statement.inner.clone()).into()
+    }
 }

data/ext/polars/src/lazyframe/mod.rs

@@ -1,5 +1,5 @@
 use magnus::{IntoValue, RArray, RHash, TryConvert, Value};
-use polars::io::RowIndex;
+use polars::io::{HiveOptions, RowIndex};
 use polars::lazy::frame::LazyFrame;
 use polars::prelude::*;
 use std::cell::RefCell;
@@ -100,13 +100,13 @@ impl RbLazyFrame {
         let row_index = Option::<(String, IdxSize)>::try_convert(arguments[17])?;
         let try_parse_dates = bool::try_convert(arguments[18])?;
         let eol_char = String::try_convert(arguments[19])?;
+        let truncate_ragged_lines = bool::try_convert(arguments[20])?;
         // end arguments
 
         let null_values = null_values.map(|w| w.0);
         let quote_char = quote_char.map(|s| s.as_bytes()[0]);
         let separator = separator.as_bytes()[0];
         let eol_char = eol_char.as_bytes()[0];
-
         let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
 
         let overwrite_dtype = overwrite_dtype.map(|overwrite_dtype| {
@@ -115,6 +115,7 @@ impl RbLazyFrame {
                 .map(|(name, dtype)| Field::new(&name, dtype.0))
                 .collect::<Schema>()
         });
+
         let r = LazyCsvReader::new(path)
             .with_infer_schema_length(infer_schema_length)
             .with_separator(separator)
@@ -124,6 +125,7 @@ impl RbLazyFrame {
             .with_n_rows(n_rows)
             .with_cache(cache)
             .with_dtype_overwrite(overwrite_dtype.as_ref())
+            // TODO add with_schema
             .low_memory(low_memory)
             .with_comment_prefix(comment_prefix.as_deref())
             .with_quote_char(quote_char)
@@ -133,7 +135,9 @@ impl RbLazyFrame {
             .with_encoding(encoding.0)
             .with_row_index(row_index)
             .with_try_parse_dates(try_parse_dates)
-            .with_null_values(null_values);
+            .with_null_values(null_values)
+            // TODO add with_missing_is_null
+            .truncate_ragged_lines(truncate_ragged_lines);
 
         if let Some(_lambda) = with_schema_modify {
             todo!();
@@ -144,7 +148,8 @@ impl RbLazyFrame {
 
     #[allow(clippy::too_many_arguments)]
     pub fn new_from_parquet(
-        path: String,
+        path: Option<PathBuf>,
+        paths: Vec<PathBuf>,
         n_rows: Option<usize>,
         cache: bool,
         parallel: Wrap<ParallelStrategy>,
@@ -153,21 +158,43 @@ impl RbLazyFrame {
         low_memory: bool,
         use_statistics: bool,
         hive_partitioning: bool,
+        hive_schema: Option<Wrap<Schema>>,
     ) -> RbResult<Self> {
+        let parallel = parallel.0;
+        let hive_schema = hive_schema.map(|s| Arc::new(s.0));
+
+        let first_path = if let Some(path) = &path {
+            path
+        } else {
+            paths
+                .first()
+                .ok_or_else(|| RbValueError::new_err("expected a path argument".to_string()))?
+        };
+
         let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
+        let hive_options = HiveOptions {
+            enabled: hive_partitioning,
+            schema: hive_schema,
+        };
+
         let args = ScanArgsParquet {
             n_rows,
             cache,
-            parallel: parallel.0,
+            parallel,
             rechunk,
             row_index,
             low_memory,
-            // TODO support cloud options
             cloud_options: None,
             use_statistics,
-            hive_partitioning,
+            hive_options,
         };
-        let lf = LazyFrame::scan_parquet(path, args).map_err(RbPolarsErr::from)?;
+
+        let lf = if path.is_some() {
+            LazyFrame::scan_parquet(first_path, args)
+        } else {
+            LazyFrame::scan_parquet_files(Arc::from(paths), args)
+        }
+        .map_err(RbPolarsErr::from)?;
         Ok(lf.into())
     }
 
@@ -185,7 +212,8 @@ impl RbLazyFrame {
             cache,
             rechunk,
             row_index,
-            memmap: memory_map,
+            memory_map,
+            cloud_options: None,
         };
         let lf = LazyFrame::scan_ipc(path, args).map_err(RbPolarsErr::from)?;
         Ok(lf.into())
@@ -242,17 +270,18 @@ impl RbLazyFrame {
     pub fn sort(
         &self,
         by_column: String,
-        reverse: bool,
+        descending: bool,
         nulls_last: bool,
         maintain_order: bool,
+        multithreaded: bool,
     ) -> Self {
         let ldf = self.ldf.clone();
         ldf.sort(
-            &by_column,
-            SortOptions {
-                descending: reverse,
+            [&by_column],
+            SortMultipleOptions {
+                descending: vec![descending],
                 nulls_last,
-                multithreaded: true,
+                multithreaded,
                 maintain_order,
             },
         )
@@ -261,15 +290,24 @@ impl RbLazyFrame {
 
     pub fn sort_by_exprs(
         &self,
-        by_column: RArray,
-        reverse: Vec<bool>,
+        by: RArray,
+        descending: Vec<bool>,
         nulls_last: bool,
         maintain_order: bool,
+        multithreaded: bool,
     ) -> RbResult<Self> {
         let ldf = self.ldf.clone();
-        let exprs = rb_exprs_to_exprs(by_column)?;
+        let exprs = rb_exprs_to_exprs(by)?;
         Ok(ldf
-            .sort_by_exprs(exprs, reverse, nulls_last, maintain_order)
+            .sort_by_exprs(
+                exprs,
+                SortMultipleOptions {
+                    descending,
+                    nulls_last,
+                    maintain_order,
+                    multithreaded,
+                },
+            )
             .into())
     }
 
@@ -326,6 +364,7 @@ impl RbLazyFrame {
         Ok(())
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn sink_csv(
         &self,
         path: PathBuf,
@@ -427,7 +466,7 @@ impl RbLazyFrame {
         let closed_window = closed.0;
         let ldf = self.ldf.clone();
         let by = rb_exprs_to_exprs(by)?;
-        let lazy_gb = ldf.group_by_rolling(
+        let lazy_gb = ldf.rolling(
             index_column.inner.clone(),
             by,
             RollingGroupOptions {

data/ext/polars/src/lib.rs

@@ -23,7 +23,7 @@ use error::{RbPolarsErr, RbTypeError, RbValueError};
 use expr::rb_exprs_to_exprs;
 use expr::RbExpr;
 use functions::string_cache::RbStringCacheHolder;
-use functions::whenthen::{RbThen, RbWhen};
+use functions::whenthen::{RbChainedThen, RbChainedWhen, RbThen, RbWhen};
 use lazyframe::RbLazyFrame;
 use lazygroupby::RbLazyGroupBy;
 use magnus::{define_module, function, method, prelude::*, Error, Ruby};
@@ -74,7 +74,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("row_tuple", method!(RbDataFrame::row_tuple, 1))?;
     class.define_method("row_tuples", method!(RbDataFrame::row_tuples, 0))?;
     class.define_method("to_numo", method!(RbDataFrame::to_numo, 0))?;
-    class.define_method("write_parquet", method!(RbDataFrame::write_parquet, 5))?;
+    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))?;
@@ -213,7 +213,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("arg_min", method!(RbExpr::arg_min, 0))?;
     class.define_method("search_sorted", method!(RbExpr::search_sorted, 2))?;
     class.define_method("gather", method!(RbExpr::gather, 1))?;
-    class.define_method("sort_by", method!(RbExpr::sort_by, 2))?;
+    class.define_method("sort_by", method!(RbExpr::sort_by, 5))?;
     class.define_method("backward_fill", method!(RbExpr::backward_fill, 1))?;
     class.define_method("forward_fill", method!(RbExpr::forward_fill, 1))?;
     class.define_method("shift", method!(RbExpr::shift, 2))?;
@@ -312,7 +312,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("arr_reverse", method!(RbExpr::arr_reverse, 0))?;
     class.define_method("arr_arg_min", method!(RbExpr::arr_arg_min, 0))?;
     class.define_method("arr_arg_max", method!(RbExpr::arr_arg_max, 0))?;
-    class.define_method("arr_get", method!(RbExpr::arr_get, 1))?;
+    class.define_method("arr_get", method!(RbExpr::arr_get, 2))?;
     class.define_method("arr_join", method!(RbExpr::arr_join, 2))?;
     class.define_method("arr_contains", method!(RbExpr::arr_contains, 1))?;
     class.define_method("arr_count_matches", method!(RbExpr::arr_count_matches, 1))?;
@@ -406,7 +406,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("dt_cast_time_unit", method!(RbExpr::dt_cast_time_unit, 1))?;
     class.define_method(
         "dt_replace_time_zone",
-        method!(RbExpr::dt_replace_time_zone, 2),
+        method!(RbExpr::dt_replace_time_zone, 3),
     )?;
     class.define_method("dt_truncate", method!(RbExpr::dt_truncate, 2))?;
     class.define_method("dt_month_start", method!(RbExpr::dt_month_start, 0))?;
@@ -448,7 +448,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("list_sort", method!(RbExpr::list_sort, 1))?;
     class.define_method("list_reverse", method!(RbExpr::list_reverse, 0))?;
     class.define_method("list_unique", method!(RbExpr::list_unique, 1))?;
-    class.define_method("list_get", method!(RbExpr::list_get, 1))?;
+    class.define_method("list_get", method!(RbExpr::list_get, 2))?;
     class.define_method("list_join", method!(RbExpr::list_join, 2))?;
     class.define_method("list_arg_min", method!(RbExpr::list_arg_min, 0))?;
     class.define_method("list_arg_max", method!(RbExpr::list_arg_max, 0))?;
@@ -554,7 +554,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     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, 2))?;
+    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))?;
@@ -689,7 +689,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_singleton_method("new_from_csv", function!(RbLazyFrame::new_from_csv, -1))?;
     class.define_singleton_method(
         "new_from_parquet",
-        function!(RbLazyFrame::new_from_parquet, 9),
+        function!(RbLazyFrame::new_from_parquet, 11),
     )?;
     class.define_singleton_method("new_from_ipc", function!(RbLazyFrame::new_from_ipc, 6))?;
     class.define_method("write_json", method!(RbLazyFrame::write_json, 1))?;
@@ -702,8 +702,8 @@ fn init(ruby: &Ruby) -> RbResult<()> {
         "optimization_toggle",
         method!(RbLazyFrame::optimization_toggle, 9),
     )?;
-    class.define_method("sort", method!(RbLazyFrame::sort, 4))?;
-    class.define_method("sort_by_exprs", method!(RbLazyFrame::sort_by_exprs, 4))?;
+    class.define_method("sort", method!(RbLazyFrame::sort, 5))?;
+    class.define_method("sort_by_exprs", method!(RbLazyFrame::sort_by_exprs, 5))?;
     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, 7))?;
@@ -835,7 +835,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("mul", method!(RbSeries::mul, 1))?;
     class.define_method("div", method!(RbSeries::div, 1))?;
     class.define_method("rem", method!(RbSeries::rem, 1))?;
-    class.define_method("sort", method!(RbSeries::sort, 2))?;
+    class.define_method("sort", method!(RbSeries::sort, 3))?;
     class.define_method("value_counts", method!(RbSeries::value_counts, 1))?;
     class.define_method("any", method!(RbSeries::any, 1))?;
     class.define_method("all", method!(RbSeries::all, 1))?;
@@ -1032,11 +1032,20 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     // 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))?;
+    class.define_method("then", method!(RbWhen::then, 1))?;
 
-    let class = module.define_class("RbWhenThen", ruby.class_object())?;
-    class.define_method("otherwise", method!(RbThen::overwise, 1))?;
+    let class = module.define_class("RbThen", ruby.class_object())?;
+    class.define_method("when", method!(RbThen::when, 1))?;
+    class.define_method("otherwise", method!(RbThen::otherwise, 1))?;
+
+    let class = module.define_class("RbChainedWhen", ruby.class_object())?;
+    class.define_method("then", method!(RbChainedWhen::then, 1))?;
+
+    let class = module.define_class("RbChainedThen", ruby.class_object())?;
+    class.define_method("when", method!(RbChainedThen::when, 1))?;
+    class.define_method("otherwise", method!(RbChainedThen::otherwise, 1))?;
 
     // sql
     let class = module.define_class("RbSQLContext", ruby.class_object())?;