Added errors for pyarrow < 2, support for TokenSpan null values, additional testing

Author: BryanCutler

text_extensions_for_pandas/array/arrow_conversion.py

@@ -20,12 +20,13 @@
 #
 # Provide Arrow compatible classes for serializing to pyarrow.
 #
+from distutils.version import LooseVersion
 
 import numpy as np
 import pyarrow as pa
 
 from text_extensions_for_pandas.array.span import SpanArray
-from text_extensions_for_pandas.array.token_span import TokenSpanArray
+from text_extensions_for_pandas.array.token_span import TokenSpanArray, _EMPTY_SPAN_ARRAY_SINGLETON
 from text_extensions_for_pandas.array.tensor import TensorArray
 from text_extensions_for_pandas.array.string_table import StringTable
 
@@ -108,6 +109,9 @@ def span_to_arrow(char_span: SpanArray) -> pa.ExtensionArray:
     :param char_span: A SpanArray to be converted
     :return: pyarrow.ExtensionArray containing Span data
     """
+    if LooseVersion(pa.__version__) < LooseVersion("2.0.0"):
+        raise NotImplementedError("Arrow serialization for SpanArray is not supported with "
+                                  "PyArrow versions < 2.0.0")
     # Create array for begins, ends
     begins_array = pa.array(char_span.begin)
     ends_array = pa.array(char_span.end)
@@ -130,9 +134,14 @@ def arrow_to_span(extension_array: pa.ExtensionArray) -> SpanArray:
     Convert a pyarrow.ExtensionArray with type ArrowSpanType to
     a SpanArray.
 
+    ..NOTE: Only supported with PyArrow >= 2.0.0
+
     :param extension_array: pyarrow.ExtensionArray with type ArrowSpanType
     :return: SpanArray
     """
+    if LooseVersion(pa.__version__) < LooseVersion("2.0.0"):
+        raise NotImplementedError("Arrow serialization for SpanArray is not supported with "
+                                  "PyArrow versions < 2.0.0")
     if isinstance(extension_array, pa.ChunkedArray):
         if extension_array.num_chunks > 1:
             raise ValueError("Only pyarrow.Array with a single chunk is supported")
@@ -175,18 +184,28 @@ def token_span_to_arrow(token_span: TokenSpanArray) -> pa.ExtensionArray:
     :param token_span: A TokenSpanArray to be converted
     :return: pyarrow.ExtensionArray containing TokenSpan data
     """
+    if LooseVersion(pa.__version__) < LooseVersion("2.0.0"):
+        raise NotImplementedError("Arrow serialization for TokenSpanArray is not supported with "
+                                  "PyArrow versions < 2.0.0")
     # Create arrays for begins/ends
     token_begins_array = pa.array(token_span.begin_token)
     token_ends_array = pa.array(token_span.end_token)
 
+    # Filter out any empty SpanArrays
+    non_null_tokens = token_span.tokens[~token_span.isna()]
+    assert len(non_null_tokens) > 0
+
     # Get either single document as a list or use a list of all if multiple docs
-    assert len(token_span.tokens) > 0
-    if all([token is token_span.tokens[0] for token in token_span.tokens]):
-        tokens_arrays = [token_span.tokens[0]]
-        tokens_indices = pa.array([0] * len(token_span.tokens))
+    if all([token is non_null_tokens[0] for token in non_null_tokens]):
+        tokens_arrays = [non_null_tokens[0]]
+        tokens_indices = pa.array([0] * len(token_span.tokens), mask=token_span.isna())
     else:
-        tokens_arrays = token_span.tokens
-        tokens_indices = pa.array(range(len(tokens_arrays)))
+        raise NotImplementedError("TokenSpan Multi-doc serialization not yet implemented due to "
+                                  "ArrowNotImplementedError: Concat with dictionary unification NYI")
+        tokens_arrays = non_null_tokens
+        tokens_indices = np.zeros_like(token_span.tokens)
+        tokens_indices[~token_span.isna()] = range(len(tokens_arrays))
+        tokens_indices = pa.array(tokens_indices, mask=token_span.isna())
 
     # Convert each token SpanArray to Arrow and get as raw storage
     arrow_tokens_arrays = [span_to_arrow(sa).storage for sa in tokens_arrays]
@@ -217,6 +236,9 @@ def arrow_to_token_span(extension_array: pa.ExtensionArray) -> TokenSpanArray:
     :param extension_array: pyarrow.ExtensionArray with type ArrowTokenSpanType
     :return: TokenSpanArray
     """
