This is an automated email from the ASF dual-hosted git repository.
jorisvandenbossche pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push:
new 373812d0aa ARROW-18280: [C++][Python] Support slicing to end in
list_slice kernel (#14749)
373812d0aa is described below
commit 373812d0aa2c8b1238ad8ed13c8a173fc2fef710
Author: Miles Granger <[email protected]>
AuthorDate: Tue Dec 6 10:41:07 2022 +0100
ARROW-18280: [C++][Python] Support slicing to end in list_slice kernel
(#14749)
Will close [ARROW-18280](https://issues.apache.org/jira/browse/ARROW-18280)
Lead-authored-by: Miles Granger <[email protected]>
Co-authored-by: Will Jones <[email protected]>
Signed-off-by: Joris Van den Bossche <[email protected]>
---
cpp/src/arrow/compute/kernels/scalar_nested.cc | 45 +++++++++++++++-------
.../arrow/compute/kernels/scalar_nested_test.cc | 37 ++++++++++++++----
python/pyarrow/tests/test_compute.py | 31 +++++++--------
3 files changed, 75 insertions(+), 38 deletions(-)
diff --git a/cpp/src/arrow/compute/kernels/scalar_nested.cc
b/cpp/src/arrow/compute/kernels/scalar_nested.cc
index 490ea73b16..1934212a08 100644
--- a/cpp/src/arrow/compute/kernels/scalar_nested.cc
+++ b/cpp/src/arrow/compute/kernels/scalar_nested.cc
@@ -106,14 +106,7 @@ struct ListSlice {
const auto opts = OptionsWrapper<ListSliceOptions>::Get(ctx);
// Invariants
- if (!opts.stop.has_value()) {
- // TODO(ARROW-18280): Support slicing to arbitrary end
- // For variable size list, this would be the largest difference in
offsets
- // For fixed size list, this would be the fixed size.
- return Status::NotImplemented(
- "Slicing to end not yet implemented, please set `stop` parameter.");
- }
- if (opts.start < 0 || opts.start >= opts.stop.value()) {
+ if (opts.start < 0 || (opts.stop.has_value() && opts.start >=
opts.stop.value())) {
// TODO(ARROW-18281): support start == stop which should give empty lists
return Status::Invalid("`start`(", opts.start,
") should be greater than 0 and smaller than
`stop`(",
@@ -130,9 +123,24 @@ struct ListSlice {
list_type->id() == arrow::Type::FIXED_SIZE_LIST);
std::unique_ptr<ArrayBuilder> builder;
+ // should have been checked in resolver
+ // if stop not set, then cannot return fixed size list without input being
fixed size
+ // list b/c we cannot determine the max list element in type resolving.
+ DCHECK(opts.stop.has_value() ||
+ (!opts.stop.has_value() && (!return_fixed_size_list ||
+ list_type->id() ==
arrow::Type::FIXED_SIZE_LIST)));
+
// construct array values
if (return_fixed_size_list) {
- const auto length = bit_util::CeilDiv(opts.stop.value() - opts.start,
opts.step);
+ int32_t stop;
+ if (opts.stop.has_value()) {
+ stop = static_cast<int32_t>(opts.stop.value());
+ } else {
+ DCHECK_EQ(list_type->id(), arrow::Type::FIXED_SIZE_LIST);
+ stop = reinterpret_cast<const
FixedSizeListType*>(list_type)->list_size();
+ }
+ const auto size = std::max(stop - static_cast<int32_t>(opts.start), 0);
+ const auto length = bit_util::CeilDiv(size, opts.step);
RETURN_NOT_OK(MakeBuilder(ctx->memory_pool(),
fixed_size_list(value_type,
static_cast<int32_t>(length)),
&builder));
@@ -217,7 +225,9 @@ struct ListSlice {
auto cursor = offset;
RETURN_NOT_OK(list_builder->Append());
- while (cursor < offset + ((opts->stop.value() - opts->start))) {
+ const auto size = opts->stop.has_value() ? (opts->stop.value() -
opts->start)
+ : ((next_offset - opts->start) -
offset);
+ while (cursor < offset + size) {
if (cursor + opts->start >= next_offset) {
if constexpr (!std::is_same_v<BuilderType, FixedSizeListBuilder>) {
break; // don't pad nulls for variable sized list output
@@ -241,14 +251,23 @@ Result<TypeHolder> MakeListSliceResolve(KernelContext*
ctx,
const auto return_fixed_size_list =
opts.return_fixed_size_list.value_or(list_type->id() ==
Type::FIXED_SIZE_LIST);
if (return_fixed_size_list) {
+ int32_t stop;
if (!opts.stop.has_value()) {
- return Status::NotImplemented(
- "Unable to produce FixedSizeListArray without `stop` being set.");
+ if (list_type->id() == Type::FIXED_SIZE_LIST) {
+ stop = checked_cast<const FixedSizeListType*>(list_type)->list_size();
+ } else {
+ return Status::NotImplemented(
+ "Unable to produce FixedSizeListArray from non-FixedSizeListArray
without "
+ "`stop` being set.");
+ }
+ } else {
+ stop = static_cast<int32_t>(opts.stop.value());
}
+ const auto size = std::max(static_cast<int32_t>(stop - opts.start), 0);
if (opts.step < 1) {
return Status::Invalid("`step` must be >= 1, got: ", opts.step);
}
- const auto length = bit_util::CeilDiv(opts.stop.value() - opts.start,
opts.step);
+ const auto length = bit_util::CeilDiv(size, opts.step);
return fixed_size_list(value_type, static_cast<int32_t>(length));
} else {
// Returning large list if that's what we got in and didn't ask for fixed
size
diff --git a/cpp/src/arrow/compute/kernels/scalar_nested_test.cc
b/cpp/src/arrow/compute/kernels/scalar_nested_test.cc
index 8b655658a3..a72ec99620 100644
--- a/cpp/src/arrow/compute/kernels/scalar_nested_test.cc
+++ b/cpp/src/arrow/compute/kernels/scalar_nested_test.cc
@@ -135,6 +135,11 @@ TEST(TestScalarNested, ListSliceVariableOutput) {
expected = ArrayFromJSON(list(value_type), "[[3], [], [], null]");
CheckScalarUnary("list_slice", input, expected, &args);
+ args.start = 1;
+ args.stop = std::nullopt;
+ expected = ArrayFromJSON(list(value_type), "[[2, 3], [5], [], null]");
+ CheckScalarUnary("list_slice", input, expected, &args);
+
args.start = 0;
args.stop = 4;
args.step = 2;
@@ -174,6 +179,27 @@ TEST(TestScalarNested, ListSliceFixedOutput) {
"[[3, null], [null, null], [null, null],
null]");
CheckScalarUnary("list_slice", input, expected, &args);
+ args.start = 1;
+ args.stop = std::nullopt;
+ expected = ArrayFromJSON(fixed_size_list(value_type, 2),
+ "[[2, 3], [5, null], [null, null], null]");
+ if (input->type()->id() == Type::FIXED_SIZE_LIST) {
+ CheckScalarUnary("list_slice", input, expected, &args);
+ } else {
+ EXPECT_RAISES_WITH_MESSAGE_THAT(
+ NotImplemented,
+ ::testing::HasSubstr("Unable to produce FixedSizeListArray from "
+ "non-FixedSizeListArray without `stop` being
set."),
+ CallFunction("list_slice", {input}, &args));
+ }
+
+ args.start = 3;
+ args.stop = std::nullopt;
+ expected = ArrayFromJSON(fixed_size_list(value_type, 0), "[[], [], [],
null]");
+ if (input->type()->id() == Type::FIXED_SIZE_LIST) {
+ CheckScalarUnary("list_slice", input, expected, &args);
+ }
+
args.start = 0;
args.stop = 4;
args.step = 2;
@@ -276,19 +302,14 @@ TEST(TestScalarNested, ListSliceBadParameters) {
::testing::HasSubstr(
"`start`(1) should be greater than 0 and smaller than `stop`(1)"),
CallFunction("list_slice", {input}, &args));
- // stop not set and FixedSizeList requested
+ // stop not set and FixedSizeList requested with variable sized input
args.stop = std::nullopt;
EXPECT_RAISES_WITH_MESSAGE_THAT(
NotImplemented,
- ::testing::HasSubstr("NotImplemented: Unable to produce
FixedSizeListArray without "
+ ::testing::HasSubstr("NotImplemented: Unable to produce
FixedSizeListArray from "
+ "non-FixedSizeListArray without "
"`stop` being set."),
CallFunction("list_slice", {input}, &args));
- // stop not set and ListArray requested
- args.stop = std::nullopt;
- args.return_fixed_size_list = false;
- EXPECT_RAISES_WITH_MESSAGE_THAT(
- NotImplemented, ::testing::HasSubstr("Slicing to end not yet
implemented"),
- CallFunction("list_slice", {input}, &args));
// Catch step must be >= 1
args.start = 0;
args.stop = 2;
diff --git a/python/pyarrow/tests/test_compute.py
b/python/pyarrow/tests/test_compute.py
index 70439ac1b7..64114521a8 100644
--- a/python/pyarrow/tests/test_compute.py
+++ b/python/pyarrow/tests/test_compute.py
@@ -2950,6 +2950,7 @@ def test_cast_table_raises():
@pytest.mark.parametrize("start,stop,expected", (
+ (0, None, [[1, 2, 3], [4, 5, None], [6, None, None], None]),
(0, 1, [[1], [4], [6], None]),
(0, 2, [[1, 2], [4, 5], [6, None], None]),
(1, 2, [[2], [5], [None], None]),
@@ -2966,13 +2967,22 @@ def test_list_slice_output_fixed(start, stop, step,
expected, value_type,
else:
arr = pa.array([[1, 2, 3], [4, 5], [6], None],
pa.list_(pa.int8())).cast(list_type(value_type()))
- result = pc.list_slice(arr, start, stop, step, return_fixed_size_list=True)
- pylist = result.cast(pa.list_(pa.int8(),
- result.type.list_size)).to_pylist()
- assert pylist == [e[::step] if e else e for e in expected]
+
+ args = arr, start, stop, step, True
+ if stop is None and list_type != "fixed":
+ msg = ("Unable to produce FixedSizeListArray from "
+ "non-FixedSizeListArray without `stop` being set.")
+ with pytest.raises(pa.ArrowNotImplementedError, match=msg):
+ pc.list_slice(*args)
+ else:
+ result = pc.list_slice(*args)
+ pylist = result.cast(pa.list_(pa.int8(),
+ result.type.list_size)).to_pylist()
+ assert pylist == [e[::step] if e else e for e in expected]
@pytest.mark.parametrize("start,stop", (
+ (0, None,),
(0, 1,),
(0, 2,),
(1, 2,),
@@ -3034,19 +3044,6 @@ def test_list_slice_bad_parameters():
with pytest.raises(pa.ArrowInvalid, match=msg):
pc.list_slice(arr, 0, 0) # start == stop?
- # TODO(ARROW-18280): support stop == None; slice to end
- # This fails first at resolve, b/c it doesn't now how big the
- # resulting FixedSizeListArray item size will be
- msg = "Unable to produce FixedSizeListArray without `stop`"
- with pytest.raises(NotImplementedError, match=msg):
- pc.list_slice(arr, 0, return_fixed_size_list=True)
-
- # cont. This fails inside of kernel function; resolver doesn't
- # need to know the item size for ListArray.
- msg = "Slicing to end not yet implemented*"
- with pytest.raises(NotImplementedError, match=msg):
- pc.list_slice(arr, 0, return_fixed_size_list=False)
-
# Step not >= 1
msg = "`step` must be >= 1, got: "
with pytest.raises(pa.ArrowInvalid, match=msg + "0"):
