Merge pull request #91 from TileDB-Inc/sethshelnutt/ch2178/expose-driver-option-to-enable-stats-printing

 Add support for enabling and dumping TileDB Stats

Author: Shelnutt2

apis/java/src/main/java/io/tiledb/libvcfnative/LibVCFNative.c

@@ -559,4 +559,52 @@ Java_io_tiledb_libvcfnative_LibVCFNative_tiledb_1vcf_1reader_1get_1last_1error_1
   jstring result = (*env)->NewStringUTF(env, msg);
   tiledb_vcf_error_free(&error);
   return result;
+}
+
+JNIEXPORT jint JNICALL
+Java_io_tiledb_libvcfnative_LibVCFNative_tiledb_1vcf_1reader_1set_1tiledb_1stats_1enabled(
+    JNIEnv* env, jclass self, jlong readerPtr, jboolean statsEnabled) {
+  (void)self;
+  tiledb_vcf_reader_t* reader = (tiledb_vcf_reader_t*)readerPtr;
+  if (reader == 0) {
+    return TILEDB_VCF_ERR;
+  }
+
+  const bool stats_enabled = statsEnabled ? true : false;
+
+  int32_t rc = tiledb_vcf_reader_set_tiledb_stats_enabled(reader, stats_enabled);
+
+  return rc;
+}
+
+JNIEXPORT jboolean JNICALL
+Java_io_tiledb_libvcfnative_LibVCFNative_tiledb_1vcf_1reader_1tiledb_1stats_1enabled(
+    JNIEnv* env, jclass self, jlong readerPtr) {
+  (void)self;
+  tiledb_vcf_reader_t* reader = (tiledb_vcf_reader_t*)readerPtr;
+  if (reader == 0) {
+    return TILEDB_VCF_ERR;
+  }
+
+  bool stats_enabled;
+
+  int rc = tiledb_vcf_reader_get_tiledb_stats_enabled(reader, &stats_enabled);
+
+  return stats_enabled;
+}
+
+JNIEXPORT jstring JNICALL
+Java_io_tiledb_libvcfnative_LibVCFNative_tiledb_1vcf_1reader_1tiledb_1stats(
+    JNIEnv* env, jclass self, jlong readerPtr) {
+  (void)self;
+  tiledb_vcf_reader_t* reader = (tiledb_vcf_reader_t*)readerPtr;
+  if (reader == 0) {
+    return NULL;
+  }
+
+  char *stats;
+  tiledb_vcf_reader_get_tiledb_stats(reader, &stats);
+
+  jstring result = (*env)->NewStringUTF(env, stats);
+  return result;
 }

apis/java/src/main/java/io/tiledb/libvcfnative/LibVCFNative.h

@@ -87,4 +87,11 @@ public static final native int tiledb_vcf_reader_get_dataset_version(
       long readerPtr, int[] version);
 
   public static final native String tiledb_vcf_reader_get_last_error_message(long readerPtr);
+
+  public static final native int tiledb_vcf_reader_set_tiledb_stats_enabled(
+      long readerPtr, boolean statsEnabled);
+
+  public static final native boolean tiledb_vcf_reader_tiledb_stats_enabled(long readerPtr);
+
+  public static final native String tiledb_vcf_reader_tiledb_stats(long readerPtr);
 }

apis/java/src/main/java/io/tiledb/libvcfnative/LibVCFNative.java

@@ -355,6 +355,27 @@ public AttributeTypeInfo getAttributeDatatype(String attribute) {
     }
   }
 