+    if LooseVersion(pa.__version__) < LooseVersion("2.0.0"):
+        raise NotImplementedError("Arrow serialization for TokenSpanArray is not supported with "
+                                  "PyArrow versions < 2.0.0")
     if isinstance(extension_array, pa.ChunkedArray):
         if extension_array.num_chunks > 1:
             raise ValueError("Only pyarrow.Array with a single chunk is supported")
@@ -252,7 +274,8 @@ def arrow_to_token_span(extension_array: pa.ExtensionArray) -> TokenSpanArray:
         tokens_arrays.append(tokens_array)
 
     # Map the token indices to the actual token SpanArray for each element in the TokenSpanArray
-    tokens = [tokens_arrays[i.as_py()] for i in tokens_indices]
+    tokens = [_EMPTY_SPAN_ARRAY_SINGLETON if i is None else tokens_arrays[i]
+              for i in tokens_indices.to_pylist()]
 
     # Zero-copy convert arrays to numpy
     token_begins = token_begins_array.to_numpy()

text_extensions_for_pandas/array/test_token_span.py

@@ -14,13 +14,15 @@
 #
 
 import pandas as pd
+from distutils.version import LooseVersion
 import os
 import tempfile
 import unittest
 # noinspection PyPackageRequirements
 import pytest
 
 from pandas.tests.extension import base
+import pyarrow as pa
 
 from text_extensions_for_pandas.array.test_span import ArrayTestBase
 from text_extensions_for_pandas.array.span import *
@@ -365,6 +367,8 @@ def test_as_frame(self):
         self.assertEqual(len(df), len(arr))
 
 
+@pytest.mark.skipif(LooseVersion(pa.__version__) < LooseVersion("2.0.0"),
+                    reason="Nested dictionaries only supported in Arrow >= 2.0.0")
 class TokenSpanArrayIOTests(ArrayTestBase):
 
     def do_roundtrip(self, df):
@@ -383,7 +387,7 @@ def test_feather(self):
         self.do_roundtrip(df1)
 
         # More token spans than tokens
-        """ts2 = TokenSpanArray(toks, [0, 1, 2, 3, 0, 2, 0], [1, 2, 3, 4, 2, 4, 4])
+        ts2 = TokenSpanArray(toks, [0, 1, 2, 3, 0, 2, 0], [1, 2, 3, 4, 2, 4, 4])
         df2 = pd.DataFrame({"ts2": ts2})
         self.do_roundtrip(df2)
 
@@ -404,7 +408,35 @@ def test_feather(self):
 
         # All columns together, TokenSpan arrays padded as needed
         df = pd.concat([df1, df2, df3, df4], axis=1)
-        self.do_roundtrip(df)"""
+        self.do_roundtrip(df)
+
+    @pytest.mark.skip(reason="ArrowNotImplementedError: Concat with dictionary unification NYI")
+    def test_feather_multi_doc(self):
+        toks = self._make_spans_of_tokens()
+        arr = TokenSpanArray(toks, np.arange(len(toks)), np.arange(len(toks)) + 1)
+        df1 = pd.DataFrame({'TokenSpan': arr})
+
+        toks = SpanArray(
+            "Have at it.", np.array([0, 5, 8]), np.array([4, 7, 11])
+        )
+        arr = TokenSpanArray(toks, np.arange(len(toks)), np.arange(len(toks)) + 1)
+        df2 = pd.DataFrame({'TokenSpan': arr})
+
+        df = pd.concat([df1, df2], ignore_index=True)
+        self.assertFalse(df["TokenSpan"].array.is_single_document)
+        self.do_roundtrip(df)
+
+    @pytest.mark.skip(reason="ArrowNotImplementedError: Writing DictionaryArray with nested dictionary type not yet supported")
+    def test_parquet(self):
+        toks = self._make_spans_of_tokens()
+        arr = TokenSpanArray(toks, np.arange(len(toks)), np.arange(len(toks)) + 1)
+        df = pd.DataFrame({'TokenSpan': arr})
+
+        with tempfile.TemporaryDirectory() as dirpath:
+            filename = os.path.join(dirpath, "token_span_array_test.parquet")
+            df.to_parquet(filename)
+            df_read = pd.read_parquet(filename)
+            pd.testing.assert_frame_equal(df, df_read)
 
 
 @pytest.fixture