com.aliyun.odps.table.arrow.writers.ArrowCompressVectorUnloader Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of odps-sdk-table-api Show documentation
The newest version!
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package com.aliyun.odps.table.arrow.writers;

import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.TypeLayout;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.VectorUnloader;
import org.apache.arrow.vector.compression.CompressionCodec;
import org.apache.arrow.vector.compression.CompressionUtil;
import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;

import java.util.ArrayList;
import java.util.List;

/**
 * Helper class that handles converting a {@link VectorSchemaRoot}
 * to a {@link ArrowRecordBatch}.
 */
class ArrowCompressVectorUnloader extends VectorUnloader {

    private final VectorSchemaRoot root;
    private final boolean includeNullCount;
    private final CompressionCodec codec;
    private final boolean alignBuffers;

    ArrowCompressVectorUnloader(
            VectorSchemaRoot root, boolean includeNullCount, CompressionCodec codec, boolean alignBuffers) {
        super(root, includeNullCount, codec, alignBuffers);
        this.root = root;
        this.includeNullCount = includeNullCount;
        this.codec = codec;
        this.alignBuffers = alignBuffers;
    }

    public ArrowRecordBatch getRecordBatch() {
        List nodes = new ArrayList<>();
        List buffers = new ArrayList<>();
        for (FieldVector vector : root.getFieldVectors()) {
            appendNodes(vector, nodes, buffers);
        }
        return new ArrowCompressRecordBatch(
                root.getRowCount(), nodes, buffers, CompressionUtil.createBodyCompression(codec), alignBuffers);
    }

    private void appendNodes(FieldVector vector, List nodes, List buffers) {
        nodes.add(new ArrowFieldNode(vector.getValueCount(), includeNullCount ? vector.getNullCount() : -1));
        List fieldBuffers = vector.getFieldBuffers();
        int expectedBufferCount = TypeLayout.getTypeBufferCount(vector.getField().getType());
        if (fieldBuffers.size() != expectedBufferCount) {
            throw new IllegalArgumentException(String.format(
                    "wrong number of buffers for field %s in vector %s. found: %s",
                    vector.getField(), vector.getClass().getSimpleName(), fieldBuffers));
        }
        for (ArrowBuf buf : fieldBuffers) {
            // If the codec is CompressionCodec, it will return a new buffer, and also close the input buffer.
            // In that case, we need to retain the input buffer still to avoid modifying the source VectorSchemaRoot.
            buf.getReferenceManager().retain();
            buffers.add(codec.compress(vector.getAllocator(), buf));
        }
        for (FieldVector child : vector.getChildrenFromFields()) {
            appendNodes(child, nodes, buffers);
        }
    }
}