EIC code displayed by LXR master/​include/​arrow/​acero/​hash_join_dict.h	Source navigation Diff markup Identifier search General search
	Version: master test Revert   Change

File indexing completed on 2025-08-28 08:26:53

 // 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.
 
 #pragma once
 
 #include <memory>
 #include <unordered_map>
 
 #include "arrow/acero/schema_util.h"
 #include "arrow/compute/exec.h"
 #include "arrow/compute/row/row_encoder_internal.h"
 #include "arrow/result.h"
 #include "arrow/status.h"
 #include "arrow/type.h"
 
 // This file contains hash join logic related to handling of dictionary encoded key
 // columns.
 //
 // A key column from probe side of the join can be matched against a key column from build
 // side of the join, as long as the underlying value types are equal. That means that:
 // - both scalars and arrays can be used and even mixed in the same column
 // - dictionary column can be matched against non-dictionary column if underlying value
 // types are equal
 // - dictionary column can be matched against dictionary column with a different index
 // type, and potentially using a different dictionary, if underlying value types are equal
 //
 // We currently require in hash join that for all dictionary encoded columns, the same
 // dictionary is used in all input exec batches.
 //
 // In order to allow matching columns with different dictionaries, different dictionary
 // index types, and dictionary key against non-dictionary key, internally comparisons will
 // be evaluated after remapping values on both sides of the join to a common
 // representation (which will be called "unified representation"). This common
 // representation is a column of int32() type (not a dictionary column). It represents an
 // index in the unified dictionary computed for the (only) dictionary present on build
 // side (an empty dictionary is still created for an empty build side). Null value is
 // always represented in this common representation as null int32 value, unified
 // dictionary will never contain a null value (so there is no ambiguity of representing
 // nulls as either index to a null entry in the dictionary or null index).
 //
 // Unified dictionary represents values present on build side. There may be values on
 // probe side that are not present in it. All such values, that are not null, are mapped
 // in the common representation to a special constant kMissingValueId.
 //
 
 namespace arrow {
 
 using compute::ExecBatch;
 using compute::ExecContext;
 using compute::internal::RowEncoder;
 
 namespace acero {
 
 /// Helper class with operations that are stateless and common to processing of dictionary
 /// keys on both build and probe side.
 class HashJoinDictUtil {
  public:
   // Null values in unified representation are always represented as null that has
   // corresponding integer set to this constant
   static constexpr int32_t kNullId = 0;
   // Constant representing a value, that is not null, missing on the build side, in
   // unified representation.
   static constexpr int32_t kMissingValueId = -1;
 
   // Check if data types of corresponding pair of key column on build and probe side are
   // compatible
   static bool KeyDataTypesValid(const std::shared_ptr<DataType>& probe_data_type,
                                 const std::shared_ptr<DataType>& build_data_type);
 
   // Input must be dictionary array or dictionary scalar.
   // A precomputed and provided here lookup table in the form of int32() array will be
   // used to remap input indices to unified representation.
   //
   static Result<std::shared_ptr<ArrayData>> IndexRemapUsingLUT(
       ExecContext* ctx, const Datum& indices, int64_t batch_length,
       const std::shared_ptr<ArrayData>& map_array,
       const std::shared_ptr<DataType>& data_type);
 
   // Return int32() array that contains indices of input dictionary array or scalar after
   // type casting.
   static Result<std::shared_ptr<ArrayData>> ConvertToInt32(
       const std::shared_ptr<DataType>& from_type, const Datum& input,
       int64_t batch_length, ExecContext* ctx);
 
   // Return an array that contains elements of input int32() array after casting to a
   // given integer type. This is used for mapping unified representation stored in the
   // hash table on build side back to original input data type of hash join, when
   // outputting hash join results to parent exec node.
   //
   static Result<std::shared_ptr<ArrayData>> ConvertFromInt32(
       const std::shared_ptr<DataType>& to_type, const Datum& input, int64_t batch_length,
       ExecContext* ctx);
 
   // Return dictionary referenced in either dictionary array or dictionary scalar
   static std::shared_ptr<Array> ExtractDictionary(const Datum& data);
 };
 
 /// Implements processing of dictionary arrays/scalars in key columns on the build side of
 /// a hash join.
 /// Each instance of this class corresponds to a single column and stores and
 /// processes only the information related to that column.
 /// Const methods are thread-safe, non-const methods are not (the caller must make sure
 /// that only one thread at any time will access them).
 ///
 class HashJoinDictBuild {
  public:
   // Returns true if the key column (described in input by its data type) requires any
   // pre- or post-processing related to handling dictionaries.
   //
   static bool KeyNeedsProcessing(const std::shared_ptr<DataType>& build_data_type) {
     return (build_data_type->id() == Type::DICTIONARY);
   }
 
   // Data type of unified representation
   static std::shared_ptr<DataType> DataTypeAfterRemapping() { return int32(); }
 
   // Should be called only once in hash join, before processing any build or probe
   // batches.
   //
   // Takes a pointer to the dictionary for a corresponding key column on the build side as
   // an input. If the build side is empty, it still needs to be called, but with
   // dictionary pointer set to null.
   //
   // Currently it is required that all input batches on build side share the same
   // dictionary. For each input batch during its pre-processing, dictionary will be
   // checked and error will be returned if it is different then the one provided in the
   // call to this method.
   //
   // Unifies the dictionary. The order of the values is still preserved.
   // Null and duplicate entries are removed. If the dictionary is already unified, its
   // copy will be produced and stored within this class.
   //
   // Prepares the mapping from ids within original dictionary to the ids in the resulting
   // dictionary. This is used later on to pre-process (map to unified representation) key
   // column on build side.
   //
   // Prepares the reverse mapping (in the form of hash table) from values to the ids in
   // the resulting dictionary. This will be used later on to pre-process (map to unified
   // representation) key column on probe side. Values on probe side that are not present
   // in the original dictionary will be mapped to a special constant kMissingValueId. The
   // exception is made for nulls, which get always mapped to nulls (both when null is
   // represented as a dictionary id pointing to a null and a null dictionary id).
   //
   Status Init(ExecContext* ctx, std::shared_ptr<Array> dictionary,
               std::shared_ptr<DataType> index_type, std::shared_ptr<DataType> value_type);
 
   // Remap array or scalar values into unified representation (array of int32()).
   // Outputs kMissingValueId if input value is not found in the unified dictionary.
   // Outputs null for null input value (with corresponding data set to kNullId).
   //
   Result<std::shared_ptr<ArrayData>> RemapInputValues(ExecContext* ctx,
                                                       const Datum& values,
                                                       int64_t batch_length) const;
 
   // Remap dictionary array or dictionary scalar on build side to unified representation.
   // Dictionary referenced in the input must match the dictionary that was
   // given during initialization.
   // The output is a dictionary array that references unified dictionary.
   //
   Result<std::shared_ptr<ArrayData>> RemapInput(
       ExecContext* ctx, const Datum& indices, int64_t batch_length,
       const std::shared_ptr<DataType>& data_type) const;
 
   // Outputs dictionary array referencing unified dictionary, given an array with 32-bit
   // ids.
   // Used to post-process values looked up in a hash table on build side of the hash join
   // before outputting to the parent exec node.
   //
   Result<std::shared_ptr<ArrayData>> RemapOutput(const ArrayData& indices32Bit,
                                                  ExecContext* ctx) const;
 
   // Release shared pointers and memory
   void CleanUp();
 
  private:
   // Data type of dictionary ids for the input dictionary on build side
   std::shared_ptr<DataType> index_type_;
   // Data type of values for the input dictionary on build side
   std::shared_ptr<DataType> value_type_;
   // Mapping from (encoded as string) values to the ids in unified dictionary
   std::unordered_map<std::string, int32_t> hash_table_;
   // Mapping from input dictionary ids to unified dictionary ids
   std::shared_ptr<ArrayData> remapped_ids_;
   // Input dictionary
   std::shared_ptr<Array> dictionary_;
   // Unified dictionary
   std::shared_ptr<ArrayData> unified_dictionary_;
 };
 
 /// Implements processing of dictionary arrays/scalars in key columns on the probe side of
 /// a hash join.
 /// Each instance of this class corresponds to a single column and stores and
 /// processes only the information related to that column.
 /// It is not thread-safe - every participating thread should use its own instance of
 /// this class.
 ///
 class HashJoinDictProbe {
  public:
   static bool KeyNeedsProcessing(const std::shared_ptr<DataType>& probe_data_type,
                                  const std::shared_ptr<DataType>& build_data_type);
 
   // Data type of the result of remapping input key column.
   //
   // The result of remapping is what is used in hash join for matching keys on build and
   // probe side. The exact data types may be different, as described below, and therefore
   // a common representation is needed for simplifying comparisons of pairs of keys on
   // both sides.
   //
   // We support matching key that is of non-dictionary type with key that is of dictionary
   // type, as long as the underlying value types are equal. We support matching when both
   // keys are of dictionary type, regardless whether underlying dictionary index types are
   // the same or not.
   //
   static std::shared_ptr<DataType> DataTypeAfterRemapping(
       const std::shared_ptr<DataType>& build_data_type);
 
   // Should only be called if KeyNeedsProcessing method returns true for a pair of
   // corresponding key columns from build and probe side.
   // Converts values in order to match the common representation for
   // both build and probe side used in hash table comparison.
   // Supports arrays and scalars as input.
   // Argument opt_build_side should be null if dictionary key on probe side is matched
   // with non-dictionary key on build side.
   //
   Result<std::shared_ptr<ArrayData>> RemapInput(
       const HashJoinDictBuild* opt_build_side, const Datum& data, int64_t batch_length,
       const std::shared_ptr<DataType>& probe_data_type,
       const std::shared_ptr<DataType>& build_data_type, ExecContext* ctx);
 
   void CleanUp();
 
  private:
   // May be null if probe side key is non-dictionary. Otherwise it is used to verify that
   // only a single dictionary is referenced in exec batch on probe side of hash join.
   std::shared_ptr<Array> dictionary_;
   // Mapping from dictionary on probe side of hash join (if it is used) to unified
   // representation.
   std::shared_ptr<ArrayData> remapped_ids_;
   // Encoder of key columns that uses unified representation instead of original data type
   // for key columns that need to use it (have dictionaries on either side of the join).
   RowEncoder encoder_;
 };
 
 // Encapsulates dictionary handling logic for build side of hash join.
 //
 class HashJoinDictBuildMulti {
  public:
   Status Init(const SchemaProjectionMaps<HashJoinProjection>& proj_map,
               const ExecBatch* opt_non_empty_batch, ExecContext* ctx);
   static void InitEncoder(const SchemaProjectionMaps<HashJoinProjection>& proj_map,
                           RowEncoder* encoder, ExecContext* ctx);
   Status EncodeBatch(size_t thread_index,
                      const SchemaProjectionMaps<HashJoinProjection>& proj_map,
                      const ExecBatch& batch, RowEncoder* encoder, ExecContext* ctx) const;
   Status PostDecode(const SchemaProjectionMaps<HashJoinProjection>& proj_map,
                     ExecBatch* decoded_key_batch, ExecContext* ctx);
   const HashJoinDictBuild& get_dict_build(int icol) const { return remap_imp_[icol]; }
 
  private:
   std::vector<bool> needs_remap_;
   std::vector<HashJoinDictBuild> remap_imp_;
 };
 
 // Encapsulates dictionary handling logic for probe side of hash join
 //
 class HashJoinDictProbeMulti {
  public:
   void Init(size_t num_threads);
   bool BatchRemapNeeded(size_t thread_index,
                         const SchemaProjectionMaps<HashJoinProjection>& proj_map_probe,
                         const SchemaProjectionMaps<HashJoinProjection>& proj_map_build,
                         ExecContext* ctx);
   Status EncodeBatch(size_t thread_index,
                      const SchemaProjectionMaps<HashJoinProjection>& proj_map_probe,
                      const SchemaProjectionMaps<HashJoinProjection>& proj_map_build,
                      const HashJoinDictBuildMulti& dict_build, const ExecBatch& batch,
                      RowEncoder** out_encoder, ExecBatch* opt_out_key_batch,
                      ExecContext* ctx);
 
  private:
   void InitLocalStateIfNeeded(
       size_t thread_index, const SchemaProjectionMaps<HashJoinProjection>& proj_map_probe,
       const SchemaProjectionMaps<HashJoinProjection>& proj_map_build, ExecContext* ctx);
   static void InitEncoder(const SchemaProjectionMaps<HashJoinProjection>& proj_map_probe,
                           const SchemaProjectionMaps<HashJoinProjection>& proj_map_build,
                           RowEncoder* encoder, ExecContext* ctx);
   struct ThreadLocalState {
     bool is_initialized;
     // Whether any key column needs remapping (because of dictionaries used) before doing
     // join hash table lookups
     bool any_needs_remap;
     // Whether each key column needs remapping before doing join hash table lookups
     std::vector<bool> needs_remap;
     std::vector<HashJoinDictProbe> remap_imp;
     // Encoder of key columns that uses unified representation instead of original data
     // type for key columns that need to use it (have dictionaries on either side of the
     // join).
     RowEncoder post_remap_encoder;
   };
   std::vector<ThreadLocalState> local_states_;
 };
 
 }  // namespace acero
 }  // namespace arrow

This page was automatically generated by the 2.3.7 LXR engine. The LXR team