feat: add support for pyarrow arrays as input

bindings/python/Cargo.toml

@@ -10,7 +10,7 @@ crate-type = ["cdylib"]
 
 [dependencies]
 rayon = "1.10"
-serde = { version = "1.0", features = [ "rc", "derive" ]}
+serde = { version = "1.0", features = ["rc", "derive"] }
 serde_json = "1.0"
 libc = "0.2"
 env_logger = "0.11"
@@ -19,6 +19,9 @@ numpy = "0.21"
 ndarray = "0.15"
 onig = { version = "6.4", default-features = false }
 itertools = "0.12"
+arrow = { git = "https://github.com/apache/arrow-rs", branch = "master", features = [
+    "pyarrow",
+], optional = true }
 
 [dependencies.tokenizers]
 path = "../../tokenizers"
@@ -29,3 +32,4 @@ pyo3 = { version = "0.21", features = ["auto-initialize"] }
 
 [features]
 defaut = ["pyo3/extension-module"]
+pyarrow = ["arrow"]

bindings/python/pyproject.toml

@@ -47,7 +47,7 @@ build-backend = "maturin"
 python-source = "py_src"
 module-name = "tokenizers.tokenizers"
 bindings = 'pyo3'
-features = ["pyo3/extension-module"]
+features = ["pyo3/extension-module", "pyarrow"]
 
 [tool.black]
 line-length = 119

bindings/python/src/tokenizer.rs

@@ -1,6 +1,11 @@
 use std::collections::{hash_map::DefaultHasher, HashMap};
 use std::hash::{Hash, Hasher};
 
+#[cfg(feature = "pyarrow")]
+use std::borrow::Cow;
+#[cfg(feature = "pyarrow")]
+use std::slice;
+
 use numpy::{npyffi, PyArray1};
 use pyo3::class::basic::CompareOp;
 use pyo3::exceptions;
@@ -260,13 +265,48 @@ impl PyAddedToken {
     }
 }
 
+#[cfg(feature = "pyarrow")]
+struct PyArrowScalarStringInput<'s>(Cow<'s, str>);
+#[cfg(feature = "pyarrow")]
+impl<'s> FromPyObject<'s> for PyArrowScalarStringInput<'s> {
+    fn extract(ob: &'s PyAny) -> PyResult<Self> {
+        let pyarrow = PyModule::import_bound(ob.py(), "pyarrow").map(Bound::into_gil_ref)?;
+        let str_scalar_class = pyarrow.getattr("StringScalar")?;
+        let large_str_scalar_class = pyarrow.getattr("LargeStringScalar")?;
+        if ob.is_exact_instance(str_scalar_class) || ob.is_exact_instance(large_str_scalar_class) {
+            let buf = ob.call_method0("as_buffer")?;
+            let addr = buf.getattr("address")?.extract::<usize>()?;
+            let size = buf.getattr("size")?.extract::<usize>()?;
+
+            // SAFETY address is valid because it's from the StringScalar buffer
+            let buf_slice = unsafe { slice::from_raw_parts::<'s>(addr as *const u8, size) };
+            let x = String::from_utf8_lossy(buf_slice);
+            Ok(Self(x))
+        } else {
+            let err =
+                exceptions::PyTypeError::new_err("TextInputSequence must be pyarrow.StringScalar");
+            Err(err)
+        }
+    }
+}
+#[cfg(feature = "pyarrow")]
+impl<'s> From<PyArrowScalarStringInput<'s>> for tk::InputSequence<'s> {
+    fn from(s: PyArrowScalarStringInput<'s>) -> Self {
+        s.0.into()
+    }
+}
+
 struct TextInputSequence<'s>(tk::InputSequence<'s>);
 impl<'s> FromPyObject<'s> for TextInputSequence<'s> {
     fn extract(ob: &'s PyAny) -> PyResult<Self> {
-        let err = exceptions::PyTypeError::new_err("TextInputSequence must be str");
         if let Ok(s) = ob.downcast::<PyString>() {
             Ok(Self(s.to_string_lossy().into()))
         } else {
+            #[cfg(feature = "pyarrow")]
+            if let Ok(s) = ob.extract::<PyArrowScalarStringInput>() {
+                return Ok(Self(s.0.into()));
+            }
+            let err = exceptions::PyTypeError::new_err("TextInputSequence must be str");
             Err(err)
         }
     }
