Merge pull request #2832 from guj/pyAPI_query

Added python query API

bindings/Python/CMakeLists.txt

@@ -6,6 +6,7 @@ Python_add_library(adios2_py MODULE
   py11Attribute.cpp
   py11Engine.cpp
   py11Operator.cpp
+  py11Query.cpp
   py11File.cpp py11File.tcc
   py11glue.cpp
 )

bindings/Python/py11Engine.h

@@ -30,6 +30,7 @@ class IO; // friend
 class Engine
 {
     friend class IO;
+    friend class Query;
 
 public:
     struct Info

bindings/Python/py11Query.cpp

@@ -0,0 +1,45 @@
+/*
+ * Distributed under the OSI-approved Apache License, Version 2.0.  See
+ * accompanying file Copyright.txt for details.
+ *
+ * py11Query.h :
+ *
+ *  Created on: August 5, 2021
+ *      Author: Junmin Gu (jgu@lbl.gov)
+ */
+
+#include "py11Query.h"
+
+namespace adios2
+{
+namespace py11
+{
+
+Query::Query(adios2::query::Worker *qw) : m_QueryWorker(qw) {}
+
+Query::Query(std::string queryFile, Engine reader)
+{
+    adios2::query::Worker *m =
+        adios2::query::GetWorker(queryFile, reader.m_Engine);
+    if (m == nullptr)
+        throw std::invalid_argument("ERROR: unable to construct query. ");
+    m_QueryWorker = std::make_shared<adios2::query::Worker>(std::move(*m));
+    delete m;
+}
+
+Query::operator bool() const noexcept
+{
+    return (m_QueryWorker == nullptr) ? false : true;
+}
+
+std::vector<Box<Dims>> Query::GetResult()
+{
+    // std::cout<<"Do something"<<std::endl;
+    adios2::Box<adios2::Dims> empty; // look into all data
+    std::vector<Box<Dims>> touched_blocks;
+    m_QueryWorker->GetResultCoverage(empty, touched_blocks);
+    return touched_blocks;
+}
+
+} // py11
+} // adios2

bindings/Python/py11Query.h

@@ -0,0 +1,48 @@
+/*
+ * Distributed under the OSI-approved Apache License, Version 2.0.  See
+ * accompanying file Copyright.txt for details.
+ *
+ * py11Query.h :
+ *
+ *  Created on: August 5, 2021
+ *      Author: Junmin Gu (jgu@lbl.gov)
+ */
+
+#ifndef ADIOS2_BINDINGS_PYTHON_QUERY_H_
+#define ADIOS2_BINDINGS_PYTHON_QUERY_H_
+
+#include <pybind11/numpy.h>
+
+//#include "adios2/toolkit/query/Query.h"
+#include "adios2/toolkit/query/Worker.h"
+#include "py11Engine.h"
+
+namespace adios2
+{
+namespace py11
+{
+
+class Engine;
+
+class Query
+{
+
+public:
+    Query(std::string queryFile, Engine reader);
+    ~Query() = default;
+
+    explicit operator bool() const noexcept;
+
+    std::vector<Box<Dims>> GetResult();
+    // const Box< Dims > & refinedSelectionIfAny,
+    // std::vector< Box< Dims > > &touched_blocks
+
+private:
+    Query(adios2::query::Worker *qw);
+    std::shared_ptr<adios2::query::Worker> m_QueryWorker;
+};
+
+} // end namespace py11
+} // end namespace adios2
+
+#endif /* BINDINGS_PYTHON_PYQUERY_H_ */

bindings/Python/py11glue.cpp

@@ -26,6 +26,7 @@
 #include "py11File.h"
 #include "py11IO.h"
 #include "py11Operator.h"
+#include "py11Query.h"
 #include "py11Variable.h"
 
 #if ADIOS2_USE_MPI
@@ -364,6 +365,25 @@ PYBIND11_MODULE(ADIOS2_PYTHON_MODULE_NAME, m)
         .def("RemoveAttribute", &adios2::py11::IO::RemoveAttribute)
         .def("RemoveAllAttributes", &adios2::py11::IO::RemoveAllAttributes);
 
+    pybind11::class_<adios2::py11::Query>(m, "Query")
+        .def("__nonzero__",
+             [](const adios2::py11::Query &query) {
+                 const bool opBool = query ? true : false;
+                 return opBool;
+             })
+        // Python 3
+        .def("__bool__",
+             [](const adios2::py11::Query &query) {
+                 const bool opBool = query ? true : false;
+                 return opBool;
+             })
+        .def(
+            pybind11::init<const std::string &, const adios2::py11::Engine &>(),
+            "adios2 query construction, a xml query File and a read engine",
+            pybind11::arg("queryFile"), pybind11::arg("reader") = true)
+
+        .def("GetResult", &adios2::py11::Query::GetResult);
+
     pybind11::class_<adios2::py11::Variable>(m, "Variable")
         // Python 2
         .def("__nonzero__",
@@ -522,7 +542,6 @@ PYBIND11_MODULE(ADIOS2_PYTHON_MODULE_NAME, m)
              [](const adios2::py11::File &stream) { return stream; })
         .def("__exit__",
              [](adios2::py11::File &stream, pybind11::args) { stream.Close(); })
-
         .def("__iter__", [](adios2::py11::File &stream) { return stream; },
              pybind11::keep_alive<0, 1>())
         .def("__next__",

source/adios2/toolkit/query/Query.cpp

@@ -247,12 +247,14 @@ bool QueryVar::IsSelectionValid(adios2::Dims &shape) const
         return false; // different dimension
     }
 
+    /*
     for (size_t i = 0; i < shape.size(); i++)
     {
         if ((m_Selection.first[i] > shape[i]) ||
             (m_Selection.second[i] > shape[i]))
             return false;
     }
+    */
     return true;
 }
 

