35
35
% %% end
36
36
% %%
37
37
% %% The result currently uses a non-standard magic prefix to prevent the KCL from
38
- % %% deaggregating the record automatically. To use compression, use
39
- % %% `kpl_agg:finish/2` with ` true` as the second argument , which uses another
38
+ % %% deaggregating the record automatically. To use compression, instantiate the
39
+ % %% aggregator using kpl_agg:new( true) , which uses another
40
40
% %% non-standard magic prefix.
41
41
% %%
42
42
% %% @end
45
45
-module (kpl_agg ).
46
46
47
47
% % API
48
- -export ([new /0 , count /1 , size_bytes /1 , finish /1 , finish /2 , add /2 , add_all /2 ]).
48
+ -export ([new /0 , new /1 , count /1 , size_bytes /1 , finish /1 , add /2 , add_all /2 ]).
49
49
50
50
-define (MD5_DIGEST_BYTES , 16 ).
51
51
% % From http://docs.aws.amazon.com/kinesis/latest/APIReference/API_PutRecord.html:
57
57
% % A set of keys, mapping each key to a unique index.
58
58
-record (keyset , {
59
59
rev_keys = [] :: list (binary ()), % % list of known keys in reverse order
60
- key_to_index = maps :new () :: map () % % maps each known key to a 0-based index
60
+ rev_keys_length = 0 :: non_neg_integer (), % % length of the rev_keys list
61
+ key_to_index = maps :new () :: map () % % maps each known key to a 0-based index
61
62
}).
62
63
63
64
% % Internal state of a record aggregator. It stores an aggregated record that
76
77
explicit_hash_keyset = # keyset {} :: # keyset {},
77
78
78
79
% % List if user records added so far, in reverse order.
79
- rev_records = [] :: [# 'Record' {}]
80
+ rev_records = [] :: [# 'Record' {}],
81
+
82
+ should_deflate = false
80
83
}).
81
84
82
85
85
88
% %%===================================================================
86
89
87
90
new () ->
88
- # state {}.
91
+ new (false ).
92
+ new (ShouldDeflate ) ->
93
+ # state {should_deflate = ShouldDeflate }.
89
94
90
95
count (# state {num_user_records = Num } = _State ) ->
91
96
Num .
@@ -101,15 +106,12 @@ size_bytes(#state{agg_size_bytes = Size,
101
106
end )
102
107
+ byte_size (kpl_agg_pb :encode_msg (# 'AggregatedRecord' {})).
103
108
104
- finish (# state {num_user_records = 0 } = State , _ ) ->
109
+ finish (# state {num_user_records = 0 } = State ) ->
105
110
{undefined , State };
106
111
107
- finish (# state {agg_partition_key = AggPK , agg_explicit_hash_key = AggEHK } = State , ShouldDeflate ) ->
112
+ finish (# state {agg_partition_key = AggPK , agg_explicit_hash_key = AggEHK , should_deflate = ShouldDeflate } = State ) ->
108
113
AggRecord = {AggPK , serialize_data (State , ShouldDeflate ), AggEHK },
109
- {AggRecord , new ()}.
110
-
111
- finish (State ) ->
112
- finish (State , false ).
114
+ {AggRecord , new (ShouldDeflate )}.
113
115
114
116
115
117
add (State , {PartitionKey , Data } = _Record ) ->
@@ -277,23 +279,23 @@ is_key(Key, #keyset{key_to_index = KeyToIndex} = _KeySet) ->
277
279
278
280
get_or_add_key (undefined , KeySet ) ->
279
281
{undefined , KeySet };
280
- get_or_add_key (Key , # keyset {rev_keys = RevKeys , key_to_index = KeyToIndex } = KeySet ) ->
282
+ get_or_add_key (Key , # keyset {rev_keys = RevKeys , rev_keys_length = Length , key_to_index = KeyToIndex } = KeySet ) ->
281
283
case maps :get (Key , KeyToIndex , not_found ) of
282
284
not_found ->
283
- Index = length (RevKeys ),
284
285
NewKeySet = KeySet # keyset {
285
286
rev_keys = [Key | RevKeys ],
286
- key_to_index = maps :put (Key , Index , KeyToIndex )
287
+ rev_keys_length = Length + 1 ,
288
+ key_to_index = maps :put (Key , Length , KeyToIndex )
287
289
},
288
- {Index , NewKeySet };
290
+ {Length , NewKeySet };
289
291
Index ->
290
292
{Index , KeySet }
291
293
end .
292
294
293
295
294
- potential_index (Key , # keyset {rev_keys = RevKeys , key_to_index = KeyToIndex } = _KeySet ) ->
296
+ potential_index (Key , # keyset {rev_keys_length = Length , key_to_index = KeyToIndex } = _KeySet ) ->
295
297
case maps :get (Key , KeyToIndex , not_found ) of
296
- not_found -> length ( RevKeys ) ;
298
+ not_found -> Length ;
297
299
Index -> Index
298
300
end .
299
301
@@ -443,9 +445,9 @@ full_record_test() ->
443
445
444
446
445
447
deflate_test () ->
446
- Agg0 = new (),
448
+ Agg0 = new (true ),
447
449
{undefined , Agg1 } = add (Agg0 , {<<" pk1" >>, <<" data1" >>, <<" ehk1" >>}),
448
- {{_ , Data , _ }, _ } = finish (Agg1 , true ),
450
+ {{_ , Data , _ }, _ } = finish (Agg1 ),
449
451
<<Magic :4 /binary , Deflated /binary >> = Data ,
450
452
? assertEqual (? KPL_AGG_MAGIC_DEFLATED , Magic ),
451
453
Inflated = zlib :uncompress (Deflated ),
0 commit comments