temp

Author: yl-lisen

src/Interpreters/Streaming/AsofHashJoin.cpp

@@ -0,0 +1,99 @@
+#pragma once
+
+#include <Interpreters/Streaming/HashJoin.h>
+
+namespace DB
+{
+namespace Streaming
+{
+AsofHashJoin::AsofHashJoin(
+    std::shared_ptr<TableJoin> table_join_,
+    JoinStreamDescriptionPtr left_join_stream_desc_,
+    JoinStreamDescriptionPtr right_join_stream_desc_)
+    : HashJoin(table_join_, left_join_stream_desc_, right_join_stream_desc_)
+    , asof_type(*table_join->getAsofType())
+    , asof_inequality(table_join->getAsofInequality())
+{
+}
+
+void AsofHashJoin::joinLeftBlock(Block & block)
+{
+    doJoinBlockWithHashTable<true>(block, hash_blocks);
+}
+
+void AsofHashJoin::insertRightBlock(Block block)
+{
+    /// FIXME, there are quite some block copies
+    /// FIXME, all_key_columns shall hold shared_ptr to columns instead of raw ptr
+    /// then we can update `source_block` in place
+    /// key columns are from source `block`
+    ColumnRawPtrMap all_key_columns = JoinCommon::materializeColumnsInplaceMap(block, table_join->getAllNames(JoinTableSide::Right));
+
+    /// We have copy of source `block` to `block_to_save` after prepare, so `block_to_save` is good to get moved to the buffered stream data
+    Block block_to_save = prepareBlockToSave(block, right_data.buffered_data->sample_block);
+
+    /// FIXME, multiple disjuncts OR clause
+    ColumnRawPtrs key_columns;
+    const Names & key_names = table_join->getClauses().front().key_names_right;
+    key_columns.reserve(key_names.size());
+    for (const auto & name : key_names)
+        key_columns.push_back(all_key_columns[name]);
+
+    /// We will insert to the map only keys, where all components are not NULL.
+    ConstNullMapPtr null_map{};
+    ColumnPtr null_map_holder = extractNestedColumnsAndNullMap(key_columns, null_map);
+
+    /// If LEFT, RIGHT or FULL save blocks with nulls for NotJoinedBlocks
+    UInt8 save_nullmap = 0;
+    if (isRightOrFull(table_join->kind()) && null_map)
+    {
+        /// Save rows with NULL keys
+        for (size_t i = 0; !save_nullmap && i < null_map->size(); ++i)
+            save_nullmap |= (*null_map)[i];
+    }
+
+    /// Add `block_to_save` to target stream data
+    /// Note `block_to_save` may be empty for cases in which the query doesn't care other non-key columns.
+    /// For example, SELECT count() FROM stream_a JOIN stream_b ON i=ii;
+    auto start_row = buffered_hash_data->addOrConcatDataBlock(std::move(block_to_save));
+    auto rows = buffered_hash_data->lastDataBlock().rows();
+
+    switch (hash_method_type)
+    {
+#define M(TYPE) \
+    case HashType::TYPE: \
+        return insertFromBlockImplType< \
+            Strictness::Asof, \
+            typename KeyGetterForType<HashType::TYPE, std::remove_reference_t<decltype(*(buffered_hash_data->maps->TYPE))>>::Type>( \
+            join, \
+            *(buffered_hash_data->maps->TYPE), \
+            rows, \
+            key_columns, \
+            key_sizes[0], \
+            &buffered_hash_data->blocks, \
+            start_row, \
+            null_map, \
+            buffered_hash_data->pool); \
+        break;
+        APPLY_FOR_HASH_KEY_VARIANTS(M)
+#undef M
+    }
+    insertFromBlockImpl<strictness_>(
+                hash_method_type,
+                map,
+                rows,
+                key_columns,
+                key_sizes[0],
+                &target_hash_blocks->blocks,
+                start_row,
+                null_map,
+                target_hash_blocks->pool);
+
+    if (save_nullmap)
+        /// FIXME, we will need account the allocated bytes for null_map_holder / not_joined_map as well
+        buffered_hash_data->blocks_nullmaps.emplace_back(&buffered_hash_data->lastDataBlock(), null_map_holder);
+
+    checkLimits();
+}
+}
+}

src/Interpreters/Streaming/AsofHashJoin.h

