Fast list prefix update (#28)

* updated prefix to apply to start range of first process

* fix slow listing bug when using prefixes

* add comment explanation for prefix removal

* updated to include safety for None input

* add verbose test output

* update fake_gcs

* reduce e2e object count

* fix removeprefix version error

* comment explaining prefix removal

Author: jdnurme

dataflux_core/benchmarking/dataflux_client_bench.py

@@ -30,7 +30,7 @@ def parse_args():
     parser.add_argument("--bucket-file-size", type=int, default=None)
     parser.add_argument("--num-workers", type=int, default=10)
     parser.add_argument("--max-compose-bytes", type=int, default=100000000)
-    parser.add_argument("--prefix", type=str, default=None)
+    parser.add_argument("--prefix", type=str, default="")
     return parser.parse_args()
 
 

dataflux_core/benchmarking/dataflux_client_parallel_bench.py

@@ -31,7 +31,7 @@ def parse_args():
     parser.add_argument("--num-workers", type=int, default=10)
     parser.add_argument("--max-compose-bytes", type=int, default=100000000)
     parser.add_argument("--parallelization", type=int, default=20)
-    parser.add_argument("--prefix", type=str, default=None)
+    parser.add_argument("--prefix", type=str, default="")
     return parser.parse_args()
 
 

dataflux_core/benchmarking/dataflux_client_threaded_bench.py

@@ -31,7 +31,7 @@ def parse_args():
     parser.add_argument("--num-workers", type=int, default=10)
     parser.add_argument("--max-compose-bytes", type=int, default=100000000)
     parser.add_argument("--threads", type=int, default=20)
-    parser.add_argument("--prefix", type=str, default=None)
+    parser.add_argument("--prefix", type=str, default="")
     return parser.parse_args()
 
 

dataflux_core/fast_list.py

@@ -29,6 +29,22 @@
 DEFAULT_ALLOWED_CLASS = ["STANDARD"]
 
 
+def remove_prefix(text: str, prefix: str):
+    """Helper function that removes prefix from a string.
+
+    Args:
+        text: String of text to trim a prefix from.
+        prefix: String of text that will be trimmed from text.
+
+    Returns:
+        Text value with the specified prefix removed.
+    """
+    # Note that as of python 3.9 removeprefix is built into string.
+    if text.startswith(prefix):
+        return text[len(prefix) :]
+    return text
+
+
 class ListWorker(object):
     """Worker that lists a range of objects from a GCS bucket.
 
@@ -75,7 +91,7 @@ def __init__(
         client: storage.Client = None,
         skip_compose: bool = True,
         list_directory_objects: bool = False,
-        prefix: str = None,
+        prefix: str = "",
         allowed_storage_classes: list[str] = DEFAULT_ALLOWED_CLASS,
         max_retries: int = 5,
     ):
@@ -98,7 +114,7 @@ def __init__(
         self.default_alph = "a"
         self.skip_compose = skip_compose
         self.list_directory_objects = list_directory_objects
-        self.prefix = prefix
+        self.prefix = prefix if prefix else ""
         self.allowed_storage_classes = allowed_storage_classes
         self.api_call_count = 0
         self.max_retries = max_retries
@@ -163,8 +179,10 @@ def run(self) -> None:
             try:
                 list_blob_args = {
                     "max_results": self.max_results,
-                    "start_offset": self.start_range,
-                    "end_offset": self.end_range,
+                    "start_offset": self.prefix + self.start_range,
+                    "end_offset": (
+                        "" if not self.end_range else self.prefix + self.end_range
+                    ),
                 }
                 if self.prefix:
                     list_blob_args["prefix"] = self.prefix
@@ -184,7 +202,10 @@ def run(self) -> None:
                         and blob.storage_class in self.allowed_storage_classes
                     ):
                         self.results.add((blob.name, blob.size))
-                    self.start_range = blob.name
+                    # Remove the prefix from the name so that range calculations remain prefix-agnostic.
+                    # This is necessary due to the unbounded end-range when splitting string namespaces
+                    # of unknown size.
+                    self.start_range = remove_prefix(blob.name, self.prefix)
                     if i == self.max_results:
                         # Only allow work stealing when paging.
                         has_results = True
@@ -237,7 +258,7 @@ def run_list_worker(
     end_range: str,
     client: storage.Client = None,
     skip_compose: bool = True,
-    prefix: str = None,
+    prefix: str = "",
     allowed_storage_classes: list[str] = DEFAULT_ALLOWED_CLASS,
 ) -> None:
     """Helper function to execute a ListWorker.
@@ -253,7 +274,7 @@ def run_list_worker(
       unidle_queue: Multiprocessing queue pushed to when the worker has successfully stolen work.
       results_queue: Multiprocessing queue on which the worker pushes its listing results onto.
       metadata_queue: Multiprocessing queue on which the worker pushes tracking metadata.
-      start_range: Stirng start range worker will begin listing from.
+      start_range: String start range worker will begin listing from.
       end_range: String end range worker will list until.
       client: The GCS storage client. When not provided, will be derived from background auth.
       skip_compose: When true, skip listing files with the composed object prefix.
@@ -303,7 +324,7 @@ def __init__(
         bucket: str,
         sort_results: bool = False,
         skip_compose: bool = True,
-        prefix: str = None,
+        prefix: str = "",
         allowed_storage_classes: list[str] = DEFAULT_ALLOWED_CLASS,
     ):
         # The maximum number of threads utilized in the fast list operation.

dataflux_core/tests/fake_gcs.py

@@ -38,7 +38,7 @@ def list_blobs(
         max_results: int = 0,
         start_offset: str = "",
         end_offset: str = "",
-        prefix=None,
+        prefix: str = "",
     ) -> list[Blob]:
         results = []
         for name in sorted(self.blobs):
@@ -47,7 +47,7 @@ def list_blobs(
             if (not start_offset or name >= start_offset) and (
                 not end_offset or name < end_offset
             ):
-                if not prefix or name.startswith(prefix):
+                if name.startswith(prefix):
                     results.append(self.blobs[name])
         return results
 

dataflux_core/tests/test_fast_list.py

@@ -31,7 +31,7 @@ def test_single_worker(self):
                 "compose_obj_count": 1,
                 "prefix_obj_count": 0,
                 "archive_obj_count": 0,
-                "prefix": None,
+                "prefix": "",
                 "object_size": 10,
                 "directory_obj_count": 10,
                 "skip_compose": True,
@@ -45,7 +45,7 @@ def test_single_worker(self):
                 "compose_obj_count": 1,
                 "prefix_obj_count": 0,
                 "archive_obj_count": 0,
-                "prefix": None,
+                "prefix": "",
                 "object_size": 10,
                 "directory_obj_count": 0,
                 "skip_compose": False,
@@ -59,7 +59,7 @@ def test_single_worker(self):
                 "compose_obj_count": 5000,
                 "prefix_obj_count": 0,
                 "archive_obj_count": 0,
-                "prefix": None,
+                "prefix": "",
                 "object_size": 10,
                 "directory_obj_count": 0,
                 "skip_compose": True,
@@ -87,7 +87,7 @@ def test_single_worker(self):
                 "compose_obj_count": 0,
                 "prefix_obj_count": 0,
                 "archive_obj_count": 0,
-                "prefix": None,
+                "prefix": "",
                 "object_size": 10,
                 "directory_obj_count": 10,
                 "skip_compose": True,
@@ -101,7 +101,7 @@ def test_single_worker(self):
                 "compose_obj_count": 0,
                 "prefix_obj_count": 0,
                 "archive_obj_count": 1000,
-                "prefix": None,
+                "prefix": "",
                 "object_size": 10,
                 "directory_obj_count": 0,
                 "skip_compose": True,
@@ -260,9 +260,9 @@ def test_list_controller_e2e(self):
         client = fake_gcs.Client()
         bucket_name = "test_bucket"
         bucket = client.bucket(bucket_name)
-        object_count = 100000
+        object_count = 1000
         object_size = 10
-        for i in range(100000):
+        for i in range(object_count):
             bucket._add_file(str(i), "aaaaaaaaaa")
         controller = fast_list.ListingController(1, "", bucket_name, True)
         controller.client = client

kokoro/build.sh

@@ -32,7 +32,7 @@ function install_requirements() {
 
 function run_unit_tests() {
     echo Running unit tests.
-    python -m pytest dataflux_core/tests --junit-xml="${KOKORO_ARTIFACTS_DIR}/unit_tests/sponge_log.xml"
+    python -m pytest dataflux_core/tests -vvv --junit-xml="${KOKORO_ARTIFACTS_DIR}/unit_tests/sponge_log.xml"
 }
 
 install_requirements