license import { TensorContainer , TensorContainerArray , TensorContainerObject } from '@tensorflow/tfjs-core' ; import { iteratorFromConcatenated , iteratorFromConcatenatedFunction , iteratorFromFunction , iteratorFromIncrementing , iteratorFromItems , iteratorFromZipped , LazyIterator , ZipMismatchMode } from './lazy_iterator' ; export class TestIntegerIterator extends LazyIterator < number > { currentIndex = 0 ; data : number [ ] ; constructor ( protected readonly length = 100 ) { super ( ) ; this . data = Array . from ( { length } , ( v , k ) => k ) ; } summary ( ) { return `TestIntegers` ; } async next ( ) { if ( this . currentIndex >= this . length ) { return { value : null , done : true } ; } const result = this . data [ this . currentIndex ] ; this . currentIndex ++ ; if ( Math . random ( ) < 0.1 ) { await new Promise ( res => setTimeout ( res , 1 ) ) ; } return { value : result , done : false } ; } } describe ( 'LazyIterator' , ( ) => { it ( 'collects all stream elements into an array' , async ( ) => { const readIterator = new TestIntegerIterator ( ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 100 ) ; } ) ; it ( 'reads chunks in order' , async ( ) => { const readIterator = new TestIntegerIterator ( ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 100 ) ; for ( let i = 0 ; i < 100 ; i ++ ) { expect ( result [ i ] ) . toEqual ( i ) ; } } ) ; it ( 'filters elements' , async ( ) => { const readIterator = new TestIntegerIterator ( ) . filter ( x => x % 2 === 0 ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 50 ) ; for ( let i = 0 ; i < 50 ; i ++ ) { expect ( result [ i ] ) . toEqual ( 2 * i ) ; } } ) ; it ( 'maps elements' , async ( ) => { const readIterator = new TestIntegerIterator ( ) . map ( x => `item ${ x }` ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 100 ) ; for ( let i = 0 ; i < 100 ; i ++ ) { expect ( result [ i ] ) . toEqual ( `item ${ i }` ) ; } } ) ; it ( 'maps elements through an async function' , async ( ) => { const readIterator = new TestIntegerIterator ( ) . mapAsync ( async ( x ) => { if ( Math . random ( ) < 0.1 ) { await new Promise ( res => setTimeout ( res , 1 ) ) ; } return Promise . resolve ( `item ${ x }` ) ; } ) ; const result = await readIterator . prefetch ( 200 ) . toArray ( ) ; expect ( result . length ) . toEqual ( 100 ) ; for ( let i = 0 ; i < 100 ; i ++ ) { expect ( result [ i ] ) . toEqual ( `item ${ i }` ) ; } } ) ; it ( 'flatmaps simple elements' , async ( ) => { const readStream = new TestIntegerIterator ( ) . flatmap ( x => [ `item ${ x } A` , `item ${ x } B` , `item ${ x } C` ] ) ; const result = await readStream . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 300 ) ; for ( let i = 0 ; i < 100 ; i ++ ) { expect ( result [ 3 * i + 0 ] ) . toEqual ( `item ${ i } A` ) ; expect ( result [ 3 * i + 1 ] ) . toEqual ( `item ${ i } B` ) ; expect ( result [ 3 * i + 2 ] ) . toEqual ( `item ${ i } C` ) ; } } ) ; it ( 'flatmap flattens object elements but not their contents' , async ( ) => { const readStream = new TestIntegerIterator ( ) . flatmap ( x => [ { foo : `foo ${ x } A` , bar : `bar ${ x } A` } , { foo : `foo ${ x } B` , bar : `bar ${ x } B` } , { foo : `foo ${ x } C` , bar : `bar ${ x } C` } , ] ) ; const result = await readStream . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 300 ) ; for ( let i = 0 ; i < 100 ; i ++ ) { expect ( result [ 3 * i + 0 ] ) . toEqual ( { foo : `foo ${ i } A` , bar : `bar ${ i } A` } ) ; expect ( result [ 3 * i + 1 ] ) . toEqual ( { foo : `foo ${ i } B` , bar : `bar ${ i } B` } ) ; expect ( result [ 3 * i + 2 ] ) . toEqual ( { foo : `foo ${ i } C` , bar : `bar ${ i } C` } ) ; } } ) ; it ( 'flatmap flattens array elements but not their contents' , async ( ) => { const readStream = new TestIntegerIterator ( ) . flatmap ( x => [ [ `foo ${ x } A` , `bar ${ x } A` ] , [ `foo ${ x } B` , `bar ${ x } B` ] , [ `foo ${ x } C` , `bar ${ x } C` ] , ] ) ; const result = await readStream . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 300 ) ; for ( let i = 0 ; i < 100 ; i ++ ) { expect ( result [ 3 * i + 0 ] ) . toEqual ( [ `foo ${ i } A` , `bar ${ i } A` ] ) ; expect ( result [ 3 * i + 1 ] ) . toEqual ( [ `foo ${ i } B` , `bar ${ i } B` ] ) ; expect ( result [ 3 * i + 2 ] ) . toEqual ( [ `foo ${ i } C` , `bar ${ i } C` ] ) ; } } ) ; it ( 'batches elements to a row-major representation' , async ( ) => { const readIterator = new TestIntegerIterator ( ) . rowMajorBatch ( 8 ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 13 ) ; for ( let i = 0 ; i < 12 ; i ++ ) { expect ( result [ i ] ) . toEqual ( Array . from ( { length : 8 } , ( v , k ) => ( i * 8 ) + k ) ) ; } expect ( result [ 12 ] ) . toEqual ( [ 96 , 97 , 98 , 99 ] ) ; } ) ; it ( 'batches elements to a column-major representation' , async ( ) => { const readIterator = new TestIntegerIterator ( ) . columnMajorBatch ( 8 ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 13 ) ; for ( let i = 0 ; i < 12 ; i ++ ) { expect ( result [ i ] ) . toEqual ( Array . from ( { length : 8 } , ( v , k ) => ( i * 8 ) + k ) ) ; } expect ( result [ 12 ] ) . toEqual ( [ 96 , 97 , 98 , 99 ] ) ; } ) ; it ( 'can be limited to a certain number of elements' , async ( ) => { const readIterator = new TestIntegerIterator ( ) . take ( 8 ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result ) . toEqual ( [ 0 , 1 , 2 , 3 , 4 , 5 , 6 , 7 ] ) ; } ) ; it ( 'is unaltered by a negative or undefined take() count.' , async ( ) => { const baseIterator = new TestIntegerIterator ( ) ; const readIterator = baseIterator . take ( - 1 ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result ) . toEqual ( baseIterator . data ) ; const baseIterator2 = new TestIntegerIterator ( ) ; const readIterator2 = baseIterator2 . take ( undefined ) ; const result2 = await readIterator2 . toArrayForTest ( ) ; expect ( result2 ) . toEqual ( baseIterator2 . data ) ; } ) ; it ( 'can skip a certain number of elements' , async ( ) => { const readIterator = new TestIntegerIterator ( ) . skip ( 88 ) . take ( 8 ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result ) . toEqual ( [ 88 , 89 , 90 , 91 , 92 , 93 , 94 , 95 ] ) ; } ) ; it ( 'is unaltered by a negative or undefined skip() count.' , async ( ) => { const baseIterator = new TestIntegerIterator ( ) ; const readIterator = baseIterator . skip ( - 1 ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result ) . toEqual ( baseIterator . data ) ; const baseIterator2 = new TestIntegerIterator ( ) ; const readIterator2 = baseIterator2 . skip ( undefined ) ; const result2 = await readIterator2 . toArrayForTest ( ) ; expect ( result2 ) . toEqual ( baseIterator2 . data ) ; } ) ; it ( 'can selectively ignore upstream errors' , async ( ) => { const readIterator = new TestIntegerIterator ( ) . take ( 10 ) . map ( x => { if ( x % 2 === 0 ) { throw new Error ( 'Oh no, an even number!' ) ; } return x ; } ) ; const errorIgnoringIterator = readIterator . handleErrors ( ( e ) => true ) ; const result = await errorIgnoringIterator . toArrayForTest ( ) ; expect ( result ) . toEqual ( [ 1 , 3 , 5 , 7 , 9 ] ) ; } ) ; it ( 'can terminate cleanly based on upstream errors' , async ( ) => { const readIterator = new TestIntegerIterator ( ) . map ( x => { if ( x % 2 === 0 ) { throw new Error ( `Oh no, an even number: ${ x }` ) ; } return x ; } ) ; const errorHandlingIterator = readIterator . handleErrors ( ( e ) => e . message !== 'Oh no, an even number: 10' ) ; const result = await errorHandlingIterator . toArrayForTest ( ) ; expect ( result ) . toEqual ( [ 1 , 3 , 5 , 7 , 9 ] ) ; } ) ; it ( 'can selectively propagate upstream errors' , async ( done ) => { const readIterator = new TestIntegerIterator ( ) . map ( x => { if ( x % 2 === 0 ) { throw new Error ( `Oh no, an even number: ${ x }` ) ; } return x ; } ) ; const errorHandlingIterator = readIterator . handleErrors ( ( e ) => { if ( e . message === 'Oh no, an even number: 2' ) { throw e ; } return true ; } ) ; try { await errorHandlingIterator . toArray ( ) ; done . fail ( ) ; } catch ( e ) { expect ( e . message ) . toEqual ( 'Oh no, an even number: 2' ) ; done ( ) ; } } ) ; it ( 'can be forced to execute serially' , async ( ) => { class ParallelFilterIterator < T > extends LazyIterator < T > { constructor ( protected upstream , protected predicate ) { super ( ) ; } summary ( ) { return `${ this . upstream . summary } -> ParallelFilterIterator` ; } async next ( ) { while ( true ) { const item = await this . upstream . next ( ) ; if ( item . done || this . predicate ( item . value ) ) { return item ; } } } } const newReadIterator = ( ) => new ParallelFilterIterator ( ( new TestIntegerIterator ( ) ) . take ( 10 ) , x => { return x % 2 === 0 ; } ) ; const badResult = await newReadIterator ( ) . toArrayForTest ( ) ; expect ( badResult ) . toEqual ( [ 0 ] ) ; const goodResult = await newReadIterator ( ) . serial ( ) . toArrayForTest ( ) ; expect ( goodResult ) . toEqual ( [ 0 , 2 , 4 , 6 , 8 ] ) ; } ) ; it ( 'can be created from an array' , async ( ) => { const readIterator = iteratorFromItems ( [ 1 , 2 , 3 , 4 , 5 , 6 ] ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result ) . toEqual ( [ 1 , 2 , 3 , 4 , 5 , 6 ] ) ; } ) ; it ( 'can be created from a function' , async ( ) => { let i = - 1 ; const func = ( ) => ++ i < 7 ? { value : i , done : false } : { value : null , done : true } ; const readIterator = iteratorFromFunction ( func ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result ) . toEqual ( [ 0 , 1 , 2 , 3 , 4 , 5 , 6 ] ) ; } ) ; it ( 'can be created with incrementing integers' , async ( ) => { const readIterator = iteratorFromIncrementing ( 0 ) . take ( 7 ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result ) . toEqual ( [ 0 , 1 , 2 , 3 , 4 , 5 , 6 ] ) ; } ) ; it ( 'can be concatenated' , async ( ) => { const a = iteratorFromItems ( [ 1 , 2 , 3 ] ) ; const b = iteratorFromItems ( [ 4 , 5 , 6 ] ) ; const readIterator = a . concatenate ( b ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result ) . toEqual ( [ 1 , 2 , 3 , 4 , 5 , 6 ] ) ; } ) ; it ( 'can be concatenated after skipping' , async ( ) => { const a = iteratorFromItems ( [ 1 , 2 , 3 ] ) ; const b = iteratorFromItems ( [ 4 , 5 , 6 ] ) ; const readIterator = a . skip ( 1 ) . concatenate ( b ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result ) . toEqual ( [ 2 , 3 , 4 , 5 , 6 ] ) ; } ) ; it ( 'can be created by concatenating streams' , async ( ) => { const a = new TestIntegerIterator ( ) ; const b = new TestIntegerIterator ( ) ; const readIterator = iteratorFromConcatenated ( iteratorFromItems ( [ a , b ] ) ) ; const result = await readIterator . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 200 ) ; } ) ; it ( 'can be created by concatenating streams from a function' , async ( ) => { const readIterator = iteratorFromConcatenatedFunction ( ( ) => ( { value : new TestIntegerIterator ( ) , done : false } ) , 3 ) ; const expectedResult = [ ] ; for ( let i = 0 ; i < 3 ; i ++ ) { for ( let j = 0 ; j < 100 ; j ++ ) { expectedResult [ i * 100 + j ] = j ; } } const result = await readIterator . toArrayForTest ( ) ; expect ( result ) . toEqual ( expectedResult ) ; } ) ; it ( 'can be created by zipping an array of streams' , async ( ) => { const a = new TestIntegerIterator ( ) ; const b = new TestIntegerIterator ( ) . map ( x => x * 10 ) ; const c = new TestIntegerIterator ( ) . map ( x => `string ${ x }` ) ; const readStream = iteratorFromZipped ( [ a , b , c ] ) ; const result = await readStream . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 100 ) ; for ( const e of result ) { const ee = e as TensorContainerArray ; expect ( ee [ 1 ] ) . toEqual ( ee [ 0 ] as number * 10 ) ; expect ( ee [ 2 ] ) . toEqual ( `string ${ ee [ 0 ] }` ) ; } } ) ; it ( 'can be created by zipping a dict of streams' , async ( ) => { const a = new TestIntegerIterator ( ) ; const b = new TestIntegerIterator ( ) . map ( x => x * 10 ) ; const c = new TestIntegerIterator ( ) . map ( x => `string ${ x }` ) ; const readStream = iteratorFromZipped ( { a , b , c } ) ; const result = await readStream . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 100 ) ; for ( const e of result ) { const ee = e as TensorContainerObject ; expect ( ee [ 'b' ] ) . toEqual ( ee [ 'a' ] as number * 10 ) ; expect ( ee [ 'c' ] ) . toEqual ( `string ${ ee [ 'a' ] }` ) ; } } ) ; it ( 'can be created by zipping a nested structure of streams' , async ( ) => { const a = new TestIntegerIterator ( ) . map ( x => ( { 'a' : x , 'constant' : 12 } ) ) ; const b = new TestIntegerIterator ( ) . map ( x => ( { 'b' : x * 10 , 'array' : [ x * 100 , x * 200 ] } ) ) ; const c = new TestIntegerIterator ( ) . map ( x => ( { 'c' : `string ${ x }` } ) ) ; const readStream = iteratorFromZipped ( [ a , b , c ] ) ; const result = await readStream . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 100 ) ; for ( const e of result ) { const ee = e as TensorContainerArray ; const aa = ee [ 0 ] as TensorContainerObject ; const bb = ee [ 1 ] as TensorContainerObject ; const cc = ee [ 2 ] as TensorContainerObject ; expect ( aa [ 'constant' ] ) . toEqual ( 12 ) ; expect ( bb [ 'b' ] ) . toEqual ( aa [ 'a' ] as number * 10 ) ; expect ( bb [ 'array' ] ) . toEqual ( [ aa [ 'a' ] as number * 100 , aa [ 'a' ] as number * 200 ] ) ; expect ( cc [ 'c' ] ) . toEqual ( `string ${ aa [ 'a' ] }` ) ; } } ) ; it ( 'zip requires streams of the same length by default' , async ( done ) => { try { const a = new TestIntegerIterator ( 10 ) ; const b = new TestIntegerIterator ( 3 ) ; const c = new TestIntegerIterator ( 2 ) ; const readStream = iteratorFromZipped ( [ a , b , c ] ) ; await readStream . toArray ( ) ; done . fail ( ) ; } catch ( error ) { expect ( error . message ) . toBe ( 'Zipped streams should have the same length. ' + 'Mismatched at element 2.' ) ; done ( ) ; } } ) ; it ( 'zip can be told to terminate when the shortest stream terminates' , async ( ) => { const a = new TestIntegerIterator ( 10 ) ; const b = new TestIntegerIterator ( 3 ) ; const c = new TestIntegerIterator ( 2 ) ; const readStream = iteratorFromZipped ( [ a , b , c ] , ZipMismatchMode . SHORTEST ) ; const result = await readStream . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 2 ) ; } ) ; it ( 'zip can be told to terminate when the longest stream terminates' , async ( ) => { const a = new TestIntegerIterator ( 10 ) ; const b = new TestIntegerIterator ( 3 ) ; const c = new TestIntegerIterator ( 2 ) ; const readStream = iteratorFromZipped ( [ a , b , c ] , ZipMismatchMode . LONGEST ) ; const result = await readStream . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 10 ) ; expect ( result [ 9 ] ) . toEqual ( [ 9 , null , null ] ) ; } ) ; /**
     * This test demonstrates behavior that is intrinsic to the tf.data zip()
     * API, but that may not be what users ultimately want when zipping dicts.
     * This may merit a convenience function (e.g., maybe flatZip()).
     */ it ( 'zipping TensorContainer streams requires manual merge' , async ( ) => { function naiveMerge ( xs ) { const result = { } ; for ( const x of xs ) { Object . assign ( result , x ) ; } return result ; } const a = new TestIntegerIterator ( ) . map ( x => ( { 'a' : x } ) ) ; const b = new TestIntegerIterator ( ) . map ( x => ( { 'b' : x * 10 } ) ) ; const c = new TestIntegerIterator ( ) . map ( x => ( { 'c' : `string ${ x }` } ) ) ; const zippedStream = iteratorFromZipped ( [ a , b , c ] ) ; const readStream = zippedStream . map ( e => naiveMerge ( e as TensorContainerArray ) ) ; const result = await readStream . toArrayForTest ( ) ; expect ( result . length ) . toEqual ( 100 ) ; for ( const e of result ) { const ee = e as TensorContainerObject ; expect ( ee [ 'b' ] ) . toEqual ( ee [ 'a' ] as number * 10 ) ; expect ( ee [ 'c' ] ) . toEqual ( `string ${ ee [ 'a' ] }` ) ; } } ) ; } ) ;