simplify and refactor perf_test example script

airbyte/cli.py

@@ -76,6 +76,7 @@
  """For example, --config='{password: "SECRET:MY_PASSWORD"}'."""
 )
 
+
 def _resolve_config(
  config: str,
 ) -> dict[str, Any]:

examples/run_perf_test_reads.py

@@ -2,80 +2,77 @@
 """
 Simple script to get performance profile of read throughput.
 
-This script accepts a single argument `-e=SCALE` as a power of 10.
+This script accepts a single argument `-n=NUM_RECORDS` with record count
+provided as a regular number or in scientific notation.
 
--e=2 is equivalent to 500 records.
--e=3 is equivalent to 5_000 records.
--e=4 is equivalent to 50_000 records.
--e=5 is equivalent to 500_000 records.
--e=6 is equivalent to 5_000_000 records.
+When providing in scientific notation:
 
-Use smaller values of `e` (2-3) to understand read and overhead costs.
-Use larger values of `e` (4-5) to understand write throughput at scale.
+-n=5e2 is equivalent to 500 records.
+-n=5e3 is equivalent to 5_000 records.
+-n=5e4 is equivalent to 50_000 records.
+-n=5e5 is equivalent to 500_000 records.
+-n=5e6 is equivalent to 5_000_000 records.
 
 For performance profiling, use `viztracer` to generate a flamegraph:
 ```
-poetry run viztracer --open -- ./examples/run_perf_test_reads.py -e=3
-poetry run viztracer --open -- ./examples/run_perf_test_reads.py -e=5
+poetry run viztracer --open -- ./examples/run_perf_test_reads.py -n=1e3
+poetry run viztracer --open -- ./examples/run_perf_test_reads.py -n=1e5
 ```
 
 To run without profiling, prefix script name with `poetry run python`:
+
 ```
 # Run with 5_000 records
-poetry run python ./examples/run_perf_test_reads.py -e=3
+poetry run python ./examples/run_perf_test_reads.py -n=1e3
 # Run with 500_000 records
-poetry run python ./examples/run_perf_test_reads.py -e=5
+poetry run python ./examples/run_perf_test_reads.py -n=1e5
 
 # Load 5_000 records to Snowflake
-poetry run python ./examples/run_perf_test_reads.py -e=3 --cache=snowflake
+poetry run python ./examples/run_perf_test_reads.py -n=1e3 --cache=snowflake
 
 # Load 5_000 records to BigQuery
-poetry run python ./examples/run_perf_test_reads.py -e=3 --cache=bigquery
+poetry run python ./examples/run_perf_test_reads.py -n=1e3 --cache=bigquery
 ```
 
 You can also use this script to test destination load performance:
 
 ```bash
 # Load 5_000 records to BigQuery
-poetry run python ./examples/run_perf_test_reads.py -e=5 --destination=e2e
+poetry run python ./examples/run_perf_test_reads.py -n=1e3 --destination=e2e
 ```
 
 Testing raw PyAirbyte throughput with and without caching:
 
 ```bash
 # Test raw PyAirbyte throughput with caching (Source->Cache):
-poetry run python ./examples/run_perf_test_reads.py -e=5
+poetry run python ./examples/run_perf_test_reads.py -n=1e3
 # Test raw PyAirbyte throughput without caching (Source->Destination):
-poetry run python ./examples/run_perf_test_reads.py -e=5 --destination=e2e --no-cache
+poetry run python ./examples/run_perf_test_reads.py -n=1e3 --destination=e2e --no-cache
 ```
 
 Testing Python CDK throughput:
 
 ```bash
-# Test max throughput:
-poetry run python ./examples/run_perf_test_reads.py -n=2400000 --source=hardcoded --destination=e2e
+# Test max throughput with 2.4 million records:
+poetry run python ./examples/run_perf_test_reads.py -n=2.4e6 --source=hardcoded --destination=e2e
+
 # Analyze tracing data:
-poetry run viztracer --open -- ./examples/run_perf_test_reads.py -e=3 --source=hardcoded --destination=e2e
+poetry run viztracer --open -- ./examples/run_perf_test_reads.py -n=1e3 --source=hardcoded --destination=e2e
 ```
-
-
-Note:
-- The Faker stream ('purchases') is assumed to be 220 bytes, meaning 4_500 records is
- approximately 1 MB. Based on this: 25K records/second is approximately 5.5 MB/s.
-- The E2E stream is assumed to be 180 bytes, meaning 5_500 records is
- approximately 1 MB. Based on this: 40K records/second is approximately 7.2 MB/s
- and 60K records/second is approximately 10.5 MB/s.
 """
 
 from __future__ import annotations
 
 import argparse
 import tempfile
