| 286 | } |
| 287 | |
| 288 | Result<std::shared_ptr<ArrayData>> HashJoinDictBuild::RemapInputValues( |
| 289 | ExecContext* ctx, const Datum& values, int64_t batch_length) const { |
| 290 | // Initialize encoder |
| 291 | // |
| 292 | RowEncoder encoder; |
| 293 | std::vector<TypeHolder> encoder_types = {value_type_}; |
| 294 | encoder.Init(encoder_types, ctx); |
| 295 | |
| 296 | // Encode all |
| 297 | // |
| 298 | ARROW_DCHECK(values.is_array() || values.is_scalar()); |
| 299 | bool is_scalar = values.is_scalar(); |
| 300 | int64_t encoded_length = is_scalar ? 1 : batch_length; |
| 301 | ExecBatch batch({values}, encoded_length); |
| 302 | RETURN_NOT_OK(encoder.EncodeAndAppend(ExecSpan(batch))); |
| 303 | |
| 304 | // Allocate output buffers |
| 305 | // |
| 306 | ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Buffer> non_nulls_buf, |
| 307 | AllocateBitmap(batch_length, ctx->memory_pool())); |
| 308 | ARROW_ASSIGN_OR_RAISE( |
| 309 | std::shared_ptr<Buffer> ids_buf, |
| 310 | AllocateBuffer(batch_length * sizeof(int32_t), ctx->memory_pool())); |
| 311 | uint8_t* non_nulls = non_nulls_buf->mutable_data(); |
| 312 | int32_t* ids = reinterpret_cast<int32_t*>(ids_buf->mutable_data()); |
| 313 | memset(non_nulls, 0xff, bit_util::BytesForBits(batch_length)); |
| 314 | |
| 315 | // Populate output buffers (for scalar only the first entry is populated) |
| 316 | // |
| 317 | for (int64_t i = 0; i < encoded_length; ++i) { |
| 318 | std::string str = encoder.encoded_row(static_cast<int32_t>(i)); |
| 319 | if (KeyEncoder::IsNull(reinterpret_cast<const uint8_t*>(str.data()))) { |
| 320 | // Map nulls to nulls |
| 321 | bit_util::ClearBit(non_nulls, i); |
| 322 | ids[i] = HashJoinDictUtil::kNullId; |
| 323 | } else { |
| 324 | auto iter = hash_table_.find(str); |
| 325 | if (iter == hash_table_.end()) { |
| 326 | ids[i] = HashJoinDictUtil::kMissingValueId; |
| 327 | } else { |
| 328 | ids[i] = iter->second; |
| 329 | } |
| 330 | } |
| 331 | } |
| 332 | |
| 333 | // Generate array of repeated values for scalar input |
| 334 | // |
| 335 | if (is_scalar) { |
| 336 | if (!bit_util::GetBit(non_nulls, 0)) { |
| 337 | memset(non_nulls, 0, bit_util::BytesForBits(batch_length)); |
| 338 | } |
| 339 | for (int64_t i = 1; i < batch_length; ++i) { |
| 340 | ids[i] = ids[0]; |
| 341 | } |
| 342 | } |
| 343 | |
| 344 | return ArrayData::Make(DataTypeAfterRemapping(), batch_length, |
| 345 | {std::move(non_nulls_buf), std::move(ids_buf)}); |
no test coverage detected