コード例 #1
0
ファイル: scanner.go プロジェクト: TuneLab/parquet-go
// Scan reads the next page inside the column chunk. returns false if no more data pages
// are present or if an error occurred.
func (s *scanner) Scan() bool {
	var (
		header thrift.PageHeader
	)

	if s.err != nil {
		return false
	}

	s.dictionary = nil
	s.dataPage = nil
	s.indexPage = nil

	err := header.Read(s.r)
	if err != nil {
		if strings.HasSuffix(err.Error(), "EOF") { // FIXME: find a better way to detect io.EOF
			s.setErr(io.EOF)
			return false
		}
		s.setErr(fmt.Errorf("column scanner: could not read chunk header: %s", err))
		return false
	}

	// setup reader
	r := io.LimitReader(s.r, int64(header.CompressedPageSize))
	r, err = s.compressionReader(r, &header)
	if err != nil {
		s.setErr(err)
		return false
	}

	// read the page
	if err := s.readPage(r, &header); err != nil {
		s.setErr(err)
		return false
	}

	// check if we consumed all the data from the limit reader as a safe guard
	if n, err := io.Copy(ioutil.Discard, r); err != nil {
		if err == io.EOF {
			return true
		}
		s.setErr(err)
		return false
	} else if n > 0 {
		err := fmt.Errorf("not all the data was consumed for page %s", header.GetType())
		s.setErr(err)
		return false
	}

	return true
}
コード例 #2
0
ファイル: scanner.go プロジェクト: TuneLab/parquet-go
// returns a reader for the right compression
func (s *scanner) compressionReader(r io.Reader, header *thrift.PageHeader) (io.Reader, error) {
	switch s.codec {
	case thrift.CompressionCodec_GZIP:
		r, err := gzip.NewReader(r)
		if err != nil {
			return nil, fmt.Errorf("could not create gzip reader:%s", err)
		}
		b, err := ioutil.ReadAll(r)
		if err != nil {
			return nil, fmt.Errorf("could not read gzip reader:%s", err)
		}
		if err := r.Close(); err != nil {
			log.Println("WARNING error closing gzip reader:%s", err)
		}
		return bytes.NewReader(b), nil

	case thrift.CompressionCodec_LZO:
		// https://github.com/rasky/go-lzo/blob/master/decompress.go#L149			s.r = r
		return nil, fmt.Errorf("NYI")

	case thrift.CompressionCodec_SNAPPY:
		src, err := ioutil.ReadAll(r)
		if err != nil {
			return nil, fmt.Errorf("could not create gzip reader:%s", err)
		}

		out := make([]byte, int(header.GetUncompressedPageSize()))
		out, err = snappy.Decode(out, src)
		if err != nil {
			return nil, fmt.Errorf("could not create gzip reader:%s", err)
		}

		return bytes.NewReader(out), nil

	case thrift.CompressionCodec_UNCOMPRESSED:
		// use the same reader
		return r, nil

	default:
		return nil, fmt.Errorf("unknown compression format %s", s.codec)
	}
}
コード例 #3
0
ファイル: scanner.go プロジェクト: TuneLab/parquet-go
func (s *scanner) readPage(r io.Reader, header *thrift.PageHeader) error {

	switch header.GetType() {

	case thrift.PageType_INDEX_PAGE:
		if !header.IsSetIndexPageHeader() {
			return nil
		}

		s.indexPage = NewIndexPage(header.GetIndexPageHeader())
		// TODO read indexPage
		return nil

	case thrift.PageType_DICTIONARY_PAGE:
		if !header.IsSetDictionaryPageHeader() {
			return fmt.Errorf("bad file format:DictionaryPageHeader flag was not set")
		}
		dictHeader := header.GetDictionaryPageHeader()
		s.dictionary = NewDictionaryPage(s.schema, dictHeader)
		return s.dictionary.Decode(r)

	case thrift.PageType_DATA_PAGE_V2:
		panic("nyi")

	case thrift.PageType_DATA_PAGE:
		s.totalRead += int(header.GetDataPageHeader().GetNumValues())

		if !header.IsSetDataPageHeader() {
			return fmt.Errorf("bad file format: DataPageHeader flag was not set")
		}
		s.dataPage = NewDataPage(s.schema, header.GetDataPageHeader())
		return s.dataPage.ReadAll(r)

	default:
		return fmt.Errorf("unknown PageHeader.PageType: %s", header.GetType())
	}
}