github.com/vc42/parquet-go@v0.0.0-20240320194221-1a9adb5f23f5/encoding/delta/byte_array_amd64.go (about)

     1  //go:build !purego
     2  
     3  package delta
     4  
     5  import (
     6  	"github.com/vc42/parquet-go/encoding/plain"
     7  	"golang.org/x/sys/cpu"
     8  )
     9  
    10  //go:noescape
    11  func validatePrefixAndSuffixLengthValuesAVX2(prefix, suffix []int32, maxLength int) (totalPrefixLength, totalSuffixLength int, ok bool)
    12  
    13  func validatePrefixAndSuffixLengthValues(prefix, suffix []int32, maxLength int) (totalPrefixLength, totalSuffixLength int, err error) {
    14  	if cpu.X86.HasAVX2 {
    15  		totalPrefixLength, totalSuffixLength, ok := validatePrefixAndSuffixLengthValuesAVX2(prefix, suffix, maxLength)
    16  		if ok {
    17  			return totalPrefixLength, totalSuffixLength, nil
    18  		}
    19  	}
    20  
    21  	lastValueLength := 0
    22  
    23  	for i := range prefix {
    24  		p := int(prefix[i])
    25  		n := int(suffix[i])
    26  		if p < 0 {
    27  			err = errInvalidNegativePrefixLength(p)
    28  			return
    29  		}
    30  		if n < 0 {
    31  			err = errInvalidNegativeValueLength(n)
    32  			return
    33  		}
    34  		if p > lastValueLength {
    35  			err = errPrefixLengthOutOfBounds(p, lastValueLength)
    36  			return
    37  		}
    38  		totalPrefixLength += p
    39  		totalSuffixLength += n
    40  		lastValueLength = p + n
    41  	}
    42  
    43  	if totalSuffixLength > maxLength {
    44  		err = errValueLengthOutOfBounds(totalSuffixLength, maxLength)
    45  		return
    46  	}
    47  
    48  	return totalPrefixLength, totalSuffixLength, nil
    49  }
    50  
    51  //go:noescape
    52  func decodeByteArrayAVX2(dst, src []byte, prefix, suffix []int32) int
    53  
    54  func decodeByteArray(dst, src []byte, prefix, suffix []int32) ([]byte, error) {
    55  	totalPrefixLength, totalSuffixLength, err := validatePrefixAndSuffixLengthValues(prefix, suffix, len(src))
    56  	if err != nil {
    57  		return dst, err
    58  	}
    59  
    60  	totalLength := plain.ByteArrayLengthSize*len(prefix) + totalPrefixLength + totalSuffixLength
    61  	dst = resizeNoMemclr(dst, totalLength+padding)
    62  
    63  	_ = prefix[:len(suffix)]
    64  	_ = suffix[:len(prefix)]
    65  
    66  	var lastValue []byte
    67  	var i int
    68  	var j int
    69  
    70  	if cpu.X86.HasAVX2 && len(src) > padding {
    71  		k := len(suffix)
    72  		n := 0
    73  
    74  		for k > 0 && n < padding {
    75  			k--
    76  			n += int(suffix[k])
    77  		}
    78  
    79  		if k > 0 && n >= padding {
    80  			i = decodeByteArrayAVX2(dst, src, prefix[:k], suffix[:k])
    81  			j = len(src) - n
    82  			lastValue = dst[i-(int(prefix[k-1])+int(suffix[k-1])):]
    83  			prefix = prefix[k:]
    84  			suffix = suffix[k:]
    85  		}
    86  	}
    87  
    88  	for k := range prefix {
    89  		p := int(prefix[k])
    90  		n := int(suffix[k])
    91  		plain.PutByteArrayLength(dst[i:], p+n)
    92  		i += plain.ByteArrayLengthSize
    93  		k := i
    94  		i += copy(dst[i:], lastValue[:p])
    95  		i += copy(dst[i:], src[j:j+n])
    96  		j += n
    97  		lastValue = dst[k:]
    98  	}
    99  
   100  	return dst[:totalLength], nil
   101  }
   102  
   103  //go:noescape
   104  func decodeFixedLenByteArrayAVX2(dst, src []byte, prefix, suffix []int32) int
   105  
   106  //go:noescape
   107  func decodeFixedLenByteArrayAVX2x128bits(dst, src []byte, prefix, suffix []int32) int
   108  
   109  func decodeFixedLenByteArray(dst, src []byte, size int, prefix, suffix []int32) ([]byte, error) {
   110  	totalPrefixLength, totalSuffixLength, err := validatePrefixAndSuffixLengthValues(prefix, suffix, len(src))
   111  	if err != nil {
   112  		return dst, err
   113  	}
   114  
   115  	totalLength := totalPrefixLength + totalSuffixLength
   116  	dst = resizeNoMemclr(dst, totalLength+padding)
   117  
   118  	_ = prefix[:len(suffix)]
   119  	_ = suffix[:len(prefix)]
   120  
   121  	var lastValue []byte
   122  	var i int
   123  	var j int
   124  
   125  	if cpu.X86.HasAVX2 && len(src) > padding {
   126  		k := len(suffix)
   127  		n := 0
   128  
   129  		for k > 0 && n < padding {
   130  			k--
   131  			n += int(suffix[k])
   132  		}
   133  
   134  		if k > 0 && n >= padding {
   135  			if size == 16 {
   136  				i = decodeFixedLenByteArrayAVX2x128bits(dst, src, prefix[:k], suffix[:k])
   137  			} else {
   138  				i = decodeFixedLenByteArrayAVX2(dst, src, prefix[:k], suffix[:k])
   139  			}
   140  			j = len(src) - n
   141  			prefix = prefix[k:]
   142  			suffix = suffix[k:]
   143  			if i >= size {
   144  				lastValue = dst[i-size:]
   145  			}
   146  		}
   147  	}
   148  
   149  	for k := range prefix {
   150  		p := int(prefix[k])
   151  		n := int(suffix[k])
   152  		k := i
   153  		i += copy(dst[i:], lastValue[:p])
   154  		i += copy(dst[i:], src[j:j+n])
   155  		j += n
   156  		lastValue = dst[k:]
   157  	}
   158  
   159  	return dst[:totalLength], nil
   160  }