Adding expr pl.len() (pola-rs#258)

Adding expr pl.len() to close pola-rs#257

---------

Co-authored-by: Cory Grinstead <[email protected]>

Author: Bidek56

__tests__/expr.test.ts

@@ -517,6 +517,26 @@ describe("expr", () => {
     const actual = df.select(col("a").last().as("last"));
     expect(actual).toFrameEqual(expected);
   });
+  test("len", () => {
+    const df = pl.DataFrame({
+      a: [1, 2, 3, 3, 3],
+      b: ["a", "a", "b", "a", "a"],
+    });
+    let actual = df.select(pl.len());
+    let expected = pl.DataFrame({ len: [5] });
+    expect(actual).toFrameEqual(expected);
+
+    actual = df.withColumn(pl.len());
+    expected = df.withColumn(pl.lit(5).alias("len"));
+    expect(actual).toFrameEqual(expected);
+
+    actual = df.withColumn(pl.intRange(pl.len()).alias("index"));
+    expected = df.withColumn(pl.Series("index", [0, 1, 2, 3, 4], pl.Int64));
+    expect(actual).toFrameEqual(expected);
+
+    actual = df.groupBy("b").agg(pl.len());
+    expect(actual.shape).toEqual({ height: 2, width: 2 });
+  });
   test("list", () => {
     const df = pl.DataFrame({
       a: ["a", "b", "c"],

biome.json

@@ -1,5 +1,5 @@
 {
-  "$schema": "https://biomejs.dev/schemas/1.0.0/schema.json",
+  "$schema": "https://biomejs.dev/schemas/1.8.3/schema.json",
   "organizeImports": {
     "enabled": false
   },

polars/index.ts

@@ -72,6 +72,7 @@ export namespace pl {
   export import groups = lazy.groups;
   export import head = lazy.head;
   export import last = lazy.last;
+  export import len = lazy.len;
   export import mean = lazy.mean;
   export import median = lazy.median;
   export import nUnique = lazy.nUnique;

polars/lazy/functions.ts

@@ -208,6 +208,11 @@ export function intRange(
   if (typeof opts?.low === "number") {
     return intRange(opts.low, opts.high, opts.step, opts.dtype, opts.eager);
   }
+  // if expression like pl.len() passed
+  if (high === undefined || high === null) {
+    high = opts;
+    opts = 0;
+  }
   const low = exprToLitOrExpr(opts, false);
   high = exprToLitOrExpr(high, false);
   if (eager) {
@@ -472,7 +477,51 @@ export function head(column: Series | ExprOrString, n?): Series | Expr {
   }
   return exprToLitOrExpr(column, false).head(n);
 }
-
+/** Return the number of elements in the column.  
+  This is similar to `COUNT(*)` in SQL.
+
+  @return Expr - Expression of data type :class:`UInt32`.
+  @example
+  ```
+  >>> const df = pl.DataFrame(
+  ...     {
+  ...         "a": [1, 2, None],
+  ...         "b": [3, None, None],
+  ...         "c": ["foo", "bar", "foo"],
+  ...     }
+  ... )
+  >>> df.select(pl.len())
+  shape: (1, 1)
+  ┌─────┐
+  │ len │
+  │ --- │
+  │ u32 │
+  ╞═════╡
+  │ 3   │
+  └─────┘
+  ```
+  Generate an index column by using `len` in conjunction with :func:`intRange`.
+
+  ```
+  >>> df.select(
+  ...     pl.intRange(pl.len(), dtype=pl.UInt32).alias("index"),
+  ...     pl.all(),
+  ... )
+  shape: (3, 4)
+  ┌───────┬──────┬──────┬─────┐
+  │ index ┆ a    ┆ b    ┆ c   │
+  │ ---   ┆ ---  ┆ ---  ┆ --- │
+  │ u32   ┆ i64  ┆ i64  ┆ str │
+  ╞═══════╪══════╪══════╪═════╡
+  │ 0     ┆ 1    ┆ 3    ┆ foo │
+  │ 1     ┆ 2    ┆ null ┆ bar │
+  │ 2     ┆ null ┆ null ┆ foo │
+  └───────┴──────┴──────┴─────┘
+  ```
+*/
+export function len(): any {
+  return _Expr(pli.len());
+}
 /** Get the last value. */
 export function last(column: ExprOrString | Series): any {
   if (Series.isSeries(column)) {

src/lazy/dataframe.rs

@@ -490,9 +490,9 @@ impl JsLazyFrame {
             .into()
     }
     #[napi(catch_unwind)]
-    pub fn slice(&self, offset: i64, len: u32) -> JsLazyFrame {
+    pub fn slice(&self, offset: i64, length: u32) -> JsLazyFrame {
         let ldf = self.ldf.clone();
-        ldf.slice(offset, len).into()
+        ldf.slice(offset, length).into()
     }
     #[napi(catch_unwind)]
     pub fn tail(&self, n: u32) -> JsLazyFrame {

src/lazy/dsl.rs

@@ -1719,6 +1719,11 @@ pub fn spearman_rank_corr(
     polars::lazy::dsl::spearman_rank_corr(a.0, b.0, ddof, propagate_nans).into()
 }
 
+#[napi(catch_unwind)]
+pub fn len() -> JsExpr {
+    polars::lazy::dsl::len().into()
+}
+
 #[napi(catch_unwind)]
 pub fn cov(a: Wrap<Expr>, b: Wrap<Expr>, ddof: u8) -> JsExpr {
     polars::lazy::dsl::cov(a.0, b.0, ddof).into()