Skip to content

Commit d38b3b1

Browse files
daniel-adam-tfsdaniel-adam-tfs
committed
Fix byte-stream split decoding on big endian machines
1 parent 8f338f1 commit d38b3b1

4 files changed

Lines changed: 116 additions & 71 deletions

File tree

‎parquet/internal/encoding/byte_stream_split_big_endian.go‎

Lines changed: 47 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222
"fmt"
2323
"unsafe"
2424

25+
"github.com/apache/arrow-go/v18/parquet"
2526
"github.com/apache/arrow-go/v18/parquet/internal/debug"
2627
)
2728

@@ -41,7 +42,8 @@ func decodeByteStreamSplitBatchWidth4InByteOrderDefault(data []byte, nValues, st
4142
out = out[:width*nValues]
4243
out32 := unsafe.Slice((*uint32)(unsafe.Pointer(&out[0])), nValues)
4344
for i := range nValues {
44-
out32[i] = uint32(s3[i]) | uint32(s2[i])<<8 | uint32(s1[i])<<16 | uint32(s0[i])<<24
45+
// Big-endian machine: put s0 as MSB, s3 as LSB
46+
out32[i] = uint32(s3[i])<<24 | uint32(s2[i])<<16 | uint32(s1[i])<<8 | uint32(s0[i])
4547
}
4648
}
4749

@@ -64,7 +66,49 @@ func decodeByteStreamSplitBatchWidth8InByteOrderDefault(data []byte, nValues, st
6466
out = out[:width*nValues]
6567
out64 := unsafe.Slice((*uint64)(unsafe.Pointer(&out[0])), nValues)
6668
for i := range nValues {
67-
out64[i] = uint64(s7[i]) | uint64(s6[i])<<8 | uint64(s5[i])<<16 | uint64(s4[i])<<24 |
68-
uint64(s3[i])<<32 | uint64(s2[i])<<40 | uint64(s1[i])<<48 | uint64(s0[i])<<56
69+
// Big-endian machine: put s0 as MSB, s7 as LSB
70+
out64[i] = uint64(s7[i])<<56 | uint64(s6[i])<<48 | uint64(s5[i])<<40 | uint64(s4[i])<<32 |
71+
uint64(s3[i])<<24 | uint64(s2[i])<<16 | uint64(s1[i])<<8 | uint64(s0[i])
72+
}
73+
}
74+
75+
// decodeByteStreamSplitBatchFLBAWidth2 decodes the batch of nValues FixedLenByteArrays of length 2 provided by 'data',
76+
// into the output slice 'out' using BYTE_STREAM_SPLIT encoding.
77+
// 'out' must have space for at least nValues slices.
78+
func decodeByteStreamSplitBatchFLBAWidth2(data []byte, nValues, stride int, out []parquet.FixedLenByteArray) {
79+
debug.Assert(len(out) >= nValues, fmt.Sprintf("not enough space in output slice for decoding, out: %d values, data: %d values", len(out), nValues))
80+
for element := 0; element < nValues; element++ {
81+
out[element][0] = data[element]
82+
out[element][1] = data[stride+element]
83+
}
84+
}
85+
86+
// decodeByteStreamSplitBatchFLBAWidth4 decodes the batch of nValues FixedLenByteArrays of length 4 provided by 'data',
87+
// into the output slice 'out' using BYTE_STREAM_SPLIT encoding.
88+
// 'out' must have space for at least nValues slices.
89+
func decodeByteStreamSplitBatchFLBAWidth4(data []byte, nValues, stride int, out []parquet.FixedLenByteArray) {
90+
debug.Assert(len(out) >= nValues, fmt.Sprintf("not enough space in output slice for decoding, out: %d values, data: %d values", len(out), nValues))
91+
for element := 0; element < nValues; element++ {
92+
out[element][0] = data[element]
93+
out[element][1] = data[stride+element]
94+
out[element][2] = data[stride*2+element]
95+
out[element][3] = data[stride*3+element]
96+
}
97+
}
98+
99+
// decodeByteStreamSplitBatchFLBAWidth8 decodes the batch of nValues FixedLenByteArrays of length 8 provided by 'data',
100+
// into the output slice 'out' using BYTE_STREAM_SPLIT encoding.
101+
// 'out' must have space for at least nValues slices.
102+
func decodeByteStreamSplitBatchFLBAWidth8(data []byte, nValues, stride int, out []parquet.FixedLenByteArray) {
103+
debug.Assert(len(out) >= nValues, fmt.Sprintf("not enough space in output slice for decoding, out: %d values, data: %d values", len(out), nValues))
104+
for element := 0; element < nValues; element++ {
105+
out[element][0] = data[element]
106+
out[element][1] = data[stride+element]
107+
out[element][2] = data[stride*2+element]
108+
out[element][3] = data[stride*3+element]
109+
out[element][4] = data[stride*4+element]
110+
out[element][5] = data[stride*5+element]
111+
out[element][6] = data[stride*6+element]
112+
out[element][7] = data[stride*7+element]
69113
}
70114
}

