| 954 | } |
| 955 | |
| 956 | Status SwissTableWithKeys::Map(Input* input, bool insert_missing, const uint32_t* hashes, |
| 957 | uint8_t* match_bitvector_maybe_null, uint32_t* key_ids) { |
| 958 | arrow::util::TempVectorStack* temp_stack = input->temp_stack; |
| 959 | |
| 960 | // Split into smaller mini-batches |
| 961 | // |
| 962 | int minibatch_size = swiss_table_.minibatch_size(); |
| 963 | int num_rows_to_process = input->selection_maybe_null |
| 964 | ? input->num_selected |
| 965 | : input->batch_end_row - input->batch_start_row; |
| 966 | auto hashes_buf = arrow::util::TempVectorHolder<uint32_t>(temp_stack, minibatch_size); |
| 967 | auto match_bitvector_buf = arrow::util::TempVectorHolder<uint8_t>( |
| 968 | temp_stack, |
| 969 | static_cast<uint32_t>(bit_util::BytesForBits(minibatch_size)) + sizeof(uint64_t)); |
| 970 | for (int minibatch_start = 0; minibatch_start < num_rows_to_process;) { |
| 971 | int minibatch_size_next = |
| 972 | std::min(minibatch_size, num_rows_to_process - minibatch_start); |
| 973 | |
| 974 | // Prepare updated input buffers that represent the current minibatch. |
| 975 | // |
| 976 | Input minibatch_input(*input, minibatch_start, minibatch_size_next); |
| 977 | uint8_t* minibatch_match_bitvector = |
| 978 | insert_missing ? match_bitvector_buf.mutable_data() |
| 979 | : match_bitvector_maybe_null + minibatch_start / 8; |
| 980 | const uint32_t* minibatch_hashes; |
| 981 | if (input->selection_maybe_null) { |
| 982 | minibatch_hashes = hashes_buf.mutable_data(); |
| 983 | for (int i = 0; i < minibatch_size_next; ++i) { |
| 984 | hashes_buf.mutable_data()[i] = hashes[minibatch_input.selection_maybe_null[i]]; |
| 985 | } |
| 986 | } else { |
| 987 | minibatch_hashes = hashes + minibatch_start; |
| 988 | } |
| 989 | uint32_t* minibatch_key_ids = key_ids + minibatch_start; |
| 990 | |
| 991 | // Lookup existing keys. |
| 992 | { |
| 993 | auto slots = |
| 994 | arrow::util::TempVectorHolder<uint8_t>(temp_stack, minibatch_size_next); |
| 995 | swiss_table_.early_filter(minibatch_size_next, minibatch_hashes, |
| 996 | minibatch_match_bitvector, slots.mutable_data()); |
| 997 | swiss_table_.find(minibatch_size_next, minibatch_hashes, minibatch_match_bitvector, |
| 998 | slots.mutable_data(), minibatch_key_ids, temp_stack, equal_impl_, |
| 999 | &minibatch_input); |
| 1000 | } |
| 1001 | |
| 1002 | // Perform inserts of missing keys if required. |
| 1003 | // |
| 1004 | if (insert_missing) { |
| 1005 | auto ids_buf = |
| 1006 | arrow::util::TempVectorHolder<uint16_t>(temp_stack, minibatch_size_next); |
| 1007 | int num_ids; |
| 1008 | arrow::util::bit_util::bits_to_indexes( |
| 1009 | 0, swiss_table_.hardware_flags(), minibatch_size_next, |
| 1010 | minibatch_match_bitvector, &num_ids, ids_buf.mutable_data()); |
| 1011 | |
| 1012 | RETURN_NOT_OK(swiss_table_.map_new_keys( |
| 1013 | num_ids, ids_buf.mutable_data(), minibatch_hashes, minibatch_key_ids, |