VariantLinkColumn.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
*/
 
 
 
#ifndef COLUMNAR_VARIANT_VARIANT_LINK_COLUMN_H
#define COLUMNAR_VARIANT_VARIANT_LINK_COLUMN_H
 
#include <ColumnarCore/LinkColumn.h>
#include <ColumnarVariant/VariantDef.h>
#include <ColumnarVariant/VariantOptObjectId.h>
#include <ColumnarCore/VectorColumn.h>
 
namespace columnar
{
 
  template<ContainerIdConcept CIBase,ContainerIdConcept... CIList>
  class ObjectLink<VariantContainerId<CIBase,CIList...>,ColumnarModeXAOD> final
  {
  public:
 
    using CI = VariantContainerId<CIBase,CIList...>;
    using CM = ColumnarModeXAOD;
 
    using LinkType = ElementLink<typename CI::xAODElementLinkType>;
 
    ObjectLink (const LinkType* val_link)
      : m_link (val_link)
    {}
 
    const typename CI::xAODObjectIdType* getXAODObject () const
    {
      if (m_link->isValid())
      {
        return **m_link;
      } else
      {
        return nullptr;
      }
    }
 
    explicit operator bool () const noexcept
    {
      return m_link->isValid();
    }
 
    [[nodiscard]] bool has_value () const noexcept
    {
      return m_link->isValid();
    }
 
    [[nodiscard]] OptObjectId<CI,CM> opt_value () const
    {
      return OptObjectId<CI,CM> (getXAODObject ());
    }
 
    template<ContainerIdConcept CI2>
    [[nodiscard]] bool isContainer () const
    {
      static constexpr unsigned variantIndex = CI::template getVariantIndex<CI2>();
      static_assert (variantIndex < CI::numVariants, "invalid container id");
      return m_link->isValid() && dynamic_cast<typename CI2::xAODObjectRangeType*>(m_link->getStorableObjectPointer());
    }
 
    template<ContainerIdConcept CI2>
    [[nodiscard]] bool operator == (ObjectId<CI2,CM> id) const
    {
      static constexpr unsigned variantIndex = CI::template getVariantIndex<CI2>();
      static_assert (variantIndex < CI::numVariants, "invalid container id");
      return m_link->isValid() && (**m_link) == &id.getXAODObject();
    }
 
    [[nodiscard]] bool operator == (const ObjectLink<CI,CM>& obj) const
    {
      return m_link == obj.m_link;
    }
 
    template<ContainerIdConcept CI2>
    [[nodiscard]] OptObjectId<CI2,CM> tryGetVariant () const
    {
      static constexpr unsigned variantIndex = CI::template getVariantIndex<CI2>();
      static_assert (variantIndex < CI::numVariants, "invalid container id");
      if (m_link->isValid())
      {
        return OptObjectId<CI2,CM> (dynamic_cast<const typename CI2::xAODObjectIdType*>(**m_link));
      } else
        return OptObjectId<CI2,CM> ();
    }
 
  private:
 
    const LinkType* m_link = nullptr;
  };
  template<typename... CIList>
  std::ostream& operator<< (std::ostream& str, const ObjectLink<VariantContainerId<CIList...>,ColumnarModeXAOD>& obj)
  {
    return str << obj.getXAODObject() << "/" << obj.getXAODObject()->type();
  }
 
  // in xAOD mode we can do a fairly straightforward conversion from
  // ElementLink as the logic is inside ObjectLink
  template<ContainerIdConcept CIBase,typename... CIList>
  struct ColumnTypeTraits<ObjectLink<VariantContainerId<CIBase,CIList...>,ColumnarModeXAOD>,ColumnarModeXAOD> final
  {
    using CI = VariantContainerId<CIBase,CIList...>;
    using CM = ColumnarModeXAOD;
    using LinkType = ElementLink<typename CI::xAODElementLinkType>;
    using ColumnType = NativeColumn<LinkType>;
    static constexpr bool isNativeType = false;
    static constexpr bool useConvertInput = true;
    static constexpr bool useConvertWithDataInput = false;
    static constexpr bool useConvertOutput = true;
    static ColumnInfo& updateColumnInfo (ColumnarTool<CM>& /*columnBase*/, ColumnInfo& info) {return info;}
    static auto convertInput (const LinkType& value) {
      return ObjectLink<VariantContainerId<CIBase,CIList...>,ColumnarModeXAOD> (&value);
    };
    template<ContainerIdConcept CI2>
      requires (CI::template isValidContainer<CI2>())
    static auto convertOutput (const ObjectId<CI2,ColumnarModeXAOD>& obj) noexcept
    {
      auto* container = static_cast<const typename CIBase::xAODElementLinkType*>(obj.getXAODObjectNoexcept().container());
      return LinkType(*container, obj.getXAODObjectNoexcept().index());
    }
    using UserType = ObjectLink<VariantContainerId<CIBase,CIList...>,ColumnarModeXAOD>;
  };
 
 
 
