ARROW-12851: [Go][Parquet] Add Golang Parquet encoding package #10379
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,101 @@ | ||
| // Licensed to the Apache Software Foundation (ASF) under one | ||
| // or more contributor license agreements. See the NOTICE file | ||
| // distributed with this work for additional information | ||
| // regarding copyright ownership. The ASF licenses this file | ||
| // to you under the Apache License, Version 2.0 (the | ||
| // "License"); you may not use this file except in compliance | ||
| // with the License. You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| package encoding | ||
|
|
||
| import ( | ||
| "github.com/apache/arrow/go/arrow/bitutil" | ||
| "github.com/apache/arrow/go/parquet" | ||
| "github.com/apache/arrow/go/parquet/internal/utils" | ||
| "golang.org/x/xerrors" | ||
| ) | ||
|
|
||
| // PlainBooleanDecoder is for the Plain Encoding type, there is no | ||
| // dictionary decoding for bools. | ||
| type PlainBooleanDecoder struct { | ||
| decoder | ||
|
|
||
| bitOffset int | ||
| } | ||
|
|
||
| // Type for the PlainBooleanDecoder is parquet.Types.Boolean | ||
| func (PlainBooleanDecoder) Type() parquet.Type { | ||
| return parquet.Types.Boolean | ||
| } | ||
|
|
||
| // Decode fills out with bools decoded from the data at the current point | ||
| // or until we reach the end of the data. | ||
| // | ||
| // Returns the number of values decoded | ||
| func (dec *PlainBooleanDecoder) Decode(out []bool) (int, error) { | ||
| max := utils.MinInt(len(out), dec.nvals) | ||
|
|
||
| unalignedExtract := func(start, end, curBitOffset int) int { | ||
| i := start | ||
| for ; curBitOffset < end; i, curBitOffset = i+1, curBitOffset+1 { | ||
| out[i] = (dec.data[0] & byte(1<<curBitOffset)) != 0 | ||
| } | ||
| return i // return the number of bits we extracted | ||
| } | ||
|
|
||
| // if we aren't at a byte boundary, then get bools until we hit | ||
| // a byte boundary with the bit offset. | ||
| i := 0 | ||
| if dec.bitOffset != 0 { | ||
| i = unalignedExtract(0, 8, dec.bitOffset) | ||
| dec.bitOffset = 0 | ||
| } | ||
|
|
||
| // determine the number of full bytes worth of bits we can decode | ||
| // given the number of values we want to decode. | ||
| bitsRemain := max - i | ||
| batch := bitsRemain / 8 * 8 | ||
| if batch > 0 { // only go in here if there's at least one full byte to decode | ||
| if i > 0 { // skip our data forward if we decoded anything above | ||
| dec.data = dec.data[1:] | ||
| out = out[i:] | ||
| } | ||
| // determine the number of aligned bytes we can grab using SIMD optimized | ||
| // functions to improve performance. | ||
| alignedBytes := bitutil.BytesForBits(int64(batch)) | ||
| utils.BytesToBools(dec.data[:alignedBytes], out) | ||
| dec.data = dec.data[alignedBytes:] | ||
| out = out[alignedBytes*8:] | ||
| } | ||
|
|
||
| // grab any trailing bits now that we've got our aligned bytes. | ||
| dec.bitOffset += unalignedExtract(dec.bitOffset, bitsRemain-batch, dec.bitOffset) | ||
|
|
||
| dec.nvals -= max | ||
| return max, nil | ||
| } | ||
|
|
||
| // DecodeSpaced is like Decode except it expands the values to leave spaces for null | ||
| // as determined by the validBits bitmap. | ||
| func (dec *PlainBooleanDecoder) DecodeSpaced(out []bool, nullCount int, validBits []byte, validBitsOffset int64) (int, error) { | ||
| if nullCount > 0 { | ||
| toRead := len(out) - nullCount | ||
| valuesRead, err := dec.Decode(out[:toRead]) | ||
| if err != nil { | ||
| return 0, err | ||
| } | ||
| if valuesRead != toRead { | ||
| return valuesRead, xerrors.New("parquet: boolean decoder: number of values / definition levels read did not match") | ||
| } | ||
| return spacedExpand(out, nullCount, validBits, validBitsOffset), nil | ||
| } | ||
| return dec.Decode(out) | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,84 @@ | ||
| // Licensed to the Apache Software Foundation (ASF) under one | ||
| // or more contributor license agreements. See the NOTICE file | ||
| // distributed with this work for additional information | ||
| // regarding copyright ownership. The ASF licenses this file | ||
| // to you under the Apache License, Version 2.0 (the | ||
| // "License"); you may not use this file except in compliance | ||
| // with the License. You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| package encoding | ||
|
|
||
| import ( | ||
| "github.com/apache/arrow/go/arrow/bitutil" | ||
| "github.com/apache/arrow/go/parquet" | ||
| "github.com/apache/arrow/go/parquet/internal/utils" | ||
| ) | ||
|
|
||
| const ( | ||
| boolBufSize = 1024 | ||
| boolsInBuf = boolBufSize * 8 | ||
| ) | ||
|
|
||
| // PlainBooleanEncoder encodes bools as a bitmap as per the Plain Encoding | ||
| type PlainBooleanEncoder struct { | ||
| encoder | ||
| bitsBuffer []byte | ||
| wr utils.BitmapWriter | ||
| } | ||
|
|
||
| // Type for the PlainBooleanEncoder is parquet.Types.Boolean | ||
| func (PlainBooleanEncoder) Type() parquet.Type { | ||
| return parquet.Types.Boolean | ||
| } | ||
|
|
||
| // Put encodes the contents of in into the underlying data buffer. | ||
| func (enc *PlainBooleanEncoder) Put(in []bool) { | ||
| if enc.bitsBuffer == nil { | ||
| enc.bitsBuffer = make([]byte, boolBufSize) | ||
| } | ||
| if enc.wr == nil { | ||
| enc.wr = utils.NewBitmapWriter(enc.bitsBuffer, 0, boolsInBuf) | ||
| } | ||
|
|
||
| n := enc.wr.AppendBools(in) | ||
| for n < len(in) { | ||
| enc.wr.Finish() | ||
| enc.append(enc.bitsBuffer) | ||
| enc.wr.Reset(0, boolsInBuf) | ||
| in = in[n:] | ||
| n = enc.wr.AppendBools(in) | ||
| } | ||
| } | ||
|
|
||
| // PutSpaced will use the validBits bitmap to determine which values are nulls | ||
| // and can be left out from the slice, and the encoded without those nulls. | ||
| func (enc *PlainBooleanEncoder) PutSpaced(in []bool, validBits []byte, validBitsOffset int64) { | ||
| bufferOut := make([]bool, len(in)) | ||
| nvalid := spacedCompress(in, bufferOut, validBits, validBitsOffset) | ||
| enc.Put(bufferOut[:nvalid]) | ||
| } | ||
|
|
||
| // EstimatedDataEncodedSize returns the current number of bytes that have | ||
| // been buffered so far | ||
| func (enc *PlainBooleanEncoder) EstimatedDataEncodedSize() int64 { | ||
| return int64(enc.sink.Len() + int(bitutil.BytesForBits(enc.wr.Pos()))) | ||
| } | ||
|
|
||
| // FlushValues returns the buffered data, the responsibility is on the caller | ||
| // to release the buffer memory | ||
| func (enc *PlainBooleanEncoder) FlushValues() Buffer { | ||
| if enc.wr.Pos() > 0 { | ||
| toFlush := int(enc.wr.Pos()) | ||
| enc.append(enc.bitsBuffer[:bitutil.BytesForBits(int64(toFlush))]) | ||
| } | ||
|
|
||
| return enc.sink.Finish() | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,88 @@ | ||
| // Licensed to the Apache Software Foundation (ASF) under one | ||
| // or more contributor license agreements. See the NOTICE file | ||
| // distributed with this work for additional information | ||
| // regarding copyright ownership. The ASF licenses this file | ||
| // to you under the Apache License, Version 2.0 (the | ||
| // "License"); you may not use this file except in compliance | ||
| // with the License. You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| package encoding | ||
|
|
||
| import ( | ||
| "encoding/binary" | ||
|
|
||
| "github.com/apache/arrow/go/parquet" | ||
| "github.com/apache/arrow/go/parquet/internal/utils" | ||
| "golang.org/x/xerrors" | ||
| ) | ||
|
|
||
| // PlainByteArrayDecoder decodes a data chunk for bytearrays according to | ||
| // the plain encoding. The byte arrays will use slices to reference the | ||
| // data rather than copying it. | ||
| // | ||
| // The parquet spec defines Plain encoding for ByteArrays as a 4 byte little | ||
| // endian integer containing the length of the bytearray followed by that many | ||
| // bytes being the raw data of the byte array. | ||
| type PlainByteArrayDecoder struct { | ||
| decoder | ||
| } | ||
|
|
||
| // Type returns parquet.Types.ByteArray for this decoder | ||
| func (PlainByteArrayDecoder) Type() parquet.Type { | ||
| return parquet.Types.ByteArray | ||
| } | ||
|
|
||
| // Decode will populate the slice of bytearrays in full or until the number | ||
| // of values is consumed. | ||
| // | ||
| // Returns the number of values that were decoded. | ||
| func (pbad *PlainByteArrayDecoder) Decode(out []parquet.ByteArray) (int, error) { | ||
| max := utils.MinInt(len(out), pbad.nvals) | ||
|
|
||
| for i := 0; i < max; i++ { | ||
| // there should always be at least four bytes which is the length of the | ||
| // next value in the data. | ||
| if len(pbad.data) < 4 { | ||
| return i, xerrors.New("parquet: eof reading bytearray") | ||
| } | ||
|
|
||
| // the first 4 bytes are a little endian int32 length | ||
| byteLen := int32(binary.LittleEndian.Uint32(pbad.data[:4])) | ||
| if byteLen < 0 { | ||
| return i, xerrors.New("parquet: invalid BYTE_ARRAY value") | ||
| } | ||
|
|
||
| if int64(len(pbad.data)) < int64(byteLen)+4 { | ||
| return i, xerrors.New("parquet: eof reading bytearray") | ||
| } | ||
|
|
||
| out[i] = pbad.data[4 : byteLen+4 : byteLen+4] | ||
| pbad.data = pbad.data[byteLen+4:] | ||
| } | ||
|
|
||
| pbad.nvals -= max | ||
| return max, nil | ||
| } | ||
|
|
||
| // DecodeSpaced is like Decode, but expands the slice out to leave empty values | ||
| // where the validBits bitmap has 0s | ||
| func (pbad *PlainByteArrayDecoder) DecodeSpaced(out []parquet.ByteArray, nullCount int, validBits []byte, validBitsOffset int64) (int, error) { | ||
|
||
| toRead := len(out) - nullCount | ||
| valuesRead, err := pbad.Decode(out[:toRead]) | ||
| if err != nil { | ||
| return valuesRead, err | ||
| } | ||
| if valuesRead != toRead { | ||
| return valuesRead, xerrors.New("parquet: number of values / definition levels read did not match") | ||
| } | ||
|
|
||
| return spacedExpand(out, nullCount, validBits, validBitsOffset), nil | ||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.