Azure Cosmos DB Spark connector: Throughput control

The Spark connector allows you to communicate with Azure Cosmos DB by using Apache Spark. This article describes how to limit throughput consumption with server-side throughput buckets or client-side throughput control. Prefer throughput buckets where they're supported and you can enable the preview. Azure Cosmos DB enforces bucket limits on the server, rather than relying on client-side RU tracking and coordination.

The client-side examples in this article use global throughput control groups to coordinate request unit (RU) consumption across Spark clients. The Java SDK also supports global groups, local groups within a single client instance, and server-side throughput buckets.

Why is throughput control important?

Throughput control helps to isolate the performance needs of applications that run against a container. Throughput control limits the amount of RUs that a specific Spark client can consume.

Several advanced scenarios benefit from throughput limits:

  • Different operations and tasks have different priorities: There can be a need to prevent normal transactions from being throttled because of data ingestion or copy activities. Some operations or tasks aren't sensitive to latency and are more tolerant to being throttled than others.
  • Provide fairness/isolation to different users or tenants: An application usually has many users. Some users might send too many requests, which consume all available throughput and cause others to get throttled.
  • Load balancing of throughput between different Azure Cosmos DB clients: In some use cases, it's important to make sure all the clients get a fair (equal) share of the throughput.

Throughput control enables the capability for more granular-level RU rate limiting, as needed.

Prefer server-side throughput buckets

Throughput buckets (preview) limit the maximum percentage of a container's throughput that a workload can consume. Requests from all clients assigned to the same bucket share its limit. When the bucket exceeds that limit, Azure Cosmos DB throttles subsequent requests with HTTP status code 429 and substatus code 3212.

Prefer buckets for Spark ingestion, migration, bulk writes, queries, and change feed workloads when your account and container support them. Server-side enforcement avoids the client-side approximation and coordination limitations described later in this article. You don't need a separate throughput control metadata container.

Buckets don't reserve throughput or guarantee that their maximum allocation is available. Throughput remains shared with other workloads. Up to five buckets, with IDs from 1 to 5, are available per container.

Enable and configure a bucket for Spark

  1. Use Spark connector version 4.43.0 or later, with an artifact that matches your Spark and Scala versions. See the Spark connector setup instructions.

  2. Register for Throughput Buckets under Preview features on your subscription page in the Azure portal. Wait for approval and enablement, which can take up to an hour after approval.

  3. In Data Explorer, open your container's Scale & Settings pane, select Throughput Buckets, activate a bucket, and set its maximum throughput percentage. For the following example, activate bucket 2. See Configure throughput buckets.

  4. Add the following options to your existing Spark read or write configuration, alongside your account connection, database, and container options:

    val throughputBucketOptions = Map(
        "spark.cosmos.throughputControl.enabled" -> "true",
        "spark.cosmos.throughputControl.name" -> "SparkBucket2",
        "spark.cosmos.throughputControl.throughputBucket" -> "2"
    )
    

    Pass these options to the Spark reader or writer by using .options(throughputBucketOptions). The connector uses the group name to associate operations with the bucket; it doesn't create or rename a server-side bucket. The bucket ID selects the limit configured on the container.

  5. When you switch from client-side throughput control, remove spark.cosmos.throughputControl.targetThroughput, spark.cosmos.throughputControl.targetThroughputThreshold, and the spark.cosmos.throughputControl.globalControl.* options. Also remove spark.cosmos.throughputControl.accountEndpoint and spark.cosmos.throughputControl.accountKey if you configured a separate metadata account. Don't combine these client-side settings with throughputBucket. Keep the main account connection options.

  6. Monitor Total Requests (preview) and Total Request Units (preview) in the Azure portal, split by ThroughputBucket, to verify that your workload uses the expected bucket.

Keep the selected bucket active. Requests that specify a deleted bucket ID consume the container's overall throughput instead of being limited by that bucket. Requests without a bucket assignment use the container's default bucket, if configured; otherwise, they have no bucket-specific limit.

When to use client-side throughput control instead

Use client-side throughput control when you can't enable the preview, need an absolute RU/s target instead of a percentage, or use an unsupported account or container. Throughput buckets aren't supported for serverless accounts or containers in a shared throughput database. Requests assigned to a bucket also can't use burst capacity. For more information, see the throughput bucket limitations and frequently asked questions.

How does client-side throughput control work?

Client-side throughput control in gateway mode requires a recent version of the Spark connector. For serverless Azure Cosmos DB accounts, attempting to use targetThroughputThreshold to define a percentage results in failure. Instead, set an absolute RU/s target by using spark.cosmos.throughputControl.targetThroughput.

To configure throughput control for the Spark connector, you first create a container that defines throughput control metadata. The partition key is groupId and ttl is enabled. Here, you create this container by using Spark SQL and call it ThroughputControl:

    %sql
    CREATE TABLE IF NOT EXISTS cosmosCatalog.`database-v4`.ThroughputControl 
    USING cosmos.oltp
    OPTIONS(spark.cosmos.database = 'database-v4')
    TBLPROPERTIES(partitionKeyPath = '/groupId', autoScaleMaxThroughput = '4000', indexingPolicy = 'AllProperties', defaultTtlInSeconds = '-1');

