license import { TensorContainer , util } from '@tensorflow/tfjs-core' ; import { Dataset } from '../dataset' ; import { DataSource } from '../datasource' ; import { LazyIterator } from '../iterators/lazy_iterator' ; import { ColumnConfig , CSVConfig } from '../types' ; import { TextLineDataset } from './text_line_dataset' ; const CODE_QUOTE = '"' ; const STATE_OUT = Symbol ( 'out' ) ; const STATE_FIELD = Symbol ( 'field' ) ; const STATE_QUOTE = Symbol ( 'quote' ) ; const STATE_QUOTE_AFTER_QUOTE = Symbol ( 'quoteafterquote' ) ; const STATE_WITHIN_QUOTE_IN_QUOTE = Symbol ( 'quoteinquote' ) ; /**
 * Represents a potentially large collection of delimited text records.
 *
 * The produced `TensorContainer`s each contain one key-value pair for
 * every column of the table.  When a field is empty in the incoming data, the
 * resulting value is `undefined`, or throw error if it is required.  Values
 * that can be parsed as numbers are emitted as type `number`, other values
 * are parsed as `string`.
 *
 * The results are not batched.
 */ doc export class CSVDataset extends Dataset < TensorContainer > { base : TextLineDataset ; private hasHeader = true ; private fullColumnNames : string [ ] = null ; private columnNamesValidated = false ; private columnConfigs : { [ key ] : ColumnConfig ; } = null ; private configuredColumnsOnly = false ; private delimiter = ',' ; private delimWhitespace = false ; /**
     * Returns column names of the csv dataset. If `configuredColumnsOnly` is
     * true, return column names in `columnConfigs`. If `configuredColumnsOnly` is
     * false and `columnNames` is provided, `columnNames`. If
     * `configuredColumnsOnly` is false and `columnNames` is not provided, return
     * all column names parsed from the csv file. For example usage please go to
     * `tf.data.csv`.
     */ doc async columnNames ( ) { if ( ! this . columnNamesValidated ) { await this . setColumnNames ( ) ; } return this . configuredColumnsOnly ? Object . keys ( this . columnConfigs ) : this . fullColumnNames ; } private async setColumnNames ( ) { const columnNamesFromFile = await this . maybeReadHeaderLine ( ) ; if ( ! this . fullColumnNames && ! columnNamesFromFile ) { throw new Error ( 'Column names must be provided if there is no header line.' ) ; } else if ( this . fullColumnNames && columnNamesFromFile ) { util . assert ( columnNamesFromFile . length === this . fullColumnNames . length , ( ) => 'The length of provided columnNames (' + this . fullColumnNames . length . toString ( ) + ') does not match the length of the header line read from ' + 'file (' + columnNamesFromFile . length . toString ( ) + ').' ) ; } if ( ! this . fullColumnNames ) { this . fullColumnNames = columnNamesFromFile ; } const counts = this . fullColumnNames . reduce ( ( countAcc , name ) => { countAcc [ name ] = ( countAcc [ name ] + 1 ) || 1 ; return countAcc ; } , { } ) ; const duplicateNames = Object . keys ( counts ) . filter ( ( name ) => ( counts [ name ] > 1 ) ) ; util . assert ( duplicateNames . length === 0 , ( ) => 'Duplicate column names found: ' + duplicateNames . toString ( ) ) ; if ( this . columnConfigs ) { for ( const key of Object . keys ( this . columnConfigs ) ) { const index = this . fullColumnNames . indexOf ( key ) ; if ( index === - 1 ) { throw new Error ( 'The key "' + key + '" provided in columnConfigs does not match any of the column ' + 'names (' + this . fullColumnNames . toString ( ) + ').' ) ; } } } this . columnNamesValidated = true ; } private async maybeReadHeaderLine ( ) { if ( this . hasHeader ) { const iter = await this . base . iterator ( ) ; const firstElement = await iter . next ( ) ; if ( firstElement . done ) { throw new Error ( 'No data was found for CSV parsing.' ) ; } const firstLine = firstElement . value ; const headers = this . parseRow ( firstLine , false ) ; return headers ; } else { return null ; } } param input param csvConfig constructor ( protected readonly input , csvConfig ? ) { super ( ) ; this . base = new TextLineDataset ( input ) ; if ( ! csvConfig ) { csvConfig = { } ; } this . hasHeader = csvConfig . hasHeader === false ? false : true ; this . fullColumnNames = csvConfig . columnNames ; this . columnConfigs = csvConfig . columnConfigs ; this . configuredColumnsOnly = csvConfig . configuredColumnsOnly ; if ( csvConfig . delimWhitespace ) { util . assert ( csvConfig . delimiter == null , ( ) => 'Delimiter should not be provided when delimWhitespace is true.' ) ; this . delimWhitespace = true ; this . delimiter = ' ' ; } else { this . delimiter = csvConfig . delimiter ? csvConfig . delimiter : ',' ; } } async iterator ( ) { if ( ! this . columnNamesValidated ) { await this . setColumnNames ( ) ; } let lines = await this . base . iterator ( ) ; if ( this . hasHeader ) { lines = lines . skip ( 1 ) ; } return lines . map ( x => this . makeDataElement ( x ) ) ; } makeDataElement ( line ) { const values = this . parseRow ( line ) ; const features = { } ; const labels = { } ; for ( let i = 0 ; i < this . fullColumnNames . length ; i ++ ) { const key = this . fullColumnNames [ i ] ; const config = this . columnConfigs ? this . columnConfigs [ key ] : null ; if ( this . configuredColumnsOnly && ! config ) { continue ; } else { const value = values [ i ] ; let parsedValue = null ; if ( value === '' ) { if ( config && config . default !== undefined ) { parsedValue = config . default ; } else if ( config && ( config . required || config . isLabel ) ) { throw new Error ( `Required column ${ key } is empty in this line: ${ line }` ) ; } else { parsedValue = undefined ; } } else { const valueAsNum = Number ( value ) ; if ( isNaN ( valueAsNum ) ) { if ( config && config . dtype === 'bool' ) { parsedValue = this . getBoolean ( value ) ; } else { parsedValue = value as string ; } } else if ( ! config || ! config . dtype ) { parsedValue = valueAsNum ; } else { switch ( config . dtype ) { case 'float32' : parsedValue = valueAsNum ; break ; case 'int32' : parsedValue = Math . floor ( valueAsNum ) ; break ; case 'bool' : parsedValue = this . getBoolean ( value ) ; break ; default : parsedValue = valueAsNum ; } } } ( config && config . isLabel ) ? labels [ key ] = parsedValue : features [ key ] = parsedValue ; } } if ( Object . keys ( labels ) . length === 0 ) { return features ; } else { return { xs : features , ys : labels } ; } } private getBoolean ( value ) { if ( value === '1' || value . toLowerCase ( ) === 'true' ) { return 1 ; } else { return 0 ; } } private parseRow ( line , validateElementCount = true ) { const result = [ ] ; let readOffset = 0 ; const readLength = line . length ; let currentState = STATE_OUT ; for ( let i = 0 ; i < readLength ; i ++ ) { switch ( currentState ) { case STATE_OUT : switch ( line . charAt ( i ) ) { case CODE_QUOTE : readOffset = i + 1 ; currentState = STATE_QUOTE ; break ; case this . delimiter : readOffset = i + 1 ; if ( this . delimiter === ' ' && this . delimWhitespace ) { break ; } result . push ( '' ) ; currentState = STATE_OUT ; break ; default : currentState = STATE_FIELD ; readOffset = i ; break ; } break ; case STATE_FIELD : switch ( line . charAt ( i ) ) { case this . delimiter : result . push ( line . substring ( readOffset , i ) ) ; currentState = STATE_OUT ; readOffset = i + 1 ; break ; default : } break ; case STATE_QUOTE : switch ( line . charAt ( i ) ) { case CODE_QUOTE : currentState = STATE_QUOTE_AFTER_QUOTE ; break ; default : } break ; case STATE_QUOTE_AFTER_QUOTE : switch ( line . charAt ( i ) ) { case this . delimiter : result . push ( line . substring ( readOffset , i - 1 ) ) ; currentState = STATE_OUT ; readOffset = i + 1 ; break ; case CODE_QUOTE : currentState = STATE_QUOTE ; break ; default : currentState = STATE_WITHIN_QUOTE_IN_QUOTE ; break ; } break ; case STATE_WITHIN_QUOTE_IN_QUOTE : switch ( line . charAt ( i ) ) { case CODE_QUOTE : currentState = STATE_QUOTE ; break ; default : } break ; default : } } if ( currentState === STATE_QUOTE_AFTER_QUOTE ) { result . push ( line . substring ( readOffset , readLength - 1 ) ) ; } else { result . push ( line . substring ( readOffset ) ) ; } if ( validateElementCount && result . length !== this . fullColumnNames . length ) { throw new Error ( `Invalid row in csv file. Should have ${ this . fullColumnNames . length } elements in a row, but got ${ result }` ) ; } return result ; } }