In-memory export of VCF info flags (#613)

Author: gspowley

apis/python/tests/test_tiledbvcf.py

@@ -1503,3 +1503,24 @@ def test_gvcf_export(tmp_path):
     for test in tests:
         df = ds.read(attrs=attrs, regions=test["region"], set_af_filter="<=1.0")
         assert set(df["sample_name"].unique()) == set(test["samples"])
+
+
+def test_flag_export(tmp_path):
+    # Create the dataset
+    uri = os.path.join(tmp_path, "dataset")
+    ds = tiledbvcf.Dataset(uri, mode="w")
+    samples = [os.path.join(TESTS_INPUT_DIR, s) for s in ["small.vcf.gz"]]
+    ds.create_dataset()
+    ds.ingest_samples(samples)
+
+    # Read info flags
+    ds = tiledbvcf.Dataset(uri, mode="r")
+    df = ds.read(attrs=["pos_start", "info_DB", "info_DS"])
+    df = df.sort_values(by=["pos_start"])
+
+    # Check if flags match the expected values
+    expected_db = [1, 1, 1, 0, 0, 1]
+    assert df["info_DB"].tolist() == expected_db
+
+    expected_ds = [1, 1, 0, 0, 1, 1]
+    assert df["info_DS"].tolist() == expected_ds

apis/spark/src/test/java/io/tiledb/vcf/VCFDatasourceTest.java

@@ -873,4 +873,36 @@ public void testTimeTravel() throws java.text.ParseException {
     Assert.assertEquals(rows.get(1).getString(0), "HG01762");
     Assert.assertEquals(rows.get(2).getString(0), "HG01762");
   }
+
+  @Test
+  public void testInfoFlagExport() {
+    Dataset<Row> dfRead =
+        session()
+            .read()
+            .format("io.tiledb.vcf")
+            .option("uri", testSampleGroupURI("small.tdb", "v4"))
+            .option("ranges", "1:10000-20000")
+            .option("tiledb.vfs.num_threads", 1)
+            .option("tiledb.vcf.log_level", "DEBUG")
+            .load();
+    Dataset<Row> df = dfRead.select("contig", "posStart", "info_DB", "info_DS");
+    df.show();
+
+    // Define expected values
+    List<Integer> expectedDB = Arrays.asList(1, 1, 1, 0, 0, 1);
+    List<Integer> expectedDS = Arrays.asList(1, 1, 0, 0, 1, 1);
+
+    // Get actual values
+    List<Row> rows = df.collectAsList();
+    List<Integer> actualDB = new ArrayList<>();
+    List<Integer> actualDS = new ArrayList<>();
+    for (Row row : rows) {
+      actualDB.add(row.getInt(2));
+      actualDS.add(row.getInt(3));
+    }
+
+    // Compare
+    Assert.assertEquals(expectedDB, actualDB);
+    Assert.assertEquals(expectedDS, actualDS);
+  }
 }

apis/spark3/src/test/java/io/tiledb/vcf/VCFDatasourceTest.java

@@ -873,4 +873,36 @@ public void testTimeTravel() throws java.text.ParseException {
     Assert.assertEquals(rows.get(1).getString(0), "HG01762");
     Assert.assertEquals(rows.get(2).getString(0), "HG01762");
   }
+
+  @Test
+  public void testInfoFlagExport() {
+    Dataset<Row> dfRead =
+        session()
+            .read()
+            .format("io.tiledb.vcf")
+            .option("uri", testSampleGroupURI("small.tdb", "v4"))
+            .option("ranges", "1:10000-20000")
+            .option("tiledb.vfs.num_threads", 1)
+            .option("tiledb.vcf.log_level", "DEBUG")
+            .load();
+    Dataset<Row> df = dfRead.select("contig", "posStart", "info_DB", "info_DS");
+    df.show();
+
+    // Define expected values
+    List<Integer> expectedDB = Arrays.asList(1, 1, 1, 0, 0, 1);
+    List<Integer> expectedDS = Arrays.asList(1, 1, 0, 0, 1, 1);
+
+    // Get actual values
+    List<Row> rows = df.collectAsList();
+    List<Integer> actualDB = new ArrayList<>();
+    List<Integer> actualDS = new ArrayList<>();
+    for (Row row : rows) {
+      actualDB.add(row.getInt(2));
+      actualDS.add(row.getInt(3));
+    }
+
+    // Compare
+    Assert.assertEquals(expectedDB, actualDB);
+    Assert.assertEquals(expectedDS, actualDS);
+  }
 }

libtiledbvcf/src/c_api/tiledbvcf_enum.h

