take

google/cloud/bigquery/dtypes.py

@@ -14,10 +14,12 @@
 
 import datetime
 import operator
+from typing import Any, Sequence
 
 import numpy
 import packaging.version
 from pandas._libs import NaT
+import pandas.core.algorithms
 import pandas.core.dtypes.base
 import pandas.core.dtypes.dtypes
 import pandas.core.dtypes.generic
@@ -114,7 +116,6 @@ def copy(self):
     def repeat(self, n):
         return self.__class__(self._ndarray.repeat(n), self._dtype)
 
-
 #
 ###########################################################################
 
@@ -180,6 +181,40 @@ def _from_factorized(self, unique, original):
     def isna(self):
         return pandas.isna(self._ndarray)
 
+    def _validate_scalar(self, value):
+        if pandas.isna(value):
+            return None
+
+        if not isinstance(value, self.dtype.type):
+            raise ValueError(value)
+
+        return value
+
+    def take(
+        self,
+        indices: Sequence[int],
+        *,
+        allow_fill: bool = False,
+        fill_value: Any = None,
+    ):
+        indices = numpy.asarray(indices, dtype=numpy.intp)
+        data = self._ndarray
+        if allow_fill:
+            fill_value = self._validate_scalar(fill_value)
+            fill_value = (
+                numpy.datetime64() if fill_value is None
+                else numpy.datetime64(self._datetime(fill_value))
+            )
+            if (indices < -1).any():
+                raise ValueError(
+                    "take called with negative indexes other than -1,"
+                    " when a fill value is provided.")
+        out = data.take(indices)
+        if allow_fill:
+            out[indices == -1] = fill_value
+
+        return self.__class__(out)
+
 
 @pandas.core.dtypes.dtypes.register_extension_dtype
 class TimeDtype(_BaseDtype):

tests/unit/test_dtypes.py

@@ -307,3 +307,48 @@ def test__from_factorized(dtype):
 def test_isna(dtype):
     a = _make_one(dtype)
     assert list(a.isna()) == [False, False, True]
+
+
+@for_date_and_time
+def test__validate_scalar_invalid(dtype):
+    with pytest.raises(ValueError):
+        _make_one(dtype)._validate_scalar("bad")
+
+
+@for_date_and_time
+@pytest.mark.parametrize(
+    "allow_fill, fill_value",
+    [
+        (False, None),
+        (True, None),
+        (True, pd._libs.NaT if pd else None),
+        (True, np.NaN if pd else None), (True, 42)])
+def test_take(dtype, allow_fill, fill_value):
+    sample_values = SAMPLE_VALUES[dtype]
+    a = _cls(dtype)(sample_values)
+    if allow_fill:
+        if fill_value == 42:
+            fill_value = expected_fill = (
+                datetime.date(1971, 4, 2) if dtype == 'date'
+                else datetime.time(0, 42, 42, 424242))
+        else:
+            expected_fill = None
+        b = a.take([1, -1, 3], allow_fill=True, fill_value=fill_value)
+        expect = [sample_values[1], expected_fill, sample_values[3]]
+    else:
+        b = a.take([1, -4, 3])
+        expect = [sample_values[1], sample_values[-4], sample_values[3]]
+
+    assert list(b) == expect
+
+
+@for_date_and_time
+def test_take_bad_index(dtype):
+    # When allow_fill is set, negative indexes < -1 raise ValueError.
+    # This is based on testing with an integer series/array.
+    # The documentation isn't clear on this at all.
+    sample_values = SAMPLE_VALUES[dtype]
+    a = _cls(dtype)(sample_values)
+    with pytest.raises(ValueError):
+        a.take([1, -2, 3], allow_fill=True, fill_value=None)
+