def test_read_tensor(tmpdir): # Create and write tensor tensor data = np.random.randn(10, 4) tensor = pa.Tensor.from_numpy(data) data_size = pa.get_tensor_size(tensor) path = os.path.join(str(tmpdir), 'pyarrow-tensor-ipc-read-tensor') write_mmap = pa.create_memory_map(path, data_size) pa.write_tensor(tensor, write_mmap) # Try to read tensor read_mmap = pa.memory_map(path, mode='r') array = pa.read_tensor(read_mmap).to_numpy() np.testing.assert_equal(data, array)
def test_store_arrow_objects(self): data = np.random.randn(10, 4) # Write an arrow object. object_id = random_object_id() tensor = pa.Tensor.from_numpy(data) data_size = pa.get_tensor_size(tensor) buf = self.plasma_client.create(object_id, data_size) stream = pa.FixedSizeBufferWriter(buf) pa.write_tensor(tensor, stream) self.plasma_client.seal(object_id) # Read the arrow object. [tensor] = self.plasma_client.get_buffers([object_id]) reader = pa.BufferReader(tensor) array = pa.read_tensor(reader).to_numpy() # Assert that they are equal. np.testing.assert_equal(data, array)
def test_tensor_size(): data = np.random.randn(10, 4) tensor = pa.Tensor.from_numpy(data) assert pa.get_tensor_size(tensor) > (data.size * 8)