Skip to content

Commit

Permalink
fix for chunked array + add python tests
Browse files Browse the repository at this point in the history
  • Loading branch information
jorisvandenbossche committed Jun 12, 2024
1 parent 7890c00 commit 050fd9d
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 2 deletions.
3 changes: 1 addition & 2 deletions cpp/src/arrow/pretty_print.cc
Original file line number Diff line number Diff line change
Expand Up @@ -496,8 +496,7 @@ Status PrettyPrint(const ChunkedArray& chunked_arr, const PrettyPrintOptions& op
} else {
PrettyPrintOptions chunk_options = options;
chunk_options.indent += options.indent_size;
ArrayPrinter printer(chunk_options, sink);
RETURN_NOT_OK(printer.Print(*chunked_arr.chunk(i)));
RETURN_NOT_OK(PrettyPrint(*chunked_arr.chunk(i), chunk_options, sink));
}
}
if (!options.skip_new_lines) {
Expand Down
54 changes: 54 additions & 0 deletions python/pyarrow/tests/test_cuda.py
Original file line number Diff line number Diff line change
Expand Up @@ -792,3 +792,57 @@ def test_IPC(size):
p.start()
p.join()
assert p.exitcode == 0


def test_print_array():
batch = make_recordbatch(10)
cbuf = cuda.serialize_record_batch(batch, global_context)
cbatch = cuda.read_record_batch(cbuf, batch.schema)
arr = batch["f0"]
carr = cbatch["f0"]
assert str(carr) == str(arr)

batch = make_recordbatch(100)
cbuf = cuda.serialize_record_batch(batch, global_context)
cbatch = cuda.read_record_batch(cbuf, batch.schema)
arr = batch["f0"]
carr = cbatch["f0"]
assert str(carr) == str(arr)


def make_chunked_array(n_elements_per_chunk, n_chunks):
arrs = []
carrs = []
for _ in range(n_chunks):
batch = make_recordbatch(n_elements_per_chunk)
cbuf = cuda.serialize_record_batch(batch, global_context)
cbatch = cuda.read_record_batch(cbuf, batch.schema)
arrs.append(batch["f0"])
carrs.append(cbatch["f0"])

return pa.chunked_array(arrs), pa.chunked_array(carrs)


def test_print_chunked_array():
arr, carr = make_chunked_array(10, 3)
assert str(carr) == str(arr)

arr, carr = make_chunked_array(100, 20)
assert str(carr) == str(arr)


def test_print_record_batch():
batch = make_recordbatch(10)
cbuf = cuda.serialize_record_batch(batch, global_context)
cbatch = cuda.read_record_batch(cbuf, batch.schema)
assert str(cbatch) == str(batch)

batch = make_recordbatch(100)
cbuf = cuda.serialize_record_batch(batch, global_context)
cbatch = cuda.read_record_batch(cbuf, batch.schema)
assert str(cbatch) == str(batch)


def test_print_table():
_, table, _, ctable = make_table_cuda()
assert str(ctable) == str(table)

0 comments on commit 050fd9d

Please sign in to comment.