|
18 | 18 | #include <numeric> |
19 | 19 |
|
20 | 20 | #include "arrow/compute/row/compare_internal.h" |
| 21 | +#include "arrow/testing/generator.h" |
21 | 22 | #include "arrow/testing/gtest_util.h" |
22 | 23 | #include "arrow/testing/random.h" |
| 24 | +#include "arrow/util/bitmap_ops.h" |
23 | 25 |
|
24 | 26 | namespace arrow { |
25 | 27 | namespace compute { |
@@ -164,5 +166,141 @@ TEST(KeyCompare, CompareColumnsToRowsTempStackUsage) { |
164 | 166 | } |
165 | 167 | } |
166 | 168 |
|
| 169 | +#ifndef ARROW_VALGRIND |
| 170 | +// Compare columns to rows at offsets over 2GB within a row table. |
| 171 | +// Certain AVX2 instructions may behave unexpectedly causing troubles like GH-41813. |
| 172 | +TEST(KeyCompare, CompareColumnsToRowsLarge) { |
| 173 | + if constexpr (sizeof(void*) == 4) { |
| 174 | + GTEST_SKIP() << "Test only works on 64-bit platforms"; |
| 175 | + } |
| 176 | + |
| 177 | + // The idea of this case is to create a row table using several fixed length columns and |
| 178 | + // one var length column (so the row is hence var length and has offset buffer), with |
| 179 | + // the overall data size exceeding 2GB. Then compare each row with itself. |
| 180 | + constexpr int64_t two_gb = 2ll * 1024ll * 1024ll * 1024ll; |
| 181 | + // The compare function requires the row id of the left column to be uint16_t, hence the |
| 182 | + // number of rows. |
| 183 | + constexpr int64_t num_rows = std::numeric_limits<uint16_t>::max() + 1; |
| 184 | + const std::vector<std::shared_ptr<DataType>> fixed_length_types{uint64(), uint32()}; |
| 185 | + // The var length column should be a little smaller than 2GB to workaround the capacity |
| 186 | + // limitation in the var length builder. |
| 187 | + constexpr int32_t var_length = two_gb / num_rows - 1; |
| 188 | + auto row_size = std::accumulate(fixed_length_types.begin(), fixed_length_types.end(), |
| 189 | + static_cast<int64_t>(var_length), |
| 190 | + [](int64_t acc, const std::shared_ptr<DataType>& type) { |
| 191 | + return acc + type->byte_width(); |
| 192 | + }); |
| 193 | + // The overall size should be larger than 2GB. |
| 194 | + ASSERT_GT(row_size * num_rows, two_gb); |
| 195 | + |
| 196 | + MemoryPool* pool = default_memory_pool(); |
| 197 | + |
| 198 | + // The left side columns. |
| 199 | + std::vector<KeyColumnArray> columns_left; |
| 200 | + ExecBatch batch_left; |
| 201 | + { |
| 202 | + std::vector<Datum> values; |
| 203 | + |
| 204 | + // Several fixed length arrays containing random content. |
| 205 | + for (const auto& type : fixed_length_types) { |
| 206 | + ASSERT_OK_AND_ASSIGN(auto value, ::arrow::gen::Random(type)->Generate(num_rows)); |
| 207 | + values.push_back(std::move(value)); |
| 208 | + } |
| 209 | + // A var length array containing 'X' repeated var_length times. |
| 210 | + ASSERT_OK_AND_ASSIGN(auto value_var_length, |
| 211 | + ::arrow::gen::Constant( |
| 212 | + std::make_shared<BinaryScalar>(std::string(var_length, 'X'))) |
| 213 | + ->Generate(num_rows)); |
| 214 | + values.push_back(std::move(value_var_length)); |
| 215 | + |
| 216 | + batch_left = ExecBatch(std::move(values), num_rows); |
| 217 | + ASSERT_OK(ColumnArraysFromExecBatch(batch_left, &columns_left)); |
| 218 | + } |
| 219 | + |
| 220 | + // The right side row table. |
| 221 | + RowTableImpl row_table_right; |
| 222 | + { |
| 223 | + // Encode the row table with the left columns. |
| 224 | + std::vector<KeyColumnMetadata> column_metadatas; |
| 225 | + ASSERT_OK(ColumnMetadatasFromExecBatch(batch_left, &column_metadatas)); |
| 226 | + RowTableMetadata table_metadata; |
| 227 | + table_metadata.FromColumnMetadataVector(column_metadatas, sizeof(uint64_t), |
| 228 | + sizeof(uint64_t)); |
| 229 | + ASSERT_OK(row_table_right.Init(pool, table_metadata)); |
| 230 | + std::vector<uint16_t> row_ids(num_rows); |
| 231 | + std::iota(row_ids.begin(), row_ids.end(), 0); |
| 232 | + RowTableEncoder row_encoder; |
| 233 | + row_encoder.Init(column_metadatas, sizeof(uint64_t), sizeof(uint64_t)); |
| 234 | + row_encoder.PrepareEncodeSelected(0, num_rows, columns_left); |
| 235 | + ASSERT_OK(row_encoder.EncodeSelected( |
| 236 | + &row_table_right, static_cast<uint32_t>(num_rows), row_ids.data())); |
| 237 | + |
| 238 | + // The row table must contain an offset buffer. |
| 239 | + ASSERT_NE(row_table_right.offsets(), NULLPTR); |
| 240 | + // The whole point of this test. |
| 241 | + ASSERT_GT(row_table_right.offsets()[num_rows - 1], two_gb); |
| 242 | + } |
| 243 | + |
| 244 | + // The rows to compare. |
| 245 | + std::vector<uint32_t> row_ids_to_compare(num_rows); |
| 246 | + std::iota(row_ids_to_compare.begin(), row_ids_to_compare.end(), 0); |
| 247 | + |
| 248 | + TempVectorStack stack; |
| 249 | + ASSERT_OK(stack.Init(pool, KeyCompare::CompareColumnsToRowsTempStackUsage(num_rows))); |
| 250 | + LightContext ctx{CpuInfo::GetInstance()->hardware_flags(), &stack}; |
| 251 | + |
| 252 | + { |
| 253 | + // No selection, output no match row ids. |
| 254 | + uint32_t num_rows_no_match; |
| 255 | + std::vector<uint16_t> row_ids_out(num_rows); |
| 256 | + KeyCompare::CompareColumnsToRows(num_rows, /*sel_left_maybe_null=*/NULLPTR, |
| 257 | + row_ids_to_compare.data(), &ctx, &num_rows_no_match, |
| 258 | + row_ids_out.data(), columns_left, row_table_right, |
| 259 | + /*are_cols_in_encoding_order=*/true, |
| 260 | + /*out_match_bitvector_maybe_null=*/NULLPTR); |
| 261 | + ASSERT_EQ(num_rows_no_match, 0); |
| 262 | + } |
| 263 | + |
| 264 | + { |
| 265 | + // No selection, output match bit vector. |
| 266 | + std::vector<uint8_t> match_bitvector(BytesForBits(num_rows)); |
| 267 | + KeyCompare::CompareColumnsToRows( |
| 268 | + num_rows, /*sel_left_maybe_null=*/NULLPTR, row_ids_to_compare.data(), &ctx, |
| 269 | + /*out_num_rows=*/NULLPTR, /*out_sel_left_maybe_same=*/NULLPTR, columns_left, |
| 270 | + row_table_right, |
| 271 | + /*are_cols_in_encoding_order=*/true, match_bitvector.data()); |
| 272 | + ASSERT_EQ(arrow::internal::CountSetBits(match_bitvector.data(), 0, num_rows), |
| 273 | + num_rows); |
| 274 | + } |
| 275 | + |
| 276 | + std::vector<uint16_t> selection_left(num_rows); |
| 277 | + std::iota(selection_left.begin(), selection_left.end(), 0); |
| 278 | + |
| 279 | + { |
| 280 | + // With selection, output no match row ids. |
| 281 | + uint32_t num_rows_no_match; |
| 282 | + std::vector<uint16_t> row_ids_out(num_rows); |
| 283 | + KeyCompare::CompareColumnsToRows(num_rows, selection_left.data(), |
| 284 | + row_ids_to_compare.data(), &ctx, &num_rows_no_match, |
| 285 | + row_ids_out.data(), columns_left, row_table_right, |
| 286 | + /*are_cols_in_encoding_order=*/true, |
| 287 | + /*out_match_bitvector_maybe_null=*/NULLPTR); |
| 288 | + ASSERT_EQ(num_rows_no_match, 0); |
| 289 | + } |
| 290 | + |
| 291 | + { |
| 292 | + // With selection, output match bit vector. |
| 293 | + std::vector<uint8_t> match_bitvector(BytesForBits(num_rows)); |
| 294 | + KeyCompare::CompareColumnsToRows( |
| 295 | + num_rows, selection_left.data(), row_ids_to_compare.data(), &ctx, |
| 296 | + /*out_num_rows=*/NULLPTR, /*out_sel_left_maybe_same=*/NULLPTR, columns_left, |
| 297 | + row_table_right, |
| 298 | + /*are_cols_in_encoding_order=*/true, match_bitvector.data()); |
| 299 | + ASSERT_EQ(arrow::internal::CountSetBits(match_bitvector.data(), 0, num_rows), |
| 300 | + num_rows); |
| 301 | + } |
| 302 | +} |
| 303 | +#endif // ARROW_VALGRIND |
| 304 | + |
167 | 305 | } // namespace compute |
168 | 306 | } // namespace arrow |
0 commit comments