chore: replace todo with error (arrow-udf#34)

datafuse-extras · Jun 20, 2024 · dfa1390 · dfa1390
1 parent d7399a4
commit dfa1390
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 26 deletions.
diff --git a/arrow-udf-js-deno/src/deno_arrow.rs b/arrow-udf-js-deno/src/deno_arrow.rs
@@ -744,7 +744,6 @@ impl Converter {
                 }
                 Ok(Arc::new(builder.finish()))
             }
-            DataType::FixedSizeBinary(_) => todo!(),
             DataType::Utf8 => match field
                 .metadata()
                 .get(self.arrow_extension_key.as_str())
@@ -917,8 +916,6 @@ impl Converter {
                     Some(nulls),
                 )))
             }
-            DataType::FixedSizeList(_, _) => todo!(),
-            DataType::LargeList(_) => todo!(),
             DataType::Struct(fields) => {
                 let mut arrays = Vec::with_capacity(fields.len());
                 for field in fields {
@@ -950,13 +947,7 @@ impl Converter {
                     Some(nulls),
                 )))
             }
-            DataType::Union(_, _) => todo!(),
-            DataType::Dictionary(_, _) => todo!(),
-            DataType::Decimal128(_, _) => todo!(),
-            DataType::Decimal256(_, _) => todo!(),
-            DataType::Map(_, _) => todo!(),
-            DataType::RunEndEncoded(_, _) => todo!(),
-            _ => todo!(),
+            other => Err(anyhow::anyhow!("Unimplemented datatype {}", other)),
         }
     }
 
@@ -1207,7 +1198,6 @@ impl Converter {
                     }
                 }
             }
-            DataType::FixedSizeBinary(_) => todo!(),
             DataType::Utf8 => {
                 let array = array.as_any().downcast_ref::<StringArray>().unwrap();
                 let string =
@@ -1290,8 +1280,6 @@ impl Converter {
                     }
                 }
             }
-            DataType::FixedSizeList(_, _) => todo!(),
-            DataType::LargeList(_) => todo!(),
             DataType::Struct(fields) => {
                 let array = array.as_any().downcast_ref::<StructArray>().unwrap();
                 let object = v8::Object::new(scope);
@@ -1306,13 +1294,7 @@ impl Converter {
                 }
                 Ok(object.into())
             }
-            DataType::Union(_, _) => todo!(),
-            DataType::Dictionary(_, _) => todo!(),
-            DataType::Decimal128(_, _) => todo!(),
-            DataType::Decimal256(_, _) => todo!(),
-            DataType::Map(_, _) => todo!(),
-            DataType::RunEndEncoded(_, _) => todo!(),
-            _ => todo!(),
+            other => Err(anyhow::anyhow!("Unimplemented datatype {}", other)),
         }
     }
 }
diff --git a/arrow-udf-js/src/jsarrow.rs b/arrow-udf-js/src/jsarrow.rs
@@ -275,6 +275,30 @@ impl Converter {
                     }
                 }
             }
+            // large list
+            DataType::LargeList(inner) => {
+                let array = array.as_any().downcast_ref::<LargeListArray>().unwrap();
+                let list = array.value(i);
+                match inner.data_type() {
+                    DataType::Int8 => get_typed_array!(Int8Array, ctx, list),
+                    DataType::Int16 => get_typed_array!(Int16Array, ctx, list),
+                    DataType::Int32 => get_typed_array!(Int32Array, ctx, list),
+                    DataType::Int64 => get_typed_array!(Int64Array, ctx, list),
+                    DataType::UInt8 => get_typed_array!(UInt8Array, ctx, list),
+                    DataType::UInt16 => get_typed_array!(UInt16Array, ctx, list),
+                    DataType::UInt32 => get_typed_array!(UInt32Array, ctx, list),
+                    DataType::UInt64 => get_typed_array!(UInt64Array, ctx, list),
+                    DataType::Float32 => get_typed_array!(Float32Array, ctx, list),
+                    DataType::Float64 => get_typed_array!(Float64Array, ctx, list),
+                    _ => {
+                        let mut values = Vec::with_capacity(list.len());
+                        for j in 0..list.len() {
+                            values.push(self.get_jsvalue(ctx, field, list.as_ref(), j)?);
+                        }
+                        values.into_js(ctx)
+                    }
+                }
+            }
             DataType::Struct(fields) => {
                 let array = array.as_any().downcast_ref::<StructArray>().unwrap();
                 let object = Object::new(ctx.clone())?;
@@ -284,7 +308,7 @@ impl Converter {
                 }
                 Ok(object.into_value())
             }
