@@ -181,65 +181,28 @@ func readHandleAndLocalPathList(br *binaryReader) ([]HandleAndLocalPath, error)
181181// Type 14 (V2) includes a separate checkpointId field.
182182func readChangelogStateHandle (br * binaryReader , kind KeyedStateHandleType , parseFull bool ) (KeyedStateHandle , error ) {
183183 _ = parseFull
184- isV2 := kind == KeyedStateHandleChangelogV2
185-
186- startKeyGroup , err := br .ReadInt32 ()
187- if err != nil {
188- return nil , fmt .Errorf ("read changelog start key group: %w" , err )
189- }
190- count , err := br .ReadInt32 ()
191- if err != nil {
192- return nil , fmt .Errorf ("read changelog key group count: %w" , err )
193- }
194- checkpointedSize , err := br .ReadInt64 ()
184+ startKeyGroup , numKeyGroups , checkpointedSize , err := readChangelogHeader (br )
195185 if err != nil {
196- return nil , fmt . Errorf ( "read changelog checkpointed size: %w" , err )
186+ return nil , err
197187 }
198188
199- materializedCount , err := br . ReadInt32 ( )
189+ materialized , err := readChangelogKeyedStateHandles ( br , "materialized" )
200190 if err != nil {
201- return nil , fmt .Errorf ("read changelog materialized count: %w" , err )
202- }
203- if materializedCount < 0 {
204- return nil , fmt .Errorf ("changelog materialized count negative: %d" , materializedCount )
205- }
206- materialized := make ([]KeyedStateHandle , 0 , materializedCount )
207- for i := int32 (0 ); i < materializedCount ; i ++ {
208- handle , err := readKeyedStateHandle (br , true )
209- if err != nil {
210- return nil , fmt .Errorf ("read changelog materialized handle: %w" , err )
211- }
212- if handle != nil {
213- materialized = append (materialized , handle )
214- }
191+ return nil , err
215192 }
216193
217- nonMaterializedCount , err := br . ReadInt32 ( )
194+ nonMaterialized , err := readChangelogKeyedStateHandles ( br , "non materialized" )
218195 if err != nil {
219- return nil , fmt .Errorf ("read changelog non materialized count: %w" , err )
220- }
221- if nonMaterializedCount < 0 {
222- return nil , fmt .Errorf ("changelog non materialized count negative: %d" , nonMaterializedCount )
223- }
224- nonMaterialized := make ([]KeyedStateHandle , 0 , nonMaterializedCount )
225- for i := int32 (0 ); i < nonMaterializedCount ; i ++ {
226- handle , err := readKeyedStateHandle (br , true )
227- if err != nil {
228- return nil , fmt .Errorf ("read changelog non materialized handle: %w" , err )
229- }
230- if handle != nil {
231- nonMaterialized = append (nonMaterialized , handle )
232- }
196+ return nil , err
233197 }
234198
235199 materializationID , err := br .ReadInt64 ()
236200 if err != nil {
237201 return nil , fmt .Errorf ("read changelog materialization id: %w" , err )
238202 }
239203
240- // checkpointId is only present in V2 (type 14); for V1, use materializationID
241204 checkpointID := materializationID
242- if isV2 {
205+ if kind == KeyedStateHandleChangelogV2 {
243206 checkpointID , err = br .ReadInt64 ()
244207 if err != nil {
245208 return nil , fmt .Errorf ("read changelog checkpoint id: %w" , err )
@@ -254,7 +217,7 @@ func readChangelogStateHandle(br *binaryReader, kind KeyedStateHandleType, parse
254217 return ChangelogStateHandle {
255218 Type : kind ,
256219 StartKeyGroup : startKeyGroup ,
257- NumKeyGroups : count ,
220+ NumKeyGroups : numKeyGroups ,
258221 CheckpointedSize : checkpointedSize ,
259222 Materialized : materialized ,
260223 NonMaterialized : nonMaterialized ,
@@ -264,6 +227,46 @@ func readChangelogStateHandle(br *binaryReader, kind KeyedStateHandleType, parse
264227 }, nil
265228}
266229
230+ func readChangelogHeader (br * binaryReader ) (startKeyGroup int32 , numKeyGroups int32 , checkpointedSize int64 , err error ) {
231+ startKeyGroup , err = br .ReadInt32 ()
232+ if err != nil {
233+ return 0 , 0 , 0 , fmt .Errorf ("read changelog start key group: %w" , err )
234+ }
235+ count , err := br .ReadInt32 ()
236+ if err != nil {
237+ return 0 , 0 , 0 , fmt .Errorf ("read changelog key group count: %w" , err )
238+ }
239+ checkpointedSize , err = br .ReadInt64 ()
240+ if err != nil {
241+ return 0 , 0 , 0 , fmt .Errorf ("read changelog checkpointed size: %w" , err )
242+ }
243+
244+ return startKeyGroup , count , checkpointedSize , nil
245+ }
246+
247+ func readChangelogKeyedStateHandles (br * binaryReader , label string ) ([]KeyedStateHandle , error ) {
248+ count , err := br .ReadInt32 ()
249+ if err != nil {
250+ return nil , fmt .Errorf ("read changelog %s count: %w" , label , err )
251+ }
252+ if count < 0 {
253+ return nil , fmt .Errorf ("changelog %s count negative: %d" , label , count )
254+ }
255+
256+ handles := make ([]KeyedStateHandle , 0 , count )
257+ for i := int32 (0 ); i < count ; i ++ {
258+ handle , err := readKeyedStateHandle (br , true )
259+ if err != nil {
260+ return nil , fmt .Errorf ("read changelog %s handle: %w" , label , err )
261+ }
262+ if handle != nil {
263+ handles = append (handles , handle )
264+ }
265+ }
266+
267+ return handles , nil
268+ }
269+
267270// readChangelogByteIncrementHandle parses in-memory changelog increments.
268271func readChangelogByteIncrementHandle (br * binaryReader , kind KeyedStateHandleType , parseFull bool ) (KeyedStateHandle , error ) {
269272 startKeyGroup , err := br .ReadInt32 ()
@@ -304,6 +307,7 @@ func readChangelogByteIncrementHandle(br *binaryReader, kind KeyedStateHandleTyp
304307 if _ , err := br .ReadBytes (int (length )); err != nil {
305308 return nil , fmt .Errorf ("read changelog byte data: %w" , err )
306309 }
310+
307311 continue
308312 }
309313 data , err := br .ReadBytes (int (length ))
0 commit comments