Skip to content

Commit

Permalink
Add arrow_typeof function
Browse files Browse the repository at this point in the history
  • Loading branch information
judahrand committed Mar 4, 2024
1 parent 76ecf56 commit cd9a04a
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 0 deletions.
16 changes: 16 additions & 0 deletions datafusion/tests/test_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
)
6 changes: 6 additions & 0 deletions src/functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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(())
}

0 comments on commit cd9a04a

Please sign in to comment.