@@ -9,8 +9,23 @@ namespace Streaming
 class AsofHashJoin final : public HashJoin
 {
 public:
-    using HashJoin::HashJoin;
+    AsofHashJoin(
+        std::shared_ptr<TableJoin> table_join_,
+        JoinStreamDescriptionPtr left_join_stream_desc_,
+        JoinStreamDescriptionPtr right_join_stream_desc_);
+
     HashJoinType type() const override { return HashJoinType::Asof; }
+
+    void joinLeftBlock(Block & block) override;
+    void insertRightBlock(Block block) override;
+
+private:
+    TypeIndex asof_type;
+    ASOFJoinInequality asof_inequality;
+
+    using DataBlock = LightChunkWithTimestamp;
+    using BufferedAsofHashData = BufferedHashData<DataBlock, AsofRowRefs<DataBlock>>;
+    SERDE std::unique_ptr<BufferedAsofHashData> buffered_hash_data;
 };
 
 }

src/Interpreters/Streaming/HashJoin.h

@@ -255,8 +255,8 @@ class HashJoin : public IHashJoin
 
 private:
     /// For non-bidirectional hash join
-    void insertRightBlock(Block right_block);
-    void joinLeftBlock(Block & left_block);
+    virtual void insertRightBlock(Block right_block);
+    virtual void joinLeftBlock(Block & left_block);
 
     /// For bidirectional hash join
     /// There are 2 blocks returned : joined block via parameter and retracted block via returned-value if there is

src/Interpreters/Streaming/joinData.h

@@ -229,5 +229,68 @@ SERDE struct BufferedStreamData
 using BucketBlocks = std::vector<BufferedStreamData::BucketBlock>;
 
 using BufferedStreamDataPtr = std::unique_ptr<BufferedStreamData>;
+
+template <typename JoinDataBlock, typename MappedRowRef>
+struct BufferedHashData
+{
+    using JoinDataBlockList = RefCountDataBlockList<JoinDataBlock>;
+    using JoinDataBlockRawPtr = const JoinDataBlock *;
+    using BlockNullmapList = std::deque<std::pair<JoinDataBlockRawPtr, ColumnPtr>>;
+    using Map = HashMapsTemplate<MappedRowRef>;
+
+    BufferedHashData(size_t data_block_size, CachedBlockMetrics & metrics);
+
+    ~BufferedHashData();
+
+    template <typename RowRefHandler>
+    void insert(JoinDataBlock && block, RowRefHandler && handler)
+    {
+        /// Add block into `blocks`
+        auto start_row = blocks.pushBackOrConcat(std::move(block));
+
+        /// Insert MappedRowRef into hash map (which is referenced to `blocks`)
+        const IColumn * asof_column [[maybe_unused]] = nullptr;
+        if constexpr (is_range_asof_join || is_asof_join)
+            asof_column = key_columns.back();
+
+        auto key_columns
+        if constexpr (is_asof_join)
+        {
+            auto key_column_copy = key_columns;
+            auto key_size_copy = key_sizes;
+            key_column_copy.pop_back();
+            key_size_copy.pop_back();
+            return KeyGetter(key_column_copy, key_size_copy, nullptr);
+        }
+        else
+            return KeyGetter(key_columns, key_sizes, nullptr);
+
+        auto key_getter = createKeyGetter < KeyGetter, is_range_asof_join || is_asof_join > (key_columns, key_sizes);
+
+        for (size_t i = start_row; i < rows; ++i)
+        {
+            auto emplace_result = key_getter.emplaceKey(map, i, pool);
+            handler(emplace_result.getMapped(), emplace_result.isInserted());
+        }
+    }
+
+    const JoinDataBlock & lastDataBlock() const { return blocks.lastDataBlock(); }
+
+    HashMapSizes hashMapSizes(const HashJoin * hash_join) const;
+
+    /// Buffered data
+    JoinDataBlockList blocks;
+
+    /// Additional data - strings for string keys and continuation elements of single-linked lists of references to rows.
+    Arena pool;
+
+    /// Hash maps variants are attached to original source blocks, and will be garbage collected
+    /// automatically along with the source blocks. Hence put it here instead of BufferedStreamData
+    Map map;
+
+    /// FIXME: support nullable type
+    // BlockNullmapList blocks_nullmaps; /// Nullmaps for blocks of "right" table (if needed)
+};
+
 }
 }