Skip to content

Commit b1e56a2

Browse files
pcmoritzwesm
pcmoritz
authored and
wesm
committed
ARROW-1453: [C++/Python] Support non-contiguous Tensors in WriteTensor
Author: Philipp Moritz <pcmoritz@gmail.com> Author: Wes McKinney <wes.mckinney@twosigma.com> Closes apache#1034 from pcmoritz/write-tensor and squashes the following commits: 7b2ae97 [Wes McKinney] Fix off by one error 084b6b9 [Wes McKinney] Remove outdated comment f7be231 [Wes McKinney] Copy strided tensor data into a scratch buffer before writing to output stream e4be4af [Philipp Moritz] fix flake8 linting 3576840 [Philipp Moritz] fixes and more tests fc44321 [Philipp Moritz] fix python test 57c7ace [Philipp Moritz] fix linting 23d991f [Philipp Moritz] add fixes and tests 357aa73 [Philipp Moritz] fix linting 0e22ebf [Philipp Moritz] implement writing strided tensors
1 parent 0e0da74 commit b1e56a2

5 files changed

Lines changed: 132 additions & 38 deletions

File tree

‎cpp/src/arrow/compare.cc‎

Lines changed: 40 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -694,6 +694,29 @@ Status ArrayRangeEquals(const Array& left, const Array& right, int64_t left_star
694694
return Status::OK();
695695
}
696696

697+
bool StridedTensorContentEquals(int dim_index, int64_t left_offset, int64_t right_offset,
698+
int elem_size, const Tensor& left, const Tensor& right) {
699+
if (dim_index == left.ndim() - 1) {
700+
for (int64_t i = 0; i < left.shape()[dim_index]; ++i) {
701+
if (memcmp(left.raw_data() + left_offset + i * left.strides()[dim_index],
702+
right.raw_data() + right_offset + i * right.strides()[dim_index],
703+
elem_size) != 0) {
704+
return false;
705+
}
706+
}
707+
return true;
708+
}
709+
for (int64_t i = 0; i < left.shape()[dim_index]; ++i) {
710+
if (!StridedTensorContentEquals(dim_index + 1, left_offset, right_offset, elem_size,
711+
left, right)) {
712+
return false;
713+
}
714+
left_offset += left.strides()[dim_index];
715+
right_offset += right.strides()[dim_index];
716+
}
717+
return true;
718+
}
719+
697720
Status TensorEquals(const Tensor& left, const Tensor& right, bool* are_equal) {
698721
// The arrays are the same object
699722
if (&left == &right) {
@@ -704,19 +727,25 @@ Status TensorEquals(const Tensor& left, const Tensor& right, bool* are_equal) {
704727
*are_equal = true;
705728
} else {
706729
if (!left.is_contiguous() || !right.is_contiguous()) {
707-
return Status::NotImplemented(
708-
"Comparison not implemented for non-contiguous tensors");
709-
}
710-
711-
const auto& size_meta = dynamic_cast<const FixedWidthType&>(*left.type());
712-
const int byte_width = size_meta.bit_width() / CHAR_BIT;
713-
DCHECK_GT(byte_width, 0);
730+
const auto& shape = left.shape();
731+
if (shape != right.shape()) {
732+
*are_equal = false;
733+
return Status::OK();
734+
}
735+
const auto& type = static_cast<const FixedWidthType&>(*left.type());
736+
*are_equal = StridedTensorContentEquals(0, 0, 0, type.bit_width() / 8, left, right);
737+
return Status::OK();
738+
} else {
739+
const auto& size_meta = dynamic_cast<const FixedWidthType&>(*left.type());
740+
const int byte_width = size_meta.bit_width() / CHAR_BIT;
741+
DCHECK_GT(byte_width, 0);
714742

715-
const uint8_t* left_data = left.data()->data();
716-
const uint8_t* right_data = right.data()->data();
743+
const uint8_t* left_data = left.data()->data();
744+
const uint8_t* right_data = right.data()->data();
717745

718-
*are_equal =
719-
memcmp(left_data, right_data, static_cast<size_t>(byte_width * left.size())) == 0;
746+
*are_equal = memcmp(left_data, right_data,
747+
static_cast<size_t>(byte_width * left.size())) == 0;
748+
}
720749
}
721750
return Status::OK();
722751
}

‎cpp/src/arrow/ipc/ipc-read-write-test.cc‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -739,11 +739,7 @@ TEST_F(TestTensorRoundTrip, NonContiguous) {
739739
auto data = test::GetBufferFromVector(values);
740740
Tensor tensor(int64(), data, {4, 3}, {48, 16});
741741

742-
int32_t metadata_length;
743-
int64_t body_length;
744-
ASSERT_OK(mmap_->Seek(0));
745-
ASSERT_RAISES(Invalid,
746-
WriteTensor(tensor, mmap_.get(), &metadata_length, &body_length));
742+
CheckTensorRoundTrip(tensor);
747743
}
748744

