Skip to content

Commit

Permalink
refactor: always use Self or RResult<Self> in Rust impl for `Da…
Browse files Browse the repository at this point in the history
…taFrame`, `LazyFrame` and `Expr` (#1098)
  • Loading branch information
etiennebacher authored May 18, 2024
1 parent 4ec1cc1 commit d497abb
Show file tree
Hide file tree
Showing 4 changed files with 28 additions and 38 deletions.
24 changes: 12 additions & 12 deletions src/rust/src/lazy/dataframe.rs
Original file line number Diff line number Diff line change
Expand Up @@ -260,7 +260,7 @@ impl RPolarsLazyFrame {
Ok(out.into())
}

fn shift(&self, periods: Robj) -> Result<Self, String> {
fn shift(&self, periods: Robj) -> RResult<Self> {
Ok(self.clone().0.shift(robj_to!(i64, periods)?).into())
}

Expand All @@ -276,27 +276,27 @@ impl RPolarsLazyFrame {
self.0.clone().reverse().into()
}

fn drop(&self, columns: Robj) -> Result<RPolarsLazyFrame, String> {
fn drop(&self, columns: Robj) -> RResult<Self> {
Ok(self.0.clone().drop(robj_to!(Vec, String, columns)?).into())
}

fn fill_nan(&self, fill_value: Robj) -> Result<Self, String> {
fn fill_nan(&self, fill_value: Robj) -> RResult<Self> {
Ok(self
.0
.clone()
.fill_nan(robj_to!(Expr, fill_value)?.0)
.into())
}

fn fill_null(&self, fill_value: Robj) -> Result<Self, String> {
fn fill_null(&self, fill_value: Robj) -> RResult<Self> {
Ok(self
.0
.clone()
.fill_null(robj_to!(Expr, fill_value)?.0)
.into())
}

fn slice(&self, offset: Robj, length: Robj) -> Result<RPolarsLazyFrame, String> {
fn slice(&self, offset: Robj, length: Robj) -> RResult<Self> {
Ok(RPolarsLazyFrame(self.0.clone().slice(
robj_to!(i64, offset)?,
robj_to!(Option, u32, length)?.unwrap_or(u32::MAX),
Expand Down Expand Up @@ -327,7 +327,7 @@ impl RPolarsLazyFrame {
Ok(RPolarsLazyFrame(self.clone().0.select_seq(exprs)))
}

fn tail(&self, n: Robj) -> Result<RPolarsLazyFrame, String> {
fn tail(&self, n: Robj) -> RResult<Self> {
Ok(RPolarsLazyFrame(self.0.clone().tail(robj_to!(u32, n)?)))
}

Expand All @@ -346,7 +346,7 @@ impl RPolarsLazyFrame {
Ok(out.into())
}

fn unique(&self, subset: Robj, keep: Robj, maintain_order: Robj) -> RResult<RPolarsLazyFrame> {
fn unique(&self, subset: Robj, keep: Robj, maintain_order: Robj) -> RResult<Self> {
let ke = robj_to!(UniqueKeepStrategy, keep)?;
let maintain_order = robj_to!(bool, maintain_order)?;
let subset = robj_to!(Option, Vec, String, subset)?;
Expand Down Expand Up @@ -450,7 +450,7 @@ impl RPolarsLazyFrame {
suffix: Robj,
allow_parallel: Robj,
force_parallel: Robj,
) -> RResult<RPolarsLazyFrame> {
) -> RResult<Self> {
Ok(RPolarsLazyFrame(
self.0
.clone()
Expand Down Expand Up @@ -512,7 +512,7 @@ impl RPolarsLazyFrame {
value_name: Robj,
variable_name: Robj,
streamable: Robj,
) -> Result<Self, String> {
) -> RResult<Self> {
let args = MeltArgs {
id_vars: strings_to_smartstrings(robj_to!(Vec, String, id_vars)?),
value_vars: strings_to_smartstrings(robj_to!(Vec, String, value_vars)?),
Expand All @@ -523,7 +523,7 @@ impl RPolarsLazyFrame {
Ok(self.0.clone().melt(args).into())
}

fn rename(&self, existing: Robj, new: Robj) -> Result<RPolarsLazyFrame, String> {
fn rename(&self, existing: Robj, new: Robj) -> RResult<Self> {
Ok(self
.0
.clone()
Expand Down Expand Up @@ -611,7 +611,7 @@ impl RPolarsLazyFrame {
profile_with_r_func_support(self.0.clone()).map(|(r, p)| list!(result = r, profile = p))
}

fn explode(&self, dotdotdot: Robj) -> RResult<RPolarsLazyFrame> {
fn explode(&self, dotdotdot: Robj) -> RResult<Self> {
Ok(self
.0
.clone()
Expand Down Expand Up @@ -735,7 +735,7 @@ impl RPolarsLazyGroupBy {
)
}

fn agg(&self, exprs: Robj) -> Result<RPolarsLazyFrame, String> {
fn agg(&self, exprs: Robj) -> RResult<RPolarsLazyFrame> {
Ok(RPolarsLazyFrame(
self.lgb.clone().agg(robj_to!(VecPLExprColNamed, exprs)?),
))
Expand Down
32 changes: 11 additions & 21 deletions src/rust/src/lazy/dsl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2195,7 +2195,7 @@ impl RPolarsExpr {
.into())
}

pub fn str_split(&self, by: Robj, inclusive: Robj) -> Result<RPolarsExpr, String> {
pub fn str_split(&self, by: Robj, inclusive: Robj) -> RResult<RPolarsExpr> {
let by = robj_to!(PLExpr, by)?;
let inclusive = robj_to!(bool, inclusive)?;
if inclusive {
Expand All @@ -2205,12 +2205,7 @@ impl RPolarsExpr {
}
}

pub fn str_split_exact(
&self,
by: Robj,
n: Robj,
inclusive: Robj,
) -> Result<RPolarsExpr, String> {
pub fn str_split_exact(&self, by: Robj, n: Robj, inclusive: Robj) -> RResult<RPolarsExpr> {
let by = robj_to!(PLExpr, by)?;
let n = robj_to!(usize, n)?;
let inclusive = robj_to!(bool, inclusive)?;
Expand All @@ -2222,7 +2217,7 @@ impl RPolarsExpr {
.into())
}

pub fn str_splitn(&self, by: Robj, n: Robj) -> Result<RPolarsExpr, String> {
pub fn str_splitn(&self, by: Robj, n: Robj) -> RResult<RPolarsExpr> {
Ok(self
.0
.clone()
Expand All @@ -2237,7 +2232,7 @@ impl RPolarsExpr {
value: Robj,
literal: Robj,
n: Robj,
) -> Result<RPolarsExpr, String> {
) -> RResult<RPolarsExpr> {
let pat = robj_to!(PLExpr, pat)?;
let value = robj_to!(PLExpr, value)?;
let literal = robj_to!(bool, literal)?;
Expand All @@ -2250,26 +2245,21 @@ impl RPolarsExpr {
.into())
}

pub fn str_replace_all(
&self,
pat: Robj,
value: Robj,
literal: Robj,
) -> Result<RPolarsExpr, String> {
pub fn str_replace_all(&self, pat: Robj, value: Robj, literal: Robj) -> RResult<RPolarsExpr> {
let pat = robj_to!(PLExpr, pat)?;
let value = robj_to!(PLExpr, value)?;
let literal = robj_to!(bool, literal)?;
Ok(self.0.clone().str().replace_all(pat, value, literal).into())
}

pub fn str_slice(&self, offset: Robj, length: Robj) -> Result<RPolarsExpr, String> {
pub fn str_slice(&self, offset: Robj, length: Robj) -> RResult<RPolarsExpr> {
let offset = robj_to!(PLExprCol, offset)?;
let length = robj_to!(PLExprCol, length)?;

Ok(self.clone().0.str().slice(offset, length).into())
}

pub fn str_explode(&self) -> Result<RPolarsExpr, String> {
pub fn str_explode(&self) -> RResult<RPolarsExpr> {
Ok(self.0.clone().str().explode().into())
}

Expand Down Expand Up @@ -2373,17 +2363,17 @@ impl RPolarsExpr {
self.0.clone().binary().base64_encode().into()
}

pub fn bin_hex_decode(&self, strict: Robj) -> Result<RPolarsExpr, String> {
pub fn bin_hex_decode(&self, strict: Robj) -> RResult<RPolarsExpr> {
let strict = robj_to!(bool, strict)?;
Ok(self.0.clone().binary().hex_decode(strict).into())
}

pub fn bin_base64_decode(&self, strict: Robj) -> Result<RPolarsExpr, String> {
pub fn bin_base64_decode(&self, strict: Robj) -> RResult<RPolarsExpr> {
let strict = robj_to!(bool, strict)?;
Ok(self.0.clone().binary().base64_decode(strict).into())
}

pub fn struct_field_by_name(&self, name: Robj) -> Result<RPolarsExpr, String> {
pub fn struct_field_by_name(&self, name: Robj) -> RResult<RPolarsExpr> {
Ok(self
.0
.clone()
Expand All @@ -2396,7 +2386,7 @@ impl RPolarsExpr {
// self.0.clone().struct_().field_by_index(index).into()
// }

pub fn struct_rename_fields(&self, names: Robj) -> Result<RPolarsExpr, String> {
pub fn struct_rename_fields(&self, names: Robj) -> RResult<RPolarsExpr> {
let string_vec: Vec<String> = robj_to!(Vec, String, names)?;
Ok(self.0.clone().struct_().rename_fields(string_vec).into())
}
Expand Down
8 changes: 4 additions & 4 deletions src/rust/src/rdataframe/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -306,19 +306,19 @@ impl RPolarsDataFrame {
RPolarsSeries(self.0.drop_in_place(names).unwrap())
}

pub fn select(&self, exprs: Robj) -> RResult<RPolarsDataFrame> {
pub fn select(&self, exprs: Robj) -> RResult<Self> {
self.lazy().select(exprs)?.collect()
}

pub fn select_seq(&self, exprs: Robj) -> RResult<RPolarsDataFrame> {
pub fn select_seq(&self, exprs: Robj) -> RResult<Self> {
self.lazy().select_seq(exprs)?.collect()
}

pub fn with_columns(&self, exprs: Robj) -> RResult<RPolarsDataFrame> {
pub fn with_columns(&self, exprs: Robj) -> RResult<Self> {
self.lazy().with_columns(exprs)?.collect()
}

pub fn with_columns_seq(&self, exprs: Robj) -> RResult<RPolarsDataFrame> {
pub fn with_columns_seq(&self, exprs: Robj) -> RResult<Self> {
self.lazy().with_columns_seq(exprs)?.collect()
}

Expand Down
2 changes: 1 addition & 1 deletion src/rust/src/series.rs
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ impl From<&RPolarsExpr> for pl::PolarsResult<RPolarsSeries> {
#[extendr]
impl RPolarsSeries {
//utility methods
pub fn new(name: Robj, values: Robj) -> RResult<RPolarsSeries> {
pub fn new(name: Robj, values: Robj) -> RResult<Self> {
robjname2series(values, robj_to!(str, name)?)
.map_err(polars_to_rpolars_err)
.map(RPolarsSeries)
Expand Down

0 comments on commit d497abb

Please sign in to comment.