The preceding example creates a container with autoscale. If you prefer standard provisioning, you can replace autoScaleMaxThroughput with manualThroughput.

The partition key must be defined as /groupId and ttl must be enabled for the throughput control feature to work.

Within the Spark configuration of a specific application, you can then specify parameters for the workload. The following example sets throughput control as enabled. The example defines a throughput control group name parameter and a targetThroughputThreshold parameter. You also define the database and container parameters in which the throughput control group is maintained:

    "spark.cosmos.throughputControl.enabled" -> "true",
    "spark.cosmos.throughputControl.name" -> "SourceContainerThroughputControl",
    "spark.cosmos.throughputControl.targetThroughputThreshold" -> "0.95", 
    "spark.cosmos.throughputControl.globalControl.database" -> "database-v4", 
    "spark.cosmos.throughputControl.globalControl.container" -> "ThroughputControl"

Microsoft Entra ID authentication and throughput control: If you use Microsoft Entra ID (formerly Azure Active Directory) authentication with the Spark connector, extra role-based access control (RBAC) configuration might be required depending on which throughput control parameter you use:

  • targetThroughputThreshold: You must assign a custom role that includes the Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers/throughputSettings/read action to your Microsoft Entra identity. The connector needs to read the container's provisioned throughput to calculate the target RU based on the threshold percentage.
  • targetThroughput: If you use Spark connector version earlier than 4.47.0, you must also assign a role with the Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers/throughputSettings/read action. Starting from version 4.47.0, this permission is no longer required when you use targetThroughput.

For more information on configuring custom roles with specific actions, see Configure role-based access control for Azure Cosmos DB.

In the preceding example, the targetThroughputThreshold parameter is defined as 0.95. Rate limiting occurs (and requests are retried) when clients consume more than 95 percent (+/- 5-10 percent) of the throughput allocated to the container. This configuration is stored as a document in the throughput container, which looks like this example:

    {
        "id": "ZGF0YWJhc2UtdjQvY3VzdG9tZXIvU291cmNlQ29udGFpbmVyVGhyb3VnaHB1dENvbnRyb2w.info",
        "groupId": "database-v4/customer/SourceContainerThroughputControl.config",
        "targetThroughput": "",
        "targetThroughputThreshold": "0.95",
        "isDefault": true,
        "_rid": "EHcYAPolTiABAAAAAAAAAA==",
        "_self": "dbs/EHcYAA==/colls/EHcYAPolTiA=/docs/EHcYAPolTiABAAAAAAAAAA==/",
        "_etag": "\"2101ea83-0000-1100-0000-627503dd0000\"",
        "_attachments": "attachments/",
        "_ts": 1651835869
    }

Warning

Client-side throughput control doesn't calculate RU consumption before an operation. It tracks the charge from the response header after the operation completes. A single operation can therefore exceed the configured limit if its RU charge is higher than that limit. This control is approximate, not a strict cap or a guarantee of available throughput. Set the limit higher than the cost of any single operation in the group, or prefer server-side throughput buckets where supported.

When you read via query or change feed, you should configure the page size in spark.cosmos.read.maxItemCount (default 1000) to be a modest amount. In this way, the client throughput control can be recalculated with higher frequency and reflected more accurately at any specific time. When you use throughput control for a write job using bulk, the number of documents executed in a single request is automatically tuned based on the throttling rate to allow the throughput control to begin as early as possible.

Warning

The targetThroughputThreshold parameter is immutable. If you change the target throughput threshold value, a new throughput control group is created. (If you use version 4.10.0 or later, it can have the same name.) You need to restart all Spark jobs that are using the group if you want to ensure that they all consume the new threshold immediately. Otherwise, they pick up the new threshold after the next restart.

For each Spark client that uses the throughput control group, a record is created in the ThroughputControl container, with a ttl of a few seconds. As a result, the documents vanish quickly if a Spark client isn't actively running anymore. Here's an example:

    {
        "id": "Zhjdieidjojdook3osk3okso3ksp3ospojsp92939j3299p3oj93pjp93jsps939pkp9ks39kp9339skp",
        "groupId": "database-v4/customer/SourceContainerThroughputControl.config",
        "_etag": "\"1782728-w98999w-ww9998w9-99990000\"",
        "ttl": 10,
        "initializeTime": "2022-06-26T02:24:40.054Z",
        "loadFactor": 0.97636377638898,
        "allocatedThroughput": 484.89444487847,
        "_rid": "EHcYAPolTiABAAAAAAAAAA==",
        "_self": "dbs/EHcYAA==/colls/EHcYAPolTiA=/docs/EHcYAPolTiABAAAAAAAAAA==/",
        "_etag": "\"2101ea83-0000-1100-0000-627503dd0000\"",
        "_attachments": "attachments/",
        "_ts": 1651835869
    }

In each client record, the loadFactor attribute represents the load on the specific client, relative to other clients in the throughput control group. The allocatedThroughput attribute shows how many RUs are currently allocated to this client. The Spark connector adjusts allocated throughput for each client based on its load. In this way, each client gets a share of the throughput available that's proportional to its load. The clients aim to keep their combined consumption within the group's allocation, subject to the client-side approximation described earlier.