def testFillNullLists(
self, list_array, value_type, fill_with, expected, list_type_factory):
actual = array_util.FillNullLists(
pa.array(list_array, type=list_type_factory(value_type)),
pa.array(fill_with, type=value_type))
self.assertTrue(
actual.equals(pa.array(expected, type=list_type_factory(value_type))),
"{} vs {}".format(actual, expected))
def testValueTypeDoesNotEqualFillType(self):
with self.assertRaisesRegex(RuntimeError, "to be of the same type"):
array_util.FillNullLists(pa.array([[1]]), pa.array(["a"]))
def testNonListArray(self):
with self.assertRaisesRegex(RuntimeError, "Unimplemented"):
array_util.FillNullLists(pa.array([1, 2, 3]), pa.array([4]))
def testNonListArray(self):
with self.assertRaisesRegex(RuntimeError, "UNIMPLEMENTED"):
array_util.FillNullLists(pa.array([1, 2, 3]), pa.array([4]))
def GetTensor(self, record_batch: pa.RecordBatch,
produce_eager_tensors: bool) -> Union[np.ndarray, tf.Tensor]:
column = record_batch.column(self._column_index)
column = array_util.FillNullLists(column, self._default_fill)
return self._ListArrayToTensor(column, produce_eager_tensors)
def testFillNullLists(self, list_array, fill_with, expected):
actual = array_util.FillNullLists(list_array, fill_with)
self.assertTrue(actual.equals(expected),
"{} vs {}".format(actual, expected))
