github.com/apache/arrow/go/v7@v7.0.1/parquet/internal/utils/bit_writer.go (about) 1 // Licensed to the Apache Software Foundation (ASF) under one 2 // or more contributor license agreements. See the NOTICE file 3 // distributed with this work for additional information 4 // regarding copyright ownership. The ASF licenses this file 5 // to you under the Apache License, Version 2.0 (the 6 // "License"); you may not use this file except in compliance 7 // with the License. You may obtain a copy of the License at 8 // 9 // http://www.apache.org/licenses/LICENSE-2.0 10 // 11 // Unless required by applicable law or agreed to in writing, software 12 // distributed under the License is distributed on an "AS IS" BASIS, 13 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 // See the License for the specific language governing permissions and 15 // limitations under the License. 16 17 package utils 18 19 import ( 20 "encoding/binary" 21 "io" 22 "log" 23 24 "github.com/apache/arrow/go/v7/arrow/bitutil" 25 ) 26 27 // WriterAtBuffer is a convenience struct for providing a WriteAt function 28 // to a byte slice for use with things that want an io.WriterAt 29 type WriterAtBuffer struct { 30 buf []byte 31 } 32 33 // NewWriterAtBuffer returns an object which fulfills the io.WriterAt interface 34 // by taking ownership of the passed in slice. 35 func NewWriterAtBuffer(buf []byte) WriterAtWithLen { 36 return &WriterAtBuffer{buf} 37 } 38 39 // Len returns the length of the underlying byte slice. 40 func (w *WriterAtBuffer) Len() int { 41 return len(w.buf) 42 } 43 44 // WriteAt fulfills the io.WriterAt interface to write len(p) bytes from p 45 // to the underlying byte slice starting at offset off. It returns the number 46 // of bytes written from p (0 <= n <= len(p)) and any error encountered. 47 func (w *WriterAtBuffer) WriteAt(p []byte, off int64) (n int, err error) { 48 if off > int64(len(w.buf)) { 49 return 0, io.ErrUnexpectedEOF 50 } 51 52 n = copy(w.buf[off:], p) 53 if n < len(p) { 54 err = io.ErrUnexpectedEOF 55 } 56 return 57 } 58 59 // WriterAtWithLen is an interface for an io.WriterAt with a Len function 60 type WriterAtWithLen interface { 61 io.WriterAt 62 Len() int 63 } 64 65 // BitWriter is a utility for writing values of specific bit widths to a stream 66 // using a uint64 as a buffer to build up between flushing for efficiency. 67 type BitWriter struct { 68 wr io.WriterAt 69 buffer uint64 70 byteoffset int 71 bitoffset uint 72 raw [8]byte 73 } 74 75 // NewBitWriter initializes a new bit writer to write to the passed in interface 76 // using WriteAt to write the appropriate offsets and values. 77 func NewBitWriter(w io.WriterAt) *BitWriter { 78 return &BitWriter{wr: w} 79 } 80 81 // ReserveBytes reserves the next aligned nbytes, skipping them and returning 82 // the offset to use with WriteAt to write to those reserved bytes. Used for 83 // RLE encoding to fill in the indicators after encoding. 84 func (b *BitWriter) ReserveBytes(nbytes int) int { 85 b.Flush(true) 86 ret := b.byteoffset 87 b.byteoffset += nbytes 88 return ret 89 } 90 91 // WriteAt fulfills the io.WriterAt interface to write len(p) bytes from p 92 // to the underlying byte slice starting at offset off. It returns the number 93 // of bytes written from p (0 <= n <= len(p)) and any error encountered. 94 // This allows writing full bytes directly to the underlying writer. 95 func (b *BitWriter) WriteAt(val []byte, off int64) (int, error) { 96 return b.wr.WriteAt(val, off) 97 } 98 99 // Written returns the number of bytes that have been written to the BitWriter, 100 // not how many bytes have been flushed. Use Flush to ensure that all data is flushed 101 // to the underlying writer. 102 func (b *BitWriter) Written() int { 103 return b.byteoffset + int(bitutil.BytesForBits(int64(b.bitoffset))) 104 } 105 106 // WriteValue writes the value v using nbits to pack it, returning false if it fails 107 // for some reason. 108 func (b *BitWriter) WriteValue(v uint64, nbits uint) error { 109 b.buffer |= v << b.bitoffset 110 b.bitoffset += nbits 111 112 if b.bitoffset >= 64 { 113 binary.LittleEndian.PutUint64(b.raw[:], b.buffer) 114 if _, err := b.wr.WriteAt(b.raw[:], int64(b.byteoffset)); err != nil { 115 return err 116 } 117 b.buffer = 0 118 b.byteoffset += 8 119 b.bitoffset -= 64 120 b.buffer = v >> (nbits - b.bitoffset) 121 } 122 return nil 123 } 124 125 // Flush will flush any buffered data to the underlying writer, pass true if 126 // the next write should be byte-aligned after this flush. 127 func (b *BitWriter) Flush(align bool) { 128 var nbytes int64 129 if b.bitoffset > 0 { 130 nbytes = bitutil.BytesForBits(int64(b.bitoffset)) 131 binary.LittleEndian.PutUint64(b.raw[:], b.buffer) 132 b.wr.WriteAt(b.raw[:nbytes], int64(b.byteoffset)) 133 } 134 135 if align { 136 b.buffer = 0 137 b.byteoffset += int(nbytes) 138 b.bitoffset = 0 139 } 140 } 141 142 // WriteAligned writes the value val as a little endian value in exactly nbytes 143 // byte-aligned to the underlying writer, flushing via Flush(true) before writing nbytes 144 // without buffering. 145 func (b *BitWriter) WriteAligned(val uint64, nbytes int) bool { 146 b.Flush(true) 147 binary.LittleEndian.PutUint64(b.raw[:], val) 148 if _, err := b.wr.WriteAt(b.raw[:nbytes], int64(b.byteoffset)); err != nil { 149 log.Println(err) 150 return false 151 } 152 b.byteoffset += nbytes 153 return true 154 } 155 156 // WriteVlqInt writes v as a vlq encoded integer byte-aligned to the underlying writer 157 // without buffering. 158 func (b *BitWriter) WriteVlqInt(v uint64) bool { 159 b.Flush(true) 160 var buf [binary.MaxVarintLen64]byte 161 nbytes := binary.PutUvarint(buf[:], v) 162 if _, err := b.wr.WriteAt(buf[:nbytes], int64(b.byteoffset)); err != nil { 163 log.Println(err) 164 return false 165 } 166 b.byteoffset += nbytes 167 return true 168 } 169 170 // WriteZigZagVlqInt writes a zigzag encoded integer byte-aligned to the underlying writer 171 // without buffering. 172 func (b *BitWriter) WriteZigZagVlqInt(v int64) bool { 173 return b.WriteVlqInt(uint64((v << 1) ^ (v >> 63))) 174 } 175 176 // Clear resets the writer so that subsequent writes will start from offset 0, 177 // allowing reuse of the underlying buffer and writer. 178 func (b *BitWriter) Clear() { 179 b.byteoffset = 0 180 b.bitoffset = 0 181 b.buffer = 0 182 }