Add a batch_mode flag to VCF scalable queries

src/tiledb/cloud/vcf/query.py

@@ -267,6 +267,7 @@ def build_read_dag(
     namespace: Optional[str] = None,
     resource_class: Optional[str] = None,
     verbose: bool = False,
+    batch_mode: bool = False,
 ) -> Tuple[tiledb.cloud.dag.DAG, tiledb.cloud.dag.Node]:
     """
     Build the DAG for a distributed read on a TileDB-VCF dataset.
@@ -289,6 +290,7 @@ def build_read_dag(
     :param namespace: TileDB-Cloud namespace, defaults to None
     :param resource_class: TileDB-Cloud resource class for UDFs, defaults to None
     :param verbose: verbose logging, defaults to False
+    :param batch_mode: run the query with batch UDFs, defaults to False
     :return: DAG and result Node
     """
 
@@ -338,10 +340,13 @@ def build_read_dag(
     logger.debug("num_sample_partitions=%d", num_sample_partitions)
     logger.debug("num_region_partitions=%d", num_region_partitions)
 
+    mode = tiledb.cloud.dag.Mode.BATCH if batch_mode else tiledb.cloud.dag.Mode.REALTIME
+
     dag = tiledb.cloud.dag.DAG(
         namespace=namespace,
         name="VCF-Distributed-Query",
         max_workers=max_workers,
+        mode=mode,
     )
 
     # If `regions` is a Delayed object, we set the parent nodes to `dag` so the
@@ -351,6 +356,12 @@ def build_read_dag(
     ):
         regions._DelayedBase__set_all_parent_nodes_same_dag(dag)
 
+    result_format = (
+        tiledb.cloud.UDFResultType.NATIVE
+        if batch_mode
+        else tiledb.cloud.UDFResultType.ARROW
+    )
+
     tables = []
     for region in range(num_region_partitions):
         for sample in range(num_sample_partitions):
@@ -373,12 +384,14 @@ def build_read_dag(
                     name=f"VCF Query - Region {region+1}/{num_region_partitions},"
                     f" Sample {sample+1}/{num_sample_partitions}",
                     resource_class=resource_class,
-                    result_format=tiledb.cloud.UDFResultType.ARROW,
+                    result_format=result_format,
                 )
             )
 
     if len(tables) > 1:
-        table = dag.submit_local(
+        submit = dag.submit if batch_mode else dag.submit_local
+
+        table = submit(
             concat_tables_udf,
             tables,
             config=config,
@@ -429,6 +442,7 @@ def read(
     namespace: Optional[str] = None,
     resource_class: Optional[str] = None,
     verbose: bool = False,
+    batch_mode: bool = False,
 ) -> pa.Table:
     """
     Run a distributed read on a TileDB-VCF dataset.
@@ -451,6 +465,7 @@ def read(
     :param namespace: TileDB-Cloud namespace, defaults to None
     :param resource_class: TileDB-Cloud resource class for UDFs, defaults to None
     :param verbose: verbose logging, defaults to False
+    :param batch_mode: run the query with batch UDFs, defaults to False
     :return: Arrow table containing the query results
     """
 
@@ -471,6 +486,7 @@ def read(
         namespace=namespace,
         resource_class=resource_class,
         verbose=verbose,
+        batch_mode=batch_mode,
     )
 
     run_dag(dag, debug=verbose)

tests/vcf/test_vcf.py

@@ -1,17 +1,26 @@
+import cloudpickle
 import numpy as np
 import pytest
 
 import tiledb.cloud.vcf as vcf
 import tiledb.cloud.vcf.vcf_toolbox as vtb
 
+# Pickle the vcf module by value, so tests run on the latest code.
+cloudpickle.register_pickle_by_value(vcf)
+
+
+# Run VCF tests with:
+#   pytest -m vcf --run-vcf -n 8
+
 
 @vtb.df_transform
 def filter_vcf(df, *, filter=None):
     return df.query(filter)
 
 
 @pytest.mark.vcf
-def test_vcf_transform():
+@pytest.mark.parametrize("batch_mode", [False, True])
+def test_vcf_transform(batch_mode):
     vcf_uri = "tiledb://TileDB-Inc/vcf-1kg-dragen-v376"
 
     regions = [
@@ -39,6 +48,7 @@ def test_vcf_transform():
         regions=regions,
         samples="NA12878",
         transform_result=filter_vcf(filter=filter),
+        batch_mode=batch_mode,
     )
 
     assert vcf_table.num_rows == 336