Skip to content

Commit

Permalink
[python-package] Add tests for passing Arrow arrays with empty chunks (
Browse files Browse the repository at this point in the history
  • Loading branch information
borchero authored Dec 10, 2023
1 parent 1548b42 commit 522f0f0
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 13 deletions.
2 changes: 2 additions & 0 deletions include/LightGBM/arrow.h
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,7 @@ class ArrowChunkedArray {
const struct ArrowSchema* schema) {
chunks_.reserve(n_chunks);
for (auto k = 0; k < n_chunks; ++k) {
if (chunks[k].length == 0) continue;
chunks_.push_back(&chunks[k]);
}
schema_ = schema;
Expand Down Expand Up @@ -220,6 +221,7 @@ class ArrowTable {
std::vector<const ArrowArray*> children_chunks;
children_chunks.reserve(n_chunks);
for (int64_t k = 0; k < n_chunks; ++k) {
if (chunks[k].length == 0) continue;
children_chunks.push_back(chunks[k].children[j]);
}
columns_.emplace_back(children_chunks, schema->children[j]);
Expand Down
38 changes: 25 additions & 13 deletions tests/python_package_test/test_arrow.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,18 +30,19 @@
]


def generate_simple_arrow_table() -> pa.Table:
def generate_simple_arrow_table(empty_chunks: bool = False) -> pa.Table:
c: list[list[int]] = [[]] if empty_chunks else []
columns = [
pa.chunked_array([[1, 2, 3, 4, 5]], type=pa.uint8()),
pa.chunked_array([[1, 2, 3, 4, 5]], type=pa.int8()),
pa.chunked_array([[1, 2, 3, 4, 5]], type=pa.uint16()),
pa.chunked_array([[1, 2, 3, 4, 5]], type=pa.int16()),
pa.chunked_array([[1, 2, 3, 4, 5]], type=pa.uint32()),
pa.chunked_array([[1, 2, 3, 4, 5]], type=pa.int32()),
pa.chunked_array([[1, 2, 3, 4, 5]], type=pa.uint64()),
pa.chunked_array([[1, 2, 3, 4, 5]], type=pa.int64()),
pa.chunked_array([[1, 2, 3, 4, 5]], type=pa.float32()),
pa.chunked_array([[1, 2, 3, 4, 5]], type=pa.float64()),
pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.uint8()),
pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.int8()),
pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.uint16()),
pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.int16()),
pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.uint32()),
pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.int32()),
pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.uint64()),
pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.int64()),
pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.float32()),
pa.chunked_array(c + [[1, 2, 3]] + c + [[4, 5]] + c, type=pa.float64()),
]
return pa.Table.from_arrays(columns, names=[f"col_{i}" for i in range(len(columns))])

Expand Down Expand Up @@ -104,6 +105,7 @@ def dummy_dataset_params() -> Dict[str, Any]:
("arrow_table_fn", "dataset_params"),
[ # Use lambda functions here to minimize memory consumption
(lambda: generate_simple_arrow_table(), dummy_dataset_params()),
(lambda: generate_simple_arrow_table(empty_chunks=True), dummy_dataset_params()),
(lambda: generate_dummy_arrow_table(), dummy_dataset_params()),
(lambda: generate_nullable_arrow_table(), dummy_dataset_params()),
(lambda: generate_random_arrow_table(3, 1000, 42), {}),
Expand Down Expand Up @@ -160,7 +162,12 @@ def test_dataset_construct_fields_fuzzy():

@pytest.mark.parametrize(
["array_type", "label_data"],
[(pa.array, [0, 1, 0, 0, 1]), (pa.chunked_array, [[0], [1, 0, 0, 1]])],
[
(pa.array, [0, 1, 0, 0, 1]),
(pa.chunked_array, [[0], [1, 0, 0, 1]]),
(pa.chunked_array, [[], [0], [1, 0, 0, 1]]),
(pa.chunked_array, [[0], [], [1, 0], [], [], [0, 1], []]),
],
)
@pytest.mark.parametrize("arrow_type", _INTEGER_TYPES + _FLOAT_TYPES)
def test_dataset_construct_labels(array_type, label_data, arrow_type):
Expand All @@ -187,7 +194,12 @@ def test_dataset_construct_weights_none():

@pytest.mark.parametrize(
["array_type", "weight_data"],
[(pa.array, [3, 0.7, 1.5, 0.5, 0.1]), (pa.chunked_array, [[3], [0.7, 1.5, 0.5, 0.1]])],
[
(pa.array, [3, 0.7, 1.5, 0.5, 0.1]),
(pa.chunked_array, [[3], [0.7, 1.5, 0.5, 0.1]]),
(pa.chunked_array, [[], [3], [0.7, 1.5, 0.5, 0.1]]),
(pa.chunked_array, [[3], [0.7], [], [], [1.5, 0.5, 0.1], []]),
],
)
@pytest.mark.parametrize("arrow_type", _FLOAT_TYPES)
def test_dataset_construct_weights(array_type, weight_data, arrow_type):
Expand Down

0 comments on commit 522f0f0

Please sign in to comment.