greenplumn CKHeap 源码

greenplumn CKHeap 代码

文件路径：/src/backend/gporca/libgpopt/include/gpopt/base/CKHeap.h

//---------------------------------------------------------------------------
//	Greenplum Database
//	Copyright (C) 2020 VMware, Inc.
//
//	@filename:
//		CKHeap.h
//
//	@doc:
//		Optimizer binary heap utility functions
//---------------------------------------------------------------------------
#ifndef GPOPT_CKHeap_H
#define GPOPT_CKHeap_H

#include "gpopt/base/CColRef.h"

namespace gpopt
{
using namespace gpos;

// a heap keeping the k lowest-cost objects in an array of class A
// A is a CDynamicPtrArray
// E is the entry type of the array and it has a method CDouble DCost()
// See https://en.wikipedia.org/wiki/Binary_heap for details
// (with the added feature of returning only the top k).
template <class A, class E>
class CKHeap : public CRefCount
{
private:
	A *m_topk;
	CMemoryPool *m_mp;
	ULONG m_k;
	BOOL m_is_heapified;
	ULONG m_num_returned;

	// the parent index is (ix-1)/2, except for 0
	ULONG
	parent(ULONG ix)
	{
		return (0 < ix ? (ix - 1) / 2 : m_topk->Size());
	}

	// children are at indexes 2*ix + 1 and 2*ix + 2
	ULONG
	left_child(ULONG ix)
	{
		return 2 * ix + 1;
	}
	ULONG
	right_child(ULONG ix)
	{
		return 2 * ix + 2;
	}

	// does the parent/child exist?
	BOOL
	exists(ULONG ix)
	{
		return ix < m_topk->Size();
	}
	// cost of an entry (this class implements a Min-Heap)
	CDouble
	cost(ULONG ix)
	{
		return (*m_topk)[ix]->GetCostForHeap();
	}

	// push node ix in the tree down into its child tree as much as needed
	void
	HeapifyDown(ULONG ix)
	{
		ULONG left_child_ix = left_child(ix);
		ULONG right_child_ix = right_child(ix);
		ULONG min_element_ix = ix;

		if (exists(left_child_ix) && cost(left_child_ix) < cost(ix))
		{
			// left child is better than parent, it becomes the new candidate
			min_element_ix = left_child_ix;
		}

		if (exists(right_child_ix) &&
			cost(right_child_ix) < cost(min_element_ix))
		{
			// right child is better than min(parent, left child)
			min_element_ix = right_child_ix;
		}

		if (min_element_ix != ix)
		{
			// make the lowest of { parent, left child, right child } the new root
			m_topk->Swap(ix, min_element_ix);
			HeapifyDown(min_element_ix);
		}
	}

	// pull node ix in the tree up as much as needed
	void
	HeapifyUp(ULONG ix)
	{
		ULONG parent_ix = parent(ix);

		if (!exists(parent_ix))
		{
			return;
		}

		if (cost(ix) < cost(parent_ix))
		{
			m_topk->Swap(ix, parent_ix);
			HeapifyUp(parent_ix);
		}
	}

	// Convert the array into a heap, heapify-down all interior nodes of the tree, bottom-up.
	// Note that we keep all the entries, not just the top k, since our k-heaps are short-lived.
	// You can only retrieve the top k with RemoveBestElement(), though.
	void
	Heapify()
	{
		// the parent of the last node is the last node in the tree that is a parent
		ULONG start_ix = parent(m_topk->Size() - 1);

		// now work our way up to the root, calling HeapifyDown
		for (ULONG ix = start_ix; exists(ix); ix--)
		{
			HeapifyDown(ix);
		}

		m_is_heapified = true;
	}

public:
	CKHeap(CMemoryPool *mp, ULONG k)
		: m_mp(mp), m_k(k), m_is_heapified(false), m_num_returned(0)
	{
		m_topk = GPOS_NEW(m_mp) A(m_mp);
	}

	~CKHeap() override
	{
		m_topk->Release();
	}

	void
	Insert(E *elem)
	{
		GPOS_ASSERT(nullptr != elem);
		// since the cost may change as we find more expressions in the group,
		// we just append to the array now and heapify at the end
		GPOS_ASSERT(!m_is_heapified);
		m_topk->Append(elem);

		// this is dead code at the moment, but other users might want to
		// heapify and then insert additional items
		if (m_is_heapified)
		{
			HeapifyUp(m_topk->Size() - 1);
		}
	}

	// remove the next of the top k elements, sorted ascending by cost
	E *
	RemoveBestElement()
	{
		if (0 == m_topk->Size() || m_k <= m_num_returned)
		{
			return nullptr;
		}

		m_num_returned++;

		return RemoveNextElement();
	}

	// remove the next best element, without the top k limit
	E *
	RemoveNextElement()
	{
		if (0 == m_topk->Size())
		{
			return nullptr;
		}

		if (!m_is_heapified)
		{
			Heapify();
		}

		// we want to remove and return the root of the tree, which is the best element

		// first, swap the root with the last element in the array
		m_topk->Swap(0, m_topk->Size() - 1);

		// now remove the new last element, which is the real root
		E *result = m_topk->RemoveLast();

		// then push the new first element down to the correct place
		HeapifyDown(0);

		return result;
	}

	ULONG
	Size()
	{
		return m_topk->Size();
	}

	BOOL
	IsLimitExceeded()
	{
		return m_topk->Size() + m_num_returned > m_k;
	}

	void
	Clear()
	{
		m_topk->Clear();
		m_is_heapified = false;
		m_num_returned = 0;
	}

};	// class CKHeap

}  // namespace gpopt

#endif	// !GPOPT_CKHeap_H

// EOF

相关信息

greenplumn 源码目录

相关文章

greenplumn CAutoOptCtxt 源码

greenplumn CCTEInfo 源码

greenplumn CCTEMap 源码

greenplumn CCTEReq 源码

greenplumn CCastUtils 源码

greenplumn CColConstraintsArrayMapper 源码

greenplumn CColConstraintsHashMapper 源码

greenplumn CColRef 源码

greenplumn CColRefComputed 源码

greenplumn CColRefSet 源码