1package vnext23import (4 "bytes"5 "encoding/binary"6 "fmt"7 "hash/crc32"8 "math"9)1011const (12 blockMagic = "ATLASPK\x00"13 blockFraming = uint16(1)14 blockHeaderSize = 8015 directorySize = 6416 columnFlagOptional = uint16(1)17 maxBlockRows = 5_000_00018 maxBlockColumns = 4_09619 maxBlockValues = 10_000_00020)2122// Block is a decoded typed block. SchemaHash identifies the embedded schema23// used to write it but does not decide whether a reader may open it.24type Block struct {25 SchemaHash [32]byte26 Table Table27}2829type encodedColumn struct {30 field Field31 nulls []byte32 data []byte33 crc uint3234 nullAt uint6435 dataAt uint6436}3738// EncodeBlock writes a deterministic, self-framed hybrid column block.39func EncodeBlock(schemaHash [32]byte, table Table) ([]byte, error) {40 table = canonicalTable(table)41 if table.Rows > maxBlockRows || len(table.Columns) > maxBlockColumns || table.Rows > 0 && len(table.Columns) > maxBlockValues/table.Rows {42 return nil, fmt.Errorf("table has %d rows; maximum is %d", table.Rows, maxBlockRows)43 }44 if err := table.validate(); err != nil {45 return nil, err46 }47 columns, err := encodeColumns(table)48 if err != nil {49 return nil, err50 }51 payloadAt := uint64(blockHeaderSize + len(columns)*directorySize)52 assignOffsets(columns, payloadAt)53 out := make([]byte, payloadAt)54 writeBlockHeader(out, schemaHash, table, payloadAt)55 for index, column := range columns {56 writeDirectory(out[blockHeaderSize+index*directorySize:], column)57 out = append(out, column.nulls...)58 out = append(out, column.data...)59 }60 return out, nil61}6263func encodeColumns(table Table) ([]encodedColumn, error) {64 columns := make([]encodedColumn, len(table.Columns))65 for index, column := range table.Columns {66 nulls := presenceBitmap(column.Values)67 data, err := encodeValues(column.Field.Kind, column.Values)68 if err != nil {69 return nil, fmt.Errorf("encode column %s: %w", column.Field.ID, err)70 }71 checksum := crc32.NewIEEE()72 _, _ = checksum.Write(nulls)73 _, _ = checksum.Write(data)74 columns[index] = encodedColumn{field: column.Field, nulls: nulls, data: data, crc: checksum.Sum32()}75 }76 return columns, nil77}7879func assignOffsets(columns []encodedColumn, at uint64) {80 for index := range columns {81 columns[index].nullAt = at82 at += uint64(len(columns[index].nulls))83 columns[index].dataAt = at84 at += uint64(len(columns[index].data))85 }86}8788func writeBlockHeader(out []byte, schemaHash [32]byte, table Table, payloadAt uint64) {89 copy(out[:8], blockMagic)90 binary.LittleEndian.PutUint16(out[8:10], blockFraming)91 binary.LittleEndian.PutUint32(out[12:16], uint32(table.Rows))92 binary.LittleEndian.PutUint32(out[16:20], uint32(len(table.Columns)))93 copy(out[20:52], schemaHash[:])94 copy(out[52:68], table.TypeID[:])95 binary.LittleEndian.PutUint32(out[68:72], directorySize)96 binary.LittleEndian.PutUint64(out[72:80], payloadAt)97}9899func writeDirectory(out []byte, column encodedColumn) {100 copy(out[:16], column.field.ID[:])101 out[16] = byte(column.field.Kind)102 out[17] = byte(column.field.Kind)103 if column.field.Optional {104 binary.LittleEndian.PutUint16(out[18:20], columnFlagOptional)105 }106 binary.LittleEndian.PutUint64(out[20:28], column.nullAt)107 binary.LittleEndian.PutUint64(out[28:36], uint64(len(column.nulls)))108 binary.LittleEndian.PutUint64(out[36:44], column.dataAt)109 binary.LittleEndian.PutUint64(out[44:52], uint64(len(column.data)))110 binary.LittleEndian.PutUint32(out[52:56], column.crc)111}112113func presenceBitmap(values []Value) []byte {114 bitmap := make([]byte, (len(values)+7)/8)115 for row, value := range values {116 if value.Kind != KindInvalid {117 bitmap[row/8] |= 1 << (row % 8)118 }119 }120 return bitmap121}122123func encodeValues(kind Kind, values []Value) ([]byte, error) {124 switch kind {125 case KindBool:126 return encodeBools(values), nil127 case KindInt64:128 return encodeInt64s(values), nil129 case KindFloat64:130 return encodeFloat64s(values), nil131 case KindString, KindBytes:132 return encodeVariable(kind, values)133 case KindID:134 return encodeIDs(values), nil135 default:136 return nil, fmt.Errorf("unsupported wire kind %d", kind)137 }138}139140func encodeBools(values []Value) []byte {141 out := make([]byte, len(values))142 for index, value := range values {143 if value.Kind != KindInvalid && value.Bool {144 out[index] = 1145 }146 }147 return out148}149150func encodeInt64s(values []Value) []byte {151 out := make([]byte, len(values)*8)152 for index, value := range values {153 binary.LittleEndian.PutUint64(out[index*8:], uint64(value.Int64))154 }155 return out156}157158func encodeFloat64s(values []Value) []byte {159 out := make([]byte, len(values)*8)160 for index, value := range values {161 binary.LittleEndian.PutUint64(out[index*8:], math.Float64bits(value.Float64))162 }163 return out164}165166func encodeIDs(values []Value) []byte {167 out := make([]byte, len(values)*16)168 for index, value := range values {169 copy(out[index*16:], value.ID[:])170 }171 return out172}173174func encodeVariable(kind Kind, values []Value) ([]byte, error) {175 var payload bytes.Buffer176 offsets := make([]uint32, len(values)+1)177 for index, value := range values {178 offsets[index] = uint32(payload.Len())179 if value.Kind == KindInvalid {180 continue181 }182 if kind == KindString {183 _, _ = payload.WriteString(value.String)184 } else {185 _, _ = payload.Write(value.Bytes)186 }187 if uint64(payload.Len()) > math.MaxUint32 {188 return nil, fmt.Errorf("variable column exceeds 4 GiB")189 }190 }191 offsets[len(values)] = uint32(payload.Len())192 out := make([]byte, len(offsets)*4+payload.Len())193 for index, offset := range offsets {194 binary.LittleEndian.PutUint32(out[index*4:], offset)195 }196 copy(out[len(offsets)*4:], payload.Bytes())197 return out, nil198}199200// DecodeBlock validates and opens one block without consulting an application201// schema. Unknown fields therefore remain ordinary columns.202func DecodeBlock(data []byte) (Block, error) {203 header, err := readBlockHeader(data)204 if err != nil {205 return Block{}, err206 }207 table := Table{TypeID: header.typeID, Rows: header.rows, Columns: make([]Column, header.columns)}208 for index := 0; index < header.columns; index++ {209 at := blockHeaderSize + index*directorySize210 column, err := decodeColumn(data, data[at:at+directorySize], header.rows)211 if err != nil {212 return Block{}, fmt.Errorf("decode column %d: %w", index, err)213 }214 table.Columns[index] = column215 }216 if err := table.validate(); err != nil {217 return Block{}, err218 }219 return Block{SchemaHash: header.schemaHash, Table: canonicalTable(table)}, nil220}221222type decodedHeader struct {223 rows int224 columns int225 schemaHash [32]byte226 typeID ID227}228229func readBlockHeader(data []byte) (decodedHeader, error) {230 if len(data) < blockHeaderSize || string(data[:8]) != blockMagic {231 return decodedHeader{}, fmt.Errorf("typed block has no Atlas magic")232 }233 if framing := binary.LittleEndian.Uint16(data[8:10]); framing != blockFraming {234 return decodedHeader{}, fmt.Errorf("typed block framing %d is not supported", framing)235 }236 columns := int(binary.LittleEndian.Uint32(data[16:20]))237 directoryEnd := blockHeaderSize + columns*directorySize238 if columns < 0 || directoryEnd < blockHeaderSize || directoryEnd > len(data) {239 return decodedHeader{}, fmt.Errorf("typed block directory is truncated")240 }241 if binary.LittleEndian.Uint32(data[68:72]) != directorySize || binary.LittleEndian.Uint64(data[72:80]) != uint64(directoryEnd) {242 return decodedHeader{}, fmt.Errorf("typed block directory framing is invalid")243 }244 header := decodedHeader{rows: int(binary.LittleEndian.Uint32(data[12:16])), columns: columns}245 if header.rows > maxBlockRows || header.columns > maxBlockColumns || header.columns == 0 && header.rows != 0 ||246 header.rows > 0 && header.columns > maxBlockValues/header.rows {247 return decodedHeader{}, fmt.Errorf("typed block row count is invalid")248 }249 copy(header.schemaHash[:], data[20:52])250 copy(header.typeID[:], data[52:68])251 return header, nil252}253254func decodeColumn(data, directory []byte, rows int) (Column, error) {255 var fieldID ID256 copy(fieldID[:], directory[:16])257 kind := Kind(directory[16])258 if !kind.valid() || directory[17] != byte(kind) {259 return Column{}, fmt.Errorf("column has unsupported encoding")260 }261 optional := binary.LittleEndian.Uint16(directory[18:20])&columnFlagOptional != 0262 nulls, err := blockSlice(data, directory[20:28], directory[28:36])263 if err != nil || len(nulls) != (rows+7)/8 {264 return Column{}, fmt.Errorf("column presence bitmap is invalid")265 }266 payload, err := blockSlice(data, directory[36:44], directory[44:52])267 if err != nil {268 return Column{}, fmt.Errorf("column payload is invalid")269 }270 if checksum(nulls, payload) != binary.LittleEndian.Uint32(directory[52:56]) {271 return Column{}, fmt.Errorf("column checksum does not match")272 }273 values, err := decodeValues(kind, rows, nulls, payload)274 if err != nil {275 return Column{}, err276 }277 return Column{Field: Field{ID: fieldID, Kind: kind, Optional: optional}, Values: values}, nil278}279280func blockSlice(data, offsetBytes, lengthBytes []byte) ([]byte, error) {281 offset := binary.LittleEndian.Uint64(offsetBytes)282 length := binary.LittleEndian.Uint64(lengthBytes)283 if offset > uint64(len(data)) || length > uint64(len(data))-offset {284 return nil, fmt.Errorf("range exceeds block")285 }286 return data[offset : offset+length], nil287}288289func checksum(parts ...[]byte) uint32 {290 hash := crc32.NewIEEE()291 for _, part := range parts {292 _, _ = hash.Write(part)293 }294 return hash.Sum32()295}296297func decodeValues(kind Kind, rows int, present, payload []byte) ([]Value, error) {298 width := 0299 switch kind {300 case KindBool:301 width = 1302 case KindInt64, KindFloat64:303 width = 8304 case KindID:305 width = 16306 case KindString, KindBytes:307 if rows == math.MaxInt || len(payload) < (rows+1)*4 {308 return nil, fmt.Errorf("variable column offsets are truncated")309 }310 default:311 return nil, fmt.Errorf("unsupported wire kind %d", kind)312 }313 if width != 0 && (rows > math.MaxInt/width || len(payload) != rows*width) {314 return nil, fmt.Errorf("fixed-width column has wrong length")315 }316 values := make([]Value, rows)317 switch kind {318 case KindBool:319 return decodeBools(values, present, payload)320 case KindInt64:321 return decodeInt64s(values, present, payload)322 case KindFloat64:323 return decodeFloat64s(values, present, payload)324 case KindString, KindBytes:325 return decodeVariable(kind, values, present, payload)326 case KindID:327 return decodeIDs(values, present, payload)328 default:329 return nil, fmt.Errorf("unsupported wire kind %d", kind)330 }331}332333func isPresent(bitmap []byte, row int) bool { return bitmap[row/8]&(1<<(row%8)) != 0 }334335func decodeBools(values []Value, present, payload []byte) ([]Value, error) {336 if len(payload) != len(values) {337 return nil, fmt.Errorf("boolean column has wrong length")338 }339 for row := range values {340 if isPresent(present, row) {341 values[row] = BoolValue(payload[row] != 0)342 }343 }344 return values, nil345}346347func decodeInt64s(values []Value, present, payload []byte) ([]Value, error) {348 if len(payload) != len(values)*8 {349 return nil, fmt.Errorf("int64 column has wrong length")350 }351 for row := range values {352 if isPresent(present, row) {353 values[row] = Int64Value(int64(binary.LittleEndian.Uint64(payload[row*8:])))354 }355 }356 return values, nil357}358359func decodeFloat64s(values []Value, present, payload []byte) ([]Value, error) {360 if len(payload) != len(values)*8 {361 return nil, fmt.Errorf("float64 column has wrong length")362 }363 for row := range values {364 if isPresent(present, row) {365 values[row] = Float64Value(math.Float64frombits(binary.LittleEndian.Uint64(payload[row*8:])))366 }367 }368 return values, nil369}370371func decodeIDs(values []Value, present, payload []byte) ([]Value, error) {372 if len(payload) != len(values)*16 {373 return nil, fmt.Errorf("ID column has wrong length")374 }375 for row := range values {376 if isPresent(present, row) {377 var id ID378 copy(id[:], payload[row*16:(row+1)*16])379 values[row] = IDValue(id)380 }381 }382 return values, nil383}384385func decodeVariable(kind Kind, values []Value, present, payload []byte) ([]Value, error) {386 offsetBytes := (len(values) + 1) * 4387 if len(payload) < offsetBytes {388 return nil, fmt.Errorf("variable column offsets are truncated")389 }390 content := payload[offsetBytes:]391 previous := uint32(0)392 for row := range values {393 start := binary.LittleEndian.Uint32(payload[row*4:])394 end := binary.LittleEndian.Uint32(payload[(row+1)*4:])395 if start < previous || end < start || uint64(end) > uint64(len(content)) {396 return nil, fmt.Errorf("variable column offsets are invalid at row %d", row)397 }398 previous = end399 if !isPresent(present, row) {400 continue401 }402 if kind == KindString {403 values[row] = StringValue(string(content[start:end]))404 } else {405 values[row] = BytesValue(content[start:end])406 }407 }408 return values, nil409}