@@ -42,6 +42,8 @@ TILEDB_VCF_READ_STATUS_ENUM(FAILED) = 0,
     TILEDB_VCF_ATTR_DATATYPE_ENUM(INT32) = 2,
     /** 32-bit floating-point  */
     TILEDB_VCF_ATTR_DATATYPE_ENUM(FLOAT32) = 3,
+    /** Info Flag */
+    TILEDB_VCF_ATTR_DATATYPE_ENUM(FLAG) = 4,
 #endif
 
 #ifdef TILEDB_VCF_CHECKSUM_TYPE_ENUM

libtiledbvcf/src/read/in_memory_exporter.cc

@@ -30,6 +30,27 @@
 namespace tiledb {
 namespace vcf {
 
+// Declare static variables
+std::set<std::string> InMemoryExporter::info_flags_;
+
+void InMemoryExporter::add_flag(const std::string& field) {
+  // Add info_ to field if it is missing
+  std::string info_field = field;
+  if (field.substr(0, 5) != "info_") {
+    info_field = "info_" + field;
+  }
+  info_flags_.insert(info_field);
+}
+
+bool InMemoryExporter::is_flag(const std::string& field) {
+  // Add info_ to field if it is missing
+  std::string info_field = field;
+  if (field.substr(0, 5) != "info_") {
+    info_field = "info_" + field;
+  }
+  return info_flags_.count(info_field) > 0;
+}
+
 void InMemoryExporter::set_buffer_values(
     const std::string& attribute, void* buff, int64_t buff_size) {
   if (buff == nullptr) {
@@ -644,6 +665,18 @@ void InMemoryExporter::attribute_datatype(
       break;
   }
 
+  if (*datatype == AttrDatatype::FLAG) {
+    // Add this attribute to the set of info flags.
+    add_flag(attribute);
+
+    // Set buffer attributes for flags and use INT32 for the datatype.
+    *var_len = false;
+    *nullable = false;
+    *list = false;
+    *datatype = AttrDatatype::INT32;
+    return;
+  }
+
   *var_len = !fixed_len_attr(attribute);
   *nullable = nullable_attr(attribute);
   *list = var_len_list_attr(attribute);
@@ -672,7 +705,7 @@ AttrDatatype InMemoryExporter::get_info_fmt_datatype(
                         dataset->fmt_field_type(field_name, hdr);
   switch (htslib_type) {
     case BCF_HT_FLAG:
-      return AttrDatatype::INT32;
+      return AttrDatatype::FLAG;
     case BCF_HT_STR:
       return AttrDatatype::CHAR;
     case BCF_HT_INT:
@@ -964,6 +997,10 @@ bool InMemoryExporter::copy_info_fmt_value(
     for (unsigned i = 0; i < nelts; i++)
       decoded[i] = bcf_gt_allele(genotype[i]);
     return copy_cell(dest, decoded, nelts * sizeof(int), nelts, hdr);
+  } else if (is_flag(field_name)) {
+    // If this is a flag, convert the src pointer to a true or false value.
+    int flag = src == nullptr ? 0 : 1;
+    return copy_cell(dest, &flag, sizeof(int), 1, hdr);
   } else {
     return copy_cell(dest, src, nbytes, nelts, hdr);
   }

libtiledbvcf/src/read/in_memory_exporter.h

@@ -255,10 +255,23 @@ class InMemoryExporter : public Exporter {
   /** Reusable string buffer for temp results. */
   std::string str_buff_;
 
+  /**
+   * Set of info field names with type flag.
+   *
+   * NOTE: Only info fields are allowed to have type flag.
+   */
+  static std::set<std::string> info_flags_;
+
   /* ********************************* */
   /*           PRIVATE METHODS         */
   /* ********************************* */
 
+  /** Adds a field to the set of info flags. */
+  static void add_flag(const std::string& field);
+
+  /** Returns true if the given field is an info flag. */
+  static bool is_flag(const std::string& field);
+
   /**
    * Returns the ExportableAttribute corresponding to the given attribute name.
    */

libtiledbvcf/src/write/writer_worker_v4.cc

@@ -464,7 +464,13 @@ void WriterWorkerV4::buffer_info_field(
   buff->append(&type, sizeof(int));
   buff->append(&num_vals, sizeof(int));
   if (val->dst) {
-    buff->append(val->dst, num_vals * utils::bcf_type_size(type));
+    if (type == BCF_HT_FLAG) {
+      // Write a dummy value for flags.
+      int flag = 1;
+      buff->append(&flag, num_vals * utils::bcf_type_size(type));
+    } else {
+      buff->append(val->dst, num_vals * utils::bcf_type_size(type));
+    }
   } else {
     // val->dst can be NULL if the only INFO value is a flag
     assert(num_vals == 1);