  template<ContainerIdConcept CIBase,ContainerIdConcept... CIList>
  class ObjectLink<VariantContainerId<CIBase,CIList...>,ColumnarModeArray> final
  {
  public:
 
    using CI = VariantContainerId<CIBase,CIList...>;
    using CM = ColumnarModeArray;
 
    ObjectLink (typename CM::LinkIndexType val_link, const typename CM::LinkKeyType* val_keys, void** val_data)
      : m_link (val_link), m_keys (val_keys), m_data (val_data)
    {}
 
    const typename CI::xAODObjectIdType* getXAODObject () const
    {
      throw std::logic_error ("can't call xAOD function in columnar mode");
    }
 
    explicit operator bool () const noexcept
    {
      return m_link != invalidObjectIndex;
    }
 
    [[nodiscard]] bool has_value () const noexcept
    {
      return m_link != invalidObjectIndex;
    }
 
    [[nodiscard]] OptObjectId<CI,CM> opt_value () const
    {
      if (m_link == CM::invalidLinkValue)
        return OptObjectId<CI,CM> ();
      const auto key = getLinkKey();
      for (unsigned i = 0; i < CI::numVariants; ++ i)
      {
        if (m_keys[i] == key)
          return OptObjectId<CI,CM> (m_data, i, getLinkIndex());
      }
      // not sure whether we should throw here or return a nullopt, but
      // throwing is probably the safer default. if it becomes an issue
      // we can revisit this.
      throw std::runtime_error ("link key does not match any known container: " + std::to_string(key));
    }
 
    template<ContainerIdConcept CI2>
    [[nodiscard]] bool isContainer () const
    {
      static constexpr unsigned variantIndex = CI::template getVariantIndex<CI2>();
      static_assert (variantIndex < CI::numVariants, "invalid container id");
      return getLinkKey() == m_keys[variantIndex];
    }
 
    template<ContainerIdConcept CI2>
    [[nodiscard]] bool operator == (ObjectId<CI2,CM> id) const
    {
      static constexpr unsigned variantIndex = CI::template getVariantIndex<CI2>();
      static_assert (variantIndex < CI::numVariants, "invalid container id");
      return getLinkIndex() == id.getIndex() && getLinkKey() == m_keys[variantIndex];
    }
 
    [[nodiscard]] bool operator == (const ObjectLink<CI,CM>& obj) const
    {
      if (m_keys == obj.m_keys)
        return m_link == obj.m_link;
      if (getLinkIndex() != obj.getLinkIndex())
        return false;
      const auto thisKey = getLinkKey();
      const auto thatKey = obj.getLinkKey();
      for (unsigned i = 0; i < CI::numVariants; ++ i)
      {
        if (m_keys[i] == thisKey)
        {
          if (obj.m_keys[i] == thatKey)
            return true;
          else
            return false;
        } else if (obj.m_keys[i] == thatKey)
          return false;
      }
      return false;
    }
 
    template<ContainerIdConcept CI2>
    [[nodiscard]] OptObjectId<CI2,CM> tryGetVariant () const
    {
      static constexpr unsigned variantIndex = CI::template getVariantIndex<CI2>();
      static_assert (variantIndex < CI::numVariants, "invalid container id");
      if (getLinkKey() == m_keys[variantIndex])
        return OptObjectId<CI2,CM> (m_data, getLinkIndex());
      else
        return OptObjectId<CI2,CM> ();
    }
 
    [[nodiscard]] auto getLinkIndex () const noexcept
    {
      return CM::getLinkIndex (m_link);
    }
 
    [[nodiscard]] auto getLinkKey () const noexcept
    {
      return CM::getLinkKey (m_link);
    }
 
  private:
 
    typename CM::LinkIndexType m_link = 0;
    const CM::LinkKeyType* m_keys = nullptr;
    void** m_data = nullptr;
  };
  template<typename... CIList>
  std::ostream& operator<< (std::ostream& str, const ObjectLink<VariantContainerId<CIList...>,ColumnarModeArray>& obj)
  {
    return str << obj.getLinkKey() << "/" << obj.getLinkIndex();
  }
 
 
 
  // in external mode we need to use a vector column, as well as an
  // extra column to contain our keys in order.
  template<ContainerIdConcept CI,ColumnAccessMode CAM,ContainerIdConcept CIBase,ContainerIdConcept... CIList>
  class AccessorTemplate<CI,ObjectLink<VariantContainerId<CIBase,CIList...>,ColumnarModeArray>,CAM,ColumnarModeArray> final
  {
  public:
 
