Standalone tool for translating Parquet files into CSV

Author: iagaponenko

src/partition/CMakeLists.txt

@@ -48,6 +48,7 @@ partition_apps(
     sph-estimate-stats
     sph-htm-index
     sph-layout
+    sph-parq2csv
     sph-partition-matches
     sph-partition
 )

src/partition/sph-parq2csv.cc

@@ -22,20 +22,23 @@
 #include <iomanip>
 #include <iostream>
 #include <fstream>
-#include <map>
 #include <memory>
+#include <set>
 #include <stdexcept>
 #include <string>
 #include <vector>
 
+// Third party headers
 #include "boost/program_options.hpp"
+#include "nlohmann/json.hpp"
 
 #include "partition/ParquetInterface.h"
 
 namespace po = boost::program_options;
 namespace part = lsst::partition;
 
 using namespace std;
+using json = nlohmann::json;
 
 namespace {
 
@@ -64,15 +67,15 @@ class CommandLineParams {
                            "Max size (MB) of RAM allocated to the process.");
         desc.add_options()("buf-size-mb", po::value<int>()->default_value(maxBuffSizeMB),
                            "Buffers size (MB) for translating batches.");
-        desc.add_options()("parq-file", po::value<vector<string>>(), "Input file to be translated.");
-        desc.add_options()("coldef-file", po::value<vector<string>>(),
-                           "Input file with the names of columns to be extracted.");
-        desc.add_options()("csv-file", po::value<vector<string>>(), "Output file to be written.");
+        desc.add_options()("parq", po::value<vector<string>>(), "Input file to be translated.");
+        desc.add_options()("config", po::value<vector<string>>(),
+                           "Input JSON file with definition of columns to be extracted.");
+        desc.add_options()("csv", po::value<vector<string>>(), "Output file to be written.");
 
         po::positional_options_description p;
-        p.add("parq-file", 1);
-        p.add("coldef-file", 1);
-        p.add("csv-file", 1);
+        p.add("parq", 1);
+        p.add("config", 1);
+        p.add("csv", 1);
 
         po::variables_map vm;
         po::store(po::command_line_parser(argc, argv).options(desc).positional(p).run(), vm);
@@ -82,17 +85,17 @@ class CommandLineParams {
             cout << desc << "\n";
             return false;
         }
-        parqFileName = vm.count("parq-file") ? vm["parq-file"].as<vector<string>>().front() : string();
-        coldefFileName = vm.count("coldef-file") ? vm["coldef-file"].as<vector<string>>().front() : string();
-        csvFileName = vm.count("csv-file") ? vm["csv-file"].as<vector<string>>().front() : string();
+        parqFileName = vm.count("parq") ? vm["parq"].as<vector<string>>().front() : string();
+        configFileName = vm.count("config") ? vm["config"].as<vector<string>>().front() : string();
+        csvFileName = vm.count("csv") ? vm["csv"].as<vector<string>>().front() : string();
 
-        if (parqFileName.empty() || coldefFileName.empty() || csvFileName.empty()) {
+        if (parqFileName.empty() || configFileName.empty() || csvFileName.empty()) {
             throw runtime_error("The names of all required files must be provided.");
         }
         if (csvFileName == parqFileName) {
             throw runtime_error("Input and output file names must be different.");
         }
-        _parseColdDefFile();
+        _parseConfigFile();
 
         if (vm.count("max-mem-alloc-mb")) {
             maxMemAllocatedMB = vm["max-mem-alloc-mb"].as<int>();
@@ -114,11 +117,11 @@ class CommandLineParams {
     // Values of the parsed parameters aere stored in the data members defined below.
 
     string parqFileName;
-    string coldefFileName;
+    string configFileName;
     string csvFileName;
 
     vector<string> columns;
-    map<string, string> optionalColumnDefs;
+    set<string> optionalColumns;
 
     int maxMemAllocatedMB = 3000;
     int maxBuffSizeMB = 16;
@@ -129,19 +132,44 @@ class CommandLineParams {
     bool verbose = false;
 
 private:
-    void _parseColdDefFile() {
-        columns.clear();
-        ifstream columnsFile(coldefFileName);
-        if (!columnsFile) {
-            throw runtime_error("Error while opening the columns file.");
+    void _parseConfigFile() {
+        ifstream file(configFileName, ios_base::in);
+        if (!file.good()) throw invalid_argument("Failed to open file: '" + configFileName + "'");
+        json config;
+        try {
+            file >> config;
+        } catch (...) {
+            throw runtime_error("Config file: '" + configFileName + "' doesn't have a valid JSON payload");
         }
-        string column;
-        while (columnsFile >> column) {
-            columns.push_back(column);
+        if (!config.is_object()) {
+            throw invalid_argument("Config file: '" + configFileName + "' is not a valid JSON object");
         }
+        if (!config.contains("columns")) {
+            throw runtime_error("The JSON file must contain a 'columns' key.");
+        }
+        if (!config["columns"].is_array()) {
+            throw runtime_error("The 'columns' key must contain an array.");
+        }
+        columns = config["columns"].get<vector<string>>();
         if (columns.empty()) {
             throw runtime_error("No columns to be extracted.");
         }
+        optionalColumns.clear();
+        if (config.contains("optional")) {
+            if (!config["optional"].is_array()) {
+                throw runtime_error("The 'optional' key must contain an object.");
+            }
+            for (auto const& column : config["optional"].get<vector<string>>()) {
+                optionalColumns.insert(column);
+            }
+        }
+        // All optional columns must be defined in the 'columns' array.
+        for (auto const& name : optionalColumns) {
+            if (find(columns.begin(), columns.end(), name) == columns.end()) {
+                throw runtime_error("The optional column '" + name +
+                                    "' is not defined in the 'columns' array.");
+            }
+        }
     }
 };
 
@@ -163,23 +191,24 @@ int main(int argc, char const* const* argv) {
             cout << "Translating '" << params.parqFileName << "' into '" << params.csvFileName << "'" << endl;
         }
         part::ParquetFile parqFile(params.parqFileName, params.maxMemAllocatedMB);
-        if (parqFile.setupBatchReader(maxBuffSizeBytes) != arrow::Status::OK()) {
-            throw runtime_error("Error while setting up the batch reader.");
-        }
+        parqFile.setupBatchReader(maxBuffSizeBytes);
+
         ofstream csvFile(params.csvFileName, ios::out | ios::binary);
         if (!csvFile) {
             throw runtime_error("Error while opening the output file.");
         }
         while (true) {
-            auto status = parqFile.readNextBatch_Table2CSV(buf.get(), buffSize, params.columns,
-                                                           params.optionalColumnDefs, params.nullStr,
-                                                           params.delimStr);
-            if ((status != arrow::Status::OK()) || (buffSize == 0)) break;
-            if (params.verbose) {
-                cout << "Writing " << setw(9) << buffSize << " bytes" << endl;
+            bool const hasMore =
+                    parqFile.readNextBatch_Table2CSV(buf.get(), buffSize, params.columns,
+                                                     params.optionalColumns, params.nullStr, params.delimStr);
+            if (buffSize > 0) {
+                if (params.verbose) {
+                    cout << "Writing " << setw(9) << buffSize << " bytes" << endl;
+                }
+                csvFile.write((char*)(buf.get()), buffSize);
+                numBytesWritten += buffSize;
             }
-            csvFile.write((char*)(buf.get()), buffSize);
-            numBytesWritten += buffSize;
+            if (!hasMore || (buffSize == 0)) break;
         }
         csvFile.close();
         if (params.verbose) {