From e5649c97d4bb139197f81aac1b1ed255cf76fa27 Mon Sep 17 00:00:00 2001
From: ongchi <ongchi@users.noreply.github.com>
Date: Tue, 9 Jan 2024 21:08:36 +0800
Subject: [PATCH] Add array_slice and list_slice

---
 datafusion/tests/test_functions.py | 8 ++++++++
 src/functions.rs                   | 4 ++++
 2 files changed, 12 insertions(+)

diff --git a/datafusion/tests/test_functions.py b/datafusion/tests/test_functions.py
index eb717dba..7e772589 100644
--- a/datafusion/tests/test_functions.py
+++ b/datafusion/tests/test_functions.py
@@ -424,6 +424,14 @@ def py_arr_replace(arr, from_, to, n=None):
             f.list_replace_all(col, literal(3.0), literal(4.0)),
             lambda: [py_arr_replace(arr, 3.0, 4.0) for arr in data],
         ],
+        [
+            f.array_slice(col, literal(2), literal(4)),
+            lambda: [arr[1:4] for arr in data],
+        ],
+        [
+            f.list_slice(col, literal(-1), literal(2)),
+            lambda: [arr[-1:2] for arr in data],
+        ],
     ]
 
     for stmt, py_expr in test_items:
diff --git a/src/functions.rs b/src/functions.rs
index c7943752..e3c485a3 100644
--- a/src/functions.rs
+++ b/src/functions.rs
@@ -406,6 +406,8 @@ scalar_function!(array_replace_n, ArrayReplaceN);
 scalar_function!(list_replace_n, ArrayReplaceN);
 scalar_function!(array_replace_all, ArrayReplaceAll);
 scalar_function!(list_replace_all, ArrayReplaceAll);
+scalar_function!(array_slice, ArraySlice);
+scalar_function!(list_slice, ArraySlice);
 
 aggregate_function!(approx_distinct, ApproxDistinct);
 aggregate_function!(approx_median, ApproxMedian);
@@ -646,6 +648,8 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(list_replace_n))?;
     m.add_wrapped(wrap_pyfunction!(array_replace_all))?;
     m.add_wrapped(wrap_pyfunction!(list_replace_all))?;
+    m.add_wrapped(wrap_pyfunction!(array_slice))?;
+    m.add_wrapped(wrap_pyfunction!(list_slice))?;
 
     Ok(())
 }