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))