github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/store/val/tuple_builder.go (about)

     1  // Copyright 2021 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package val
    16  
    17  import (
    18  	"time"
    19  
    20  	"github.com/dolthub/go-mysql-server/sql/analyzer/analyzererrors"
    21  	"github.com/shopspring/decimal"
    22  
    23  	"github.com/dolthub/dolt/go/store/hash"
    24  	"github.com/dolthub/dolt/go/store/pool"
    25  )
    26  
    27  const (
    28  	builderBufferSize = 128
    29  )
    30  
    31  // OrdinalMapping is a mapping from one field ordering to another.
    32  // It's used to construct index tuples from another index's tuples.
    33  type OrdinalMapping []int
    34  
    35  // NewIdentityOrdinalMapping returns a new OrdinalMapping that maps every ordinal to itself.
    36  func NewIdentityOrdinalMapping(size int) OrdinalMapping {
    37  	newMapping := make(OrdinalMapping, size)
    38  	for i := 0; i < size; i++ {
    39  		newMapping[i] = i
    40  	}
    41  	return newMapping
    42  }
    43  
    44  // MapOrdinal returns the ordinal of the field in the source tuple that maps to the |to| ordinal in the destination tuple.
    45  func (om OrdinalMapping) MapOrdinal(to int) (from int) {
    46  	from = om[to]
    47  	return
    48  }
    49  
    50  // IsIdentityMapping returns true if this mapping is the identity mapping (i.e. every position is mapped
    51  // to the same position and no columns are reordered).
    52  func (om OrdinalMapping) IsIdentityMapping() bool {
    53  	for i, mapping := range om {
    54  		if i != mapping {
    55  			return false
    56  		}
    57  	}
    58  	return true
    59  }
    60  
    61  type TupleBuilder struct {
    62  	Desc   TupleDesc
    63  	fields [][]byte
    64  	buf    []byte
    65  	pos    ByteSize
    66  }
    67  
    68  func NewTupleBuilder(desc TupleDesc) *TupleBuilder {
    69  	return &TupleBuilder{
    70  		Desc:   desc,
    71  		fields: make([][]byte, len(desc.Types)),
    72  		buf:    make([]byte, builderBufferSize),
    73  	}
    74  }
    75  
    76  // Build materializes a Tuple from the fields written to the TupleBuilder.
    77  func (tb *TupleBuilder) Build(pool pool.BuffPool) (tup Tuple) {
    78  	for i, typ := range tb.Desc.Types {
    79  		if !typ.Nullable && tb.fields[i] == nil {
    80  			panic("cannot write NULL to non-NULL field")
    81  		}
    82  	}
    83  	return tb.BuildPermissive(pool)
    84  }
    85  
    86  // BuildPermissive materializes a Tuple from the fields
    87  // written to the TupleBuilder without validating nullability.
    88  func (tb *TupleBuilder) BuildPermissive(pool pool.BuffPool) (tup Tuple) {
    89  	values := tb.fields[:tb.Desc.Count()]
    90  	tup = NewTuple(pool, values...)
    91  	tb.Recycle()
    92  	return
    93  }
    94  
    95  // BuildPrefix materializes a prefix Tuple from the first |k| fields written to the TupleBuilder.
    96  func (tb *TupleBuilder) BuildPrefix(pool pool.BuffPool, k int) (tup Tuple) {
    97  	for i, typ := range tb.Desc.Types[:k] {
    98  		if !typ.Nullable && tb.fields[i] == nil {
    99  			panic("cannot write NULL to non-NULL field")
   100  		}
   101  	}
   102  	values := tb.fields[:k]
   103  	tup = NewTuple(pool, values...)
   104  	tb.Recycle()
   105  	return
   106  }
   107  
   108  // BuildPrefixNoRecycle materializes a prefix Tuple from the first |k| fields
   109  // but does not call Recycle.
   110  func (tb *TupleBuilder) BuildPrefixNoRecycle(pool pool.BuffPool, k int) (tup Tuple) {
   111  	for i, typ := range tb.Desc.Types[:k] {
   112  		if !typ.Nullable && tb.fields[i] == nil {
   113  			panic("cannot write NULL to non-NULL field")
   114  		}
   115  	}
   116  	values := tb.fields[:k]
   117  	tup = NewTuple(pool, values...)
   118  	return
   119  }
   120  
   121  // Recycle resets the TupleBuilder so it can build a new Tuple.
   122  func (tb *TupleBuilder) Recycle() {
   123  	for i := 0; i < tb.Desc.Count(); i++ {
   124  		tb.fields[i] = nil
   125  	}
   126  	tb.pos = 0
   127  }
   128  
   129  // PutBool writes a bool to the ith field of the Tuple being built.
   130  func (tb *TupleBuilder) PutBool(i int, v bool) {
   131  	tb.Desc.expectEncoding(i, Int8Enc)
   132  	tb.ensureCapacity(int8Size)
   133  	tb.fields[i] = tb.buf[tb.pos : tb.pos+int8Size]
   134  	writeBool(tb.fields[i], v)
   135  	tb.pos += int8Size
   136  }
   137  
   138  // PutInt8 writes an int8 to the ith field of the Tuple being built.
   139  func (tb *TupleBuilder) PutInt8(i int, v int8) {
   140  	tb.Desc.expectEncoding(i, Int8Enc)
   141  	tb.ensureCapacity(int8Size)
   142  	tb.fields[i] = tb.buf[tb.pos : tb.pos+int8Size]
   143  	writeInt8(tb.fields[i], v)
   144  	tb.pos += int8Size
   145  }
   146  
   147  // PutUint8 writes a uint8 to the ith field of the Tuple being built.
   148  func (tb *TupleBuilder) PutUint8(i int, v uint8) {
   149  	tb.Desc.expectEncoding(i, Uint8Enc)
   150  	tb.ensureCapacity(uint8Size)
   151  	tb.fields[i] = tb.buf[tb.pos : tb.pos+uint8Size]
   152  	writeUint8(tb.fields[i], v)
   153  	tb.pos += uint8Size
   154  }
   155  
   156  // PutInt16 writes an int16 to the ith field of the Tuple being built.
   157  func (tb *TupleBuilder) PutInt16(i int, v int16) {
   158  	tb.Desc.expectEncoding(i, Int16Enc)
   159  	tb.ensureCapacity(int16Size)
   160  	tb.fields[i] = tb.buf[tb.pos : tb.pos+int16Size]
   161  	writeInt16(tb.fields[i], v)
   162  	tb.pos += int16Size
   163  }
   164  
   165  // PutUint16 writes a uint16 to the ith field of the Tuple being built.
   166  func (tb *TupleBuilder) PutUint16(i int, v uint16) {
   167  	tb.Desc.expectEncoding(i, Uint16Enc)
   168  	tb.ensureCapacity(uint16Size)
   169  	tb.fields[i] = tb.buf[tb.pos : tb.pos+uint16Size]
   170  	WriteUint16(tb.fields[i], v)
   171  	tb.pos += uint16Size
   172  }
   173  
   174  // PutInt32 writes an int32 to the ith field of the Tuple being built.
   175  func (tb *TupleBuilder) PutInt32(i int, v int32) {
   176  	tb.Desc.expectEncoding(i, Int32Enc)
   177  	tb.ensureCapacity(int32Size)
   178  	tb.fields[i] = tb.buf[tb.pos : tb.pos+int32Size]
   179  	writeInt32(tb.fields[i], v)
   180  	tb.pos += int32Size
   181  }
   182  
   183  // PutUint32 writes a uint32 to the ith field of the Tuple being built.
   184  func (tb *TupleBuilder) PutUint32(i int, v uint32) {
   185  	tb.Desc.expectEncoding(i, Uint32Enc)
   186  	tb.ensureCapacity(uint32Size)
   187  	tb.fields[i] = tb.buf[tb.pos : tb.pos+uint32Size]
   188  	writeUint32(tb.fields[i], v)
   189  	tb.pos += uint32Size
   190  }
   191  
   192  // PutInt64 writes an int64 to the ith field of the Tuple being built.
   193  func (tb *TupleBuilder) PutInt64(i int, v int64) {
   194  	tb.Desc.expectEncoding(i, Int64Enc)
   195  	tb.ensureCapacity(int64Size)
   196  	tb.fields[i] = tb.buf[tb.pos : tb.pos+int64Size]
   197  	writeInt64(tb.fields[i], v)
   198  	tb.pos += int64Size
   199  }
   200  
   201  // PutUint64 writes a uint64 to the ith field of the Tuple being built.
   202  func (tb *TupleBuilder) PutUint64(i int, v uint64) {
   203  	tb.Desc.expectEncoding(i, Uint64Enc)
   204  	tb.ensureCapacity(uint64Size)
   205  	tb.fields[i] = tb.buf[tb.pos : tb.pos+uint64Size]
   206  	writeUint64(tb.fields[i], v)
   207  	tb.pos += uint64Size
   208  }
   209  
   210  // PutFloat32 writes a float32 to the ith field of the Tuple being built.
   211  func (tb *TupleBuilder) PutFloat32(i int, v float32) {
   212  	tb.Desc.expectEncoding(i, Float32Enc)
   213  	tb.ensureCapacity(float32Size)
   214  	tb.fields[i] = tb.buf[tb.pos : tb.pos+float32Size]
   215  	writeFloat32(tb.fields[i], v)
   216  	tb.pos += float32Size
   217  }
   218  
   219  // PutFloat64 writes a float64 to the ith field of the Tuple being built.
   220  func (tb *TupleBuilder) PutFloat64(i int, v float64) {
   221  	tb.Desc.expectEncoding(i, Float64Enc)
   222  	tb.ensureCapacity(float64Size)
   223  	tb.fields[i] = tb.buf[tb.pos : tb.pos+float64Size]
   224  	writeFloat64(tb.fields[i], v)
   225  	tb.pos += float64Size
   226  }
   227  
   228  func (tb *TupleBuilder) PutBit(i int, v uint64) {
   229  	tb.Desc.expectEncoding(i, Bit64Enc)
   230  	tb.ensureCapacity(bit64Size)
   231  	tb.fields[i] = tb.buf[tb.pos : tb.pos+bit64Size]
   232  	writeBit64(tb.fields[i], v)
   233  	tb.pos += bit64Size
   234  }
   235  
   236  func (tb *TupleBuilder) PutDecimal(i int, v decimal.Decimal) {
   237  	tb.Desc.expectEncoding(i, DecimalEnc)
   238  	sz := sizeOfDecimal(v)
   239  	tb.ensureCapacity(sz)
   240  	tb.fields[i] = tb.buf[tb.pos : tb.pos+sz]
   241  	writeDecimal(tb.fields[i], v)
   242  	tb.pos += sz
   243  }
   244  
   245  // PutYear writes an int16-encoded year to the ith field of the Tuple being built.
   246  func (tb *TupleBuilder) PutYear(i int, v int16) {
   247  	tb.Desc.expectEncoding(i, YearEnc)
   248  	tb.ensureCapacity(yearSize)
   249  	tb.fields[i] = tb.buf[tb.pos : tb.pos+yearSize]
   250  	writeYear(tb.fields[i], v)
   251  	tb.pos += int16Size
   252  }
   253  
   254  func (tb *TupleBuilder) PutDate(i int, v time.Time) {
   255  	tb.Desc.expectEncoding(i, DateEnc)
   256  	tb.ensureCapacity(dateSize)
   257  	tb.fields[i] = tb.buf[tb.pos : tb.pos+dateSize]
   258  	writeDate(tb.fields[i], v)
   259  	tb.pos += dateSize
   260  }
   261  
   262  // PutSqlTime writes a string to the ith field of the Tuple being built.
   263  func (tb *TupleBuilder) PutSqlTime(i int, v int64) {
   264  	tb.Desc.expectEncoding(i, TimeEnc)
   265  	tb.ensureCapacity(timeSize)
   266  	tb.fields[i] = tb.buf[tb.pos : tb.pos+timeSize]
   267  	writeTime(tb.fields[i], v)
   268  	tb.pos += timeSize
   269  }
   270  
   271  func (tb *TupleBuilder) PutDatetime(i int, v time.Time) {
   272  	tb.Desc.expectEncoding(i, DatetimeEnc)
   273  	tb.ensureCapacity(datetimeSize)
   274  	tb.fields[i] = tb.buf[tb.pos : tb.pos+datetimeSize]
   275  	writeDatetime(tb.fields[i], v)
   276  	tb.pos += datetimeSize
   277  }
   278  
   279  func (tb *TupleBuilder) PutEnum(i int, v uint16) {
   280  	tb.Desc.expectEncoding(i, EnumEnc)
   281  	tb.ensureCapacity(enumSize)
   282  	tb.fields[i] = tb.buf[tb.pos : tb.pos+enumSize]
   283  	writeEnum(tb.fields[i], v)
   284  	tb.pos += enumSize
   285  }
   286  
   287  func (tb *TupleBuilder) PutSet(i int, v uint64) {
   288  	tb.Desc.expectEncoding(i, SetEnc)
   289  	tb.ensureCapacity(setSize)
   290  	tb.fields[i] = tb.buf[tb.pos : tb.pos+setSize]
   291  	writeSet(tb.fields[i], v)
   292  	tb.pos += setSize
   293  }
   294  
   295  // PutString writes a string to the ith field of the Tuple being built.
   296  func (tb *TupleBuilder) PutString(i int, v string) error {
   297  	tb.Desc.expectEncoding(i, StringEnc)
   298  	sz := ByteSize(len(v)) + 1
   299  	offSz := 0
   300  	if i > 0 {
   301  		offSz = 2 * int(uint16Size)
   302  	}
   303  	if int(tb.pos)+len(v)+offSz > int(MaxTupleDataSize) {
   304  		return analyzererrors.ErrInvalidRowLength.New(MaxTupleDataSize, int(tb.pos)+len(v)+int(offsetsSize(i)))
   305  	}
   306  	tb.ensureCapacity(sz)
   307  	tb.fields[i] = tb.buf[tb.pos : tb.pos+sz]
   308  	writeString(tb.fields[i], v)
   309  	tb.pos += sz
   310  	return nil
   311  }
   312  
   313  // PutByteString writes a []byte to the ith field of the Tuple being built.
   314  func (tb *TupleBuilder) PutByteString(i int, v []byte) {
   315  	tb.Desc.expectEncoding(i, ByteStringEnc)
   316  	sz := ByteSize(len(v)) + 1
   317  	tb.ensureCapacity(sz)
   318  	tb.fields[i] = tb.buf[tb.pos : tb.pos+sz]
   319  	writeByteString(tb.fields[i], v)
   320  	tb.pos += sz
   321  }
   322  
   323  // PutJSON writes a []byte to the ith field of the Tuple being built.
   324  func (tb *TupleBuilder) PutJSON(i int, v []byte) {
   325  	tb.Desc.expectEncoding(i, JSONEnc)
   326  	sz := ByteSize(len(v)) + 1
   327  	tb.ensureCapacity(sz)
   328  	tb.fields[i] = tb.buf[tb.pos : tb.pos+sz]
   329  	writeByteString(tb.fields[i], v)
   330  	tb.pos += sz
   331  }
   332  
   333  // PutGeometry writes a []byte to the ith field of the Tuple being built.
   334  func (tb *TupleBuilder) PutGeometry(i int, v []byte) {
   335  	tb.Desc.expectEncoding(i, GeometryEnc)
   336  	sz := ByteSize(len(v)) + 1
   337  	tb.ensureCapacity(sz)
   338  	tb.fields[i] = tb.buf[tb.pos : tb.pos+sz]
   339  	writeByteString(tb.fields[i], v)
   340  	tb.pos += sz
   341  }
   342  
   343  // PutGeometryAddr writes a Geometry's address ref to the ith field
   344  func (tb *TupleBuilder) PutGeometryAddr(i int, v hash.Hash) {
   345  	tb.Desc.expectEncoding(i, GeomAddrEnc)
   346  	tb.ensureCapacity(hash.ByteLen)
   347  	tb.putAddr(i, v)
   348  }
   349  
   350  // PutHash128 writes a hash128 to the ith field of the Tuple being built.
   351  func (tb *TupleBuilder) PutHash128(i int, v []byte) {
   352  	tb.Desc.expectEncoding(i, Hash128Enc)
   353  	tb.ensureCapacity(hash128Size)
   354  	tb.fields[i] = tb.buf[tb.pos : tb.pos+hash128Size]
   355  	writeHash128(tb.fields[i], v)
   356  	tb.pos += hash128Size
   357  }
   358  
   359  // PutExtended writes a []byte to the ith field of the Tuple being built.
   360  func (tb *TupleBuilder) PutExtended(i int, v []byte) {
   361  	tb.Desc.expectEncoding(i, ExtendedEnc)
   362  	sz := ByteSize(len(v))
   363  	tb.ensureCapacity(sz)
   364  	tb.fields[i] = tb.buf[tb.pos : tb.pos+sz]
   365  	writeExtended(tb.Desc.Handlers[i], tb.fields[i], v)
   366  	tb.pos += sz
   367  }
   368  
   369  // PutExtendedAddr writes a []byte to the ith field of the Tuple being built.
   370  func (tb *TupleBuilder) PutExtendedAddr(i int, v hash.Hash) {
   371  	tb.Desc.expectEncoding(i, ExtendedAddrEnc)
   372  	tb.ensureCapacity(hash.ByteLen)
   373  	tb.putAddr(i, v)
   374  }
   375  
   376  // PutRaw writes a []byte to the ith field of the Tuple being built.
   377  func (tb *TupleBuilder) PutRaw(i int, buf []byte) {
   378  	if buf == nil {
   379  		// todo(andy): does it make sense to
   380  		//  allow/expect nulls here?
   381  		return
   382  	}
   383  	sz := ByteSize(len(buf))
   384  	tb.ensureCapacity(sz)
   385  	tb.fields[i] = tb.buf[tb.pos : tb.pos+sz]
   386  	writeRaw(tb.fields[i], buf)
   387  	tb.pos += sz
   388  }
   389  
   390  // PutCommitAddr writes a commit's address ref to the ith field
   391  // of the Tuple being built.
   392  func (tb *TupleBuilder) PutCommitAddr(i int, v hash.Hash) {
   393  	tb.Desc.expectEncoding(i, CommitAddrEnc)
   394  	tb.ensureCapacity(hash.ByteLen)
   395  	tb.putAddr(i, v)
   396  }
   397  
   398  // PutBytesAddr writes a blob's address ref to the ith field
   399  // of the Tuple being built.
   400  func (tb *TupleBuilder) PutBytesAddr(i int, v hash.Hash) {
   401  	tb.Desc.expectEncoding(i, BytesAddrEnc)
   402  	tb.ensureCapacity(hash.ByteLen)
   403  	tb.putAddr(i, v)
   404  }
   405  
   406  // PutStringAddr writes a string's address ref to the ith field
   407  // of the Tuple being built.
   408  func (tb *TupleBuilder) PutStringAddr(i int, v hash.Hash) {
   409  	tb.Desc.expectEncoding(i, StringAddrEnc)
   410  	tb.ensureCapacity(hash.ByteLen)
   411  	tb.putAddr(i, v)
   412  }
   413  
   414  // PutJSONAddr writes a JSON string's address ref to the ith field
   415  // of the Tuple being built.
   416  func (tb *TupleBuilder) PutJSONAddr(i int, v hash.Hash) {
   417  	tb.Desc.expectEncoding(i, JSONAddrEnc)
   418  	tb.ensureCapacity(hash.ByteLen)
   419  	tb.putAddr(i, v)
   420  }
   421  
   422  func (tb *TupleBuilder) putAddr(i int, v hash.Hash) {
   423  	tb.fields[i] = tb.buf[tb.pos : tb.pos+hash.ByteLen]
   424  	writeAddr(tb.fields[i], v[:])
   425  	tb.pos += hash.ByteLen
   426  }
   427  
   428  func (tb *TupleBuilder) ensureCapacity(sz ByteSize) {
   429  	need := int(tb.pos+sz) - len(tb.buf)
   430  	if need > 0 {
   431  		for i := 0; i < need; i++ {
   432  			tb.buf = append(tb.buf, byte(0))
   433  		}
   434  	}
   435  }
   436  
   437  // PutCell writes a Cell to the ith field of the Tuple being built.
   438  func (tb *TupleBuilder) PutCell(i int, v Cell) {
   439  	tb.Desc.expectEncoding(i, CellEnc)
   440  	tb.ensureCapacity(cellSize)
   441  	tb.fields[i] = tb.buf[tb.pos : tb.pos+cellSize]
   442  	writeCell(tb.fields[i], v)
   443  	tb.pos += cellSize
   444  }