ceph/src/arrow/java/vector/src/main/java/org/apache/arrow/vector/Decimal256Vector.java

   1 /*
   2  * Licensed to the Apache Software Foundation (ASF) under one or more
   3  * contributor license agreements.  See the NOTICE file distributed with
   4  * this work for additional information regarding copyright ownership.
   5  * The ASF licenses this file to You under the Apache License, Version 2.0
   6  * (the "License"); you may not use this file except in compliance with
   7  * the License.  You may obtain a copy of the License at
   8  *
   9  *    http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  */
  17
  18 package org.apache.arrow.vector;
  19
  20 import static org.apache.arrow.vector.NullCheckingForGet.NULL_CHECKING_ENABLED;
  21
  22 import java.math.BigDecimal;
  23 import java.nio.ByteOrder;
  24
  25 import org.apache.arrow.memory.ArrowBuf;
  26 import org.apache.arrow.memory.BufferAllocator;
  27 import org.apache.arrow.vector.complex.impl.Decimal256ReaderImpl;
  28 import org.apache.arrow.vector.complex.reader.FieldReader;
  29 import org.apache.arrow.vector.holders.Decimal256Holder;
  30 import org.apache.arrow.vector.holders.NullableDecimal256Holder;
  31 import org.apache.arrow.vector.types.Types.MinorType;
  32 import org.apache.arrow.vector.types.pojo.ArrowType;
  33 import org.apache.arrow.vector.types.pojo.Field;
  34 import org.apache.arrow.vector.types.pojo.FieldType;
  35 import org.apache.arrow.vector.util.DecimalUtility;
  36 import org.apache.arrow.vector.util.TransferPair;
  37
  38 import io.netty.util.internal.PlatformDependent;
  39
  40 /**
  41  * Decimal256Vector implements a fixed width vector (32 bytes) of
  42  * decimal values which could be null. A validity buffer (bit vector) is
  43  * maintained to track which elements in the vector are null.
  44  */
  45 public final class Decimal256Vector extends BaseFixedWidthVector {
  46   public static final byte TYPE_WIDTH = 32;
  47   private static final boolean LITTLE_ENDIAN = ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN;
  48   private final FieldReader reader;
  49
  50   private final int precision;
  51   private final int scale;
  52
  53   /**
  54    * Instantiate a Decimal256Vector. This doesn't allocate any memory for
  55    * the data in vector.
  56    *
  57    * @param name name of the vector
  58    * @param allocator allocator for memory management.
  59    */
  60   public Decimal256Vector(String name, BufferAllocator allocator,
  61                                int precision, int scale) {
  62     this(name, FieldType.nullable(new ArrowType.Decimal(precision, scale, /*bitWidth=*/TYPE_WIDTH * 8)), allocator);
  63   }
  64
  65   /**
  66    * Instantiate a Decimal256Vector. This doesn't allocate any memory for
  67    * the data in vector.
  68    *
  69    * @param name name of the vector
  70    * @param fieldType type of Field materialized by this vector
  71    * @param allocator allocator for memory management.
  72    */
  73   public Decimal256Vector(String name, FieldType fieldType, BufferAllocator allocator) {
  74     this(new Field(name, fieldType, null), allocator);
  75   }
  76
  77   /**
  78    * Instantiate a Decimal256Vector. This doesn't allocate any memory for
  79    * the data in vector.
  80    *
  81    * @param field field materialized by this vector
  82    * @param allocator allocator for memory management.
  83    */
  84   public Decimal256Vector(Field field, BufferAllocator allocator) {
  85     super(field, allocator, TYPE_WIDTH);
  86     ArrowType.Decimal arrowType = (ArrowType.Decimal) field.getFieldType().getType();
  87     reader = new Decimal256ReaderImpl(Decimal256Vector.this);
  88     this.precision = arrowType.getPrecision();
  89     this.scale = arrowType.getScale();
  90   }
  91
  92   /**
  93    * Get a reader that supports reading values from this vector.
  94    *
  95    * @return Field Reader for this vector
  96    */
  97   @Override
  98   public FieldReader getReader() {
  99     return reader;
 100   }
 101
 102   /**
 103    * Get minor type for this vector. The vector holds values belonging
 104    * to a particular type.
 105    *
 106    * @return {@link org.apache.arrow.vector.types.Types.MinorType}
 107    */
 108   @Override
 109   public MinorType getMinorType() {
 110     return MinorType.DECIMAL256;
 111   }
 112
 113
 114   /*----------------------------------------------------------------*
 115    |                                                                |
 116    |          vector value retrieval methods                        |
 117    |                                                                |
 118    *----------------------------------------------------------------*/
 119
 120
 121   /**
 122    * Get the element at the given index from the vector.
 123    *
 124    * @param index   position of element
 125    * @return element at given index
 126    */
 127   public ArrowBuf get(int index) throws IllegalStateException {
 128     if (NULL_CHECKING_ENABLED && isSet(index) == 0) {
 129       throw new IllegalStateException("Value at index is null");
 130     }
 131     return valueBuffer.slice((long) index * TYPE_WIDTH, TYPE_WIDTH);
 132   }
 133
 134   /**
 135    * Get the element at the given index from the vector and
 136    * sets the state in holder. If element at given index
 137    * is null, holder.isSet will be zero.
 138    *
 139    * @param index   position of element
 140    */
 141   public void get(int index, NullableDecimal256Holder holder) {
 142     if (isSet(index) == 0) {
 143       holder.isSet = 0;
 144       return;
 145     }
 146     holder.isSet = 1;
 147     holder.buffer = valueBuffer;
 148     holder.precision = precision;
 149     holder.scale = scale;
 150     holder.start = ((long) index) * TYPE_WIDTH;
 151   }
 152
 153   /**
 154    * Same as {@link #get(int)}.
 155    *
 156    * @param index   position of element
 157    * @return element at given index
 158    */
 159   public BigDecimal getObject(int index) {
 160     if (isSet(index) == 0) {
 161       return null;
 162     } else {
 163       return DecimalUtility.getBigDecimalFromArrowBuf(valueBuffer, index, scale, TYPE_WIDTH);
 164     }
 165   }
 166
 167   /**
 168    * Return precision for the decimal value.
 169    */
 170   public int getPrecision() {
 171     return precision;
 172   }
 173
 174   /**
 175    * Return scale for the decimal value.
 176    */
 177   public int getScale() {
 178     return scale;
 179   }
 180
 181
 182   /*----------------------------------------------------------------*
 183    |                                                                |
 184    |          vector value setter methods                           |
 185    |                                                                |
 186    *----------------------------------------------------------------*/
 187
 188
 189   /**
 190    * Set the element at the given index to the given value.
 191    *
 192    * @param index    position of element
 193    * @param buffer   ArrowBuf containing decimal value.
 194    */
 195   public void set(int index, ArrowBuf buffer) {
 196     BitVectorHelper.setBit(validityBuffer, index);
 197     valueBuffer.setBytes((long) index * TYPE_WIDTH, buffer, 0, TYPE_WIDTH);
 198   }
 199
 200   /**
 201    * Set the decimal element at given index to the provided array of bytes.
 202    * Decimal256 is now implemented as Native Endian. This API allows the user
 203    * to pass a decimal value in the form of byte array in BE byte order.
 204    *
 205    * <p>Consumers of Arrow code can use this API instead of first swapping
 206    * the source bytes (doing a write and read) and then finally writing to
 207    * ArrowBuf of decimal vector.
 208    *
 209    * <p>This method takes care of adding the necessary padding if the length
 210    * of byte array is less then 32 (length of decimal type).
 211    *
 212    * @param index position of element
 213    * @param value array of bytes containing decimal in big endian byte order.
 214    */
 215   public void setBigEndian(int index, byte[] value) {
 216     BitVectorHelper.setBit(validityBuffer, index);
 217     final int length = value.length;
 218
 219     // do the bound check.
 220     valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH);
 221
 222     long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH;
 223     if (length == 0) {
 224       PlatformDependent.setMemory(outAddress, Decimal256Vector.TYPE_WIDTH, (byte) 0);
 225       return;
 226     }
 227     if (LITTLE_ENDIAN) {
 228       // swap bytes to convert BE to LE
 229       for (int byteIdx = 0; byteIdx < length; ++byteIdx) {
 230         PlatformDependent.putByte(outAddress + byteIdx, value[length - 1 - byteIdx]);
 231       }
 232
 233       if (length == TYPE_WIDTH) {
 234         return;
 235       }
 236
 237       if (length < TYPE_WIDTH) {
 238         // sign extend
 239         final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00);
 240         PlatformDependent.setMemory(outAddress + length, Decimal256Vector.TYPE_WIDTH - length, pad);
 241         return;
 242       }
 243     } else {
 244       if (length <= TYPE_WIDTH) {
 245         // copy data from value to outAddress
 246         PlatformDependent.copyMemory(value, 0, outAddress + Decimal256Vector.TYPE_WIDTH - length, length);
 247         // sign extend
 248         final byte pad = (byte) (value[0] < 0 ? 0xFF : 0x00);
 249         PlatformDependent.setMemory(outAddress, Decimal256Vector.TYPE_WIDTH - length, pad);
 250         return;
 251       }
 252     }
 253     throw new IllegalArgumentException(
 254         "Invalid decimal value length. Valid length in [1 - 32], got " + length);
 255   }
 256
 257   /**
 258    * Set the element at the given index to the given value.
 259    *
 260    * @param index    position of element
 261    * @param start    start index of data in the buffer
 262    * @param buffer   ArrowBuf containing decimal value.
 263    */
 264   public void set(int index, long start, ArrowBuf buffer) {
 265     BitVectorHelper.setBit(validityBuffer, index);
 266     valueBuffer.setBytes((long) index * TYPE_WIDTH, buffer, start, TYPE_WIDTH);
 267   }
 268
 269   /**
 270    * Sets the element at given index using the buffer whose size maybe <= 32 bytes.
 271    * @param index index to write the decimal to
 272    * @param start start of value in the buffer
 273    * @param buffer contains the decimal in native endian bytes
 274    * @param length length of the value in the buffer
 275    */
 276   public void setSafe(int index, long start, ArrowBuf buffer, int length) {
 277     handleSafe(index);
 278     BitVectorHelper.setBit(validityBuffer, index);
 279
 280     // do the bound checks.
 281     buffer.checkBytes(start, start + length);
 282     valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH);
 283
 284     long inAddress = buffer.memoryAddress() + start;
 285     long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH;
 286     if (LITTLE_ENDIAN) {
 287       PlatformDependent.copyMemory(inAddress, outAddress, length);
 288       // sign extend
 289       if (length < TYPE_WIDTH) {
 290         byte msb = PlatformDependent.getByte(inAddress + length - 1);
 291         final byte pad = (byte) (msb < 0 ? 0xFF : 0x00);
 292         PlatformDependent.setMemory(outAddress + length, Decimal256Vector.TYPE_WIDTH - length, pad);
 293       }
 294     } else {
 295       PlatformDependent.copyMemory(inAddress, outAddress + Decimal256Vector.TYPE_WIDTH - length, length);
 296       // sign extend
 297       if (length < TYPE_WIDTH) {
 298         byte msb = PlatformDependent.getByte(inAddress);
 299         final byte pad = (byte) (msb < 0 ? 0xFF : 0x00);
 300         PlatformDependent.setMemory(outAddress, Decimal256Vector.TYPE_WIDTH - length, pad);
 301       }
 302     }
 303   }
 304
 305
 306   /**
 307    * Sets the element at given index using the buffer whose size maybe <= 32 bytes.
 308    * @param index index to write the decimal to
 309    * @param start start of value in the buffer
 310    * @param buffer contains the decimal in big endian bytes
 311    * @param length length of the value in the buffer
 312    */
 313   public void setBigEndianSafe(int index, long start, ArrowBuf buffer, int length) {
 314     handleSafe(index);
 315     BitVectorHelper.setBit(validityBuffer, index);
 316
 317     // do the bound checks.
 318     buffer.checkBytes(start, start + length);
 319     valueBuffer.checkBytes((long) index * TYPE_WIDTH, (long) (index + 1) * TYPE_WIDTH);
 320
 321     // not using buffer.getByte() to avoid boundary checks for every byte.
 322     long inAddress = buffer.memoryAddress() + start;
 323     long outAddress = valueBuffer.memoryAddress() + (long) index * TYPE_WIDTH;
 324     if (LITTLE_ENDIAN) {
 325       // swap bytes to convert BE to LE
 326       for (int byteIdx = 0; byteIdx < length; ++byteIdx) {
 327         byte val = PlatformDependent.getByte((inAddress + length - 1) - byteIdx);
 328         PlatformDependent.putByte(outAddress + byteIdx, val);
 329       }
 330       // sign extend
 331       if (length < 32) {
 332         byte msb = PlatformDependent.getByte(inAddress);
 333         final byte pad = (byte) (msb < 0 ? 0xFF : 0x00);
 334         PlatformDependent.setMemory(outAddress + length, Decimal256Vector.TYPE_WIDTH - length, pad);
 335       }
 336     } else {
 337       PlatformDependent.copyMemory(inAddress, outAddress + Decimal256Vector.TYPE_WIDTH - length, length);
 338       // sign extend
 339       if (length < TYPE_WIDTH) {
 340         byte msb = PlatformDependent.getByte(inAddress);
 341         final byte pad = (byte) (msb < 0 ? 0xFF : 0x00);
 342         PlatformDependent.setMemory(outAddress, Decimal256Vector.TYPE_WIDTH - length, pad);
 343       }
 344     }
 345   }
 346
 347   /**
 348    * Set the element at the given index to the given value.
 349    *
 350    * @param index   position of element
 351    * @param value   BigDecimal containing decimal value.
 352    */
 353   public void set(int index, BigDecimal value) {
 354     BitVectorHelper.setBit(validityBuffer, index);
 355     DecimalUtility.checkPrecisionAndScale(value, precision, scale);
 356     DecimalUtility.writeBigDecimalToArrowBuf(value, valueBuffer, index, TYPE_WIDTH);
 357   }
 358
 359   /**
 360    * Set the element at the given index to the given value.
 361    *
 362    * @param index   position of element
 363    * @param value   long value.
 364    */
 365   public void set(int index, long value) {
 366     BitVectorHelper.setBit(validityBuffer, index);
 367     DecimalUtility.writeLongToArrowBuf(value, valueBuffer, index, TYPE_WIDTH);
 368   }
 369
 370   /**
 371    * Set the element at the given index to the value set in data holder.
 372    * If the value in holder is not indicated as set, element in the
 373    * at the given index will be null.
 374    *
 375    * @param index   position of element
 376    * @param holder  nullable data holder for value of element
 377    */
 378   public void set(int index, NullableDecimal256Holder holder) throws IllegalArgumentException {
 379     if (holder.isSet < 0) {
 380       throw new IllegalArgumentException();
 381     } else if (holder.isSet > 0) {
 382       BitVectorHelper.setBit(validityBuffer, index);
 383       valueBuffer.setBytes((long) index * TYPE_WIDTH, holder.buffer, holder.start, TYPE_WIDTH);
 384     } else {
 385       BitVectorHelper.unsetBit(validityBuffer, index);
 386     }
 387   }
 388
 389   /**
 390    * Set the element at the given index to the value set in data holder.
 391    *
 392    * @param index   position of element
 393    * @param holder  data holder for value of element
 394    */
 395   public void set(int index, Decimal256Holder holder) {
 396     BitVectorHelper.setBit(validityBuffer, index);
 397     valueBuffer.setBytes((long) index * TYPE_WIDTH, holder.buffer, holder.start, TYPE_WIDTH);
 398   }
 399
 400   /**
 401    * Same as {@link #set(int, ArrowBuf)} except that it handles the
 402    * case when index is greater than or equal to existing
 403    * value capacity {@link #getValueCapacity()}.
 404    *
 405    * @param index   position of element
 406    * @param buffer  ArrowBuf containing decimal value.
 407    */
 408   public void setSafe(int index, ArrowBuf buffer) {
 409     handleSafe(index);
 410     set(index, buffer);
 411   }
 412
 413   /**
 414    * Same as {@link #setBigEndian(int, byte[])} except that it handles the
 415    * case when index is greater than or equal to existing
 416    * value capacity {@link #getValueCapacity()}.
 417    */
 418   public void setBigEndianSafe(int index, byte[] value) {
 419     handleSafe(index);
 420     setBigEndian(index, value);
 421   }
 422
 423   /**
 424    * Same as {@link #set(int, int, ArrowBuf)} except that it handles the
 425    * case when index is greater than or equal to existing
 426    * value capacity {@link #getValueCapacity()}.
 427    *
 428    * @param index    position of element
 429    * @param start    start index of data in the buffer
 430    * @param buffer   ArrowBuf containing decimal value.
 431    */
 432   public void setSafe(int index, long start, ArrowBuf buffer) {
 433     handleSafe(index);
 434     set(index, start, buffer);
 435   }
 436
 437   /**
 438    * Same as {@link #set(int, BigDecimal)} except that it handles the
 439    * case when index is greater than or equal to existing
 440    * value capacity {@link #getValueCapacity()}.
 441    *
 442    * @param index   position of element
 443    * @param value   BigDecimal containing decimal value.
 444    */
 445   public void setSafe(int index, BigDecimal value) {
 446     handleSafe(index);
 447     set(index, value);
 448   }
 449
 450   /**
 451    * Same as {@link #set(int, long)} except that it handles the
 452    * case when index is greater than or equal to existing
 453    * value capacity {@link #getValueCapacity()}.
 454    *
 455    * @param index   position of element
 456    * @param value   long value.
 457    */
 458   public void setSafe(int index, long value) {
 459     handleSafe(index);
 460     set(index, value);
 461   }
 462
 463   /**
 464    * Same as {@link #set(int, NullableDecimalHolder)} except that it handles the
 465    * case when index is greater than or equal to existing
 466    * value capacity {@link #getValueCapacity()}.
 467    *
 468    * @param index   position of element
 469    * @param holder  nullable data holder for value of element
 470    */
 471   public void setSafe(int index, NullableDecimal256Holder holder) throws IllegalArgumentException {
 472     handleSafe(index);
 473     set(index, holder);
 474   }
 475
 476   /**
 477    * Same as {@link #set(int, Decimal256Holder)} except that it handles the
 478    * case when index is greater than or equal to existing
 479    * value capacity {@link #getValueCapacity()}.
 480    *
 481    * @param index   position of element
 482    * @param holder  data holder for value of element
 483    */
 484   public void setSafe(int index, Decimal256Holder holder) {
 485     handleSafe(index);
 486     set(index, holder);
 487   }
 488
 489   /**
 490    * Store the given value at a particular position in the vector. isSet indicates
 491    * whether the value is NULL or not.
 492    *
 493    * @param index position of the new value
 494    * @param isSet 0 for NULL value, 1 otherwise
 495    * @param start start position of the value in the buffer
 496    * @param buffer buffer containing the value to be stored in the vector
 497    */
 498   public void set(int index, int isSet, long start, ArrowBuf buffer) {
 499     if (isSet > 0) {
 500       set(index, start, buffer);
 501     } else {
 502       BitVectorHelper.unsetBit(validityBuffer, index);
 503     }
 504   }
 505
 506   /**
 507    * Same as {@link #setSafe(int, int, int, ArrowBuf)} except that it handles
 508    * the case when the position of new value is beyond the current value
 509    * capacity of the vector.
 510    *
 511    * @param index position of the new value
 512    * @param isSet 0 for NULL value, 1 otherwise
 513    * @param start start position of the value in the buffer
 514    * @param buffer buffer containing the value to be stored in the vector
 515    */
 516   public void setSafe(int index, int isSet, long start, ArrowBuf buffer) {
 517     handleSafe(index);
 518     set(index, isSet, start, buffer);
 519   }
 520
 521   /*----------------------------------------------------------------*
 522    |                                                                |
 523    |                      vector transfer                           |
 524    |                                                                |
 525    *----------------------------------------------------------------*/
 526
 527
 528   /**
 529    * Construct a TransferPair comprising of this and a target vector of
 530    * the same type.
 531    *
 532    * @param ref name of the target vector
 533    * @param allocator allocator for the target vector
 534    * @return {@link TransferPair}
 535    */
 536   @Override
 537   public TransferPair getTransferPair(String ref, BufferAllocator allocator) {
 538     return new TransferImpl(ref, allocator);
 539   }
 540
 541   /**
 542    * Construct a TransferPair with a desired target vector of the same type.
 543    *
 544    * @param to target vector
 545    * @return {@link TransferPair}
 546    */
 547   @Override
 548   public TransferPair makeTransferPair(ValueVector to) {
 549     return new TransferImpl((Decimal256Vector) to);
 550   }
 551
 552   private class TransferImpl implements TransferPair {
 553     Decimal256Vector to;
 554
 555     public TransferImpl(String ref, BufferAllocator allocator) {
 556       to = new Decimal256Vector(ref, allocator, Decimal256Vector.this.precision,
 557               Decimal256Vector.this.scale);
 558     }
 559
 560     public TransferImpl(Decimal256Vector to) {
 561       this.to = to;
 562     }
 563
 564     @Override
 565     public Decimal256Vector getTo() {
 566       return to;
 567     }
 568
 569     @Override
 570     public void transfer() {
 571       transferTo(to);
 572     }
 573
 574     @Override
 575     public void splitAndTransfer(int startIndex, int length) {
 576       splitAndTransferTo(startIndex, length, to);
 577     }
 578
 579     @Override
 580     public void copyValueSafe(int fromIndex, int toIndex) {
 581       to.copyFromSafe(fromIndex, toIndex, Decimal256Vector.this);
 582     }
 583   }
 584 }