    using VariantCI = VariantContainerId<CIBase,CIList...>;
    using CM = ColumnarModeArray;
    static constexpr std::array containerIdNames = {CIList::idName...};
 
    AccessorTemplate () = default;
 
    AccessorTemplate (ColumnarTool<CM>& columnBase, const std::string& name, ColumnInfo&& info = {})
    {
      std::string dataName = std::string (CI::idName) + "." + name;
      std::string keysName = std::string (CI::idName) + "." + name + ".keys";
 
      auto dataInfo = info;
      dataInfo.offsetName = CI::idName;
      dataInfo.variantLinkKeyColumn = keysName;
      dataInfo.linkTargetNames.reserve (VariantCI::numVariants);
      for (unsigned i = 0; i < VariantCI::numVariants; ++ i)
        dataInfo.linkTargetNames.emplace_back (containerIdNames[i]);
      auto keyInfo = info;
      keyInfo.accessMode = ColumnAccessMode::input;
      keyInfo.fixedDimensions.push_back (VariantCI::numVariants);
 
      m_dataData = std::make_unique<ColumnAccessorDataArray> (&m_dataIndex, &m_dataData, &typeid (typename CM::LinkIndexType), CAM);
      columnBase.addColumn (dataName, m_dataData.get(), std::move (dataInfo));
      m_keysData = std::make_unique<ColumnAccessorDataArray> (&m_keysIndex, &m_keysData, &typeid (typename CM::LinkKeyType), ColumnAccessMode::input);
      columnBase.addColumn (keysName, m_keysData.get(), std::move (keyInfo));
    }
 
    AccessorTemplate (AccessorTemplate&& that)
    {
      moveAccessor (m_dataIndex, m_dataData, that.m_dataIndex, that.m_dataData);
      moveAccessor (m_keysIndex, m_keysData, that.m_keysIndex, that.m_keysData);
    }
 
    AccessorTemplate& operator = (AccessorTemplate&& that)
    {
      if (this != &that)
      {
        moveAccessor (m_dataIndex, m_dataData, that.m_dataIndex, that.m_dataData);
        moveAccessor (m_keysIndex, m_keysData, that.m_keysIndex, that.m_keysData);
      }
      return *this;
    }
 
    auto operator () (ObjectId<CI,CM> id) const noexcept
      requires (CAM == ColumnAccessMode::input)
    {
      auto *data = static_cast<const typename CM::LinkIndexType*>(id.getData()[m_dataIndex]);
      auto *keys = static_cast<const typename CM::LinkKeyType*>(id.getData()[m_keysIndex]);
      return ObjectLink<VariantContainerId<CIBase,CIList...>,ColumnarModeArray> (data[id.getIndex()], keys, id.getData());
    }
 
    template<ContainerIdConcept CI2>
      requires (CAM == ColumnAccessMode::output && VariantCI::template isValidContainer<CI2>())
    void set (ObjectId<CI,CM> id, ObjectId<CI2,ColumnarModeArray> obj) const noexcept
    {
      auto *data = static_cast<typename CM::LinkIndexType*>(id.getData()[m_dataIndex]);
      auto *keys = static_cast<const typename CM::LinkKeyType*>(id.getData()[m_keysIndex]);
      data[id.getIndex()] = CM::mergeLinkKeyIndex (keys[VariantCI::template getVariantIndex<CI2>()], obj.getIndex());
    }
 
    [[nodiscard]] bool isAvailable (ObjectId<CI,CM> id) const noexcept
    {
      auto *data = static_cast<const typename CM::LinkIndexType*>(id.getData()[m_dataIndex]);
      return data != nullptr;
    }
 
 
 
  private:
 
    unsigned m_dataIndex = 0u;
    std::unique_ptr<ColumnAccessorDataArray> m_dataData;
    unsigned m_keysIndex = 0u;
    std::unique_ptr<ColumnAccessorDataArray> m_keysData;
  };
 
 
 
  // in external mode we need to use a vector column, as well as an
  // extra column to contain our keys in order.
  template<ContainerIdConcept CI,ContainerIdConcept CIBase,ContainerIdConcept... CIList>
  class AccessorTemplate<CI,std::vector<ObjectLink<VariantContainerId<CIBase,CIList...>,ColumnarModeArray>>,ColumnAccessMode::input,ColumnarModeArray> final
  {
  public:
 
    using VariantCI = VariantContainerId<CIBase,CIList...>;
    static constexpr ColumnAccessMode CAM = ColumnAccessMode::input;
    using CM = ColumnarModeArray;
    static constexpr std::array containerIdNames = {CIList::idName...};
 
    AccessorTemplate () = default;
 