749745
} // namespace ipc

‎cpp/src/arrow/ipc/writer.cc‎

Lines changed: 50 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -563,23 +563,62 @@ Status WriteLargeRecordBatch(const RecordBatch& batch, int64_t buffer_start_offs
563563
pool, kMaxNestingDepth, true);
564564
}
565565

566-
Status WriteTensor(const Tensor& tensor, io::OutputStream* dst, int32_t* metadata_length,
567-
int64_t* body_length) {
568-
if (!tensor.is_contiguous()) {
569-
return Status::Invalid("No support yet for writing non-contiguous tensors");
566+
static Status WriteStridedTensorData(int dim_index, int64_t offset, int elem_size,
567+
const Tensor& tensor, uint8_t* scratch_space,
568+
io::OutputStream* dst) {
569+
if (dim_index == tensor.ndim() - 1) {
570+
const uint8_t* data_ptr = tensor.raw_data() + offset;
571+
const int64_t stride = tensor.strides()[dim_index];
572+
for (int64_t i = 0; i < tensor.shape()[dim_index]; ++i) {
573+
memcpy(scratch_space + i * elem_size, data_ptr, elem_size);
574+
data_ptr += stride;
575+
}
576+
return dst->Write(scratch_space, elem_size * tensor.shape()[dim_index]);
577+
}
578+
for (int64_t i = 0; i < tensor.shape()[dim_index]; ++i) {
579+
RETURN_NOT_OK(WriteStridedTensorData(dim_index + 1, offset, elem_size, tensor,
580+
scratch_space, dst));
581+
offset += tensor.strides()[dim_index];
570582
}
583+
return Status::OK();
584+
}
571585

586+
Status WriteTensorHeader(const Tensor& tensor, io::OutputStream* dst,
587+
int32_t* metadata_length, int64_t* body_length) {
572588
RETURN_NOT_OK(AlignStreamPosition(dst));
573589
std::shared_ptr<Buffer> metadata;
574590
RETURN_NOT_OK(WriteTensorMessage(tensor, 0, &metadata));
575-
RETURN_NOT_OK(WriteMessage(*metadata, dst, metadata_length));
576-
auto data = tensor.data();
577-
if (data) {
578-
*body_length = data->size();
579-
return dst->Write(data->data(), *body_length);
591+
return WriteMessage(*metadata, dst, metadata_length);
592+
}
593+
594+
Status WriteTensor(const Tensor& tensor, io::OutputStream* dst, int32_t* metadata_length,
595+
int64_t* body_length) {
596+
if (tensor.is_contiguous()) {
597+
RETURN_NOT_OK(WriteTensorHeader(tensor, dst, metadata_length, body_length));
598+
auto data = tensor.data();
599+
if (data) {
600+
*body_length = data->size();
601+
return dst->Write(data->data(), *body_length);
602+
} else {
603+
*body_length = 0;
604+
return Status::OK();
605+
}
580606
} else {
581-
*body_length = 0;
582-
return Status::OK();
607+
Tensor dummy(tensor.type(), tensor.data(), tensor.shape());
608+
const auto& type = static_cast<const FixedWidthType&>(*tensor.type());
609+
RETURN_NOT_OK(WriteTensorHeader(dummy, dst, metadata_length, body_length));
610+
611+
const int elem_size = type.bit_width() / 8;
612+
613+
// TODO(wesm): Do we care enough about this temporary allocation to pass in
614+
// a MemoryPool to this function?
615+
std::shared_ptr<Buffer> scratch_space;
616+
RETURN_NOT_OK(AllocateBuffer(default_memory_pool(),
617+
tensor.shape()[tensor.ndim() - 1] * elem_size,
618+
&scratch_space));
619+
620+
return WriteStridedTensorData(0, 0, elem_size, tensor, scratch_space->mutable_data(),
621+
dst);
583622
}
584623
}
585624

‎python/pyarrow/tests/test_serialization.py‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -93,14 +93,17 @@ def assert_equal(obj1, obj2):
9393
np.uint8(3), np.uint32(4), np.uint64(5), np.float32(1.9),
9494
np.float64(1.9), np.zeros([100, 100]),
9595
np.random.normal(size=[100, 100]), np.array(["hi", 3]),
96-
np.array(["hi", 3], dtype=object)]
96+
np.array(["hi", 3], dtype=object),
97+
np.random.normal(size=[45, 22]).T]
98+
9799

98100
if sys.version_info >= (3, 0):
99101
PRIMITIVE_OBJECTS += [0, np.array([["hi", u"hi"], [1.3, 1]])]
100102
else:
101-
PRIMITIVE_OBJECTS += [long(42), long(1 << 62), long(0), \
102-
np.array([["hi", u"hi"], \
103-
[1.3, long(1)]])] # noqa: E501,F821
103+
PRIMITIVE_OBJECTS += [long(42), long(1 << 62), long(0),
104+
np.array([["hi", u"hi"],
105+
[1.3, long(1)]])] # noqa: E501,F821
106+
104107

105108
COMPLEX_OBJECTS = [
106109
[[[[[[[[[[[[]]]]]]]]]]]],
@@ -261,13 +264,15 @@ def test_custom_serialization(large_memory_map):
261264
for obj in CUSTOM_OBJECTS:
262265
serialization_roundtrip(obj, mmap)
263266

267+
264268
def test_numpy_serialization(large_memory_map):
265269
with pa.memory_map(large_memory_map, mode="r+") as mmap:
266270
for t in ["int8", "uint8", "int16", "uint16",
267271
"int32", "uint32", "float32", "float64"]:
268272
obj = np.random.randint(0, 10, size=(100, 100)).astype(t)
269273
serialization_roundtrip(obj, mmap)
270274

275+
271276
def test_numpy_immutable(large_memory_map):
272277
with pa.memory_map(large_memory_map, mode="r+") as mmap:
273278
obj = np.zeros([10])
@@ -278,9 +283,12 @@ def test_numpy_immutable(large_memory_map):
278283
with pytest.raises(ValueError):
279284
result[0] = 1.0
280285

286+
281287
@pytest.mark.skip(reason="extensive memory requirements")
282288
def test_arrow_limits(self):
283-
huge_memory_map = lambda temp_dir: large_memory_map(temp_dir, 100 * 1024 * 1024 * 1024)
289+
def huge_memory_map(temp_dir):
290+
return large_memory_map(temp_dir, 100 * 1024 * 1024 * 1024)
291+
284292
with pa.memory_map(huge_memory_map, mode="r+") as mmap:
285293
# Test that objects that are too large for Arrow throw a Python
286294
# exception. These tests give out of memory errors on Travis and need
@@ -294,8 +302,8 @@ def test_arrow_limits(self):
294302
l = 2 ** 29 * [["1"], 2, 3, [{"s": 4}]]
295303
serialization_roundtrip(l, mmap)
296304
del l
297-
serialization_roundtrip(l, mmap)
298305
l = 2 ** 29 * [{"s": 1}] + 2 ** 29 * [1.0]
306+
serialization_roundtrip(l, mmap)
299307
del l
300308
l = np.zeros(2 ** 25)
301309
serialization_roundtrip(l, mmap)
@@ -304,6 +312,7 @@ def test_arrow_limits(self):
304312
serialization_roundtrip(l, mmap)
305313
del l
306314

315+
307316
def test_serialization_callback_error():
308317

309318
class TempClass(object):
@@ -320,7 +329,7 @@ class TempClass(object):
320329
serialization_context.register_type(TempClass, 20*b"\x00")
321330
serialized_object = pa.serialize(TempClass(), serialization_context)
322331
deserialization_context = pa.SerializationContext()
323-
332+
324333
# Pass a Serialization Context into deserialize, but TempClass
325334
# is not registered
326335
with pytest.raises(pa.DeserializationCallbackError) as err:

‎python/pyarrow/tests/test_tensor.py‎

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -101,14 +101,35 @@ def test_tensor_ipc_roundtrip(tmpdir):
101101

102102

103103
def test_tensor_ipc_strided(tmpdir):
104-
data = np.random.randn(10, 4)
105-
tensor = pa.Tensor.from_numpy(data[::2])
104+
data1 = np.random.randn(10, 4)
105+
tensor1 = pa.Tensor.from_numpy(data1[::2])
106+
107+
data2 = np.random.randn(10, 6, 4)
108+
tensor2 = pa.Tensor.from_numpy(data2[::, ::2, ::])
106109

107110
path = os.path.join(str(tmpdir), 'pyarrow-tensor-ipc-strided')
108-
with pytest.raises(ValueError):
109-
mmap = pa.create_memory_map(path, 1024)
111+
mmap = pa.create_memory_map(path, 2048)
112+
113+
for tensor in [tensor1, tensor2]:
114+
mmap.seek(0)
110115
pa.write_tensor(tensor, mmap)
111116

117+
mmap.seek(0)
118+
result = pa.read_tensor(mmap)
119+
120+
assert result.equals(tensor)
121+
122+
123+
def test_tensor_equals():
124+
data = np.random.randn(10, 6, 4)[::, ::2, ::]
125+
tensor1 = pa.Tensor.from_numpy(data)
126+
tensor2 = pa.Tensor.from_numpy(np.ascontiguousarray(data))
127+
assert tensor1.equals(tensor2)
128+
data = data.copy()
129+
data[9, 0, 0] = 1.0
130+
tensor2 = pa.Tensor.from_numpy(np.ascontiguousarray(data))
131+
assert not tensor1.equals(tensor2)
132+
112133

113134
def test_tensor_size():
114135
data = np.random.randn(10, 4)

0 commit comments

Comments
 (0)