parquet_metadata.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <cudf/io/datasource.hpp>
#include <cudf/io/parquet_schema.hpp>
#include <cudf/io/types.hpp>
#include <cudf/table/table.hpp>
#include <cudf/utilities/default_stream.hpp>
#include <cudf/utilities/export.hpp>
#include <cudf/utilities/memory_resource.hpp>
 
#include <memory>
#include <span>
#include <string>
#include <string_view>
#include <vector>
 
namespace CUDF_EXPORT cudf {
namespace io {
using cudf::io::parquet::Type;
 
struct parquet_column_schema {
 public:
  explicit parquet_column_schema() : _cudf_type{data_type{type_id::EMPTY}} {}
 
  parquet_column_schema(std::string_view name,
                        Type type,
                        std::vector<parquet_column_schema>&& children,
                        data_type cudf_type)
    : _name{name}, _type{type}, _children{std::move(children)}, _cudf_type{cudf_type}
  {
  }
 
  [[nodiscard]] auto name() const { return _name; }
 
  [[nodiscard]] auto type() const { return _type; }
 
  [[nodiscard]] auto const& children() const& { return _children; }
 
  [[nodiscard]] auto children() && { return std::move(_children); }
 
  [[nodiscard]] auto const& child(int idx) const& { return children().at(idx); }
 
  [[nodiscard]] auto child(int idx) && { return std::move(children().at(idx)); }
 
  [[nodiscard]] auto num_children() const { return children().size(); }
 
  [[nodiscard]] auto cudf_type() const { return _cudf_type; }
 
 private:
  std::string _name;
  // 3 types available: Physical, Converted, Logical
  Type _type;  // Physical type
  std::vector<parquet_column_schema> _children;
  data_type _cudf_type;
};
 
struct parquet_schema {
 public:
  explicit parquet_schema() = default;
 
  parquet_schema(parquet_column_schema root_column_schema) : _root{std::move(root_column_schema)} {}
 
  [[nodiscard]] auto const& root() const& { return _root; }
 
  [[nodiscard]] auto root() && { return std::move(_root); }
 
 private:
  parquet_column_schema _root;
};
 
class parquet_metadata {
 public:
  using key_value_metadata = std::unordered_map<std::string, std::string>;
  using row_group_metadata = std::unordered_map<std::string, int64_t>;
  using column_chunk_metadata = std::unordered_map<std::string, std::vector<int64_t>>;
 
  explicit parquet_metadata() = default;
 
  parquet_metadata(parquet_schema schema,
                   int64_t num_rows,
                   size_type num_rowgroups,
                   std::vector<size_type> num_rowgroups_per_file,
                   key_value_metadata file_metadata,
                   std::vector<row_group_metadata> rg_metadata,
                   column_chunk_metadata column_chunk_metadata)
    : _schema{std::move(schema)},
      _num_rows{num_rows},
      _num_rowgroups{num_rowgroups},
      _num_rowgroups_per_file{std::move(num_rowgroups_per_file)},
      _file_metadata{std::move(file_metadata)},
      _rowgroup_metadata{std::move(rg_metadata)},
      _column_chunk_metadata{std::move(column_chunk_metadata)}
  {
  }
 
  [[nodiscard]] auto const& schema() const { return _schema; }
 
  [[nodiscard]] auto num_rows() const { return _num_rows; }
 
  [[nodiscard]] auto num_rowgroups() const { return _num_rowgroups; }
 
  [[nodiscard]] auto const& num_rowgroups_per_file() const { return _num_rowgroups_per_file; }
 
  [[nodiscard]] auto const& metadata() const { return _file_metadata; }
 
  [[nodiscard]] auto const& rowgroup_metadata() const { return _rowgroup_metadata; }
 
  [[nodiscard]] auto const& columnchunk_metadata() const { return _column_chunk_metadata; }
 
 private:
  parquet_schema _schema;
  int64_t _num_rows;
  size_type _num_rowgroups;
  std::vector<size_type> _num_rowgroups_per_file;
  key_value_metadata _file_metadata;
  std::vector<row_group_metadata> _rowgroup_metadata;
  column_chunk_metadata _column_chunk_metadata;
};
 
parquet_metadata read_parquet_metadata(source_info const& src_info);
 
std::vector<parquet::FileMetaData> read_parquet_footers(
  std::span<std::unique_ptr<cudf::io::datasource> const> sources);
 
std::unique_ptr<table> read_parquet_column_chunk_bounds(
  std::span<parquet::FileMetaData const> parquet_metadatas,
  std::span<std::string const> column_names,
  cuda::stream_ref stream   = cudf::get_default_stream(),
  cudf::memory_resources mr = cudf::get_current_device_resource_ref());
  // end of group
}  // namespace io
}  // namespace CUDF_EXPORT cudf