beam-sdks-java-io-astra

Apache Beam SDK to work with Astra Pipelines

How To

Installation

To use this SDK, add the following dependency to your project:

<dependency>
  <groupId>com.datastax.astra</groupId>
  <artifactId>beam-sdks-java-io-astra</artifactId>
  <version>${latest-version}</version>
</dependency>

Usage

Documentation is available in Awesome Astra with sample codes

Read Data From Astra

// Get binary from File path
byte[] scbZip = AstraSecureConnectBundleUtils
                .loadFromFilePath(options.getAstraSecureConnectBundle());

// Build a Source
AstraDbIO.Read<LanguageCode> read = AstraDbIO.<LanguageCode>read()
  .withToken(options.getAstraToken())
  .withKeyspace(options.getAstraKeyspace())
  .withSecureConnectBundle(scbZip)
  .withTable(options.getTable())
  .withCoder(SerializableCoder.of(LanguageCode.class))
  .withMapperFactoryFn(new LanguageCodeDaoMapperFactoryFn())
  .withEntity(LanguageCode.class))

Rate Limiting

The AstraDB IO connector supports rate limiting to control the rate of read and write operations against your Astra database. This helps prevent overload and ensures you stay within rate limits.

Key Concepts

Per-DoFn Rate Limiting

Due to the Beam programming model, rate limiting is applied per DoFn instance. This means:

• Each parallel worker has its own rate limiter
• You need to estimate and configure limits based on expected parallelism
• Example: If you expect 10 parallel workers and want a total rate of 4096 ops/sec, configure each worker with ~409.6 ops/sec

RateLimiter Interface

The RateLimiter interface allows for pluggable rate limiting implementations:

public interface RateLimiter extends Serializable {
  void acquire() throws InterruptedException;
  void acquire(int permits) throws InterruptedException;
  double getRate();
}

SimpleRateLimiter

The built-in SimpleRateLimiter uses a token bucket algorithm:

• Tokens are added at a constant rate
• Operations consume tokens
• If no tokens available, operations block until tokens become available
• Supports burst capacity (1 second worth of permits)

Usage Examples

Basic Read with Rate Limiting

import org.apache.beam.sdk.io.astra.db.AstraDbIO;
import org.apache.beam.sdk.io.astra.db.options.RateLimiter;
import org.apache.beam.sdk.io.astra.db.options.SimpleRateLimiter;

// Create a rate limiter: 4096 operations per second per worker
RateLimiter rateLimiter = SimpleRateLimiter.create(4096.0);

PCollection<MyEntity> results = pipeline
    .apply(AstraDbIO.<MyEntity>read()
        .withToken(token)
        .withKeyspace("my_keyspace")
        .withTable("my_table")
        .withEntity(MyEntity.class)
        .withCoder(SerializableCoder.of(MyEntity.class))
        .withRateLimiter(rateLimiter)  // Add rate limiting
        .withMinNumberOfSplits(20));

Basic Write with Rate Limiting

// Create a rate limiter: 2000 operations per second per worker
RateLimiter rateLimiter = SimpleRateLimiter.create(2000.0);

myData.apply(AstraDbIO.<MyEntity>write()
    .withToken(token)
    .withKeyspace("my_keyspace")
    .withEntity(MyEntity.class)
    .withRateLimiter(rateLimiter));  // Add rate limiting

Delete with Rate Limiting

// Create a rate limiter: 1000 operations per second per worker
RateLimiter rateLimiter = SimpleRateLimiter.create(1000.0);

myData.apply(AstraDbIO.<MyEntity>delete()
    .withToken(token)
    .withKeyspace("my_keyspace")
    .withEntity(MyEntity.class)
    .withRateLimiter(rateLimiter));  // Add rate limiting

Calculating Rate Limits for Parallel Execution

If you know your desired total throughput and expected parallelism:

// Desired total throughput: 10,000 ops/sec
// Expected parallel workers: 25
// Rate per worker: 10,000 / 25 = 400 ops/sec

RateLimiter rateLimiter = SimpleRateLimiter.create(400.0);

Custom Rate Limiter Implementation

You can implement your own rate limiter for more advanced scenarios:

public class DistributedRateLimiter implements RateLimiter {
    private static final long serialVersionUID = 1L;
    
    // Your custom implementation
    // Could use external coordination (Redis, etc.)
    
    @Override
    public void acquire() throws InterruptedException {
        // Your logic here
    }
    
    @Override
    public void acquire(int permits) throws InterruptedException {
        // Your logic here
    }
    
    @Override
    public double getRate() {
        return yourConfiguredRate;
    }
}

// Use your custom rate limiter
RateLimiter customLimiter = new DistributedRateLimiter(config);
pipeline.apply(AstraDbIO.<MyEntity>read()
    .withRateLimiter(customLimiter)
    // ... other configuration
);

Best Practices

1. Start Conservative: Begin with lower rates and increase gradually
2. Monitor Performance: Watch for throttling or timeout errors
3. Consider Burst Traffic: SimpleRateLimiter allows 1 second of burst capacity
4. Account for Parallelism: Remember that rate limits are per-worker
5. Test Thoroughly: Validate rate limiting behavior in your specific environment

Performance Considerations

• Rate limiting adds minimal overhead (microseconds per operation)
• Blocking occurs only when rate limit is exceeded
• Burst capacity allows handling of temporary spikes
• Thread-safe implementation suitable for concurrent access

Troubleshooting

Operations Taking Too Long

If operations are slower than expected:

• Check if rate limit is too restrictive
• Verify parallelism matches your configuration
• Monitor for other bottlenecks (network, database)

Rate Limit Not Being Enforced

If rate limiting seems ineffective:

• Verify rate limiter is properly configured
• Check that parallelism is as expected
• Ensure rate limiter is being passed to the transform

InterruptedException

If you see InterruptedException:

• This is normal when pipeline is cancelled
• Rate limiter properly handles interruption
• No action needed unless it occurs during normal operation

---