Changeset 19280

Timestamp:

10/23/18 18:48:24 (6 years ago)

Author:

tbretz

Message:

Implemented automatic splitting into root-trees or ascii-files.

File:

• trunk/FACT++/src/rootifysql.cc (modified) (13 diffs)

trunk/FACT++/src/rootifysql.cc

@@ -60 +60 @@
         ;
 
+    po::options_description split("Splitting options");
+    split.add_options()
+        ("split-sequence,S", vars<uint16_t>(),            "Split data sequentially into several trees/files (e.g. 1, 1, 2)")
+        ("split-quantile,Q", vars<double>(),              "Split data randomly into several trees/files (e.g. 0.5, 1)")
+        ("seed", var<uint64_t>(mt19937_64::default_seed), "Seed value in case of random split")
+        ;
+
     po::positional_options_description p;
     p.add("file", 1); // The 1st positional options (n=1)
@@ -67 +74 @@
     conf.AddOptions(ascii);
     conf.AddOptions(root);
+    conf.AddOptions(split);
     conf.SetArgumentPositions(p);
 }
@@ -114 +122 @@
         "Comments in the query-file can be placed according to the SQL standard inline "
         "/*comment*/ or introduced with # (shell script style) or -- (SQL style).\n"
+        "\n"
+        "For several purposes, it might be convenient to split the output to several "
+        "different root-trees or ascii files. This can be done using the --split-sequence (-S) "
+        "and the --split-quantile (-Q) options. If a split sequence is defined as "
+        "-S 1 -S 2 -S 1 the events are split by 1:2:1 in this sequence order. If "
+        "quantiled are given as -Q 0.5 -Q 0.6, the first tree will contain 50% of "
+        "the second one 10% and the third one 40%. The corresponding seed value can "
+        "be set with --seed.\n"
         "\n"
         "In case of success, 0 is returned, a value>0 otherwise.\n"
@@ -502 +518 @@
     //const vector<Map> mymap    = conf.Vec<Map>("map");
 
+    vector<uint16_t> split_seq   = conf.Vec<uint16_t>("split-sequence");
+    vector<double>   split_quant = conf.Vec<double>("split-quantile");
+
+    if (!split_seq.empty() && !split_quant.empty())
+        throw runtime_error("Only splitting by --split-sequence or --split-quantile is allowed.");
+
+    const size_t num_split = ::max(split_seq.size(), split_quant.size()+1);
+
+    map<size_t, size_t> split_lut;
+    for (size_t i=0; i<split_seq.size(); i++)
+    {
+        const size_t sz = split_lut.size();
+        for (size_t j=0; j<split_seq[i]; j++)
+            split_lut.emplace(j+sz, i);
+    }
+
+    for (size_t i=0; i<split_quant.size(); i++)
+        if (split_quant[i]<0 || split_quant[i]>=1)
+            throw runtime_error("Splitting quantiles must be in the range [0;1)");
+
+    for (size_t i=1; i<split_quant.size(); i++)
+    {
+        if (split_quant[i]<=split_quant[i-1])
+            throw runtime_error("Splitting quantiles must be in increasing order.");
+    }
+
     // -------------------------------------------------------------------------
 
@@ -582 +624 @@
     // waiting long for the query result
     //
+
     // I am using root here instead of boost to be
     // consistent with the access pattern by TFile
@@ -778 +821 @@
         cout << "Opening file '" << path << "' [compression=" << compression << "]...\n";
         cout << "Writing data to tree '" << tree << "'" << (nofill?" (--skipped--)":"") << endl;
+        if (num_split)
+            cout << "Splitting configured with " << num_split << " branches." << endl;
     }
 
@@ -808 +853 @@
     UInt_t cols = 0;
 
+    // IMPLEMENT FILE SPLITTING!
+    // OpenFile(tree, query)
+    // SetupColumns
+    // WriteRow
+    // CloseFile
+
+    // Ratio[3]: 50%, 20%, 30%
+    // File[x3]: root, cout, fout
+
 
     // -------------------- Configure branches of TTree ------------------------
-    TTree *ttree = new TTree(tree.c_str(), query.c_str());
+    vector<TTree*> ttree;
+
+    if (num_split==0)
+        ttree.emplace_back(new TTree(tree.c_str(), query.c_str()));
+    else
+        for (int i=0; i<num_split; i++)
+            ttree.emplace_back(new TTree((tree+"["+to_string(i)+"]").c_str(), query.c_str()));
 
     size_t skipat  = 0;
