kszucs commented on code in PR #45001: URL: https://github.com/apache/arrow/pull/45001#discussion_r3921602486
########## cpp/src/arrow/compute/kernels/scalar_hash.cc: ########## @@ -0,0 +1,471 @@ +// 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 <algorithm> +#include <utility> + +#include "arrow/array/array_base.h" +#include "arrow/array/util.h" +#include "arrow/compute/cast.h" +#include "arrow/compute/kernels/common_internal.h" +#include "arrow/compute/key_hash_internal.h" +#include "arrow/compute/light_array_internal.h" +#include "arrow/compute/registry_internal.h" +#include "arrow/compute/util.h" +#include "arrow/result.h" +#include "arrow/util/bit_run_reader.h" +#include "arrow/util/bit_util.h" +#include "arrow/util/bitmap_ops.h" + +namespace arrow { +namespace compute { +namespace internal { + +// Define symbols visible within `arrow::compute::internal` in this file; +// these symbols are not visible outside of this file. +namespace { + +// ------------------------------ +// Kernel implementations +// It is expected that HashArrowType is either UInt32Type or UInt64Type (default) + +// Free function (not dependent on ArrowType/Hasher) to avoid codegen per instantiation. +// Only called with a plain column; HashArray routes everything else (see +// NeedsRecursiveHash) elsewhere first. +Result<KeyColumnArray> ToColumnArray(const ArraySpan& array) { + KeyColumnMetadata metadata; + const uint8_t* validity_buffer = nullptr; + const uint8_t* fixed_length_buffer = nullptr; + const uint8_t* var_length_buffer = nullptr; + + if (array.GetBuffer(0) != nullptr) { + validity_buffer = array.GetBuffer(0)->data(); + } + if (array.GetBuffer(1) != nullptr) { + fixed_length_buffer = array.GetBuffer(1)->data(); + } + + auto type = array.type; + auto type_id = type->id(); + if (type_id == Type::NA) { + metadata = KeyColumnMetadata(true, 0, true); + } else if (type_id == Type::BOOL) { + metadata = KeyColumnMetadata(true, 0); + } else if (is_fixed_width(type_id)) { + metadata = KeyColumnMetadata(true, type->bit_width() / 8); + } else if (is_binary_like(type_id)) { + metadata = KeyColumnMetadata(false, sizeof(uint32_t)); + if (array.GetBuffer(2) != nullptr) { + var_length_buffer = array.GetBuffer(2)->data(); + } + } else if (is_large_binary_like(type_id)) { + metadata = KeyColumnMetadata(false, sizeof(uint64_t)); + if (array.GetBuffer(2) != nullptr) { + var_length_buffer = array.GetBuffer(2)->data(); + } + } else { + return Status::TypeError("Unsupported column data type ", type->name(), + " used with hash32/hash64 compute kernel"); + } + + return KeyColumnArray(metadata, array.length, validity_buffer, fixed_length_buffer, + var_length_buffer); +} + +// Whether HashArray must handle `type` itself rather than passing it to +// ToColumnArray/HashMultiColumn. Broader than is_nested(): EXTENSION and DICTIONARY +// aren't nested, but ToColumnArray has no case for either (and hashing a dictionary's +// raw indices would be wrong anyway). A zero-width fixed_size_binary isn't nested +// either, but ToColumnArray can only describe it as a fixed-width column of length 0 -- +// exactly how a bit-packed boolean is encoded -- so HashMultiColumn would take each +// row's hash from a bit that doesn't exist (see HashArray's dedicated branch). +bool NeedsRecursiveHash(const DataType& type) { + auto type_id = type.id(); + return type_id == Type::EXTENSION || type_id == Type::DICTIONARY || + is_nested(type_id) || + (type_id == Type::FIXED_SIZE_BINARY && type.byte_width() == 0); +} + +// Writes `array`'s own validity into `out_validity` (a fresh 0-offset bitmap), rebasing +// off array.offset. Only called for types whose validity really is a plain bitmap; union +// and run-end-encoded, which ArraySpan::IsValid computes specially, never reach here. +void WriteOwnValidity(const ArraySpan& array, uint8_t* out_validity) { + if (array.GetBuffer(0) == nullptr) { + // No bitmap: every row shares one answer, all valid except for NullType, which is + // implicitly all-null. + bit_util::SetBitsTo(out_validity, 0, array.length, + /*bits_are_set=*/array.type->id() != Type::NA); + return; + } + ::arrow::internal::CopyBitmap(array.GetBuffer(0)->data(), array.offset, array.length, + out_validity, /*dest_offset=*/0); +} + +// Folds one row's child hashes into a single hash. Seeded with CombineHashes(0, 0) rather +// than 0 just so an empty list doesn't hash to a bare 0 -- a hash-quality nicety, not a +// requirement, since a list row's validity is independent of its hash value. +template <typename c_type, typename Hasher> +c_type CombineRange(const c_type* value_hashes, int64_t start, int64_t end) { + c_type combined = Hasher::CombineHashes(0, 0); + for (int64_t j = start; j < end; j++) { + combined = Hasher::CombineHashes(combined, value_hashes[j]); + } + return combined; +} + +template <typename ArrowType, typename Hasher> +struct FastHashScalar { + using c_type = typename ArrowType::c_type; + + // Hashes the [offset, offset + length) slice of `child` into hash values plus real + // validity, always based at offset 0 whatever `child`'s own offset (callers read the + // buffers row-0-based). Only a null row's validity bit is meaningful, not its hash + // value; callers folding these into a parent hash must handle that (see + // HashListArray). + static Result<std::shared_ptr<ArrayData>> HashChild(const ArraySpan& child, + int64_t offset, int64_t length, + LightContext* hash_ctx, + ExecContext* exec_ctx) { + auto sliced = child; + sliced.SetSlice(offset, length); + auto arrow_type = TypeTraits<ArrowType>::type_singleton(); + ARROW_ASSIGN_OR_RAISE(auto buffer, AllocateBuffer(sliced.length * sizeof(c_type), + exec_ctx->memory_pool())); + ARROW_ASSIGN_OR_RAISE(auto validity, + AllocateBitmap(sliced.length, exec_ctx->memory_pool())); + ARROW_RETURN_NOT_OK(HashArray(sliced, hash_ctx, exec_ctx, + buffer->mutable_data_as<c_type>(), + validity->mutable_data())); + return ArrayData::Make(arrow_type, sliced.length, + {std::move(validity), std::move(buffer)}, kUnknownNullCount); + } + + static Status HashStructArray(const ArraySpan& array, LightContext* hash_ctx, + ExecContext* exec_ctx, c_type* out, + uint8_t* out_validity) { + // Row validity is the struct's own ANDed with every field's (in place, as + // swiss_join.cc does for multi-column nulls): an independently-null field makes the + // row invalid too (GH-17211), just like the struct row being null. + WriteOwnValidity(array, out_validity); + + if (array.child_data.empty()) { + // struct<>: HashMultiColumn needs >=1 column, so give every row one fixed hash; + // validity is already fully set above. + c_type empty_struct_hash = Hasher::CombineHashes(0, 0); + for (int64_t i = 0; i < array.length; i++) { + out[i] = empty_struct_hash; + } + return Status::OK(); + } + + std::vector<std::shared_ptr<ArrayData>> child_hashes(array.child_data.size()); + std::vector<KeyColumnArray> columns(array.child_data.size()); + for (size_t i = 0; i < array.child_data.size(); i++) { + // By reference: ArraySpan owns a child_data vector, so copying one heap-allocates. + const ArraySpan& child = array.child_data[i]; + // `child` may have its own offset independent of the struct's (see + // StructArray::GetFlattenedField): struct row r reads child row + // (child.offset + array.offset + r). + if (NeedsRecursiveHash(*child.type)) { + // StructArray::Slice() doesn't reslice child_data, so `child` may be larger + // than this slice of `array` references -- hash only the referenced range. + ARROW_ASSIGN_OR_RAISE(child_hashes[i], + HashChild(child, child.offset + array.offset, array.length, + hash_ctx, exec_ctx)); + ::arrow::internal::BitmapAnd(out_validity, 0, child_hashes[i]->buffers[0]->data(), + 0, array.length, 0, out_validity); + ARROW_ASSIGN_OR_RAISE(auto column, ToColumnArray(*child_hashes[i])); + // child_hashes[i] already covers exactly [0, array.length): no further slice. + columns[i] = column.Slice(0, array.length); + } else { + if (child.GetBuffer(0) != nullptr) { + ::arrow::internal::BitmapAnd(out_validity, 0, child.GetBuffer(0)->data(), + child.offset + array.offset, array.length, 0, + out_validity); + } else if (child.type->id() == Type::NA) { + // No bitmap, but NullType is implicitly all-null: invalidates every row. + bit_util::SetBitsTo(out_validity, 0, array.length, false); + } + ARROW_ASSIGN_OR_RAISE(auto column, ToColumnArray(child)); + columns[i] = column.Slice(child.offset + array.offset, array.length); + } + } + Hasher::HashMultiColumn(columns, hash_ctx, out); + return Status::OK(); + } + + // Handles FIXED_SIZE_LIST, LARGE_LIST, LIST, and MAP. `offsets` is null for + // FIXED_SIZE_LIST, which uses `list_size` as a constant stride instead. + template <typename OffsetT> + static Status HashListArray(const ArraySpan& array, int64_t list_size, + const OffsetT* offsets, LightContext* hash_ctx, + ExecContext* exec_ctx, c_type* out, uint8_t* out_validity) { + // The range of `values` this array actually references, as logical indices relative + // to values.offset. Needed because ArraySpan::SetSlice() doesn't reslice child_data, + // so `values` can be far larger than what this (possibly sliced) array covers. + // offsets[] are already such logical indices; FIXED_SIZE_LIST derives them from its + // constant stride instead. + int64_t rel_start = 0, rel_end = 0; + if (array.length > 0) { + if (offsets != nullptr) { + rel_start = offsets[0]; + rel_end = offsets[array.length]; + } else { + rel_start = array.offset * list_size; + rel_end = (array.offset + array.length) * list_size; + } + } + + // By reference: ArraySpan owns a child_data vector, so copying one heap-allocates. + const ArraySpan& values = array.child_data[0]; + // Element k of the result is original values row (values.offset + rel_start + k). + ARROW_ASSIGN_OR_RAISE(auto value_hashes, + HashChild(values, values.offset + rel_start, + rel_end - rel_start, hash_ctx, exec_ctx)); + // Zero the null elements' hashes: CombineRange folds values blind, and a null slot's + // bytes are undefined per the columnar spec, so otherwise a null element contributes + // whatever garbage it sat on and list<struct<f0:int32>> rows [{f0: 7}] and [null] + // (whose f0 slot also holds 7) hash alike. Filling only the gaps between runs of + // valid rows leaves the common all-valid case free. HashStructArray needs no + // equivalent: HashMultiColumn gets its fields' validity and already fixes each null + // row's contribution. + c_type* value_hash_data = value_hashes->buffers[1]->mutable_data_as<c_type>(); + int64_t valid_end = 0; + ::arrow::internal::VisitSetBitRunsVoid( + value_hashes->buffers[0]->data(), /*offset=*/0, value_hashes->length, + [&](int64_t position, int64_t run_length) { + std::fill(value_hash_data + valid_end, value_hash_data + position, c_type{0}); + valid_end = position + run_length; + }); + std::fill(value_hash_data + valid_end, value_hash_data + value_hashes->length, Review Comment: Verified, fixing. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
