diff --git a/cpp/src/arrow/compute/kernels/vector_statistics.cc b/cpp/src/arrow/compute/kernels/vector_statistics.cc index 074f2ec0a738..cc2fbc40164f 100644 --- a/cpp/src/arrow/compute/kernels/vector_statistics.cc +++ b/cpp/src/arrow/compute/kernels/vector_statistics.cc @@ -26,6 +26,7 @@ #include "arrow/compute/function.h" #include "arrow/compute/kernel.h" #include "arrow/compute/kernels/codegen_internal.h" +#include "arrow/compute/kernels/util_internal.h" #include "arrow/compute/registry.h" #include "arrow/compute/registry_internal.h" #include "arrow/result.h" @@ -71,6 +72,7 @@ struct Winsorize { // Only nulls and NaNs => return input as-is out_data->null_count = data->null_count.load(); out_data->length = data->length; + out_data->offset = data->offset; out_data->buffers = data->buffers; return Status::OK(); } @@ -127,7 +129,9 @@ struct Winsorize { DCHECK_EQ(out->buffers.size(), data.buffers.size()); out->null_count = data.null_count.load(); out->length = data.length; - out->buffers[0] = data.buffers[0]; + out->offset = 0; + ARROW_ASSIGN_OR_RAISE(out->buffers[0], + GetOrCopyNullBitmapBuffer(data, ctx->memory_pool())); ARROW_ASSIGN_OR_RAISE(out->buffers[1], ctx->Allocate(out->length * sizeof(CType))); // Avoid leaving uninitialized memory under null entries std::memset(out->buffers[1]->mutable_data(), 0, out->length * sizeof(CType)); diff --git a/cpp/src/arrow/compute/kernels/vector_statistics_test.cc b/cpp/src/arrow/compute/kernels/vector_statistics_test.cc index 97715cdaedd2..72a765be1039 100644 --- a/cpp/src/arrow/compute/kernels/vector_statistics_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_statistics_test.cc @@ -87,6 +87,45 @@ TEST_F(TestWinsorize, FloatingPoint) { } } +TEST_F(TestWinsorize, SlicedInput) { + options_.lower_limit = 0.0; + options_.upper_limit = 1.0; + auto parent = ArrayFromJSON(float64(), "[1.1, 2.2, null, 4.4, null, 6.6, 7.7, 8.8]"); + auto expected = ArrayFromJSON(float64(), "[null, 4.4, null, 6.6, 7.7]"); + CheckWinsorize(parent->Slice(2, 5), expected); + + options_.lower_limit = 0.25; + options_.upper_limit = 0.75; + auto dense = ArrayFromJSON(float64(), "[1.0, 2.0, 3.0, 44.0, 55.0, 66.0, 77.0]"); + CheckWinsorize(dense->Slice(1, 5), + ArrayFromJSON(float64(), "[3.0, 3.0, 44.0, 55.0, 55.0]")); +} + +TEST_F(TestWinsorize, SlicedChunkedInput) { + options_.lower_limit = 0.0; + options_.upper_limit = 1.0; + auto parent = ArrayFromJSON(float64(), "[1.1, 2.2, null, 4.4, null, 6.6, 7.7, 8.8]"); + auto chunked = std::make_shared( + ArrayVector{parent->Slice(2, 3), parent->Slice(5, 3)}); + auto expected = std::make_shared(ArrayVector{ + ArrayFromJSON(float64(), "[null, 4.4, null]"), + ArrayFromJSON(float64(), "[6.6, 7.7, 8.8]"), + }); + CheckWinsorize(chunked, expected); +} + +TEST_F(TestWinsorize, SlicedInputWithoutQuantiles) { + for (const auto* json_input : + {"[1, 2, null, null, null, 3]", "[1, 2, NaN, null, NaN, 3]", + "[1, 2, NaN, NaN, NaN, 3]"}) { + auto parent = ArrayFromJSON(float64(), json_input); + auto sliced = parent->Slice(2, 3); + CheckWinsorize(sliced, sliced); + auto chunked = std::make_shared(ArrayVector{sliced}); + CheckWinsorize(chunked, chunked); + } +} + TEST_F(TestWinsorize, Integral) { for (auto type : IntTypes()) { options_.lower_limit = 0.25;