data_logger.h

Go to the documentation of this file.

 
#ifndef DATA_LOGGER_H
#define DATA_LOGGER_H
 
#include <arrow/api.h>
#include <arrow/io/file.h>
#include <arrow/ipc/api.h>
#include <arrow/util/compression.h>
#include <unordered_map>
#include <variant>
#include <vector>

class DataLogger {
public:
    DataLogger() = default;

    explicit DataLogger(int64_t flush_row_count);

    virtual ~DataLogger();

    bool column_exists(const std::string& column_name);

    bool column_value_already_set(const std::string& column_name);

    void add_metadata(const std::string& key, const std::string& value);

    void add_field(const std::string& name, std::shared_ptr<arrow::DataType> type, bool ignore_existing_name = false);

    void open_file(const std::string& filename);

    void set_value(const std::string& column_name, int64_t value);

    void set_value(const std::string& column_name, int32_t value);

    void set_value(const std::string& column_name, int16_t value);

    void set_value(const std::string& column_name, int8_t value);

    void set_value(const std::string& column_name, double value);

    void set_value(const std::string& column_name, float value);

    void set_value(const std::string& column_name, const std::string& value);

    void set_value(const std::string& column_name, bool value);

    void set_value_float16(const std::string& column_name, float value);

    void save_row();

    void flush();
 
private:
    // Store Arrow float16 values as their raw 16-bit representation
    using half_float_t = arrow::NumericBuilder<arrow::HalfFloatType>::value_type;

    std::vector<std::shared_ptr<arrow::Field>> fields_;
    std::shared_ptr<arrow::Schema> schema_;
    std::shared_ptr<arrow::io::OutputStream> outfile_;
    std::shared_ptr<arrow::ipc::RecordBatchWriter> writer_;
    std::vector<std::shared_ptr<arrow::ArrayBuilder>> builders_;
    std::unordered_map<std::string, size_t> column_indices_;
    std::unordered_map<std::string, std::variant<int64_t, int32_t, int16_t, int8_t, float, double, std::string, bool, half_float_t>> row_values_;
    bool file_opened_ = false;
    std::unordered_map<std::string, std::string> user_metadata_;
    int64_t flush_row_count_ = 1024;
    int64_t buffered_row_count_ = 0;

    void check_column(const std::string& column_name);

    void reset_row();

    void initialize_builders();

    void append_current_row_to_builders();
 
    // Convert float32 -> IEEE-754 half payload (16-bit)
    static half_float_t float_to_half_bits(float f);
};
 
#endif // DATA_LOGGER_H