ceph/src/arrow/java/vector/src/main/java/org/apache/arrow/vector/complex/impl/PromotableWriter.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.complex.impl;
  19
  20 import java.math.BigDecimal;
  21
  22 import org.apache.arrow.memory.ArrowBuf;
  23 import org.apache.arrow.vector.FieldVector;
  24 import org.apache.arrow.vector.NullVector;
  25 import org.apache.arrow.vector.ValueVector;
  26 import org.apache.arrow.vector.complex.AbstractStructVector;
  27 import org.apache.arrow.vector.complex.FixedSizeListVector;
  28 import org.apache.arrow.vector.complex.LargeListVector;
  29 import org.apache.arrow.vector.complex.ListVector;
  30 import org.apache.arrow.vector.complex.MapVector;
  31 import org.apache.arrow.vector.complex.StructVector;
  32 import org.apache.arrow.vector.complex.UnionVector;
  33 import org.apache.arrow.vector.complex.writer.FieldWriter;
  34 import org.apache.arrow.vector.holders.Decimal256Holder;
  35 import org.apache.arrow.vector.holders.DecimalHolder;
  36 import org.apache.arrow.vector.types.Types.MinorType;
  37 import org.apache.arrow.vector.types.pojo.ArrowType;
  38 import org.apache.arrow.vector.types.pojo.Field;
  39 import org.apache.arrow.vector.types.pojo.FieldType;
  40 import org.apache.arrow.vector.util.TransferPair;
  41
  42 /**
  43  * This FieldWriter implementation delegates all FieldWriter API calls to an inner FieldWriter. This inner field writer
  44  * can start as a specific type, and this class will promote the writer to a UnionWriter if a call is made that the
  45  * specifically typed writer cannot handle. A new UnionVector is created, wrapping the original vector, and replaces the
  46  * original vector in the parent vector, which can be either an AbstractStructVector or a ListVector.
  47  *
  48  * <p>The writer used can either be for single elements (struct) or lists.</p>
  49  */
  50 public class PromotableWriter extends AbstractPromotableFieldWriter {
  51
  52   private final AbstractStructVector parentContainer;
  53   private final ListVector listVector;
  54   private final FixedSizeListVector fixedListVector;
  55   private final LargeListVector largeListVector;
  56   private final NullableStructWriterFactory nullableStructWriterFactory;
  57   private int position;
  58   private static final int MAX_DECIMAL_PRECISION = 38;
  59   private static final int MAX_DECIMAL256_PRECISION = 76;
  60
  61   private enum State {
  62     UNTYPED, SINGLE, UNION
  63   }
  64
  65   private MinorType type;
  66   private ValueVector vector;
  67   private UnionVector unionVector;
  68   private State state;
  69   private FieldWriter writer;
  70
  71   /**
  72    * Constructs a new instance.
  73    *
  74    * @param v The vector to write.
  75    * @param parentContainer The parent container for the vector.
  76    */
  77   public PromotableWriter(ValueVector v, AbstractStructVector parentContainer) {
  78     this(v, parentContainer, NullableStructWriterFactory.getNullableStructWriterFactoryInstance());
  79   }
  80
  81   /**
  82    * Constructs a new instance.
  83    *
  84    * @param v The vector to initialize the writer with.
  85    * @param parentContainer The parent container for the vector.
  86    * @param nullableStructWriterFactory The factory to create the delegate writer.
  87    */
  88   public PromotableWriter(
  89       ValueVector v,
  90       AbstractStructVector parentContainer,
  91       NullableStructWriterFactory nullableStructWriterFactory) {
  92     this.parentContainer = parentContainer;
  93     this.listVector = null;
  94     this.fixedListVector = null;
  95     this.largeListVector = null;
  96     this.nullableStructWriterFactory = nullableStructWriterFactory;
  97     init(v);
  98   }
  99
 100   /**
 101    * Constructs a new instance.
 102    *
 103    * @param v The vector to initialize the writer with.
 104    * @param listVector The vector that serves as a parent of v.
 105    */
 106   public PromotableWriter(ValueVector v, ListVector listVector) {
 107     this(v, listVector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance());
 108   }
 109
 110   /**
 111    * Constructs a new instance.
 112    *
 113    * @param v The vector to initialize the writer with.
 114    * @param fixedListVector The vector that serves as a parent of v.
 115    */
 116   public PromotableWriter(ValueVector v, FixedSizeListVector fixedListVector) {
 117     this(v, fixedListVector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance());
 118   }
 119
 120   /**
 121    * Constructs a new instance.
 122    *
 123    * @param v The vector to initialize the writer with.
 124    * @param largeListVector The vector that serves as a parent of v.
 125    */
 126   public PromotableWriter(ValueVector v, LargeListVector largeListVector) {
 127     this(v, largeListVector, NullableStructWriterFactory.getNullableStructWriterFactoryInstance());
 128   }
 129
 130   /**
 131    * Constructs a new instance.
 132    *
 133    * @param v The vector to initialize the writer with.
 134    * @param listVector The vector that serves as a parent of v.
 135    * @param nullableStructWriterFactory The factory to create the delegate writer.
 136    */
 137   public PromotableWriter(
 138       ValueVector v,
 139       ListVector listVector,
 140       NullableStructWriterFactory nullableStructWriterFactory) {
 141     this.listVector = listVector;
 142     this.parentContainer = null;
 143     this.fixedListVector = null;
 144     this.largeListVector = null;
 145     this.nullableStructWriterFactory = nullableStructWriterFactory;
 146     init(v);
 147   }
 148
 149   /**
 150    * Constructs a new instance.
 151    *
 152    * @param v The vector to initialize the writer with.
 153    * @param fixedListVector The vector that serves as a parent of v.
 154    * @param nullableStructWriterFactory The factory to create the delegate writer.
 155    */
 156   public PromotableWriter(
 157       ValueVector v,
 158       FixedSizeListVector fixedListVector,
 159       NullableStructWriterFactory nullableStructWriterFactory) {
 160     this.fixedListVector = fixedListVector;
 161     this.parentContainer = null;
 162     this.listVector = null;
 163     this.largeListVector = null;
 164     this.nullableStructWriterFactory = nullableStructWriterFactory;
 165     init(v);
 166   }
 167
 168   /**
 169    * Constructs a new instance.
 170    *
 171    * @param v The vector to initialize the writer with.
 172    * @param largeListVector The vector that serves as a parent of v.
 173    * @param nullableStructWriterFactory The factory to create the delegate writer.
 174    */
 175   public PromotableWriter(
 176       ValueVector v,
 177       LargeListVector largeListVector,
 178       NullableStructWriterFactory nullableStructWriterFactory) {
 179     this.largeListVector = largeListVector;
 180     this.fixedListVector = null;
 181     this.parentContainer = null;
 182     this.listVector = null;
 183     this.nullableStructWriterFactory = nullableStructWriterFactory;
 184     init(v);
 185   }
 186
 187   private void init(ValueVector v) {
 188     if (v instanceof UnionVector) {
 189       state = State.UNION;
 190       unionVector = (UnionVector) v;
 191       writer = new UnionWriter(unionVector, nullableStructWriterFactory);
 192     } else if (v instanceof NullVector) {
 193       state = State.UNTYPED;
 194     } else {
 195       setWriter(v);
 196     }
 197   }
 198
 199   @Override
 200   public void setAddVectorAsNullable(boolean nullable) {
 201     super.setAddVectorAsNullable(nullable);
 202     if (writer instanceof AbstractFieldWriter) {
 203       ((AbstractFieldWriter) writer).setAddVectorAsNullable(nullable);
 204     }
 205   }
 206
 207   private void setWriter(ValueVector v) {
 208     state = State.SINGLE;
 209     vector = v;
 210     type = v.getMinorType();
 211     switch (type) {
 212       case STRUCT:
 213         writer = nullableStructWriterFactory.build((StructVector) vector);
 214         break;
 215       case LIST:
 216         writer = new UnionListWriter((ListVector) vector, nullableStructWriterFactory);
 217         break;
 218       case MAP:
 219         writer = new UnionMapWriter((MapVector) vector);
 220         break;
 221       case UNION:
 222         writer = new UnionWriter((UnionVector) vector, nullableStructWriterFactory);
 223         break;
 224       default:
 225         writer = type.getNewFieldWriter(vector);
 226         break;
 227     }
 228   }
 229
 230   @Override
 231   public void writeNull() {
 232     FieldWriter w = getWriter();
 233     if (w != null) {
 234       w.writeNull();
 235     }
 236     setPosition(idx() + 1);
 237   }
 238
 239   @Override
 240   public void setPosition(int index) {
 241     super.setPosition(index);
 242     FieldWriter w = getWriter();
 243     if (w == null) {
 244       position = index;
 245     } else {
 246       w.setPosition(index);
 247     }
 248   }
 249
 250   @Override
 251   protected FieldWriter getWriter(MinorType type, ArrowType arrowType) {
 252     if (state == State.UNION) {
 253       if (type == MinorType.DECIMAL || type == MinorType.MAP) {
 254         ((UnionWriter) writer).getWriter(type, arrowType);
 255       } else {
 256         ((UnionWriter) writer).getWriter(type);
 257       }
 258     } else if (state == State.UNTYPED) {
 259       if (type == null) {
 260         // ???
 261         return null;
 262       }
 263       if (arrowType == null) {
 264         arrowType = type.getType();
 265       }
 266       FieldType fieldType = new FieldType(addVectorAsNullable, arrowType, null, null);
 267       ValueVector v;
 268       if (listVector != null) {
 269         v = listVector.addOrGetVector(fieldType).getVector();
 270       } else if (fixedListVector != null) {
 271         v = fixedListVector.addOrGetVector(fieldType).getVector();
 272       } else {
 273         v = largeListVector.addOrGetVector(fieldType).getVector();
 274       }
 275       v.allocateNew();
 276       setWriter(v);
 277       writer.setPosition(position);
 278     } else if (type != this.type) {
 279       promoteToUnion();
 280       if (type == MinorType.DECIMAL || type == MinorType.MAP) {
 281         ((UnionWriter) writer).getWriter(type, arrowType);
 282       } else {
 283         ((UnionWriter) writer).getWriter(type);
 284       }
 285     }
 286     return writer;
 287   }
 288
 289   @Override
 290   public boolean isEmptyStruct() {
 291     return writer.isEmptyStruct();
 292   }
 293
 294   protected FieldWriter getWriter() {
 295     return writer;
 296   }
 297
 298   private FieldWriter promoteToUnion() {
 299     String name = vector.getField().getName();
 300     TransferPair tp = vector.getTransferPair(vector.getMinorType().name().toLowerCase(), vector.getAllocator());
 301     tp.transfer();
 302     if (parentContainer != null) {
 303       // TODO allow dictionaries in complex types
 304       unionVector = parentContainer.addOrGetUnion(name);
 305       unionVector.allocateNew();
 306     } else if (listVector != null) {
 307       unionVector = listVector.promoteToUnion();
 308     } else if (fixedListVector != null) {
 309       unionVector = fixedListVector.promoteToUnion();
 310     } else if (largeListVector != null) {
 311       unionVector = largeListVector.promoteToUnion();
 312     }
 313     unionVector.addVector((FieldVector) tp.getTo());
 314     writer = new UnionWriter(unionVector, nullableStructWriterFactory);
 315     writer.setPosition(idx());
 316     for (int i = 0; i <= idx(); i++) {
 317       unionVector.setType(i, vector.getMinorType());
 318     }
 319     vector = null;
 320     state = State.UNION;
 321     return writer;
 322   }
 323
 324   @Override
 325   public void write(DecimalHolder holder) {
 326     getWriter(MinorType.DECIMAL,
 327               new ArrowType.Decimal(MAX_DECIMAL_PRECISION, holder.scale, /*bitWidth=*/128)).write(holder);
 328   }
 329
 330   @Override
 331   public void writeDecimal(long start, ArrowBuf buffer, ArrowType arrowType) {
 332     getWriter(MinorType.DECIMAL, new ArrowType.Decimal(MAX_DECIMAL_PRECISION,
 333         ((ArrowType.Decimal) arrowType).getScale(), /*bitWidth=*/128)).writeDecimal(start, buffer, arrowType);
 334   }
 335
 336   @Override
 337   public void writeDecimal(BigDecimal value) {
 338     getWriter(MinorType.DECIMAL,
 339       new ArrowType.Decimal(MAX_DECIMAL_PRECISION, value.scale(), /*bitWidth=*/128)).writeDecimal(value);
 340   }
 341
 342   @Override
 343   public void writeBigEndianBytesToDecimal(byte[] value, ArrowType arrowType) {
 344     getWriter(MinorType.DECIMAL, new ArrowType.Decimal(MAX_DECIMAL_PRECISION,
 345         ((ArrowType.Decimal) arrowType).getScale(), /*bitWidth=*/128)).writeBigEndianBytesToDecimal(value, arrowType);
 346   }
 347
 348   @Override
 349   public void write(Decimal256Holder holder) {
 350     getWriter(MinorType.DECIMAL256,
 351               new ArrowType.Decimal(MAX_DECIMAL256_PRECISION, holder.scale, /*bitWidth=*/256)).write(holder);
 352   }
 353
 354   @Override
 355   public void writeDecimal256(long start, ArrowBuf buffer, ArrowType arrowType) {
 356     getWriter(MinorType.DECIMAL256, new ArrowType.Decimal(MAX_DECIMAL256_PRECISION,
 357         ((ArrowType.Decimal) arrowType).getScale(), /*bitWidth=*/256)).writeDecimal256(start, buffer, arrowType);
 358   }
 359
 360   @Override
 361   public void writeDecimal256(BigDecimal value) {
 362     getWriter(MinorType.DECIMAL256,
 363         new ArrowType.Decimal(MAX_DECIMAL256_PRECISION, value.scale(), /*bitWidth=*/256)).writeDecimal256(value);
 364   }
 365
 366   @Override
 367   public void writeBigEndianBytesToDecimal256(byte[] value, ArrowType arrowType) {
 368     getWriter(MinorType.DECIMAL256, new ArrowType.Decimal(MAX_DECIMAL256_PRECISION,
 369         ((ArrowType.Decimal) arrowType).getScale(),
 370           /*bitWidth=*/256)).writeBigEndianBytesToDecimal256(value, arrowType);
 371   }
 372
 373
 374   @Override
 375   public void allocate() {
 376     getWriter().allocate();
 377   }
 378
 379   @Override
 380   public void clear() {
 381     getWriter().clear();
 382   }
 383
 384   @Override
 385   public Field getField() {
 386     return getWriter().getField();
 387   }
 388
 389   @Override
 390   public int getValueCapacity() {
 391     return getWriter().getValueCapacity();
 392   }
 393
 394   @Override
 395   public void close() throws Exception {
 396     getWriter().close();
 397   }
 398 }