‎parquet/internal/encoding/byte_stream_split_decode.go‎

Lines changed: 0 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@ package encoding
1818

1919
import (
2020
"fmt"
21-
"unsafe"
2221

2322
"github.com/apache/arrow-go/v18/parquet"
2423
"github.com/apache/arrow-go/v18/parquet/internal/debug"
@@ -42,57 +41,3 @@ func decodeByteStreamSplitBatchFLBA(data []byte, nValues, stride, width int, out
4241
}
4342
}
4443
}
45-
46-
// decodeByteStreamSplitBatchFLBAWidth2 decodes the batch of nValues FixedLenByteArrays of length 2 provided by 'data',
47-
// into the output slice 'out' using BYTE_STREAM_SPLIT encoding.
48-
// 'out' must have space for at least nValues slices.
49-
func decodeByteStreamSplitBatchFLBAWidth2(data []byte, nValues, stride int, out []parquet.FixedLenByteArray) {
50-
const width = 2
51-
debug.Assert(len(out) >= nValues, fmt.Sprintf("not enough space in output slice for decoding, out: %d values, data: %d values", len(out), nValues))
52-
debug.Assert(len(data) >= stride+nValues, fmt.Sprintf("not enough data for decoding, data: %d bytes, expected at least: %d bytes", len(data), stride+nValues))
53-
s0 := data[:nValues]
54-
s1 := data[stride : stride+nValues]
55-
for i := range nValues {
56-
out16 := (*uint16)(unsafe.Pointer(&out[i][0]))
57-
*out16 = uint16(s0[i]) | uint16(s1[i])<<8
58-
}
59-
}
60-
61-
// decodeByteStreamSplitBatchFLBAWidth4 decodes the batch of nValues FixedLenByteArrays of length 4 provided by 'data',
62-
// into the output slice 'out' using BYTE_STREAM_SPLIT encoding.
63-
// 'out' must have space for at least nValues slices.
64-
func decodeByteStreamSplitBatchFLBAWidth4(data []byte, nValues, stride int, out []parquet.FixedLenByteArray) {
65-
const width = 4
66-
debug.Assert(len(out) >= nValues, fmt.Sprintf("not enough space in output slice for decoding, out: %d values, data: %d values", len(out), nValues))
67-
debug.Assert(len(data) >= 3*stride+nValues, fmt.Sprintf("not enough data for decoding, data: %d bytes, expected at least: %d bytes", len(data), 3*stride+nValues))
68-
s0 := data[:nValues]
69-
s1 := data[stride : stride+nValues]
70-
s2 := data[stride*2 : stride*2+nValues]
71-
s3 := data[stride*3 : stride*3+nValues]
72-
for i := range nValues {
73-
out32 := (*uint32)(unsafe.Pointer(&out[i][0]))
74-
*out32 = uint32(s0[i]) | uint32(s1[i])<<8 | uint32(s2[i])<<16 | uint32(s3[i])<<24
75-
}
76-
}
77-
78-
// decodeByteStreamSplitBatchFLBAWidth8 decodes the batch of nValues FixedLenByteArrays of length 8 provided by 'data',
79-
// into the output slice 'out' using BYTE_STREAM_SPLIT encoding.
80-
// 'out' must have space for at least nValues slices.
81-
func decodeByteStreamSplitBatchFLBAWidth8(data []byte, nValues, stride int, out []parquet.FixedLenByteArray) {
82-
const width = 8
83-
debug.Assert(len(out) >= nValues, fmt.Sprintf("not enough space in output slice for decoding, out: %d values, data: %d values", len(out), nValues))
84-
debug.Assert(len(data) >= 7*stride+nValues, fmt.Sprintf("not enough data for decoding, data: %d bytes, expected at least: %d bytes", len(data), 7*stride+nValues))
85-
s0 := data[:nValues]
86-
s1 := data[stride : stride+nValues]
87-
s2 := data[stride*2 : stride*2+nValues]
88-
s3 := data[stride*3 : stride*3+nValues]
89-
s4 := data[stride*4 : stride*4+nValues]
90-
s5 := data[stride*5 : stride*5+nValues]
91-
s6 := data[stride*6 : stride*6+nValues]
92-
s7 := data[stride*7 : stride*7+nValues]
93-
for i := range nValues {
94-
out64 := (*uint64)(unsafe.Pointer(&out[i][0]))
95-
*out64 = uint64(s0[i]) | uint64(s1[i])<<8 | uint64(s2[i])<<16 | uint64(s3[i])<<24 |
96-
uint64(s4[i])<<32 | uint64(s5[i])<<40 | uint64(s6[i])<<48 | uint64(s7[i])<<56
97-
}
98-
}

