csv_utility.hpp

#pragma once
#include "common.hpp"
#include "csv_exceptions.hpp"
#include "csv_format.hpp"
#include "csv_reader.hpp"
#include "data_frame.hpp"
#include "data_type.hpp"
#include "string_view_stream.hpp"
 
#include <memory>
#include <sstream>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <utility>
 
namespace csv {
    struct CSVFileInfo {
        std::string filename;               
        std::vector<std::string> col_names; 
        char delim;                         
        size_t n_rows;                      
        size_t n_cols;                      
        size_t parse_worker_count;          
        internals::SpeculativeParseDiagnostics speculative_diagnostics;
    };
 
 
    inline CSVReader parse(csv::string_view in, const CSVFormat& format = CSVFormat::guess_csv()) {
        std::unique_ptr<std::istream> ss(new std::stringstream(std::string(in)));
        return CSVReader(std::move(ss), format);
    }
 
    inline CSVReader parse_unsafe(csv::string_view in, CSVFormat format = CSVFormat::guess_csv()) {
        std::unique_ptr<std::istream> stream(new internals::StringViewStream(in));
        return CSVReader(std::move(stream), format);
    }
 
    inline CSVReader parse_no_header(csv::string_view in) {
        CSVFormat format;
        format.header_row(-1);
        return parse(in, format);
    }
 
    inline CSVReader operator ""_csv(const char* in, size_t n) {
        return parse_unsafe(csv::string_view(in, n));
    }
 
    inline CSVReader operator ""_csv_no_header(const char* in, size_t n) {
        CSVFormat format;
        format.header_row(-1);
        return parse_unsafe(csv::string_view(in, n), format);
    }
 
 
    std::unordered_map<std::string, DataType> csv_data_types(CSVReader& reader);
 
    template<
        typename... ReaderArgs,
        csv::enable_if_t<std::is_constructible<CSVReader, ReaderArgs...>::value, int> = 0
    >
    inline std::unordered_map<std::string, DataType> csv_data_types(ReaderArgs&&... reader_args) {
        CSVReader reader(std::forward<ReaderArgs>(reader_args)...);
        return csv_data_types(reader);
    }
 
    template<typename State, typename Fn>
    inline void chunk_parallel_apply(
        CSVReader& reader,
        DataFrameExecutor& executor,
        std::vector<State>& states,
        Fn&& fn,
        size_t chunk_size = 50000
    ) {
        if (chunk_size == 0) {
            throw std::invalid_argument(internals::ERROR_CHUNK_PARALLEL_APPLY_ZERO);
        }
 
        std::vector<CSVRow> rows;
        while (reader.read_chunk(rows, chunk_size)) {
            DataFrame<> batch(std::move(rows));
            batch.column_parallel_apply(executor, states, std::forward<Fn>(fn));
        }
    }
 
    template<typename State, typename Fn>
    inline void chunk_parallel_apply(
        CSVReader& reader,
        std::vector<State>& states,
        Fn&& fn,
        size_t chunk_size = 50000
    ) {
        DataFrameExecutor executor;
        chunk_parallel_apply(reader, executor, states, std::forward<Fn>(fn), chunk_size);
    }
 
    inline CSVFileInfo get_file_info(const std::string& filename) {
        CSVFormat reader_format = CSVFormat::guess_csv();
 
        CSVReader reader(filename, reader_format);
        CSVFormat format = reader.get_format();
 
        std::vector<CSVRow> rows;
        while (reader.read_chunk(rows, 50000)) {}
 
        return {
            filename,
            reader.get_col_names(),
            format.get_delim(),
            reader.n_rows(),
            reader.get_col_names().size(),
            reader.parse_worker_count(),
            reader.speculative_diagnostics()
        };
    }
 
    inline std::vector<std::string> get_col_names(
        csv::string_view filename,
        const CSVFormat& format = CSVFormat::guess_csv()) {
        auto head = internals::parser::get_csv_head(filename);
        return parse_unsafe(head, format).get_col_names();
    }
 
    inline long long get_col_pos(csv::string_view filename, csv::string_view col_name,
        const CSVFormat& format = CSVFormat::guess_csv()) {
        auto col_names = get_col_names(filename, format);
        return col_names.empty() ? CSV_NOT_FOUND :
            std::distance(col_names.begin(), std::find(col_names.begin(), col_names.end(), col_name));
    }
}