+from decimal import Decimal
 from typing import TYPE_CHECKING
 
 import airbyte as ab
 from airbyte.caches import BigQueryCache, CacheBase, SnowflakeCache
+from airbyte.destinations import Destination, get_noop_destination
 from airbyte.secrets.google_gsm import GoogleGSMSecretManager
+from airbyte.sources import get_benchmark_source
 from typing_extensions import Literal
 
 if TYPE_CHECKING:
@@ -142,8 +139,11 @@ def get_cache(
 
 def get_source(
  source_alias: str,
- num_records: int,
+ num_records: int | str,
 ) -> Source:
+ if isinstance(num_records, str):
+ num_records = int(Decimal(num_records))
+
  if source_alias == "faker":
  return ab.get_source(
  "source-faker",
@@ -152,20 +152,8 @@ def get_source(
  streams=["purchases"],
  )
 
- if source_alias == "e2e":
- return ab.get_source(
- "source-e2e",
- docker_image="airbyte/source-e2e-test:latest",
- streams="*",
- config={
- "type": "BENCHMARK",
- "schema": "FIVE_STRING_COLUMNS",
- "terminationCondition": {
- "type": "MAX_RECORDS",
- "max": num_records,
- },
- },
- )
+ if source_alias in ["e2e", "benchmark"]:
+ return get_benchmark_source(num_records=num_records)
 
  if source_alias == "hardcoded":
  return ab.get_source(
@@ -180,32 +168,19 @@ def get_source(
 
 
 def get_destination(destination_type: str) -> ab.Destination:
- if destination_type == "e2e":
- return ab.get_destination(
- name="destination-e2e-test",
- config={
- "test_destination": {
- "test_destination_type": "LOGGING",
- "logging_config": {
- "logging_type": "FirstN",
- "max_entry_count": 100,
- },
- }
- },
- docker_image="airbyte/destination-e2e-test:latest",
- )
+ if destination_type in ["e2e", "noop"]:
+ return get_noop_destination()
 
  raise ValueError(f"Unknown destination type: {destination_type}") # noqa: TRY003
 
 
 def main(
- e: int | None = None,
- n: int | None = None,
+ n: int | str = "5e5",
  cache_type: Literal["duckdb", "bigquery", "snowflake", False] = "duckdb",
  source_alias: str = "e2e",
  destination_type: str | None = None,
 ) -> None:
- num_records: int = n or 5 * (10 ** (e or 3))
+ num_records = int(Decimal(n))
  cache_type = "duckdb" if cache_type is None else cache_type
 
  cache: CacheBase | Literal[False] = get_cache(
@@ -216,34 +191,31 @@ def main(
  num_records=num_records,
  )
  source.check()
+ destination: Destination | None = None
+
  if destination_type:
  destination = get_destination(destination_type=destination_type)
+
  if cache is not False:
  read_result = source.read(cache)
- if destination_type:
+ if destination:
  destination.write(read_result)
  else:
+ assert (
+ destination is not None
+ ), "Destination is required when caching is disabled."
  destination.write(source, cache=False)
 
 
 if __name__ == "__main__":
  parser = argparse.ArgumentParser(description="Run performance test reads.")
- parser.add_argument(
- "-e",
- type=int,
- help=(
- "The scale, as a power of 10."
- "Recommended values: 2-3 (500 or 5_000) for read and overhead costs, "
- " 4-6 (50K or 5MM) for write throughput. "
- "This is mutually exclusive with -n."
- ),
- )
  parser.add_argument(
  "-n",
- type=int,
+ type=str,
  help=(
  "The number of records to generate in the source. "
- "This is mutually exclusive with -e."
+ "This can be provided in scientific notation, for instance "
+ "'2.4e6' for 2.4 million and '5e5' for 500K."
  ),
  )
  parser.add_argument(
@@ -266,7 +238,12 @@ def main(
  "while the `faker` source runs natively in Python. The 'hardcoded' source is "
  "similar to the 'e2e' source, but written in Python."
  ),
- choices=["faker", "e2e", "hardcoded"],
+ choices=[
+ "benchmark",
+ "e2e",
+ "hardcoded",
+ "faker",
+ ],
  default="hardcoded",
  )
  parser.add_argument(
@@ -279,7 +256,6 @@ def main(
  args = parser.parse_args()
 
  main(
- e=args.e,
  n=args.n,
  cache_type=args.cache if not args.no_cache else False,
  source_alias=args.source,