diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc index 8f577539e34ae..090a901cb5ed4 100644 --- a/cpp/src/arrow/compute/exec.cc +++ b/cpp/src/arrow/compute/exec.cc @@ -89,15 +89,22 @@ void PrintTo(const ExecBatch& batch, std::ostream* os) { if (value.is_scalar()) { *os << "Scalar[" << value.scalar()->ToString() << "]\n"; - continue; + } else if (value.is_array() || value.is_chunked_array()) { + PrettyPrintOptions options; + options.skip_new_lines = true; + if (value.is_array()) { + auto array = value.make_array(); + *os << "Array"; + ARROW_CHECK_OK(PrettyPrint(*array, options, os)); + } else { + auto array = value.chunked_array(); + *os << "Chunked Array"; + ARROW_CHECK_OK(PrettyPrint(*array, options, os)); + } + *os << "\n"; + } else { + ARROW_DCHECK(false); } - - auto array = value.make_array(); - PrettyPrintOptions options; - options.skip_new_lines = true; - *os << "Array"; - ARROW_CHECK_OK(PrettyPrint(*array, options, os)); - *os << "\n"; } } @@ -118,8 +125,15 @@ std::string ExecBatch::ToString() const { ExecBatch ExecBatch::Slice(int64_t offset, int64_t length) const { ExecBatch out = *this; for (auto& value : out.values) { - if (value.is_scalar()) continue; - value = value.array()->Slice(offset, length); + if (value.is_scalar()) { + // keep value as is + } else if (value.is_array()) { + value = value.array()->Slice(offset, length); + } else if (value.is_chunked_array()) { + value = value.chunked_array()->Slice(offset, length); + } else { + ARROW_DCHECK(false); + } } out.length = std::min(length, this->length - offset); return out; diff --git a/cpp/src/arrow/compute/exec_test.cc b/cpp/src/arrow/compute/exec_test.cc index 7f29a673d935d..f18af71dba15b 100644 --- a/cpp/src/arrow/compute/exec_test.cc +++ b/cpp/src/arrow/compute/exec_test.cc @@ -56,6 +56,30 @@ using ::arrow::internal::BitmapEquals; using ::arrow::internal::CopyBitmap; using ::arrow::internal::CountSetBits; +TEST(ExecBatch, SliceBasics) { + int64_t length = 4, cut_length = 2, left_length = length - cut_length; + ExecBatch batch{{Int32Scalar(0), ArrayFromJSON(utf8(), R"(["a", "b", "c", "d"])"), + ChunkedArrayFromJSON(float64(), {"[1.1]", "[2.2]", "[3.3]", "[4.4]"})}, + length}; + std::vector expected_sliced{ + {{Int32Scalar(0), ArrayFromJSON(utf8(), R"(["a", "b"])"), + ChunkedArrayFromJSON(float64(), {"[1.1]", "[2.2]"})}, + cut_length}, + {{Int32Scalar(0), ArrayFromJSON(utf8(), R"(["c", "d"])"), + ChunkedArrayFromJSON(float64(), {"[3.3]", "[4.4]"})}, + left_length}}; + std::vector actual_sliced = {batch.Slice(0, cut_length), + batch.Slice(cut_length, left_length)}; + for (size_t i = 0; i < expected_sliced.size(); i++) { + ASSERT_EQ(expected_sliced[i].length, actual_sliced[i].length); + ASSERT_EQ(expected_sliced[i].values.size(), actual_sliced[i].values.size()); + for (size_t j = 0; j < expected_sliced[i].values.size(); j++) { + AssertDatumsEqual(expected_sliced[i].values[j], actual_sliced[i].values[j]); + } + ASSERT_EQ(expected_sliced[i].ToString(), actual_sliced[i].ToString()); + } +} + TEST(ExecBatch, ToRecordBatch) { auto i32_array = ArrayFromJSON(int32(), "[0, 1, 2]"); auto utf8_array = ArrayFromJSON(utf8(), R"(["a", "b", "c"])");