split up benchmarks into separate files

Author: timgraham

performance_tests/perftest/base.py

@@ -0,0 +1,113 @@
+import os
+import time
+import warnings
+from pathlib import Path
+
+from .tests import result_data
+
+if os.environ.get("FASTBENCH"):
+    NUM_ITERATIONS = 1
+    MIN_ITERATION_TIME = 5
+    MAX_ITERATION_TIME = 10
+    NUM_DOCS = 1000
+else:
+    NUM_ITERATIONS = 2
+    MIN_ITERATION_TIME = 30
+    MAX_ITERATION_TIME = 300
+    NUM_DOCS = 10000
+
+
+class Timer:
+    def __enter__(self):
+        self.start = time.monotonic()
+        return self
+
+    def __exit__(self, *args):
+        self.end = time.monotonic()
+        self.interval = self.end - self.start
+
+
+# Copied from the driver benchmarking suite.
+class PerformanceTest:
+    dataset: str
+    data_size: int
+    test_data_path = os.environ.get(
+        "DJANGO_MONGODB_PERFORMANCE_TEST_DATA_PATH",
+        Path(os.path.realpath(__file__)).parent.parent / "odm-data",
+    )
+    num_docs = NUM_DOCS
+
+    def setUp(self):
+        self.setup_time = time.monotonic()
+
+    def tearDown(self):
+        duration = time.monotonic() - self.setup_time
+        # Remove "Test" so that TestMyTestName is reported as "MyTestName".
+        name = self.__class__.__name__[4:]
+        median = self.percentile(50)
+        megabytes_per_sec = self.data_size / median / 1000000
+        print(  # noqa: T201
+            f"Completed {self.__class__.__name__} {megabytes_per_sec:.3f} MB/s, "
+            f"MEDIAN={self.percentile(50):.3f}s, "
+            f"total time={duration:.3f}s, iterations={len(self.results)}"
+        )
+        result_data.append(
+            {
+                "info": {
+                    "test_name": name,
+                },
+                "metrics": [
+                    {
+                        "name": "megabytes_per_sec",
+                        "type": "MEDIAN",
+                        "value": megabytes_per_sec,
+                        "metadata": {
+                            "improvement_direction": "up",
+                            "measurement_unit": "megabytes_per_second",
+                        },
+                    },
+                ],
+            }
+        )
+
+    def before(self):
+        pass
+
+    def do_task(self):
+        raise NotImplementedError
+
+    def after(self):
+        pass
+
+    def percentile(self, percentile):
+        if hasattr(self, "results"):
+            sorted_results = sorted(self.results)
+            percentile_index = int(len(sorted_results) * percentile / 100) - 1
+            return sorted_results[percentile_index]
+        self.fail("Test execution failed")
+        return None
+
+    def runTest(self):
+        results = []
+        start = time.monotonic()
+        i = 0
+        while True:
+            i += 1
+            self.before()
+            with Timer() as timer:
+                self.do_task()
+            self.after()
+            results.append(timer.interval)
+            duration = time.monotonic() - start
+            if duration > MIN_ITERATION_TIME and i >= NUM_ITERATIONS:
+                break
+            if duration > MAX_ITERATION_TIME:
+                with warnings.catch_warnings():
+                    warnings.simplefilter("default")
+                    warnings.warn(
+                        f"{self.__class__.__name__} timed out after {MAX_ITERATION_TIME}s, "
+                        f"completed {i}/{NUM_ITERATIONS} iterations.",
+                        stacklevel=2,
+                    )
+                break
+        self.results = results

performance_tests/perftest/large_flat_doc.py

@@ -0,0 +1,75 @@
+from pathlib import Path
+from unittest import TestCase
+
+from bson import encode, json_util
+
+from .base import PerformanceTest
+from .models import LargeFlatModel
+
+
+class LargeFlatDocTest(PerformanceTest):
+    """Parent class for large flat document tests."""
+
+    dataset = "large_doc.json"
+
+    def setUp(self):
+        super().setUp()
+        with open(  # noqa: PTH123
+            Path(self.test_data_path) / Path("flat-models") / self.dataset
+        ) as data:
+            self.document = json_util.loads(data.read())
+
+        self.data_size = len(encode(self.document)) * self.num_docs
+        self.documents = [self.document.copy() for _ in range(self.num_docs)]
+
+
+class TestLargeFlatDocCreation(LargeFlatDocTest, TestCase):
+    """Benchmark for creating a large flat document."""
+
+    def do_task(self):
+        for doc in self.documents:
+            LargeFlatModel.objects.create(**doc)
+
+    def after(self):
+        LargeFlatModel.objects.all().delete()
+
+
+class TestLargeFlatDocUpdate(LargeFlatDocTest, TestCase):
+    """Benchmark for updating a field within a large flat document."""
+
+    def setUp(self):
+        super().setUp()
+        for doc in self.documents:
+            LargeFlatModel.objects.create(**doc)
+        self.models = list(LargeFlatModel.objects.all())
+        self.data_size = len(encode({"field1": "updated_value0"})) * self.num_docs
+        self.iteration = 0
+
+    def do_task(self):
+        for model in self.models:
+            model.field1 = "updated_value" + str(self.iteration)
+            model.save()
+        self.iteration += 1
+
+    def tearDown(self):
+        super().tearDown()
+        LargeFlatModel.objects.all().delete()
+
+
+class TestLargeFlatDocFilterPkByIn(LargeFlatDocTest, TestCase):
+    """Benchmark for filtering large flat documents using the __in operator for primary keys."""
+
+    def setUp(self):
+        super().setUp()
+        models = []
+        for doc in self.documents:
+            models.append(LargeFlatModel(**doc))
+        LargeFlatModel.objects.bulk_create(models)
+        self.ids = [model.id for model in models]
+
+    def do_task(self):
+        list(LargeFlatModel.objects.filter(id__in=self.ids))
+
+    def tearDown(self):
+        super().tearDown()
+        LargeFlatModel.objects.all().delete()