    AccessorTemplate (ColumnarTool<CM>& columnBase, const std::string& name, ColumnInfo&& info = {})
    {
      std::string offsetName = std::string (CI::idName) + "." + name + ".offset";
      std::string dataName = std::string (CI::idName) + "." + name + ".data";
      std::string keysName = std::string (CI::idName) + "." + name + ".keys";
 
      auto offsetInfo = info;
      offsetInfo.offsetName = CI::idName;
      offsetInfo.isOffset = true;
      auto dataInfo = info;
      dataInfo.offsetName = offsetName;
      dataInfo.variantLinkKeyColumn = keysName;
      dataInfo.linkTargetNames.reserve (VariantCI::numVariants);
      for (unsigned i = 0; i < VariantCI::numVariants; ++ i)
        dataInfo.linkTargetNames.emplace_back (containerIdNames[i]);
      auto keyInfo = info;
      keyInfo.fixedDimensions.push_back (VariantCI::numVariants);
 
      m_offsetData = std::make_unique<ColumnAccessorDataArray> (&m_offsetIndex, &m_offsetData, &typeid (ColumnarOffsetType), ColumnAccessMode::input);
      columnBase.addColumn (offsetName, m_offsetData.get(), std::move (offsetInfo));
      m_dataData = std::make_unique<ColumnAccessorDataArray> (&m_dataIndex, &m_dataData, &typeid (typename CM::LinkIndexType), ColumnAccessMode::input);
      columnBase.addColumn (dataName, m_dataData.get(), std::move (dataInfo));
      m_keysData = std::make_unique<ColumnAccessorDataArray> (&m_keysIndex, &m_keysData, &typeid (typename CM::LinkKeyType), ColumnAccessMode::input);
      columnBase.addColumn (keysName, m_keysData.get(), std::move (keyInfo));
    }
 
    AccessorTemplate (AccessorTemplate&& that)
    {
      moveAccessor (m_offsetIndex, m_offsetData, that.m_offsetIndex, that.m_offsetData);
      moveAccessor (m_dataIndex, m_dataData, that.m_dataIndex, that.m_dataData);
      moveAccessor (m_keysIndex, m_keysData, that.m_keysIndex, that.m_keysData);
    }
 
    AccessorTemplate& operator = (AccessorTemplate&& that)
    {
      if (this != &that)
      {
        moveAccessor (m_offsetIndex, m_offsetData, that.m_offsetIndex, that.m_offsetData);
        moveAccessor (m_dataIndex, m_dataData, that.m_dataIndex, that.m_dataData);
        moveAccessor (m_keysIndex, m_keysData, that.m_keysIndex, that.m_keysData);
      }
      return *this;
    }
 
    auto operator () (ObjectId<CI,CM> id) const noexcept
    {
      auto *offset = static_cast<const ColumnarOffsetType*>(id.getData()[m_offsetIndex]);
      auto *data = static_cast<const typename CM::LinkIndexType*>(id.getData()[m_dataIndex]);
      auto *keys = static_cast<const typename CM::LinkKeyType*>(id.getData()[m_keysIndex]);
      return detail::VectorConvertView ([keys,data=id.getData()](typename CM::LinkIndexType value) {return ObjectLink<VariantContainerId<CIBase,CIList...>,ColumnarModeArray> (value, keys,data);}, std::span<const typename CM::LinkIndexType> (data + offset[id.getIndex()], offset[id.getIndex()+1]-offset[id.getIndex()]));
    }
 
  private:
 
    unsigned m_offsetIndex = 0u;
    std::unique_ptr<ColumnAccessorDataArray> m_offsetData;
    unsigned m_dataIndex = 0u;
    std::unique_ptr<ColumnAccessorDataArray> m_dataData;
    unsigned m_keysIndex = 0u;
    std::unique_ptr<ColumnAccessorDataArray> m_keysData;
  };
 
  template<ContainerIdConcept CIBase,ContainerIdConcept... CIList,typename ELT>
  struct ColumnTypeTraits<LinkCastColumn<VariantContainerId<CIBase,CIList...>,ELT>,ColumnarModeArray> final
  {
    using CI = VariantContainerId<CIBase,CIList...>;
    using CM = ColumnarModeArray;
    using ColumnType = ObjectLink<CI,CM>;
    using UserType = OptObjectId<CI,CM>;
    static constexpr bool isNativeType = false;
    static constexpr bool useConvertInput = true;
    static constexpr bool useConvertWithDataInput = false;
    static ColumnInfo& updateColumnInfo (ColumnarTool<CM>& /*columnBase*/, ColumnInfo& info) {return info;}
    static UserType convertInput (const ColumnType& value) {
      return value.opt_value();
    };
  };
}
 
#endif