From 22dde21af01cabb6213a4b97b4d3c121027673e0 Mon Sep 17 00:00:00 2001 From: zanmato1984 Date: Mon, 15 Jan 2024 00:23:55 +0800 Subject: [PATCH 1/4] Fix fsb by-word comparison exceeding buffer boundary for tail bytes --- cpp/src/arrow/compute/row/compare_internal.cc | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/compute/row/compare_internal.cc b/cpp/src/arrow/compute/row/compare_internal.cc index 7c402e7a238..807275633dc 100644 --- a/cpp/src/arrow/compute/row/compare_internal.cc +++ b/cpp/src/arrow/compute/row/compare_internal.cc @@ -224,8 +224,10 @@ void KeyCompare::CompareBinaryColumnToRow(uint32_t offset_within_row, uint64_t key_right = key_right_ptr[i]; result_or |= key_left ^ key_right; } - uint64_t key_left = util::SafeLoad(key_left_ptr + i); - uint64_t key_right = key_right_ptr[i]; + uint64_t key_left = 0; + memcpy(&key_left, key_left_ptr + i, length - num_loops_less_one * 8); + uint64_t key_right = 0; + memcpy(&key_right, key_right_ptr + i, length - num_loops_less_one * 8); result_or |= tail_mask & (key_left ^ key_right); return result_or == 0 ? 0xff : 0; }); From d9958362de25593e2f64b5caa06455c4b2c6f4c5 Mon Sep 17 00:00:00 2001 From: zanmato1984 Date: Tue, 16 Jan 2024 01:13:42 +0800 Subject: [PATCH 2/4] Add case to reproduce the issue --- cpp/src/arrow/compute/row/CMakeLists.txt | 2 + cpp/src/arrow/compute/row/compare_test.cc | 107 ++++++++++++++++++++++ 2 files changed, 109 insertions(+) create mode 100644 cpp/src/arrow/compute/row/compare_test.cc diff --git a/cpp/src/arrow/compute/row/CMakeLists.txt b/cpp/src/arrow/compute/row/CMakeLists.txt index 6ae982dbaf3..cda5bc748ca 100644 --- a/cpp/src/arrow/compute/row/CMakeLists.txt +++ b/cpp/src/arrow/compute/row/CMakeLists.txt @@ -19,3 +19,5 @@ # in a row-major order. arrow_install_all_headers("arrow/compute/row") + +add_arrow_compute_test(compare_test SOURCES compare_test.cc) diff --git a/cpp/src/arrow/compute/row/compare_test.cc b/cpp/src/arrow/compute/row/compare_test.cc new file mode 100644 index 00000000000..52b6dd39151 --- /dev/null +++ b/cpp/src/arrow/compute/row/compare_test.cc @@ -0,0 +1,107 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include + +#include "arrow/compute/row/compare_internal.h" +#include "arrow/testing/gtest_util.h" + +namespace arrow { +namespace compute { + +using namespace arrow::util; + +// Specialized case for GH-39577. +TEST(KeyCompare, CompareColumnsToRowsCuriousFSB) { + int fsb_length = 9; + MemoryPool* pool = default_memory_pool(); + TempVectorStack stack; + ASSERT_OK(stack.Init(pool, 8 * util::MiniBatch::kMiniBatchLength * sizeof(uint64_t))); + + int num_rows = 7; + auto column_right = ArrayFromJSON(fixed_size_binary(fsb_length), R"([ + "000000000", + "111111111", + "222222222", + "333333333", + "444444444", + "555555555", + "666666666"])"); + ExecBatch batch_right({column_right}, num_rows); + + std::vector column_metadatas_right; + ASSERT_OK(ColumnMetadatasFromExecBatch(batch_right, &column_metadatas_right)); + + RowTableMetadata table_metadata_right; + table_metadata_right.FromColumnMetadataVector(column_metadatas_right, sizeof(uint64_t), + sizeof(uint64_t)); + + std::vector column_arrays_right; + ASSERT_OK(ColumnArraysFromExecBatch(batch_right, &column_arrays_right)); + + RowTableImpl row_table; + ASSERT_OK(row_table.Init(pool, table_metadata_right)); + + RowTableEncoder row_encoder; + row_encoder.Init(column_metadatas_right, sizeof(uint64_t), sizeof(uint64_t)); + row_encoder.PrepareEncodeSelected(0, num_rows, column_arrays_right); + + std::vector row_ids_right(num_rows); + std::iota(row_ids_right.begin(), row_ids_right.end(), 0); + ASSERT_OK(row_encoder.EncodeSelected(&row_table, num_rows, row_ids_right.data())); + + auto column_left = ArrayFromJSON(fixed_size_binary(fsb_length), R"([ + "000000000", + "111111111", + "222222222", + "333333333", + "444444444", + "555555555", + "777777777"])"); + ExecBatch batch_left({column_left}, num_rows); + std::vector column_arrays_left; + ASSERT_OK(ColumnArraysFromExecBatch(batch_left, &column_arrays_left)); + + std::vector row_ids_left(num_rows); + std::iota(row_ids_left.begin(), row_ids_left.end(), 0); + + LightContext ctx{arrow::internal::CpuInfo::GetInstance()->hardware_flags(), &stack}; + + { + uint32_t num_rows_no_match; + std::vector row_ids_out(num_rows); + KeyCompare::CompareColumnsToRows(num_rows, NULLPTR, row_ids_left.data(), &ctx, + &num_rows_no_match, row_ids_out.data(), + column_arrays_left, row_table, true, NULLPTR); + ASSERT_EQ(num_rows_no_match, 1); + ASSERT_EQ(row_ids_out[0], 6); + } + + { + std::vector match_bitvector(arrow::bit_util::BytesForBits(num_rows)); + KeyCompare::CompareColumnsToRows(num_rows, NULLPTR, row_ids_left.data(), &ctx, + NULLPTR, NULLPTR, column_arrays_left, row_table, + true, match_bitvector.data()); + for (int i = 0; i < num_rows; ++i) { + SCOPED_TRACE(i); + ASSERT_EQ(arrow::bit_util::GetBit(match_bitvector.data(), i), i != 6); + } + } +} + +} // namespace compute +} // namespace arrow From 02f975dfd9c268ef5f6ed467338498c51a5d63b1 Mon Sep 17 00:00:00 2001 From: zanmato1984 Date: Tue, 16 Jan 2024 01:23:00 +0800 Subject: [PATCH 3/4] Fix lint --- cpp/src/arrow/compute/row/compare_test.cc | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/cpp/src/arrow/compute/row/compare_test.cc b/cpp/src/arrow/compute/row/compare_test.cc index 52b6dd39151..1d8562cd56d 100644 --- a/cpp/src/arrow/compute/row/compare_test.cc +++ b/cpp/src/arrow/compute/row/compare_test.cc @@ -23,14 +23,17 @@ namespace arrow { namespace compute { -using namespace arrow::util; +using arrow::bit_util::BytesForBits; +using arrow::internal::CpuInfo; +using arrow::util::MiniBatch; +using arrow::util::TempVectorStack; // Specialized case for GH-39577. TEST(KeyCompare, CompareColumnsToRowsCuriousFSB) { int fsb_length = 9; MemoryPool* pool = default_memory_pool(); TempVectorStack stack; - ASSERT_OK(stack.Init(pool, 8 * util::MiniBatch::kMiniBatchLength * sizeof(uint64_t))); + ASSERT_OK(stack.Init(pool, 8 * MiniBatch::kMiniBatchLength * sizeof(uint64_t))); int num_rows = 7; auto column_right = ArrayFromJSON(fixed_size_binary(fsb_length), R"([ @@ -79,7 +82,7 @@ TEST(KeyCompare, CompareColumnsToRowsCuriousFSB) { std::vector row_ids_left(num_rows); std::iota(row_ids_left.begin(), row_ids_left.end(), 0); - LightContext ctx{arrow::internal::CpuInfo::GetInstance()->hardware_flags(), &stack}; + LightContext ctx{CpuInfo::GetInstance()->hardware_flags(), &stack}; { uint32_t num_rows_no_match; @@ -92,7 +95,7 @@ TEST(KeyCompare, CompareColumnsToRowsCuriousFSB) { } { - std::vector match_bitvector(arrow::bit_util::BytesForBits(num_rows)); + std::vector match_bitvector(BytesForBits(num_rows)); KeyCompare::CompareColumnsToRows(num_rows, NULLPTR, row_ids_left.data(), &ctx, NULLPTR, NULLPTR, column_arrays_left, row_table, true, match_bitvector.data()); From f8cce17a0b57f45f705a4cbf3e9df6a714862e5a Mon Sep 17 00:00:00 2001 From: zanmato1984 Date: Tue, 16 Jan 2024 23:32:31 +0800 Subject: [PATCH 4/4] Address comments --- cpp/src/arrow/compute/CMakeLists.txt | 3 ++- cpp/src/arrow/compute/row/CMakeLists.txt | 2 -- cpp/src/arrow/compute/row/compare_internal.cc | 9 ++++----- 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/cpp/src/arrow/compute/CMakeLists.txt b/cpp/src/arrow/compute/CMakeLists.txt index 1134e0a98ae..e14d78ff6e5 100644 --- a/cpp/src/arrow/compute/CMakeLists.txt +++ b/cpp/src/arrow/compute/CMakeLists.txt @@ -89,7 +89,8 @@ add_arrow_test(internals_test kernel_test.cc light_array_test.cc registry_test.cc - key_hash_test.cc) + key_hash_test.cc + row/compare_test.cc) add_arrow_compute_test(expression_test SOURCES expression_test.cc) diff --git a/cpp/src/arrow/compute/row/CMakeLists.txt b/cpp/src/arrow/compute/row/CMakeLists.txt index cda5bc748ca..6ae982dbaf3 100644 --- a/cpp/src/arrow/compute/row/CMakeLists.txt +++ b/cpp/src/arrow/compute/row/CMakeLists.txt @@ -19,5 +19,3 @@ # in a row-major order. arrow_install_all_headers("arrow/compute/row") - -add_arrow_compute_test(compare_test SOURCES compare_test.cc) diff --git a/cpp/src/arrow/compute/row/compare_internal.cc b/cpp/src/arrow/compute/row/compare_internal.cc index 807275633dc..078a8287c71 100644 --- a/cpp/src/arrow/compute/row/compare_internal.cc +++ b/cpp/src/arrow/compute/row/compare_internal.cc @@ -208,8 +208,7 @@ void KeyCompare::CompareBinaryColumnToRow(uint32_t offset_within_row, // Non-zero length guarantees no underflow int32_t num_loops_less_one = static_cast(bit_util::CeilDiv(length, 8)) - 1; - - uint64_t tail_mask = ~0ULL >> (64 - 8 * (length - num_loops_less_one * 8)); + int32_t num_tail_bytes = length - num_loops_less_one * 8; const uint64_t* key_left_ptr = reinterpret_cast(left_base + irow_left * length); @@ -225,10 +224,10 @@ void KeyCompare::CompareBinaryColumnToRow(uint32_t offset_within_row, result_or |= key_left ^ key_right; } uint64_t key_left = 0; - memcpy(&key_left, key_left_ptr + i, length - num_loops_less_one * 8); + memcpy(&key_left, key_left_ptr + i, num_tail_bytes); uint64_t key_right = 0; - memcpy(&key_right, key_right_ptr + i, length - num_loops_less_one * 8); - result_or |= tail_mask & (key_left ^ key_right); + memcpy(&key_right, key_right_ptr + i, num_tail_bytes); + result_or |= key_left ^ key_right; return result_or == 0 ? 0xff : 0; }); }