ceph/src/arrow/java/c/src/main/java/org/apache/arrow/c/ArrayExporter.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.c;
  19
  20 import static org.apache.arrow.c.NativeUtil.NULL;
  21 import static org.apache.arrow.c.NativeUtil.addressOrNull;
  22 import static org.apache.arrow.util.Preconditions.checkNotNull;
  23
  24 import java.util.ArrayList;
  25 import java.util.List;
  26
  27 import org.apache.arrow.c.jni.JniWrapper;
  28 import org.apache.arrow.c.jni.PrivateData;
  29 import org.apache.arrow.memory.ArrowBuf;
  30 import org.apache.arrow.memory.BufferAllocator;
  31 import org.apache.arrow.vector.FieldVector;
  32 import org.apache.arrow.vector.dictionary.Dictionary;
  33 import org.apache.arrow.vector.dictionary.DictionaryProvider;
  34 import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
  35
  36 /**
  37  * Exporter for {@link ArrowArray}.
  38  */
  39 final class ArrayExporter {
  40   private final BufferAllocator allocator;
  41
  42   public ArrayExporter(BufferAllocator allocator) {
  43     this.allocator = allocator;
  44   }
  45
  46   /**
  47    * Private data structure for exported arrays.
  48    */
  49   static class ExportedArrayPrivateData implements PrivateData {
  50     ArrowBuf buffers_ptrs;
  51     List<ArrowBuf> buffers;
  52     ArrowBuf children_ptrs;
  53     List<ArrowArray> children;
  54     ArrowArray dictionary;
  55
  56     @Override
  57     public void close() {
  58       NativeUtil.closeBuffer(buffers_ptrs);
  59
  60       if (buffers != null) {
  61         for (ArrowBuf buffer : buffers) {
  62           NativeUtil.closeBuffer(buffer);
  63         }
  64       }
  65       NativeUtil.closeBuffer(children_ptrs);
  66
  67       if (children != null) {
  68         for (ArrowArray child : children) {
  69           child.close();
  70         }
  71       }
  72
  73       if (dictionary != null) {
  74         dictionary.close();
  75       }
  76     }
  77   }
  78
  79   void export(ArrowArray array, FieldVector vector, DictionaryProvider dictionaryProvider) {
  80     List<FieldVector> children = vector.getChildrenFromFields();
  81     List<ArrowBuf> buffers = vector.getFieldBuffers();
  82     int valueCount = vector.getValueCount();
  83     int nullCount = vector.getNullCount();
  84     DictionaryEncoding dictionaryEncoding = vector.getField().getDictionary();
  85
  86     ExportedArrayPrivateData data = new ExportedArrayPrivateData();
  87     try {
  88       if (children != null) {
  89         data.children = new ArrayList<>(children.size());
  90         data.children_ptrs = allocator.buffer((long) children.size() * Long.BYTES);
  91         for (int i = 0; i < children.size(); i++) {
  92           ArrowArray child = ArrowArray.allocateNew(allocator);
  93           data.children.add(child);
  94           data.children_ptrs.writeLong(child.memoryAddress());
  95         }
  96       }
  97
  98       if (buffers != null) {
  99         data.buffers = new ArrayList<>(buffers.size());
 100         data.buffers_ptrs = allocator.buffer((long) buffers.size() * Long.BYTES);
 101         for (ArrowBuf arrowBuf : buffers) {
 102           if (arrowBuf != null) {
 103             arrowBuf.getReferenceManager().retain();
 104             data.buffers_ptrs.writeLong(arrowBuf.memoryAddress());
 105           } else {
 106             data.buffers_ptrs.writeLong(NULL);
 107           }
 108           data.buffers.add(arrowBuf);
 109         }
 110       }
 111
 112       if (dictionaryEncoding != null) {
 113         Dictionary dictionary = dictionaryProvider.lookup(dictionaryEncoding.getId());
 114         checkNotNull(dictionary, "Dictionary lookup failed on export of dictionary encoded array");
 115
 116         data.dictionary = ArrowArray.allocateNew(allocator);
 117         FieldVector dictionaryVector = dictionary.getVector();
 118         export(data.dictionary, dictionaryVector, dictionaryProvider);
 119       }
 120
 121       ArrowArray.Snapshot snapshot = new ArrowArray.Snapshot();
 122       snapshot.length = valueCount;
 123       snapshot.null_count = nullCount;
 124       snapshot.offset = 0;
 125       snapshot.n_buffers = (data.buffers != null) ? data.buffers.size() : 0;
 126       snapshot.n_children = (data.children != null) ? data.children.size() : 0;
 127       snapshot.buffers = addressOrNull(data.buffers_ptrs);
 128       snapshot.children = addressOrNull(data.children_ptrs);
 129       snapshot.dictionary = addressOrNull(data.dictionary);
 130       snapshot.release = NULL;
 131       array.save(snapshot);
 132
 133       // sets release and private data
 134       JniWrapper.get().exportArray(array.memoryAddress(), data);
 135     } catch (Exception e) {
 136       data.close();
 137       throw e;
 138     }
 139
 140     // Export children
 141     if (children != null) {
 142       for (int i = 0; i < children.size(); i++) {
 143         FieldVector childVector = children.get(i);
 144         ArrowArray child = data.children.get(i);
 145         export(child, childVector, dictionaryProvider);
 146       }
 147     }
 148   }
 149 }