Partitioning Kinesis Records with AWS IoT Payload Values

Use a customer ID from an IoT topic payload as the Kinesis partition key to preserve record order for each customer.

Takahiro Iwasa

2020-04-16

3 min read

An AWS IoT topic rule can use a value from the message payload as the partition key for Kinesis records. This is useful when downstream processing must preserve the order of records for each customer. The following example uses customer_id as the partition key.

Prerequisites

Install the following tools on your machine:

AWS SAM CLI
Python 3.x

Building

The AWS IoT topic rule specifies PartitionKey: ${customer_id} (line 37), which determines how records are distributed across shards.

1
AWSTemplateFormatVersion: '2010-09-09'
2
Transform: AWS::Serverless-2016-10-31
3

4
Resources:
5
  Lambda:
6
    Type: AWS::Serverless::Function
7
    Properties:
8
      CodeUri: ./src/
9
      Events:
10
        Kinesis:
11
          Type: Kinesis
12
          Properties:
13
            BatchSize: 100
14
            BisectBatchOnFunctionError: true
15
            Enabled: true
16
            StartingPosition: LATEST
17
            Stream: !GetAtt Kinesis.Arn
18
      FunctionName: aws_iot_kinesis_partition_lambda
19
      Handler: lambda_function.lambda_handler
20
      Role: !GetAtt IamRole.Arn
21
      Runtime: python3.8
22

23
  Kinesis:
24
    Type: AWS::Kinesis::Stream
25
    Properties:
26
      Name: aws_iot_kinesis_partition_stream
27
      RetentionPeriodHours: 24
28
      ShardCount: 1
29

30
  TopicRule:
31
    Type: AWS::IoT::TopicRule
32
    Properties:
33
      RuleName: aws_iot_kinesis_partition_topic_rule
34
      TopicRulePayload:
35
        Actions:
36
          - Kinesis:
37
              PartitionKey: ${customer_id}
38
              RoleArn: !GetAtt IamRole.Arn
39
              StreamName: aws_iot_kinesis_partition_stream
40
        AwsIotSqlVersion: 2016-03-23
41
        RuleDisabled: false
42
        Sql: SELECT * FROM 'aws-iot-kinesis-partition-topic'
43

44
  IamRole:
45
    Type: AWS::IAM::Role
46
    Properties:
47
      AssumeRolePolicyDocument:
48
        Version: '2012-10-17'
49
        Statement:
50
          - Effect: Allow
51
            Principal:
52
              Service: lambda.amazonaws.com
53
            Action: sts:AssumeRole
54
          - Effect: Allow
55
            Principal:
56
              Service: iot.amazonaws.com
57
            Action: sts:AssumeRole
58
      ManagedPolicyArns:
59
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
60
      Policies:
61
        - PolicyDocument:
62
            Version: '2012-10-17'
63
            Statement:
64
              - Effect: Allow
65
                Action:
66
                  - kinesis:GetShardIterator
67
                  - kinesis:GetRecords
68
                  - kinesis:DescribeStream
69
                  - kinesis:PutRecord
70
                Resource:
71
                  - !GetAtt Kinesis.Arn
72
          PolicyName: policy
73

74
  LogGroup:
75
    Type: AWS::Logs::LogGroup
76
    Properties:
77
      LogGroupName: /aws/lambda/aws_iot_kinesis_partition_lambda
78
      RetentionInDays: 1

The Lambda function that processes the streamed records is a short Python script:

1
import json
2

3

4
def lambda_handler(event, context):
5
    for record in event['Records']:
6
        print(json.dumps(record['kinesis']))

Build and deploy the stack with SAM:

sam build
sam deploy \
  --stack-name aws-iot-kinesis-partition-lambda \
  --capabilities CAPABILITY_NAMED_IAM

Testing

Test payloads can be sent to the AWS IoT topic directly:

aws iot-data publish \
  --topic aws-iot-kinesis-partition-topic \
  --payload '{"customer_id": 1, "message": "Hello from AWS IoT"}' \
  --cli-binary-format raw-in-base64-out

Check the Lambda logs to confirm that the partitionKey matches the customer_id in your payload.

Cleaning Up

Remove the resources provisioned by this example with:

sam delete --stack-name aws-iot-kinesis-partition-lambda

Conclusion

Setting PartitionKey: ${customer_id} in the IoT topic rule sends records for the same customer to the same Kinesis shard without requiring partitioning logic in the Lambda consumer. Because Kinesis uses the partition key to select a shard, records with the same key can be processed in order by downstream consumers.

Load is distributed evenly only when traffic is reasonably balanced across customer_id values. A single high-volume customer can still concentrate traffic on one shard even when the partition key is configured correctly.

Back to posts