parquet_metadata.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2023-2024, NVIDIA CORPORATION.
 *
 * 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 <cudf/io/types.hpp>
#include <cudf/utilities/export.hpp>
 
#include <optional>
#include <string_view>
#include <variant>
#include <vector>
 
namespace CUDF_EXPORT cudf {
namespace io {
namespace parquet {
enum class TypeKind : int8_t {
  UNDEFINED_TYPE       = -1,  // Undefined for non-leaf nodes
  BOOLEAN              = 0,
  INT32                = 1,
  INT64                = 2,
  INT96                = 3,  // Deprecated
  FLOAT                = 4,
  DOUBLE               = 5,
  BYTE_ARRAY           = 6,
  FIXED_LEN_BYTE_ARRAY = 7,
};
}  // namespace parquet
 
struct parquet_column_schema {
 public:
  explicit parquet_column_schema() = default;
 
  parquet_column_schema(std::string_view name,
                        parquet::TypeKind type,
                        std::vector<parquet_column_schema> children)
    : _name{name}, _type_kind{type}, _children{std::move(children)}
  {
  }
 
  [[nodiscard]] auto name() const { return _name; }
 
  [[nodiscard]] auto type_kind() const { return _type_kind; }
 
  [[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(); }
 
 private:
  std::string _name;
  // 3 types available: Physical, Converted, Logical.
  parquet::TypeKind _type_kind;  // Physical
  std::vector<parquet_column_schema> _children;
};
 
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>;
 
  explicit parquet_metadata() = default;
 
  parquet_metadata(parquet_schema schema,
                   int64_t num_rows,
                   size_type num_rowgroups,
                   key_value_metadata file_metadata,
                   std::vector<row_group_metadata> rg_metadata)
    : _schema{std::move(schema)},
      _num_rows{num_rows},
      _num_rowgroups{num_rowgroups},
      _file_metadata{std::move(file_metadata)},
      _rowgroup_metadata{std::move(rg_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& metadata() const { return _file_metadata; }
 
  [[nodiscard]] auto const& rowgroup_metadata() const { return _rowgroup_metadata; }
 
 private:
  parquet_schema _schema;
  int64_t _num_rows;
  size_type _num_rowgroups;
  key_value_metadata _file_metadata;
  std::vector<row_group_metadata> _rowgroup_metadata;
};
 
parquet_metadata read_parquet_metadata(source_info const& src_info);
  // end of group
}  // namespace io
}  // namespace CUDF_EXPORT cudf