performance_tests/perftest/large_nested_doc.py

@@ -0,0 +1,135 @@
+from pathlib import Path
+from unittest import TestCase
+
+from bson import ObjectId, encode, json_util
+
+from .base import PerformanceTest
+from .models import (
+    IntegerEmbeddedModel,
+    LargeNestedModel,
+    StringEmbeddedModel,
+)
+
+
+class LargeNestedDocTest(PerformanceTest):
+    """Parent class for large nested document tests."""
+
+    dataset = "large_doc_nested.json"
+
+    def setUp(self):
+        super().setUp()
+        with open(  # noqa: PTH123
+            Path(self.test_data_path) / Path("nested-models") / self.dataset
+        ) as data:
+            self.document = json_util.loads(data.read())
+
+        self.data_size = len(encode(self.document)) * self.num_docs
+        self.documents = [self.document.copy() for _ in range(self.num_docs)]
+
+    def setUpData(self):
+        for doc in self.documents:
+            model = LargeNestedModel()
+            for field_name, model_data in doc.items():
+                if "array" in field_name:
+                    array_models = []
+                    for item in model_data:
+                        embedded_str_model = StringEmbeddedModel(**item)
+                        embedded_str_model.unique_field = str(ObjectId())
+                        array_models.append(embedded_str_model)
+                    setattr(model, field_name, array_models)
+                elif "embedded_str_doc" in field_name:
+                    embedded_str_model = StringEmbeddedModel(**model_data)
+                    embedded_str_model.unique_field = str(ObjectId())
+                    setattr(model, field_name, embedded_str_model)
+                else:
+                    embedded_int_model = IntegerEmbeddedModel(**model_data)
+                    setattr(model, field_name, embedded_int_model)
+            model.save()
+
+
+class TestLargeNestedDocCreation(LargeNestedDocTest, TestCase):
+    """Benchmark for creating a large nested document."""
+
+    def do_task(self):
+        for doc in self.documents:
+            model = LargeNestedModel()
+            for field_name, model_data in doc.items():
+                if "array" in field_name:
+                    array_models = []
+                    for item in model_data:
+                        embedded_str_model = StringEmbeddedModel(**item)
+                        embedded_str_model.unique_field = str(ObjectId())
+                        array_models.append(embedded_str_model)
+                    setattr(model, field_name, array_models)
+                elif "embedded_str_doc" in field_name:
+                    embedded_str_model = StringEmbeddedModel(**model_data)
+                    embedded_str_model.unique_field = str(ObjectId())
+                    setattr(model, field_name, embedded_str_model)
+                else:
+                    embedded_int_model = IntegerEmbeddedModel(**model_data)
+                    setattr(model, field_name, embedded_int_model)
+            model.save()
+
+    def after(self):
+        LargeNestedModel.objects.all().delete()
+
+
+class TestLargeNestedDocUpdate(LargeNestedDocTest, TestCase):
+    """Benchmark for updating an embedded field within a large nested document."""
+
+    def setUp(self):
+        super().setUp()
+        self.setUpData()
+        self.models = list(LargeNestedModel.objects.all())
+        self.data_size = len(encode({"field1": "updated_value0"})) * self.num_docs
+        self.iteration = 0
+
+    def do_task(self):
+        for model in self.models:
+            model.embedded_str_doc_1.field1 = "updated_value" + str(self.iteration)
+            model.save()
+        self.iteration += 1
+
+    def tearDown(self):
+        super().tearDown()
+        LargeNestedModel.objects.all().delete()
+
+
+class TestLargeNestedDocFilterById(LargeNestedDocTest, TestCase):
+    """Benchmark for filtering large nested documents by a unique field in an embedded document."""
+
+    def setUp(self):
+        super().setUp()
+        self.setUpData()
+        self.ids = [
+            model.embedded_str_doc_1.unique_field for model in list(LargeNestedModel.objects.all())
+        ]
+
+    def do_task(self):
+        for _id in self.ids:
+            list(LargeNestedModel.objects.filter(embedded_str_doc_1__unique_field=_id))
+
+    def tearDown(self):
+        super().tearDown()
+        LargeNestedModel.objects.all().delete()
+
+
+class TestLargeNestedDocFilterArray(LargeNestedDocTest, TestCase):
+    """Benchmark for filtering large nested documents using the __in operator
+    for unique values in an embedded document array."""
+
+    def setUp(self):
+        super().setUp()
+        self.setUpData()
+        self.ids = [
+            model.embedded_str_doc_array[0].unique_field
+            for model in list(LargeNestedModel.objects.all())
+        ]
+
+    def do_task(self):
+        for _id in self.ids:
+            list(LargeNestedModel.objects.filter(embedded_str_doc_array__unique_field__in=[_id]))
+
+    def tearDown(self):
+        super().tearDown()
+        LargeNestedModel.objects.all().delete()