def test_numpy_extractor(self):
pa_table = self._create_dummy_table()
extractor = NumpyArrowExtractor()
row = extractor.extract_row(pa_table)
np.testing.assert_equal(row, {"a": _COL_A[0], "b": _COL_B[0], "c": np.array(_COL_C[0])})
col = extractor.extract_column(pa_table)
np.testing.assert_equal(col, np.array(_COL_A))
batch = extractor.extract_batch(pa_table)
np.testing.assert_equal(batch, {"a": np.array(_COL_A), "b": np.array(_COL_B), "c": np.array(_COL_C)})
def test_numpy_extractor_np_array_kwargs(self):
pa_table = self._create_dummy_table().drop(["b"])
extractor = NumpyArrowExtractor(dtype=np.float16)
row = extractor.extract_row(pa_table)
self.assertEqual(row["c"].dtype, np.dtype(np.float16))
col = extractor.extract_column(pa_table)
self.assertEqual(col.dtype, np.float16)
batch = extractor.extract_batch(pa_table)
self.assertEqual(batch["a"].dtype, np.dtype(np.float16))
self.assertEqual(batch["c"].dtype, np.dtype(np.float16))
