github.com/parquet-go/parquet-go@v0.21.1-0.20240501160520-b3c3a0c3ed6f/encoding/delta/byte_array_amd64.go (about)

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