@@ -344,6 +384,26 @@ impl From<PyArrayUnicode> for tk::InputSequence<'_> {
     }
 }
 
+#[cfg(feature = "pyarrow")]
+struct PyArrowArray<'s>(Vec<Cow<'s, str>>);
+#[cfg(feature = "pyarrow")]
+impl<'s> FromPyObject<'s> for PyArrowArray<'s> {
+    fn extract(ob: &'s PyAny) -> PyResult<Self> {
+        let array = ob.extract::<Vec<&PyAny>>()?;
+        let str_array: Vec<Cow<'s, str>> = array
+            .iter()
+            .map(|item| item.extract::<PyArrowScalarStringInput>().map(|res| res.0))
+            .collect::<PyResult<Vec<_>>>()?;
+        Ok(Self(str_array))
+    }
+}
+#[cfg(feature = "pyarrow")]
+impl<'s> From<PyArrowArray<'s>> for tk::InputSequence<'s> {
+    fn from(s: PyArrowArray<'s>) -> Self {
+        s.0.into()
+    }
+}
+
 struct PyArrayStr(Vec<String>);
 impl FromPyObject<'_> for PyArrayStr {
     fn extract(ob: &PyAny) -> PyResult<Self> {
@@ -373,6 +433,10 @@ impl<'s> FromPyObject<'s> for PreTokenizedInputSequence<'s> {
         if let Ok(seq) = ob.extract::<PyArrayUnicode>() {
             return Ok(Self(seq.into()));
         }
+        #[cfg(feature = "pyarrow")]
+        if let Ok(seq) = ob.extract::<PyArrowArray>() {
+            return Ok(Self(seq.into()));
+        }
         if let Ok(seq) = ob.extract::<PyArrayStr>() {
             return Ok(Self(seq.into()));
         }

bindings/python/tests/bindings/test_tokenizer.py

@@ -1,6 +1,7 @@
 import pickle
 
 import numpy as np
+import pyarrow as pa
 import pytest
 
 from tokenizers import AddedToken, Encoding, Tokenizer
@@ -200,6 +201,11 @@ def test_pair(input, is_pretokenized=False):
         test_pair(np.array([("My name is John", "pair"), ("My name is Georges", "pair")]))
         test_pair(np.array([["My name is John", "pair"], ["My name is Georges", "pair"]]))
 
+        # Pyarrow
+        test_single(pa.array(["My name is John", "My name is Georges"]))
+        test_pair(pa.array([("My name is John", "pair"), ("My name is Georges", "pair")]))
+        test_pair(pa.array([["My name is John", "pair"], ["My name is Georges", "pair"]]))
+
         # PreTokenized inputs
 
         # Lists
@@ -266,6 +272,34 @@ def test_pair(input, is_pretokenized=False):
             True,
         )
 
+        # Pyarrow
+        test_single(
+            pa.array([["My", "name", "is", "John"], ["My", "name", "is", "Georges"]]),
+            True,
+        )
+        test_single(
+            pa.array((("My", "name", "is", "John"), ("My", "name", "is", "Georges"))),
+            True,
+        )
+        test_pair(
+            pa.array(
+                [
+                    [["My", "name", "is", "John"], ["pair"]],
+                    [["My", "name", "is", "Georges"], ["pair"]],
+                ],
+            ),
+            True,
+        )
+        test_pair(
+            pa.array(
+                (
+                    (("My", "name", "is", "John"), ("pair",)),
+                    (("My", "name", "is", "Georges"), ("pair",)),
+                ),
+            ),
+            True,
+        )
+
         # Mal formed
         with pytest.raises(TypeError, match="TextInputSequence must be str"):
             tokenizer.encode([["my", "name"]])