+  public VCFReader setStatsEnabled(boolean statsEnabled) {
+    int rc = LibVCFNative.tiledb_vcf_reader_set_tiledb_stats_enabled(this.readerPtr, statsEnabled);
+    if (rc != 0) {
+      String msg = getLastErrorMessage();
+      throw new RuntimeException("Error setting stats enabled: " + msg);
+    }
+    return this;
+  }
+
+  public boolean getStatsEnabled() {
+    boolean statsEnabled = LibVCFNative.tiledb_vcf_reader_tiledb_stats_enabled(this.readerPtr);
+
+    return statsEnabled;
+  }
+
+  public String stats() {
+    String stats = LibVCFNative.tiledb_vcf_reader_tiledb_stats(this.readerPtr);
+
+    return stats;
+  }
+
   public VCFReader resetBuffers() {
     Iterator it = buffers.entrySet().iterator();
     while (it.hasNext()) {

apis/java/src/main/java/io/tiledb/libvcfnative/VCFReader.java

@@ -181,4 +181,29 @@ public void testSetSingleBuffer() throws IOException {
 
     Assert.assertEquals(results, BED_FILE_EXPECTED_RECORDS);
   }
+
+  @Test
+  public void testSetStatsEnabled() throws IOException {
+    VCFReader reader = getVFCReader(Optional.empty(), Optional.of(constructBEDURI()));
+
+    reader.setStatsEnabled(true);
+  }
+
+  @Test
+  public void testGetStatsEnabled() throws IOException {
+    VCFReader reader = getVFCReader(Optional.empty(), Optional.of(constructBEDURI()));
+
+    Assert.assertFalse(reader.getStatsEnabled());
+    reader.setStatsEnabled(true);
+    Assert.assertTrue(reader.getStatsEnabled());
+    reader.setStatsEnabled(false);
+    Assert.assertFalse(reader.getStatsEnabled());
+  }
+
+  @Test
+  public void testStats() throws IOException {
+    VCFReader reader = getVFCReader(Optional.empty(), Optional.of(constructBEDURI()));
+    reader.setStatsEnabled(true);
+    Assert.assertNotNull(reader.stats());
+  }
 }

apis/java/src/test/java/io/tiledb/libvcfnative/VCFReaderTest.java

@@ -1 +1,2 @@
 src/tiledbvcf/version.py
+src/tiledbvcf/libhts.so.1.8

apis/python/.gitignore

@@ -29,11 +29,14 @@ PYBIND11_MODULE(libtiledbvcf, m) {
       .def("set_max_num_records", &Reader::set_max_num_records)
       .def("set_tiledb_config", &Reader::set_tiledb_config)
       .def("set_attributes", &Reader::set_attributes)
+      .def("set_tiledb_stats_enabled", &Reader::set_tiledb_stats_enabled)
       .def("read", &Reader::read)
       .def("get_buffers", &Reader::get_buffers)
       .def("get_results_arrow", &Reader::get_results_arrow)
       .def("completed", &Reader::completed)
-      .def("result_num_records", &Reader::result_num_records);
+      .def("result_num_records", &Reader::result_num_records)
+      .def("get_tiledb_stats_enabled", &Reader::get_tiledb_stats_enabled)
+      .def("get_tiledb_stats", &Reader::get_tiledb_stats);
 
   py::class_<Writer>(m, "Writer")
       .def(py::init())

apis/python/src/tiledbvcf/binding/libtiledbvcf.cc

@@ -76,6 +76,11 @@ void Reader::set_attributes(const std::vector<std::string>& attributes) {
   attributes_ = attributes;
 }
 
+void Reader::set_tiledb_stats_enabled(const bool stats_enabled) {
+  auto reader = ptr.get();
+  check_error(reader, tiledb_vcf_reader_set_tiledb_stats_enabled(reader, stats_enabled));
+}
+
 void Reader::set_samples(const std::string& samples) {
   auto reader = ptr.get();
   check_error(reader, tiledb_vcf_reader_set_samples(reader, samples.c_str()));
@@ -325,6 +330,20 @@ bool Reader::completed() {
   return status == TILEDB_VCF_COMPLETED;
 }
 
+bool Reader::get_tiledb_stats_enabled() {
+  auto reader = ptr.get();
+  bool stats_enabled;
+  check_error(reader, tiledb_vcf_reader_get_tiledb_stats_enabled(reader, &stats_enabled));
+  return stats_enabled;
+}
+
+std::string  Reader::get_tiledb_stats() {
+  auto reader = ptr.get();
+  char* stats;
+  check_error(reader, tiledb_vcf_reader_get_tiledb_stats(reader, &stats));
+  return std::string(stats);
+}
+
 py::dtype Reader::to_numpy_dtype(tiledb_vcf_attr_datatype_t datatype) {
   switch (datatype) {
     case TILEDB_VCF_CHAR:

apis/python/src/tiledbvcf/binding/reader.cc

@@ -85,6 +85,9 @@ class Reader {
   /** Sets CSV TileDB config parameters. */
   void set_tiledb_config(const std::string& config_str);
 
+  /** Sets whether internal TileDB Statistics are Enabled or Disabled*/
+  void set_tiledb_stats_enabled(const bool stats_enabled);
+
   /** Performs a blocking read operation. */
   void read();
 
@@ -106,6 +109,12 @@ class Reader {
   /** Returns true if the last read operation was complete. */
   bool completed();
 
+    /** Gets whether internal TileDB Statistics are Enabled or Disabled*/
+  bool get_tiledb_stats_enabled();
+
+  /** Fetches TileDB statistics */
+  std::string get_tiledb_stats();
+
  private:
   /** Buffer struct to hold attribute data read from the dataset. */
   struct BufferInfo {

apis/python/src/tiledbvcf/binding/reader.h

@@ -24,7 +24,7 @@
 class TileDBVCFDataset(object):
     """A handle on a TileDB-VCF dataset."""
 
-    def __init__(self, uri, mode='r', cfg=None):
+    def __init__(self, uri, mode='r', cfg=None, stats=False):
         """ Initializes a TileDB-VCF dataset for interaction.
 
         :param uri: URI of TileDB-VCF dataset
@@ -38,6 +38,7 @@ def __init__(self, uri, mode='r', cfg=None):
             self.reader = libtiledbvcf.Reader()
             self._set_read_cfg(cfg)
             self.reader.init(uri)
+            self.reader.set_tiledb_stats_enabled(stats)
         elif self.mode == 'w':
             self.writer = libtiledbvcf.Writer()
             self.writer.init(uri)
@@ -199,3 +200,12 @@ def ingest_samples(self, sample_uris=None, extra_attrs=None, checksum_type=None,
         self.writer.create_dataset()
         self.writer.register_samples()
         self.writer.ingest_samples()
+
+    def tiledb_stats(self):
+        if self.mode != 'r':
+            raise Exception('Stats can only be called for reader')
+
+        if not self.reader.get_tiledb_stats_enabled:
+            raise Exception('Stats not enabled')
+
+        return self.reader.get_tiledb_stats();

apis/python/src/tiledbvcf/dataset.py

@@ -115,6 +115,15 @@ public Optional<AWSSessionCredentialsProvider> getCredentialsProvider() {
     return Optional.empty();
   }
 
+  /** @return The log level for the VCFReader stats reporting */
+  public Optional<String> getTileDBStatsLogLevel() {
+    if (options.containsKey("tiledb_stats_log_level")) {
+      String statsLogLevel = options.get("tiledb_stats_log_level");
+      return Optional.of(statsLogLevel);
+    }
+    return Optional.empty();
+  }
+
   /** @return Optional CSV String of config parameters */
   public Optional<String> getConfigCSV() {
     return getConfigCSV(options);

apis/python/src/tiledbvcf/libhts.so.1.8

@@ -18,6 +18,7 @@
 import org.apache.arrow.vector.types.FloatingPointPrecision;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.FieldType;
+import org.apache.log4j.Level;
 import org.apache.log4j.Logger;
 import org.apache.spark.sql.execution.arrow.ArrowUtils;
 import org.apache.spark.sql.sources.v2.reader.InputPartitionReader;
@@ -78,6 +79,8 @@ public class VCFInputPartitionReader implements InputPartitionReader<ColumnarBat
   /** Stats counter: number of bytes in allocated buffers. */
   private long statsTotalBufferBytes;
 
+  private Level enableStatsLogLevel;
+
   /**
    * Creates a TileDB-VCF reader.
    *
@@ -115,6 +118,12 @@ public VCFInputPartitionReader(
     } else {
       this.samples = new String[] {};
     }
+
+    this.enableStatsLogLevel = Level.OFF;
+    if (this.options.getTileDBStatsLogLevel().isPresent()) {
+      // If an invalid log level is set, the default is DEBUG
+      this.enableStatsLogLevel = Level.toLevel(this.options.getTileDBStatsLogLevel().get());
+    }
   }
 
   @Override
@@ -183,6 +192,10 @@ public ColumnarBatch get() {
   public void close() {
     log.info("Closing VCFReader for partition " + (this.partitionId));
 
+    if (!this.enableStatsLogLevel.equals(Level.OFF)) {
+      log.log(this.enableStatsLogLevel, this.vcfReader.stats());
+    }
+
     if (vcfReader != null) {
       vcfReader.close();
       vcfReader = null;
@@ -251,6 +264,9 @@ private void initVCFReader() {
       vcfReader.setSortRegions(sortRegions.get().booleanValue());
     }
 
+    // Enable VCFReader stats
+    if (!this.enableStatsLogLevel.equals(Level.OFF)) this.vcfReader.setStatsEnabled(true);
+
     // Set logical partition in array
     vcfReader.setRangePartition(
         rangePartitionInfo.getNumPartitions(), rangePartitionInfo.getIndex());

apis/spark/src/main/java/io/tiledb/vcf/VCFDataSourceOptions.java

@@ -10,6 +10,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
+import org.apache.log4j.Level;
 import org.apache.spark.sql.Dataset;
 import org.apache.spark.sql.Row;
 import org.apache.spark.sql.SparkSession;
@@ -43,6 +44,7 @@ private Dataset<Row> testSampleDataset() {
             .option("samples", "HG01762,HG00280")
             .option("ranges", "1:12100-13360,1:13500-17350")
             .option("tiledb.vfs.num_threads", 1)
+            .option("tiledb_stats_log_level", Level.INFO.toString())
             .load();
     return dfRead;
   }

apis/spark/src/main/java/io/tiledb/vcf/VCFInputPartitionReader.java

@@ -354,6 +354,40 @@ int32_t tiledb_vcf_reader_set_tiledb_config(
   return TILEDB_VCF_OK;
 }
 
+int32_t tiledb_vcf_reader_set_tiledb_stats_enabled(
+    tiledb_vcf_reader_t* reader, const bool stats_enabled) {
+  if (sanity_check(reader) == TILEDB_VCF_ERR)
+    return TILEDB_VCF_ERR;
+
+  if (SAVE_ERROR_CATCH(
+          reader, reader->reader_->set_tiledb_stats_enabled(stats_enabled)))
+    return TILEDB_VCF_ERR;
+
+  return TILEDB_VCF_OK;
+}
+
+int32_t tiledb_vcf_reader_get_tiledb_stats_enabled(
+    tiledb_vcf_reader_t* reader, bool* enabled) {
+  if (sanity_check(reader) == TILEDB_VCF_ERR)
+    return TILEDB_VCF_ERR;
+
+  if (SAVE_ERROR_CATCH(reader, reader->reader_->tiledb_stats_enabled(enabled)))
+    return TILEDB_VCF_ERR;
+
+  return TILEDB_VCF_OK;
+}
+
+int32_t tiledb_vcf_reader_get_tiledb_stats(
+    tiledb_vcf_reader_t* reader, char** stats) {
+  if (sanity_check(reader) == TILEDB_VCF_ERR)
+    return TILEDB_VCF_ERR;
+
+  if (SAVE_ERROR_CATCH(reader, reader->reader_->tiledb_stats(stats)))
+    return TILEDB_VCF_ERR;
+
+  return TILEDB_VCF_OK;
+}
+
 int32_t tiledb_vcf_reader_read(tiledb_vcf_reader_t* reader) {
   if (sanity_check(reader) == TILEDB_VCF_ERR)
     return TILEDB_VCF_ERR;