fix: minor fixes to list conversion (#2277)

pola-rs · Jan 6, 2022 · 8a8cabb · 8a8cabb
1 parent a5fba05
commit 8a8cabb
Show file tree

Hide file tree

Showing 3 changed files with 119 additions and 4 deletions.
diff --git a/nodejs-polars/polars/lazy/functions.ts b/nodejs-polars/polars/lazy/functions.ts
@@ -241,7 +241,7 @@ export function cov(a: ExprOrString, b: ExprOrString): Expr {
  * >>> pl.col("*").exclude(columns)
  * ```
  */
-export function exclude(columns: string[] | string)
+export function exclude(columns: string[] | string): Expr
 export function exclude(col: string, ...cols: string[]): Expr
 export function exclude(...columns): Expr {
   return col("*").exclude(columns as any);

diff --git a/nodejs-polars/src/conversion/into.rs b/nodejs-polars/src/conversion/into.rs
@@ -1,15 +1,24 @@
 use crate::dataframe::JsDataFrame;
+use crate::prelude::JsPolarsEr;
 use crate::prelude::Wrap;
 use crate::series::JsSeries;
-use napi::{CallContext, JsBigint, JsExternal, JsNumber, JsString, JsUnknown, Result};
-use polars_core::prelude::AnyValue;
+use napi::{CallContext, JsBigint, JsExternal, JsNumber, JsObject, JsString, JsUnknown, Result};
+use polars::datatypes::DataType;
+use polars_core::prelude::*;
 pub trait IntoJs<T>: Send + Sized {
     fn into_js(self, cx: &CallContext) -> T {
         self.try_into_js(cx).expect("Use 'try_into_js' instead")
     }
     fn try_into_js(self, cx: &CallContext) -> Result<T>;
 }
 
+pub trait IntoJsRef<T>: Send + Sized {
+    fn into_js_ref(&self, cx: &CallContext) -> T {
+        self.try_into_js_ref(cx).expect("Use 'try_into_js_ref' instead")
+    }
+    fn try_into_js_ref(&self, cx: &CallContext) -> Result<T>;
+}
+
 impl IntoJs<JsExternal> for JsSeries {
     fn try_into_js(self, cx: &CallContext) -> Result<JsExternal> {
         cx.env.create_external(self, None)
@@ -92,6 +101,74 @@ impl IntoJs<JsNumber> for f64 {
         cx.env.create_double(self)
     }
 }
+impl IntoJs<napi::JsBoolean> for bool {
+    fn try_into_js(self, cx: &CallContext) -> Result<napi::JsBoolean> {
+        cx.env.get_boolean(self)
+    }
+}
+
+macro_rules! into_js_chunked {
+    ($name:ident) => {
+        impl IntoJsRef<JsObject> for $name {
+            fn try_into_js_ref(&self, cx: &CallContext) -> Result<JsObject> {
+                let mut arr = cx.env.create_array()?;
+                for (i, val) in self.into_iter().enumerate() {
+                    match val {
+                        Some(v) => arr.set_element(i as u32, v.into_js(&cx))?,
+                        None => arr.set_element(i as u32, cx.env.get_null()?)?,
+                    };
+                }
+                Ok(arr)
+            }
+        }
+    };
+}
+
+impl IntoJsRef<JsObject> for DatetimeChunked {
+    fn try_into_js_ref(&self, cx: &CallContext) -> Result<JsObject> {
+        let mut arr = cx.env.create_array()?;
+        for (i, val) in self.into_iter().enumerate() {
+            match val {
+                Some(v) => {
+                    arr.set_element(i as u32, cx.env.create_date(v as f64)?)?
+                },
+                None => arr.set_element(i as u32, cx.env.get_null()?)?,
+            };
+        }
+        Ok(arr)
+    }
+}
+impl IntoJsRef<JsObject> for DateChunked {
+    fn try_into_js_ref(&self, cx: &CallContext) -> Result<JsObject> {
+        let mut arr = cx.env.create_array()?;
+        for (i, val) in self.into_iter().enumerate() {
+            match val {
+                Some(v) => {
+                    arr.set_element(i as u32, cx.env.create_date(v as f64)?)?
+                },
+                None => arr.set_element(i as u32, cx.env.get_null()?)?,
+            };
+        }
+        Ok(arr)
+    }
+}
+
+
+
+into_js_chunked!(Int8Chunked);
+into_js_chunked!(Int16Chunked);
+into_js_chunked!(Int32Chunked);
+into_js_chunked!(Int64Chunked);
+into_js_chunked!(UInt8Chunked);
+into_js_chunked!(UInt16Chunked);
+into_js_chunked!(UInt32Chunked);
+into_js_chunked!(UInt64Chunked);
+into_js_chunked!(Float32Chunked);
+into_js_chunked!(Float64Chunked);
+into_js_chunked!(BooleanChunked);
+into_js_chunked!(Utf8Chunked);
+
+
 
 impl IntoJs<JsUnknown> for Wrap<AnyValue<'_>> {
     fn try_into_js(self, cx: &CallContext) -> Result<JsUnknown> {
@@ -110,7 +187,37 @@ impl IntoJs<JsUnknown> for Wrap<AnyValue<'_>> {
             AnyValue::Float64(v) => cx.env.create_double(v).map(|v| v.into_unknown()),
             AnyValue::Date(v) => cx.env.create_date(v as f64).map(|v| v.into_unknown()),
             AnyValue::Datetime(v, _, _) => cx.env.create_date(v as f64).map(|v| v.into_unknown()),
-            AnyValue::List(v) => cx.env.to_js_value(&v).map(|v| v.into_unknown()),
+            AnyValue::List(v) => {
+                let v = match *v.dtype() {
+                    DataType::Null => {
+                        let mut arr = cx.env.create_array()?;
+                        let null_count = v.null_count();
+                        for i in 0..null_count {
+                            arr.set_element(i as u32, cx.env.get_null()?)?;
+                        }
+                        Ok(arr)
+                    }
+                    DataType::Boolean => v.bool().map_err(JsPolarsEr::from)?.try_into_js_ref(&cx),
+                    DataType::Utf8 => v.utf8().map_err(JsPolarsEr::from)?.try_into_js_ref(&cx),
+                    DataType::Float32 => v.f32().map_err(JsPolarsEr::from)?.try_into_js_ref(&cx),
+                    DataType::Float64 => v.f64().map_err(JsPolarsEr::from)?.try_into_js_ref(&cx),
+                    DataType::UInt8 => v.u8().map_err(JsPolarsEr::from)?.try_into_js_ref(&cx),
+                    DataType::UInt16 => v.u16().map_err(JsPolarsEr::from)?.try_into_js_ref(&cx),
+                    DataType::UInt32 => v.u32().map_err(JsPolarsEr::from)?.try_into_js_ref(&cx),
+                    DataType::UInt64 => v.u64().map_err(JsPolarsEr::from)?.try_into_js_ref(&cx),
+                    DataType::Int8 => v.i8().map_err(JsPolarsEr::from)?.try_into_js_ref(&cx),
+                    DataType::Int16 => v.i16().map_err(JsPolarsEr::from)?.try_into_js_ref(&cx),
+                    DataType::Int32 => v.i32().map_err(JsPolarsEr::from)?.try_into_js_ref(&cx),
+                    DataType::Int64 => v.i64().map_err(JsPolarsEr::from)?.try_into_js_ref(&cx),
+                    DataType::Date => v.date().map_err(JsPolarsEr::from)?.try_into_js_ref(&cx),
+                    DataType::Datetime(_, _) => v.datetime().map_err(JsPolarsEr::from)?.try_into_js_ref(&cx),
+
+                    _ => panic!("unsupported"),
+                };
+
+                v.map(|val| val.into_unknown())
+
+            }
             _ => cx.env.get_null().map(|v| v.into_unknown()),
         }
     }

diff --git a/nodejs-polars/src/series.rs b/nodejs-polars/src/series.rs
@@ -157,6 +157,14 @@ pub fn repeat(cx: CallContext) -> JsResult<JsExternal> {
             ca.rename(name);
             ca.into_series()
         }
+        DataType::Datetime(_, _) => {
+            let val = val.extract::<i64>().unwrap();
+            let mut ca: NoNull<Int64Chunked> = (0..n).map(|_| val).collect_trusted();
+            ca.rename(name);
+            ca.clone()
+                .into_datetime(TimeUnit::Milliseconds, None)
+                .into_series()
+        }
         dt => {
             panic!("cannot create repeat with dtype: {:?}", dt);
         }