-            t => todo!("unsupported data type: {:?}", t),
+            _other => Err(Error::Unknown),
         }
     }
 
@@ -470,7 +494,7 @@ impl Converter {
                     Some(nulls),
                 )))
             }
-            t => todo!("unsupported data type: {:?}", t),
+            other => Err(anyhow::anyhow!("Unimplemented datatype {}", other)),
         }
     }
 

diff --git a/arrow-udf-python/src/pyarrow.rs b/arrow-udf-python/src/pyarrow.rs
@@ -17,7 +17,7 @@
 use arrow_array::{array::*, builder::*};
 use arrow_buffer::OffsetBuffer;
 use arrow_schema::{DataType, Field};
-use pyo3::{types::PyAnyMethods, IntoPy, PyObject, PyResult, Python};
+use pyo3::{exceptions::PyTypeError, types::PyAnyMethods, IntoPy, PyObject, PyResult, Python};
 use std::{borrow::Cow, sync::Arc};
 
 macro_rules! get_pyobject {
@@ -185,7 +185,12 @@ impl Converter {
                 }
                 object.into()
             }
-            _ => todo!(),
+            other => {
+                return Err(PyTypeError::new_err(format!(
+                    "Unimplemented datatype {}",
+                    other
+                )))
+            }
         })
     }
 
@@ -278,6 +283,31 @@ impl Converter {
                     Some(nulls),
                 )))
             }
+            // large list
+            DataType::LargeList(inner) => {
+                // flatten lists
+                let mut flatten_values = vec![];
+                let mut offsets = Vec::<i64>::with_capacity(values.len() + 1);
+                offsets.push(0);
+                for val in values {
+                    if !val.is_none(py) {
+                        let array = val.bind(py);
+                        flatten_values.reserve(array.len()?);
+                        for elem in array.iter()? {
+                            flatten_values.push(elem?.into());
+                        }
+                    }
+                    offsets.push(flatten_values.len() as i64);
+                }
+                let values_array = self.build_array(inner, py, &flatten_values)?;
+                let nulls = values.iter().map(|v| !v.is_none(py)).collect();
+                Ok(Arc::new(LargeListArray::new(
+                    inner.clone(),
+                    OffsetBuffer::new(offsets.into()),
+                    values_array,
+                    Some(nulls),
+                )))
+            }
             DataType::Struct(fields) => {
                 let mut arrays = Vec::with_capacity(fields.len());
                 for field in fields {
@@ -301,7 +331,10 @@ impl Converter {
                     Some(nulls),
                 )))
             }
-            _ => todo!(),
+            other => Err(PyTypeError::new_err(format!(
+                "Unimplemented datatype {}",
+                other
+            ))),
         }
     }
 }
diff --git a/arrow-udf/src/ffi.rs b/arrow-udf/src/ffi.rs
@@ -181,7 +181,7 @@ fn call_table(function: TableFunction, input_bytes: &[u8]) -> Result<Box<RecordB
 
     let input = Box::new(input_batch);
     // SAFETY: The lifetime of `input` is longer than `iter`.
-    let input_ref = unsafe { std::mem::transmute(input.as_ref()) };
+    let input_ref: &RecordBatch = unsafe { std::mem::transmute(input.as_ref()) };
     let iter = function(input_ref)?;
     Ok(Box::new(RecordBatchIter {
         _input: input,