ArrowIOHelpers.h

Go to the documentation of this file.

// Copyright (c) 2002-present, OpenMS Inc. -- EKU Tuebingen, ETH Zurich, and FU Berlin
// SPDX-License-Identifier: BSD-3-Clause
//
// --------------------------------------------------------------------------
// $Maintainer: Timo Sachsenberg $
// $Authors: Timo Sachsenberg $
// --------------------------------------------------------------------------
 
#pragma once
 
#include <OpenMS/config.h>
#include <OpenMS/CONCEPT/Types.h>
#include <OpenMS/DATASTRUCTURES/StringUtils.h>
#include <OpenMS/FORMAT/MSExperimentArrowExport.h>  // for ParquetWriteConfig
 
#include <cstdint>
#include <map>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <variant>
#include <unordered_set>
#include <vector>
 
// Forward declarations
namespace arrow
{
  class Array;
  class KeyValueMetadata;
  class Table;
}
 
namespace OpenMS
{
class ConsensusMap;
class DataValue;
class MetaInfoInterface;
namespace MS1LabelState { struct Keys; }
 
namespace ArrowIOHelpers
{
  OPENMS_DLLAPI std::string generateUuidV4();
 
  OPENMS_DLLAPI Size tableRowCount(const std::shared_ptr<arrow::Table>& table);
 
  OPENMS_DLLAPI bool writeTableToParquet(
    const std::shared_ptr<arrow::Table>& table,
    const std::string& filename,
    const ParquetWriteConfig& config = ParquetWriteConfig{});
 
  struct OPENMS_DLLAPI ParquetColumn
  {
    struct Score
    {
      std::string name;
      double value = 0.0;
      std::optional<bool> higher_better; 
    };
    enum class Type { INT16, INT64, UTF8, FLOAT_LIST, UTF8_LIST, SCORE_LIST, INT32, BOOL, DOUBLE, STRUCT_LIST, FLOAT, INT32_LIST };
    struct StructField
    {
      StructField() = default;
      StructField(std::string field_name, Type field_type, bool field_nullable = true, std::vector<StructField> struct_fields = {}) :
        name(std::move(field_name)), type(field_type), nullable(field_nullable), fields(std::move(struct_fields)) {}
      std::string name;
      Type type = Type::UTF8;
      bool nullable = true;
      std::vector<StructField> fields; 
    };
    struct StructValue;
    using StructEntry = std::vector<StructValue>;
    struct OPENMS_DLLAPI StructValue
    {
      std::variant<std::monostate, std::string, Int32, bool, double, std::vector<StructEntry>> value;
      StructValue() = default; 
      StructValue(std::string v) : value(std::move(v)) {}
      StructValue(const char* v) : value(std::string(v)) {} 
      StructValue(Int32 v) : value(v) {}
      StructValue(bool v) : value(v) {}
      StructValue(double v) : value(v) {}
      StructValue(std::vector<StructEntry> v) : value(std::move(v)) {}
      template <typename T> StructValue(const std::optional<T>& v) { if (v) { *this = StructValue(*v); } }
    };
    std::string name;
    Type type = Type::INT64;
    bool nullable = true;
    std::vector<Int64> int64_values;                              
    std::vector<double> double_values;                            
    std::vector<std::string> utf8_values;                         
    std::vector<std::vector<float>> float_lists;                  
    std::vector<std::vector<Int32>> int32_lists;                  
    std::vector<std::vector<std::string>> utf8_lists;             
    std::vector<std::vector<Score>> score_lists;                  
    std::vector<StructField> struct_fields;                       
    std::vector<std::vector<StructEntry>> struct_lists;           
    std::vector<bool> null_rows;                                  
  };
 
  OPENMS_DLLAPI std::shared_ptr<arrow::Table> withColumns(
    const std::shared_ptr<arrow::Table>& table,
    const std::vector<ParquetColumn>& columns);
 
  OPENMS_DLLAPI std::shared_ptr<arrow::Table> tableOfColumns(
    const std::vector<ParquetColumn>& columns,
    Size rows);
 
  OPENMS_DLLAPI bool concatenateAndWriteToParquet(
    const std::vector<std::shared_ptr<arrow::Table>>& tables,
    const std::string& filename,
    const ParquetWriteConfig& config = ParquetWriteConfig{},
    const std::shared_ptr<const arrow::KeyValueMetadata>& metadata = nullptr);
 
  OPENMS_DLLAPI bool validateQPXTable(
    const std::shared_ptr<arrow::Table>& table,
    const std::string& filename,
    const ParquetWriteConfig& config = ParquetWriteConfig{},
    const std::shared_ptr<const arrow::KeyValueMetadata>& metadata = nullptr);
 
