github.com/parquet-go/parquet-go@v0.21.1-0.20240501160520-b3c3a0c3ed6f/scan.go (about)

     1  package parquet
     2  
     3  import "io"
     4  
     5  // ScanRowReader constructs a RowReader which exposes rows from reader until
     6  // the predicate returns false for one of the rows, or EOF is reached.
     7  func ScanRowReader(reader RowReader, predicate func(Row, int64) bool) RowReader {
     8  	return &scanRowReader{reader: reader, predicate: predicate}
     9  }
    10  
    11  type scanRowReader struct {
    12  	reader    RowReader
    13  	predicate func(Row, int64) bool
    14  	rowIndex  int64
    15  }
    16  
    17  func (s *scanRowReader) ReadRows(rows []Row) (int, error) {
    18  	if s.rowIndex < 0 {
    19  		return 0, io.EOF
    20  	}
    21  
    22  	n, err := s.reader.ReadRows(rows)
    23  
    24  	for i, row := range rows[:n] {
    25  		if !s.predicate(row, s.rowIndex) {
    26  			s.rowIndex = -1
    27  			return i, io.EOF
    28  		}
    29  		s.rowIndex++
    30  	}
    31  
    32  	return n, err
    33  }