Add arrow_typeof function

Author: judahrand

datafusion/tests/test_functions.py

@@ -760,3 +760,19 @@ def test_binary_string_functions(df):
     assert pa.array(result.column(1)).cast(pa.string()) == pa.array(
         ["Hello", "World", "!"]
     )
+
+
+def test_other_functions(df):
+    test_items = [
+        [
+            f.arrow_typeof(column("b").cast(pa.int32())),
+            lambda: [pa.int32()] * 3,
+        ]
+    ]
+
+    for stmt, py_expr in test_items:
+        query_result = df.select(stmt).collect()[0].column(0)
+        for a, b in zip(query_result, py_expr()):
+            np.testing.assert_array_almost_equal(
+                np.array(a.as_py(), dtype=float), np.array(b, dtype=float)
+            )

src/functions.rs

@@ -446,6 +446,9 @@ scalar_function!(array_slice, ArraySlice);
 scalar_function!(list_slice, ArraySlice);
 scalar_function!(flatten, Flatten);
 
+// Other Functions
+scalar_function!(arrow_typeof, ArrowTypeof);
+
 aggregate_function!(approx_distinct, ApproxDistinct);
 aggregate_function!(approx_median, ApproxMedian);
 aggregate_function!(approx_percentile_cont, ApproxPercentileCont);
@@ -689,5 +692,8 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(list_slice))?;
     m.add_wrapped(wrap_pyfunction!(flatten))?;
 
+    // Other Functions
+    m.add_wrapped(wrap_pyfunction!(arrow_typeof))?;
+
     Ok(())
 }