relational_operators/SortMergeRunOperator.cpp - incubator-retired-quickstep - Git at Google

 /**
  * 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 "relational_operators/SortMergeRunOperator.hpp"

 #include <cstddef>
 #include <cstdlib>
 #include <utility>
 #include <vector>

 #include "query_execution/QueryExecutionTypedefs.hpp"
 #include "query_execution/WorkOrderProtosContainer.hpp"
 #include "query_execution/WorkOrdersContainer.hpp"
 #include "relational_operators/SortMergeRunOperator.pb.h"
 #include "relational_operators/SortMergeRunOperatorHelpers.hpp"
 #include "relational_operators/WorkOrder.pb.h"
 #include "threading/ThreadIDBasedMap.hpp"

 #include "glog/logging.h"

 #include "tmb/id_typedefs.h"

 namespace quickstep {

 class InsertDestination;

 using merge_run_operator::Run;
 using merge_run_operator::RunMerger;
 using merge_run_operator::MergeTree;

 bool SortMergeRunOperator::getAllWorkOrders(
     WorkOrdersContainer *container,
     QueryContext *query_context,
     StorageManager *storage_manager,
     const tmb::client_id scheduler_client_id,
     tmb::MessageBus *bus) {
   if (input_relation_is_stored_) {
     // Input blocks (or runs) are from base relation. Only possible when base
     // relation is stored sorted.
     if (!started_) {
       // Initialize merge tree completely, since all input runs are known.
       merge_tree_.initializeTree(input_relation_block_ids_.size());
       started_ = true;
       initializeInputRuns();
     }
   } else {
     // Input blocks (or runs) are pipelined from the sorted run generation
     // operator.
     if (!started_ && !input_stream_done_) {
       // Initialize merge tree for first pipeline mode.
       merge_tree_.initializeForPipeline();
       started_ = true;
       initializeInputRuns();
     }
   }
   // Generate runs from merge tree.
   return generateWorkOrders(container, query_context, storage_manager, scheduler_client_id, bus);
 }

 bool SortMergeRunOperator::getAllWorkOrderProtos(WorkOrderProtosContainer *container) {
   if (input_relation_is_stored_) {
     // Input blocks (or runs) are from base relation. Only possible when base
     // relation is stored sorted.
     if (!started_) {
       // Initialize merge tree completely, since all input runs are known.
       merge_tree_.initializeTree(input_relation_block_ids_.size());
       started_ = true;
       initializeInputRuns();
     }
   } else {
     // Input blocks (or runs) are pipelined from the sorted run generation
     // operator.
     if (!started_ && !input_stream_done_) {
       // Initialize merge tree for first pipeline mode.
       merge_tree_.initializeForPipeline();
       started_ = true;
       initializeInputRuns();
     }
   }

   // Get merge jobs from merge tree.
   std::vector<MergeTree::MergeJob> jobs;
   const bool done_generating = merge_tree_.getMergeJobs(&jobs);

   for (std::vector<MergeTree::MergeJob>::size_type job_id = 0;
        job_id < jobs.size();
        ++job_id) {
     // Add work order for each merge job.
     container->addWorkOrderProto(createWorkOrderProto(&jobs[job_id]), op_index_);
   }

   return done_generating;
 }

 serialization::WorkOrder* SortMergeRunOperator::createWorkOrderProto(
     merge_run_operator::MergeTree::MergeJob *job) {
   DCHECK(job != nullptr);
   DCHECK(!job->runs.empty());

   serialization::WorkOrder *proto = new serialization::WorkOrder;
   proto->set_work_order_type(serialization::SORT_MERGE_RUN);
   proto->set_query_id(query_id_);

   proto->SetExtension(serialization::SortMergeRunWorkOrder::operator_index, op_index_);
   proto->SetExtension(serialization::SortMergeRunWorkOrder::sort_config_index, sort_config_index_);

   for (const merge_run_operator::Run &run : job->runs) {
     serialization::Run *run_proto = proto->AddExtension(serialization::SortMergeRunWorkOrder::runs);
     for (const block_id block : run) {
       run_proto->add_blocks(block);
     }
   }

   proto->SetExtension(serialization::SortMergeRunWorkOrder::top_k, top_k_);
   proto->SetExtension(serialization::SortMergeRunWorkOrder::merge_level, job->level);
   proto->SetExtension(serialization::SortMergeRunWorkOrder::relation_id,
                       job->level > 0 ? run_relation_.getID()
                                      : input_relation_.getID());
   proto->SetExtension(serialization::SortMergeRunWorkOrder::insert_destination_index,
                       job->is_final_level ? output_destination_index_
                                           : run_block_destination_index_);

   return proto;
 }

 WorkOrder *SortMergeRunOperator::createWorkOrder(
     merge_run_operator::MergeTree::MergeJob *job,
     QueryContext *query_context,
     StorageManager *storage_manager,
     const tmb::client_id scheduler_client_id,
     tmb::MessageBus *bus) {
   DCHECK(!job->runs.empty());
   DCHECK(query_context != nullptr);

   InsertDestination *output_destination =
       query_context->getInsertDestination(job->is_final_level
                                               ? output_destination_index_
                                               : run_block_destination_index_);

   // Create a work order from the merge job from merge tree.
   return new SortMergeRunWorkOrder(
       query_id_,
       query_context->getSortConfig(sort_config_index_),
       job->level > 0 ? run_relation_ : input_relation_,
       std::move(job->runs),
       top_k_,
       job->level,
       output_destination,
       storage_manager,
       op_index_,
       scheduler_client_id,
       bus);
 }

 bool SortMergeRunOperator::generateWorkOrders(
     WorkOrdersContainer *container,
     QueryContext *query_context,
     StorageManager *storage_manager,
     const tmb::client_id scheduler_client_id,
     tmb::MessageBus *bus) {
   std::vector<MergeTree::MergeJob> jobs;

   // Get merge jobs from merge tree.
   bool done_generating = merge_tree_.getMergeJobs(&jobs);

   for (std::vector<MergeTree::MergeJob>::size_type job_id = 0;
        job_id != jobs.size();
        ++job_id) {
     // Add work order for each merge job.
     container->addNormalWorkOrder(createWorkOrder(&jobs[job_id],
                                                   query_context,
                                                   storage_manager,
                                                   scheduler_client_id,
                                                   bus),
                                   op_index_);
   }

   return done_generating;
 }

 void SortMergeRunOperator::initializeInputRuns() {
   DCHECK(started_);
   if (input_relation_is_stored_) {
     // Input blocks are from base relation; add all blocks.
     merge_tree_.addInputBlocks(input_relation_block_ids_);
   } else {
     // Input blocks are pipelined; add only newly available blocks.
     if (num_input_workorders_generated_ < input_relation_block_ids_.size()) {
       std::vector<block_id> new_blocks(
           input_relation_block_ids_.begin() + num_input_workorders_generated_,
           input_relation_block_ids_.end());
       merge_tree_.addInputBlocks(new_blocks);
       num_input_workorders_generated_ += new_blocks.size();
     }
   }
 }

 void SortMergeRunOperator::doneFeedingInputBlocks(const relation_id input_relation_id) {
   if (input_relation_is_stored_) {
     return;
   }

   // Now we know all the input blocks; compute the merge tree.
   merge_tree_.initializeTree(input_relation_block_ids_.size());

   // Initialize runs for the final time.
   initializeInputRuns();

   // If the final merge was already scheduled, fix it to write to correct output
   // destinaton.
   merge_tree_.checkAndFixFinalMerge();
 }

 namespace {

 /**
  * @brief Object to store the output run of a merge step with methods to
  * serialize and deserialize from its corresponding protobuf.
  **/
 class SortMergeRunOutput {
  public:
   /**
    * @brief Constructor with move semantics for creating the object from a
    * vector of block_ids.
    *
    * @param merge_level Level of the current merge.
    * @param blocks Rvalue of the vector of blocks in the output.
    **/
   SortMergeRunOutput(std::size_t merge_level, std::vector<block_id> &&blocks)
       : merge_level_(merge_level), blocks_(std::move(blocks)) {}

   /**
    * @brief Constructor from protobuf object.
    *
    * @param proto Protobuf of merged run.
    **/
   explicit SortMergeRunOutput(serialization::SortMergeRunOutput &proto)
       : merge_level_(static_cast<std::size_t>(proto.merge_level())) {
     blocks_.reserve(proto.blocks_size());
     for (const auto block : proto.blocks()) {
       blocks_.emplace_back(block);
     }
   }

   /**
    * @brief Method to serialize the object.
    *
    * @return Pair with pointer serialized blob and size of the blob.
    **/
   std::pair<void *, std::size_t> serialize() const {
     serialization::SortMergeRunOutput msg;
     msg.set_merge_level(merge_level_);
     for (const block_id block : blocks_) {
       msg.add_blocks(block);
     }

     const std::size_t msg_size = msg.ByteSize();
     char *msg_bytes = static_cast<char *>(std::malloc(msg_size));
     CHECK(msg.SerializeToArray(msg_bytes, msg_size));
     return std::make_pair(msg_bytes, msg_size);
   }

   /**
    * @brief Get merge level.
    *
    * @return Merge level.
    **/
   std::size_t getMergeLevel() const { return merge_level_; }

   /**
    * @brief Get mutable pointer to the blocks in the run.
    *
    * @return Pointer to vector of block_ids in the run.
    **/
   std::vector<block_id>* getBlocksMutable() { return &blocks_; }

  private:
   std::size_t merge_level_;
   std::vector<block_id> blocks_;
 };

 }  // namespace

 void SortMergeRunOperator::receiveFeedbackMessage(
     const WorkOrder::FeedbackMessage &msg) {
   CHECK(SortMergeRunOperator::kRunOutputMessage == msg.type());

   // Deserialize completion message from merge run work order.
   serialization::SortMergeRunOutput run_output_proto;
   CHECK(run_output_proto.ParseFromArray(msg.payload(), msg.payload_size()));

   // Write the output run to merge tree.
   SortMergeRunOutput run_output(run_output_proto);
   merge_tree_.writeOutputRun(run_output.getMergeLevel(),
                              run_output.getBlocksMutable());
 }

 void SortMergeRunWorkOrder::execute() {
   // Merge input runs.
   merge_run_operator::RunMerger run_merger(sort_config_,
                                            std::move(input_runs_),
                                            top_k_,
                                            run_relation_,
                                            output_destination_,
                                            merge_level_,
                                            storage_manager_);
   run_merger.doMerge();

   // Serialize completion message with output run.
   SortMergeRunOutput output(run_merger.getMergeLevel(),
                             std::move(*run_merger.getOutputRunMutable()));
   std::pair<void *, std::size_t> serialized_output = output.serialize();

   // Send completion message to operator.
   FeedbackMessage msg(SortMergeRunOperator::kRunOutputMessage,
                       getQueryID(),
                       operator_index_,
                       serialized_output.first,
                       serialized_output.second);
   SendFeedbackMessage(
       bus_, ClientIDMap::Instance()->getValue(), scheduler_client_id_, msg);
 }

 }  // namespace quickstep
	/**
	* 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 "relational_operators/SortMergeRunOperator.hpp"

	#include <cstddef>
	#include <cstdlib>
	#include <utility>
	#include <vector>

	#include "query_execution/QueryExecutionTypedefs.hpp"
	#include "query_execution/WorkOrderProtosContainer.hpp"
	#include "query_execution/WorkOrdersContainer.hpp"
	#include "relational_operators/SortMergeRunOperator.pb.h"
	#include "relational_operators/SortMergeRunOperatorHelpers.hpp"
	#include "relational_operators/WorkOrder.pb.h"
	#include "threading/ThreadIDBasedMap.hpp"

	#include "glog/logging.h"

	#include "tmb/id_typedefs.h"

	namespace quickstep {

	class InsertDestination;

	using merge_run_operator::Run;
	using merge_run_operator::RunMerger;
	using merge_run_operator::MergeTree;

	bool SortMergeRunOperator::getAllWorkOrders(
	WorkOrdersContainer *container,
	QueryContext *query_context,
	StorageManager *storage_manager,
	const tmb::client_id scheduler_client_id,
	tmb::MessageBus *bus) {
	if (input_relation_is_stored_) {
	// Input blocks (or runs) are from base relation. Only possible when base
	// relation is stored sorted.
	if (!started_) {
	// Initialize merge tree completely, since all input runs are known.
	merge_tree_.initializeTree(input_relation_block_ids_.size());
	started_ = true;
	initializeInputRuns();
	}
	} else {
	// Input blocks (or runs) are pipelined from the sorted run generation
	// operator.
	if (!started_ && !input_stream_done_) {
	// Initialize merge tree for first pipeline mode.
	merge_tree_.initializeForPipeline();
	started_ = true;
	initializeInputRuns();
	}
	}
	// Generate runs from merge tree.
	return generateWorkOrders(container, query_context, storage_manager, scheduler_client_id, bus);
	}

	bool SortMergeRunOperator::getAllWorkOrderProtos(WorkOrderProtosContainer *container) {
	if (input_relation_is_stored_) {
	// Input blocks (or runs) are from base relation. Only possible when base
	// relation is stored sorted.
	if (!started_) {
	// Initialize merge tree completely, since all input runs are known.
	merge_tree_.initializeTree(input_relation_block_ids_.size());
	started_ = true;
	initializeInputRuns();
	}
	} else {
	// Input blocks (or runs) are pipelined from the sorted run generation
	// operator.
	if (!started_ && !input_stream_done_) {
	// Initialize merge tree for first pipeline mode.
	merge_tree_.initializeForPipeline();
	started_ = true;
	initializeInputRuns();
	}
	}

	// Get merge jobs from merge tree.
	std::vector<MergeTree::MergeJob> jobs;
	const bool done_generating = merge_tree_.getMergeJobs(&jobs);

	for (std::vector<MergeTree::MergeJob>::size_type job_id = 0;
	job_id < jobs.size();
	++job_id) {
	// Add work order for each merge job.
	container->addWorkOrderProto(createWorkOrderProto(&jobs[job_id]), op_index_);
	}

	return done_generating;
	}

	serialization::WorkOrder* SortMergeRunOperator::createWorkOrderProto(
	merge_run_operator::MergeTree::MergeJob *job) {
	DCHECK(job != nullptr);
	DCHECK(!job->runs.empty());

	serialization::WorkOrder *proto = new serialization::WorkOrder;
	proto->set_work_order_type(serialization::SORT_MERGE_RUN);
	proto->set_query_id(query_id_);

	proto->SetExtension(serialization::SortMergeRunWorkOrder::operator_index, op_index_);
	proto->SetExtension(serialization::SortMergeRunWorkOrder::sort_config_index, sort_config_index_);

	for (const merge_run_operator::Run &run : job->runs) {
	serialization::Run *run_proto = proto->AddExtension(serialization::SortMergeRunWorkOrder::runs);
	for (const block_id block : run) {
	run_proto->add_blocks(block);
	}
	}

	proto->SetExtension(serialization::SortMergeRunWorkOrder::top_k, top_k_);
	proto->SetExtension(serialization::SortMergeRunWorkOrder::merge_level, job->level);
	proto->SetExtension(serialization::SortMergeRunWorkOrder::relation_id,
	job->level > 0 ? run_relation_.getID()
	: input_relation_.getID());
	proto->SetExtension(serialization::SortMergeRunWorkOrder::insert_destination_index,
	job->is_final_level ? output_destination_index_
	: run_block_destination_index_);

	return proto;
	}

	WorkOrder *SortMergeRunOperator::createWorkOrder(
	merge_run_operator::MergeTree::MergeJob *job,
	QueryContext *query_context,
	StorageManager *storage_manager,
	const tmb::client_id scheduler_client_id,
	tmb::MessageBus *bus) {
	DCHECK(!job->runs.empty());
	DCHECK(query_context != nullptr);

	InsertDestination *output_destination =
	query_context->getInsertDestination(job->is_final_level
	? output_destination_index_
	: run_block_destination_index_);

	// Create a work order from the merge job from merge tree.
	return new SortMergeRunWorkOrder(
	query_id_,
	query_context->getSortConfig(sort_config_index_),
	job->level > 0 ? run_relation_ : input_relation_,
	std::move(job->runs),
	top_k_,
	job->level,
	output_destination,
	storage_manager,
	op_index_,
	scheduler_client_id,
	bus);
	}

	bool SortMergeRunOperator::generateWorkOrders(
	WorkOrdersContainer *container,
	QueryContext *query_context,
	StorageManager *storage_manager,
	const tmb::client_id scheduler_client_id,
	tmb::MessageBus *bus) {
	std::vector<MergeTree::MergeJob> jobs;

	// Get merge jobs from merge tree.
	bool done_generating = merge_tree_.getMergeJobs(&jobs);

	for (std::vector<MergeTree::MergeJob>::size_type job_id = 0;
	job_id != jobs.size();
	++job_id) {
	// Add work order for each merge job.
	container->addNormalWorkOrder(createWorkOrder(&jobs[job_id],
	query_context,
	storage_manager,
	scheduler_client_id,
	bus),
	op_index_);
	}

	return done_generating;
	}

	void SortMergeRunOperator::initializeInputRuns() {
	DCHECK(started_);
	if (input_relation_is_stored_) {
	// Input blocks are from base relation; add all blocks.
	merge_tree_.addInputBlocks(input_relation_block_ids_);
	} else {
	// Input blocks are pipelined; add only newly available blocks.
	if (num_input_workorders_generated_ < input_relation_block_ids_.size()) {
	std::vector<block_id> new_blocks(
	input_relation_block_ids_.begin() + num_input_workorders_generated_,
	input_relation_block_ids_.end());
	merge_tree_.addInputBlocks(new_blocks);
	num_input_workorders_generated_ += new_blocks.size();
	}
	}
	}

	void SortMergeRunOperator::doneFeedingInputBlocks(const relation_id input_relation_id) {
	if (input_relation_is_stored_) {
	return;
	}

	// Now we know all the input blocks; compute the merge tree.
	merge_tree_.initializeTree(input_relation_block_ids_.size());

	// Initialize runs for the final time.
	initializeInputRuns();

	// If the final merge was already scheduled, fix it to write to correct output
	// destinaton.
	merge_tree_.checkAndFixFinalMerge();
	}

	namespace {

	/**
	* @brief Object to store the output run of a merge step with methods to
	* serialize and deserialize from its corresponding protobuf.
	**/
	class SortMergeRunOutput {
	public:
	/**
	* @brief Constructor with move semantics for creating the object from a
	* vector of block_ids.
	*
	* @param merge_level Level of the current merge.
	* @param blocks Rvalue of the vector of blocks in the output.
	**/
	SortMergeRunOutput(std::size_t merge_level, std::vector<block_id> &&blocks)
	: merge_level_(merge_level), blocks_(std::move(blocks)) {}

	/**
	* @brief Constructor from protobuf object.
	*
	* @param proto Protobuf of merged run.
	**/
	explicit SortMergeRunOutput(serialization::SortMergeRunOutput &proto)
	: merge_level_(static_cast<std::size_t>(proto.merge_level())) {
	blocks_.reserve(proto.blocks_size());
	for (const auto block : proto.blocks()) {
	blocks_.emplace_back(block);
	}
	}

	/**
	* @brief Method to serialize the object.
	*
	* @return Pair with pointer serialized blob and size of the blob.
	**/
	std::pair<void *, std::size_t> serialize() const {
	serialization::SortMergeRunOutput msg;
	msg.set_merge_level(merge_level_);
	for (const block_id block : blocks_) {
	msg.add_blocks(block);
	}

	const std::size_t msg_size = msg.ByteSize();
	char msg_bytes = static_cast<char >(std::malloc(msg_size));
	CHECK(msg.SerializeToArray(msg_bytes, msg_size));
	return std::make_pair(msg_bytes, msg_size);
	}

	/**
	* @brief Get merge level.
	*
	* @return Merge level.
	**/
	std::size_t getMergeLevel() const { return merge_level_; }

	/**
	* @brief Get mutable pointer to the blocks in the run.
	*
	* @return Pointer to vector of block_ids in the run.
	**/
	std::vector<block_id>* getBlocksMutable() { return &blocks_; }

	private:
	std::size_t merge_level_;
	std::vector<block_id> blocks_;
	};

	} // namespace

	void SortMergeRunOperator::receiveFeedbackMessage(
	const WorkOrder::FeedbackMessage &msg) {
	CHECK(SortMergeRunOperator::kRunOutputMessage == msg.type());

	// Deserialize completion message from merge run work order.
	serialization::SortMergeRunOutput run_output_proto;
	CHECK(run_output_proto.ParseFromArray(msg.payload(), msg.payload_size()));

	// Write the output run to merge tree.
	SortMergeRunOutput run_output(run_output_proto);
	merge_tree_.writeOutputRun(run_output.getMergeLevel(),
	run_output.getBlocksMutable());
	}

	void SortMergeRunWorkOrder::execute() {
	// Merge input runs.
	merge_run_operator::RunMerger run_merger(sort_config_,
	std::move(input_runs_),
	top_k_,
	run_relation_,
	output_destination_,
	merge_level_,
	storage_manager_);
	run_merger.doMerge();

	// Serialize completion message with output run.
	SortMergeRunOutput output(run_merger.getMergeLevel(),
	std::move(*run_merger.getOutputRunMutable()));
	std::pair<void *, std::size_t> serialized_output = output.serialize();

	// Send completion message to operator.
	FeedbackMessage msg(SortMergeRunOperator::kRunOutputMessage,
	getQueryID(),
	operator_index_,
	serialized_output.first,
	serialized_output.second);
	SendFeedbackMessage(
	bus_, ClientIDMap::Instance()->getValue(), scheduler_client_id_, msg);
	}

	} // namespace quickstep