From cd9a04aaeb12fe821729318754447c13c62bd3ae Mon Sep 17 00:00:00 2001
From: Judah Rand <17158624+judahrand@users.noreply.github.com>
Date: Mon, 4 Mar 2024 16:25:52 +0000
Subject: [PATCH] Add `arrow_typeof` function

---
 datafusion/tests/test_functions.py | 16 ++++++++++++++++
 src/functions.rs                   |  6 ++++++
 2 files changed, 22 insertions(+)

diff --git a/datafusion/tests/test_functions.py b/datafusion/tests/test_functions.py
index 1a18fd82..6e921cea 100644
--- a/datafusion/tests/test_functions.py
+++ b/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)
+            )
diff --git a/src/functions.rs b/src/functions.rs
index 757fb31d..cca250ce 100644
--- a/src/functions.rs
+++ b/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(())
 }