  class OPENMS_DLLAPI ParquetTableStream
  {
  public:
    ParquetTableStream(std::string filename, const ParquetWriteConfig& config,
                       std::shared_ptr<const arrow::KeyValueMetadata> metadata);
    ~ParquetTableStream();
    ParquetTableStream(const ParquetTableStream&) = delete;
    ParquetTableStream& operator=(const ParquetTableStream&) = delete;
 
    bool append(const std::shared_ptr<arrow::Table>& table, int threads = 0);
    bool close(const std::map<std::string, std::string>& keys = {});
    void abandon();
    Size rows() const;
    bool failed() const;
    bool closed() const;
    const std::string& filename() const;
 
  private:
    struct Impl;
    std::unique_ptr<Impl> impl_;
  };
 
  // ---------------------------------------------------------------------------
  // QPX file metadata
  // ---------------------------------------------------------------------------
 
  OPENMS_DLLAPI std::shared_ptr<const arrow::KeyValueMetadata> qpxFileMetadata(
    const std::string& file_type,
    const ParquetWriteConfig& config = ParquetWriteConfig{},
    const std::map<std::string, std::string>& extra = {});
 
  OPENMS_DLLAPI std::string qpxScanFormat(const std::string& native_id);
 
  OPENMS_DLLAPI std::string qpxScanFormat(const std::vector<std::string>& native_ids);
 
  OPENMS_DLLAPI std::string qpxScanFormat(const std::vector<std::string>& native_ids, int threads);
 
  OPENMS_DLLAPI std::string qpxRunFileName(const std::string& ms_run_path);
 
  OPENMS_DLLAPI std::vector<Int32> qpxScanComponents(const std::string& spectrum_reference);
 
  struct OPENMS_DLLAPI QPXRunFileNameKeys
  {
    QPXRunFileNameKeys();
    UInt reference_file_name;
    UInt run_file_name;
  };
 
  OPENMS_DLLAPI std::string qpxPsmRunFileName(const MetaInfoInterface& hit,
                                              const MetaInfoInterface& identification,
                                              const std::string& resolved_run_file,
                                              const QPXRunFileNameKeys& keys);
 
  OPENMS_DLLAPI bool qpxWarnOnRunNameCollisions(
    const std::string& context,
    const std::vector<std::string>& ms_run_paths);
 
 
  OPENMS_DLLAPI std::string qpxIntensityLabel(
    const std::string& column_label,
    const std::string& channel_name);
 
  OPENMS_DLLAPI std::map<std::uint64_t, std::string> qpxIntensityLabels(
    const ConsensusMap& cmap);
 
  OPENMS_DLLAPI bool qpxIsCanonicalIntensityLabel(const std::string& label);
 
  OPENMS_DLLAPI std::vector<std::pair<std::string, std::string>> qpxCvParams(const MetaInfoInterface& hit);
 
  OPENMS_DLLAPI std::vector<std::pair<std::string, std::string>> qpxCvParams(const MetaInfoInterface& hit,
                                                                              const MS1LabelState::Keys& keys);
 
  // ---------------------------------------------------------------------------
  // Read helpers
  // ---------------------------------------------------------------------------
 
  OPENMS_DLLAPI std::shared_ptr<arrow::Array> getColumn(
    const std::shared_ptr<arrow::Table>& table,
    const std::string& name,
    bool required = true);
 
  OPENMS_DLLAPI std::string getStringValue(
    const std::shared_ptr<arrow::Array>& array,
    int64_t row);
 
  OPENMS_DLLAPI double getDoubleValue(
    const std::shared_ptr<arrow::Array>& array,
    int64_t row,
    double default_val = 0.0);
 
  OPENMS_DLLAPI float getFloatValue(
    const std::shared_ptr<arrow::Array>& array,
    int64_t row,
    float default_val = 0.0f);
 
  OPENMS_DLLAPI int32_t getInt32Value(
    const std::shared_ptr<arrow::Array>& array,
    int64_t row,
    int32_t default_val = 0);
 
  OPENMS_DLLAPI int64_t getInt64Value(
    const std::shared_ptr<arrow::Array>& array,
    int64_t row,
    int64_t default_val = 0);
 
  OPENMS_DLLAPI bool getBoolValue(
    const std::shared_ptr<arrow::Array>& array,
    int64_t row,
    bool default_val = false);
 
  OPENMS_DLLAPI bool isNull(
    const std::shared_ptr<arrow::Array>& array,
    int64_t row);
 
  OPENMS_DLLAPI std::string metaValueString(const DataValue& value);
 
  OPENMS_DLLAPI void readMetaValues(
    const std::shared_ptr<arrow::Array>& array,
    int64_t row,
    MetaInfoInterface& target,
    const std::unordered_set<std::string>& excluded_keys = {});
}
 
} // namespace OpenMS