Adding categorical ToNapiValue (#97)

* Adding categorical ToNapiValue * Removing addtional unsafe * Adding DF selection by data type --------- Co-authored-by: Darek <dchrostowski@medallia.com>
pola-rs · Aug 21, 2023 · 5f2ada0 · 5f2ada0
1 parent 534f19b
commit 5f2ada0
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 2 deletions.
diff --git a/__tests__/dataframe.test.ts b/__tests__/dataframe.test.ts
@@ -763,6 +763,34 @@ describe("dataframe", () => {
       [3, 8, "c"],
     ]);
   });
+  test("rows:categorical", () => {
+    const actual = pl
+      .DataFrame({
+        foo: ["one", "two", "one", "two", "three"],
+        flt: [1, 2, 1, 2, 3],
+      })
+      .withColumns(
+        pl.col("foo").cast(pl.Categorical).alias("cat"),
+        pl.col("flt").cast(pl.Int32).alias("int"),
+      );
+
+    const expected = [
+      ["one", 1.0, "one", 1],
+      ["two", 2.0, "two", 2],
+      ["one", 1.0, "one", 1],
+      ["two", 2.0, "two", 2],
+      ["three", 3.0, "three", 3],
+    ];
+
+    expect(actual.rows()).toEqual(expected);
+    const byType = actual.select(
+      pl.col(pl.Utf8),
+      pl.col(pl.Float64),
+      pl.col(pl.Categorical),
+      pl.col(pl.Int32),
+    );
+    expect(byType.rows()).toEqual(expected);
+  });
   test("sample:n", () => {
     const actual = pl
       .DataFrame({

diff --git a/polars/lazy/functions.ts b/polars/lazy/functions.ts
@@ -1,4 +1,5 @@
 import { Expr, _Expr, exprToLitOrExpr } from "./expr";
+import { DataType } from "../datatypes";
 import { Series } from "../series";
 import { DataFrame } from "../dataframe";
 import { ExprOrString, range, selectionToExprList } from "../utils";
@@ -97,12 +98,14 @@ import pli from "../internals/polars_internal";
  * ╰───────────┴─────╯
  * ```
  */
-export function col(col: string | string[] | Series): Expr {
+export function col(col: string | string[] | Series | DataType): Expr {
   if (Series.isSeries(col)) {
     col = col.toArray();
   }
   if (Array.isArray(col)) {
     return _Expr(pli.cols(col));
+  } else if (typeof col === "object" && Object.values(col)[0] === "DataType") {
+    return _Expr(pli.dtypeCols([col]));
   } else {
     return _Expr(pli.col(col));
   }

diff --git a/src/conversion.rs b/src/conversion.rs
@@ -124,6 +124,7 @@ impl<'a> ToNapiValue for Wrap<AnyValue<'a>> {
             AnyValue::Float32(n) => f64::to_napi_value(env, n as f64),
             AnyValue::Float64(n) => f64::to_napi_value(env, n),
             AnyValue::Utf8(s) => String::to_napi_value(env, s.to_owned()),
+            AnyValue::Utf8Owned(s) => String::to_napi_value(env, s.to_string()),
             AnyValue::Date(v) => {
                 let mut ptr = std::ptr::null_mut();
 
@@ -143,6 +144,15 @@ impl<'a> ToNapiValue for Wrap<AnyValue<'a>> {
 
                 Ok(js_value)
             }
+            AnyValue::Categorical(idx, rev, arr) => {
+                let s = if arr.is_null() {
+                    rev.get(idx)
+                } else {
+                    arr.deref_unchecked().value(idx as usize)
+                };
+                let ptr = String::to_napi_value(env, s.to_string());
+                Ok(ptr.unwrap())
+            }
             AnyValue::Duration(v, _) => i64::to_napi_value(env, v),
             AnyValue::Time(v) => i64::to_napi_value(env, v),
             AnyValue::List(ser) => Wrap::<&Series>::to_napi_value(env, Wrap(&ser)),
@@ -592,7 +602,12 @@ impl FromNapiValue for Wrap<DataType> {
         match ty {
             ValueType::Object => {
                 let obj = Object::from_napi_value(env, napi_val)?;
-                let variant = obj.get::<_, String>("variant")?.unwrap();
+                let variant = if let Some(variant) = obj.get::<_, String>("variant")? {
+                    variant
+                } else {
+                    "".into()
+                };
+
                 let dtype = match variant.as_ref() {
                     "Int8" => DataType::Int8,
                     "Int16" => DataType::Int16,