@@ -862 +922 @@
             //     name = boost::regex_replace(l[i], boost::regex(m.first), m.second);
 
-            ttree->Branch(l[i].c_str(), buf.data()+i);
+            for (auto it=ttree.begin(); it!=ttree.end(); it++)
+                it[0]->Branch(l[i].c_str(), buf.data()+i);
             cols++;
         }
@@ -879 +940 @@
     }
 
-    // -------------------------------------------------------------------------
+    // ------------------------- Open the ascii files --------------------------
+
+    vector<ofstream> fout;
+    if (!write.empty())
+    {
+        vector<string> names;
+        if (num_split==0)
+            names.emplace_back(write);
+        else
+            for (int i=0; i<num_split; i++)
+                names.emplace_back(write+"-"+to_string(i));
+
+        for (auto it=names.cbegin(); it!=names.cend(); it++)
+        {
+            fout.emplace_back(*it);
+            if (!*fout.rbegin())
+                cout << "WARNING: Writing to '" << write << "' failed: " << strerror(errno) << endl;
+        }
+    }
+
+    // ----------------------- Prepare the ascii comment -----------------------
 
     string contents;
@@ -905 +986 @@
             (copy_comments && comment && !header))
             contents += '#' + ibuf + '\n';
-
-    }
-
-    ofstream fout(write);
-    if (!write.empty() && !fout)
-        cout << "WARNING: Writing to '" << write << "' failed: " << strerror(errno) << endl;
+    }
+
+    // ----------------------- Write the ascii headers -------------------------
+
+    ostringstream htxt;
+    if (display || !fout.empty())
+        htxt << row.field_list(delimiter.c_str());
 
     if (display)
@@ -916 +998 @@
         cout << endl;
         cout << contents << endl;
-        cout << "# " << row.field_list(delimiter.c_str()) << endl;
-    }
-
-    if (!write.empty())
-    {
-        fout << contents;
-        fout << "# " << row.field_list(delimiter.c_str()) << endl;
+        cout << "# " << htxt.str() << endl;
+    }
+    for (auto ff=fout.begin(); ff!=fout.end(); ff++)
+    {
+        *ff << contents;
+        *ff << "# " << htxt.str() << endl;
     }
 
     // ---------------------- Fill TTree with DB data --------------------------
+
+    const uniform_real_distribution<double> distribution(0,1);
+    mt19937_64 generator;
+    if (conf.Has("seed"))
+        generator.seed(conf.Get<uint64_t>("seed"));
+    auto rndm = bind(distribution, generator);
+
     size_t count = 0;
     size_t skip  = 0;
     do
     {
+        size_t index = 0;
+        if (!split_lut.empty())
+            index = split_lut[count % split_lut.size()];
+        if (!split_quant.empty())
+        {
+            const float r = rndm();
+            for (; r>=split_quant[index]; index++)
+                if (index==split_quant.size())
+                    break;
+        }
+
         count++;
 
+        ostringstream rtxt;
+        if (display || !fout.empty())
+            rtxt << row.value_list(delimiter.c_str(), mysqlpp::do_nothing);
+
         if (display)
-            cout << row.value_list(delimiter.c_str(), mysqlpp::do_nothing) << '\n';
-        if (!write.empty())
-            fout << row.value_list(delimiter.c_str(), mysqlpp::do_nothing) << '\n';
+            cout << rtxt.str() << '\n';
+        if (!fout.empty())
+            fout[index] << rtxt.str() << '\n';
 
         size_t idx=0;
@@ -972 +1075 @@
 
         if (idx==row.size() && !nofill)
-            ttree->Fill();
+            ttree[index]->Fill();
 
         row = res.fetch_row();
@@ -990 +1093 @@
             cout << skip << " rows skipped due to NULL field." << endl;
 
-        cout << ttree->GetEntries() << " rows filled into tree." << endl;
-    }
-
-    ttree->Write();
+        for (int i=0; i<ttree.size(); i++)
+            cout << ttree[i]->GetEntries() << " rows filled into tree #" << i << "." << endl;
+    }
+
+    for (auto it=ttree.begin(); it!=ttree.end(); it++)
+        (*it)->Write();
     tfile.Close();