‎parquet/internal/encoding/byte_stream_split_decode_test.go‎

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,9 @@ import (
2020
"bytes"
2121
"fmt"
2222
"testing"
23+
"unsafe"
2324

25+
"github.com/apache/arrow-go/v18/internal/utils"
2426
"github.com/apache/arrow-go/v18/parquet"
2527
)
2628

@@ -54,12 +56,12 @@ func TestDecodeByteStreamSplitWidth4(t *testing.T) {
5456
out := make([]byte, width*nValues)
5557
t.Run(fmt.Sprintf("nValues=%d", nValues), func(t *testing.T) {
5658
decodeByteStreamSplitBatchWidth4InByteOrder(data, nValues, stride, out)
57-
if !bytes.Equal(out, expected) {
58-
for i := 0; i < len(expected); i++ {
59-
if out[i] != expected[i] {
60-
t.Errorf("First mismatch at index %d: got %d, want %d", i, out[i], expected[i])
61-
break
62-
}
59+
for i := 0; i < nValues; i++ {
60+
got := utils.ToLE(*(*uint32)(unsafe.Pointer(&out[i*4])))
61+
want := *(*uint32)(unsafe.Pointer(&expected[i*4]))
62+
if got != want {
63+
t.Errorf("Mismatch at index %d: got %08x, want %08x", i, got, want)
64+
break
6365
}
6466
}
6567
})
@@ -130,13 +132,12 @@ func TestDecodeByteStreamSplitWidth8(t *testing.T) {
130132
t.Run(fmt.Sprintf("nValues=%d", nValues), func(t *testing.T) {
131133
out := make([]byte, width*nValues)
132134
decodeByteStreamSplitBatchWidth8InByteOrder(data, nValues, stride, out)
133-
if !bytes.Equal(out, expected) {
134-
t.Errorf("Reference implementation produced incorrect output")
135-
for i := 0; i < len(expected); i++ {
136-
if out[i] != expected[i] {
137-
t.Errorf("First mismatch at index %d: got %d, want %d", i, out[i], expected[i])
138-
break
139-
}
135+
for i := 0; i < nValues; i++ {
136+
got := utils.ToLE(*(*uint64)(unsafe.Pointer(&out[i*8])))
137+
want := *(*uint64)(unsafe.Pointer(&expected[i*8]))
138+
if got != want {
139+
t.Errorf("Mismatch at index %d: got %016x, want %016x", i, got, want)
140+
break
140141
}
141142
}
142143
})

‎parquet/internal/encoding/byte_stream_split_little_endian.go‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222
"fmt"
2323
"unsafe"
2424

25+
"github.com/apache/arrow-go/v18/parquet"
2526
"github.com/apache/arrow-go/v18/parquet/internal/debug"
2627
)
2728

@@ -68,3 +69,57 @@ func decodeByteStreamSplitBatchWidth8InByteOrderDefault(data []byte, nValues, st
6869
uint64(s4[i])<<32 | uint64(s5[i])<<40 | uint64(s6[i])<<48 | uint64(s7[i])<<56
6970
}
7071
}
72+
73+
// decodeByteStreamSplitBatchFLBAWidth2 decodes the batch of nValues FixedLenByteArrays of length 2 provided by 'data',
74+
// into the output slice 'out' using BYTE_STREAM_SPLIT encoding.
75+
// 'out' must have space for at least nValues slices.
76+
func decodeByteStreamSplitBatchFLBAWidth2(data []byte, nValues, stride int, out []parquet.FixedLenByteArray) {
77+
const width = 2
78+
debug.Assert(len(out) >= nValues, fmt.Sprintf("not enough space in output slice for decoding, out: %d values, data: %d values", len(out), nValues))
79+
debug.Assert(len(data) >= stride+nValues, fmt.Sprintf("not enough data for decoding, data: %d bytes, expected at least: %d bytes", len(data), stride+nValues))
80+
s0 := data[:nValues]
81+
s1 := data[stride : stride+nValues]
82+
for i := range nValues {
83+
out16 := (*uint16)(unsafe.Pointer(&out[i][0]))
84+
*out16 = uint16(s0[i]) | uint16(s1[i])<<8
85+
}
86+
}
87+
88+
// decodeByteStreamSplitBatchFLBAWidth4 decodes the batch of nValues FixedLenByteArrays of length 4 provided by 'data',
89+
// into the output slice 'out' using BYTE_STREAM_SPLIT encoding.
90+
// 'out' must have space for at least nValues slices.
91+
func decodeByteStreamSplitBatchFLBAWidth4(data []byte, nValues, stride int, out []parquet.FixedLenByteArray) {
92+
const width = 4
93+
debug.Assert(len(out) >= nValues, fmt.Sprintf("not enough space in output slice for decoding, out: %d values, data: %d values", len(out), nValues))
94+
debug.Assert(len(data) >= 3*stride+nValues, fmt.Sprintf("not enough data for decoding, data: %d bytes, expected at least: %d bytes", len(data), 3*stride+nValues))
95+
s0 := data[:nValues]
96+
s1 := data[stride : stride+nValues]
97+
s2 := data[stride*2 : stride*2+nValues]
98+
s3 := data[stride*3 : stride*3+nValues]
99+
for i := range nValues {
100+
out32 := (*uint32)(unsafe.Pointer(&out[i][0]))
101+
*out32 = uint32(s0[i]) | uint32(s1[i])<<8 | uint32(s2[i])<<16 | uint32(s3[i])<<24
102+
}
103+
}
104+
105+
// decodeByteStreamSplitBatchFLBAWidth8 decodes the batch of nValues FixedLenByteArrays of length 8 provided by 'data',
106+
// into the output slice 'out' using BYTE_STREAM_SPLIT encoding.
107+
// 'out' must have space for at least nValues slices.
108+
func decodeByteStreamSplitBatchFLBAWidth8(data []byte, nValues, stride int, out []parquet.FixedLenByteArray) {
109+
const width = 8
110+
debug.Assert(len(out) >= nValues, fmt.Sprintf("not enough space in output slice for decoding, out: %d values, data: %d values", len(out), nValues))
111+
debug.Assert(len(data) >= 7*stride+nValues, fmt.Sprintf("not enough data for decoding, data: %d bytes, expected at least: %d bytes", len(data), 7*stride+nValues))
112+
s0 := data[:nValues]
113+
s1 := data[stride : stride+nValues]
114+
s2 := data[stride*2 : stride*2+nValues]
115+
s3 := data[stride*3 : stride*3+nValues]
116+
s4 := data[stride*4 : stride*4+nValues]
117+
s5 := data[stride*5 : stride*5+nValues]
118+
s6 := data[stride*6 : stride*6+nValues]
119+
s7 := data[stride*7 : stride*7+nValues]
120+
for i := range nValues {
121+
out64 := (*uint64)(unsafe.Pointer(&out[i][0]))
122+
*out64 = uint64(s0[i]) | uint64(s1[i])<<8 | uint64(s2[i])<<16 | uint64(s3[i])<<24 |
123+
uint64(s4[i])<<32 | uint64(s5[i])<<40 | uint64(s6[i])<<48 | uint64(s7[i])<<56
124+
}
125+
}

0 commit comments

Comments
 (0)