ceph/src/arrow/java/c/src/main/java/org/apache/arrow/vector/StructVectorUnloader.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 java.util.ArrayList;
  21 import java.util.List;
  22
  23 import org.apache.arrow.memory.ArrowBuf;
  24 import org.apache.arrow.vector.complex.StructVector;
  25 import org.apache.arrow.vector.compression.CompressionCodec;
  26 import org.apache.arrow.vector.compression.CompressionUtil;
  27 import org.apache.arrow.vector.compression.NoCompressionCodec;
  28 import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
  29 import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
  30
  31 /**
  32  * Helper class that handles converting a {@link StructVector} to a
  33  * {@link ArrowRecordBatch}.
  34  */
  35 public class StructVectorUnloader {
  36
  37   private final StructVector root;
  38   private final boolean includeNullCount;
  39   private final CompressionCodec codec;
  40   private final boolean alignBuffers;
  41
  42   /**
  43    * Constructs a new instance of the given struct vector.
  44    */
  45   public StructVectorUnloader(StructVector root) {
  46     this(root, true, NoCompressionCodec.INSTANCE, true);
  47   }
  48
  49   /**
  50    * Constructs a new instance.
  51    *
  52    * @param root             The struct vector to serialize to an
  53    *                         {@link ArrowRecordBatch}.
  54    * @param includeNullCount Controls whether null count is copied to the
  55    *                         {@link ArrowRecordBatch}
  56    * @param alignBuffers     Controls if buffers get aligned to 8-byte boundaries.
  57    */
  58   public StructVectorUnloader(StructVector root, boolean includeNullCount, boolean alignBuffers) {
  59     this(root, includeNullCount, NoCompressionCodec.INSTANCE, alignBuffers);
  60   }
  61
  62   /**
  63    * Constructs a new instance.
  64    *
  65    * @param root             The struct vector to serialize to an
  66    *                         {@link ArrowRecordBatch}.
  67    * @param includeNullCount Controls whether null count is copied to the
  68    *                         {@link ArrowRecordBatch}
  69    * @param codec            the codec for compressing data. If it is null, then
  70    *                         no compression is needed.
  71    * @param alignBuffers     Controls if buffers get aligned to 8-byte boundaries.
  72    */
  73   public StructVectorUnloader(StructVector root, boolean includeNullCount, CompressionCodec codec,
  74       boolean alignBuffers) {
  75     this.root = root;
  76     this.includeNullCount = includeNullCount;
  77     this.codec = codec;
  78     this.alignBuffers = alignBuffers;
  79   }
  80
  81   /**
  82    * Performs the depth first traversal of the Vectors to create an
  83    * {@link ArrowRecordBatch} suitable for serialization.
  84    */
  85   public ArrowRecordBatch getRecordBatch() {
  86     List<ArrowFieldNode> nodes = new ArrayList<>();
  87     List<ArrowBuf> buffers = new ArrayList<>();
  88     for (FieldVector vector : root.getChildrenFromFields()) {
  89       appendNodes(vector, nodes, buffers);
  90     }
  91     return new ArrowRecordBatch(root.getValueCount(), nodes, buffers, CompressionUtil.createBodyCompression(codec),
  92         alignBuffers);
  93   }
  94
  95   private void appendNodes(FieldVector vector, List<ArrowFieldNode> nodes, List<ArrowBuf> buffers) {
  96     nodes.add(new ArrowFieldNode(vector.getValueCount(), includeNullCount ? vector.getNullCount() : -1));
  97     List<ArrowBuf> fieldBuffers = vector.getFieldBuffers();
  98     int expectedBufferCount = TypeLayout.getTypeBufferCount(vector.getField().getType());
  99     if (fieldBuffers.size() != expectedBufferCount) {
 100       throw new IllegalArgumentException(String.format("wrong number of buffers for field %s in vector %s. found: %s",
 101           vector.getField(), vector.getClass().getSimpleName(), fieldBuffers));
 102     }
 103     for (ArrowBuf buf : fieldBuffers) {
 104       buffers.add(codec.compress(vector.getAllocator(), buf));
 105     }
 106     for (FieldVector child : vector.getChildrenFromFields()) {
 107       appendNodes(child, nodes, buffers);
 108     }
 109   }
 110 }