Program Listing for File table_info.hpp

↰ Return to documentation for file (python/morpheus/morpheus/_lib/include/morpheus/objects/table_info.hpp)

/*
 * SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed 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 "morpheus/export.h"
#include "morpheus/objects/data_table.hpp"
#include "morpheus/objects/dtype.hpp"
#include "morpheus/objects/table_info_data.hpp"

#include <cudf/column/column_view.hpp>  // for column_view
#include <cudf/table/table_view.hpp>
#include <cudf/types.hpp>      // for size_type
#include <pybind11/pytypes.h>  // for object

#include <memory>
#include <mutex>
#include <optional>
#include <shared_mutex>
#include <string>
#include <tuple>
#include <vector>

namespace morpheus {

struct CudfHelper;

/****** Component public implementations *******************/
/****** TableInfo******************************************/
struct MORPHEUS_EXPORT TableInfoBase
{
    const cudf::table_view& get_view() const;

    std::vector<std::string> get_index_names() const;

    std::vector<std::string> get_column_names() const;

    cudf::size_type num_indices() const;

    cudf::size_type num_columns() const;

    cudf::size_type num_rows() const;

    const cudf::column_view& get_column(cudf::size_type idx) const;

    const cudf::column_view& get_column(const std::string& column_name) const;

    bool has_sliceable_index() const;

    const std::shared_ptr<const IDataTable>& get_parent() const;

    const TableInfoData& get_data() const;

  protected:
    TableInfoBase() = default;

    TableInfoBase(std::shared_ptr<const IDataTable> parent, TableInfoData data);

    TableInfoData& get_data();

  private:
    std::shared_ptr<const IDataTable> m_parent;
    TableInfoData m_data;

    // Give access to internal m_parent and m_data for converting to cudf dataframe
    friend CudfHelper;
};

struct MORPHEUS_EXPORT TableInfo : public TableInfoBase
{
  public:
    TableInfo() = default;
    TableInfo(std::shared_ptr<const IDataTable> parent, std::shared_lock<std::shared_mutex> lock, TableInfoData data);

    TableInfo get_slice(cudf::size_type start, cudf::size_type stop, std::vector<std::string> column_names = {}) const;

  private:
    // We use a shared_lock to allow multiple `TableInfo` to be in use at the same time.
    std::shared_lock<std::shared_mutex> m_lock;
};

struct MORPHEUS_EXPORT MutableTableInfo : public TableInfoBase
{
  public:
    MutableTableInfo(std::shared_ptr<const IDataTable> parent,
                     std::unique_lock<std::shared_mutex> lock,
                     TableInfoData data);

    MutableTableInfo(MutableTableInfo&& other) = default;

    ~MutableTableInfo();

    MutableTableInfo get_slice(cudf::size_type start,
                               cudf::size_type stop,
                               std::vector<std::string> column_names = {}) &&;

    void insert_columns(const std::vector<std::tuple<std::string, morpheus::DType>>& columns);

    void insert_missing_columns(const std::vector<std::tuple<std::string, morpheus::DType>>& columns);

    std::unique_ptr<pybind11::object> checkout_obj();

    void return_obj(std::unique_ptr<pybind11::object>&& obj);

    std::optional<std::string> ensure_sliceable_index();

  private:
    // We use a unique_lock here to enforce exclusive access
    std::unique_lock<std::shared_mutex> m_lock;

    mutable int m_checked_out_ref_count{-1};
};
  // end of group
}  // namespace morpheus