source/adios2/toolkit/query/XmlWorker.cpp

@@ -196,8 +196,6 @@ void XmlWorker::ConstructQuery(QueryVar &simpleQ, const pugi::xml_node &node)
 
     if (bbNode)
     {
-        adios2::Box<adios2::Dims> box =
-            adios2::Box<adios2::Dims>({100, 100}, {200, 200});
         std::string startStr =
             adios2::helper::XMLAttribute("start", bbNode, "in query")->value();
         std::string countStr =

testing/adios2/bindings/python/TestQuery.py

@@ -0,0 +1,106 @@
+from mpi4py import MPI
+import numpy as np
+import adios2
+import sys
+import os
+import math
+import matplotlib.pyplot as plt
+
+# MPI
+comm = MPI.COMM_WORLD
+rank = comm.Get_rank()
+size = comm.Get_size()
+
+configFile = './defaultConfig.xml'
+queryFile = './sampleQuery.xml'
+dataPath = './heat.bp'
+
+def doAnalysis(reader, touched_blocks, varList):
+    print(" Step: ", reader.CurrentStep(),
+          "  num touched blocks: ", len(touched_blocks))
+    values = []
+    data = {}
+
+    for var in varList:
+        data[var] = []
+
+    if (len(touched_blocks) > 0):
+        for n in touched_blocks:
+            for var in varList:
+                values = np.zeros(n[1], dtype=np.double)
+                var.SetSelection(n)
+                reader.Get(var, values, adios2.Mode.Sync)
+                data[var].extend(values)
+                # do analysis with data here
+
+def runQuery():
+    adios = adios2.ADIOS(configFile, comm, True)
+    queryIO = adios.DeclareIO("query")
+    reader = queryIO.Open(dataPath, adios2.Mode.Read, comm)
+    w = adios2.Query(queryFile, reader)
+
+    touched_blocks = []
+
+    var = [queryIO.InquireVariable("T")]
+
+    print("Num steps: ", reader.Steps())
+
+    while (reader.BeginStep() == adios2.StepStatus.OK):
+        # say only rank 0 wants to process result
+        if (rank == 0):
+            touched_blocks = w.GetResult()
+            doAnalysis(reader, touched_blocks, var)
+
+    reader.EndStep()
+    reader.Close()
+
+def createConfigFile():
+    print(".. Writing config file to: ", configFile)
+    file1 = open(configFile, 'w')
+
+    xmlContent = ["<?xml version=\"1.0\"?>\n",
+                  "<adios-config>\n",
+                  "<io name=\"query\">\n",
+                  "  <engine type=\"BPFile\">\n",
+                  "  </engine>\n",
+                  "  <transport type=\"File\">\n",
+                  "    <parameter key=\"Library\" value=\"POSIX\"/>\n",
+                  "  </transport>\n",
+                  "</io>\n", "</adios-config>\n"]
+
+    file1.writelines(xmlContent)
+    file1.close()
+
+def createQueryFile():
+    print(".. Writing query file to: ", queryFile)
+
+    file1 = open(queryFile, 'w')
+    queryContent = [
+        "<?xml version=\"1.0\"?>\n", "<adios-query>\n",
+        "  <io name=\"query\">\n"
+        "  <var name=\"T\">\n",
+        "    <op value=\"AND\">\n",
+        "      <range  compare=\"LT\" value=\"2.7\"/>\n",
+        "      <range  compare=\"GT\" value=\"2.66\"/>\n", "    </op>\n",
+        "  </var>\n", "  </io>\n", "</adios-query>\n"
+    ]
+    file1.writelines(queryContent)
+    file1.close()
+
+
+if (os.path.exists(dataPath) is False):
+    print("Please generate data file:", dataPath,
+          " from heat transfer example first.")
+else:
+    # configFile created
+    createConfigFile()
+
+    # queryFile Generated
+    createQueryFile()
+
+    print(".. Running query against: ", dataPath)
+    runQuery()
+
+    print("Now clean up.")
+    os.remove(queryFile)
+    os.remove(configFile)

testing/adios2/bindings/python/heat.bp/profiling.json

@@ -0,0 +1,6 @@
+[
+{ "rank": 0, "start": "Thu_Aug_19_14:58:09_2021", "threads": 1, "bytes": 3964, "mkdir_mus": 186, "memcpy_mus": 0, "aggregation_mus": 0, "meta_sort_merge_mus": 542, "minmax_mus": 48, "buffering_mus": 1136, "transport_0": { "type": "File_stdio", "close_mus": 142, "write_mus": 182, "open_mus": 0 },"transport_1": { "type": "File_stdio", "close_mus": 0, "write_mus": 48, "open_mus": 0 } },
+{ "rank": 1, "start": "Thu_Aug_19_14:58:09_2021", "threads": 1, "bytes": 3888, "mkdir_mus": 193, "memcpy_mus": 0, "aggregation_mus": 0, "meta_sort_merge_mus": 65, "minmax_mus": 61, "buffering_mus": 774,  },
+{ "rank": 2, "start": "Thu_Aug_19_14:58:09_2021", "threads": 1, "bytes": 3888, "mkdir_mus": 196, "memcpy_mus": 0, "aggregation_mus": 0, "meta_sort_merge_mus": 110, "minmax_mus": 47, "buffering_mus": 768,  },
+{ "rank": 3, "start": "Thu_Aug_19_14:58:09_2021", "threads": 1, "bytes": 3888, "mkdir_mus": 197, "memcpy_mus": 0, "aggregation_mus": 0, "meta_sort_merge_mus": 100, "minmax_mus": 47, "buffering_mus": 671,  }
+]