java/vector/src/main/java/org/apache/arrow/vector/ipc/ArrowWriter.java

/*
 * 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 org.apache.arrow.vector.ipc;

import java.io.IOException;
import java.nio.channels.WritableByteChannel;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import org.apache.arrow.vector.FieldVector;
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.compression.NoCompressionCodec;
import org.apache.arrow.vector.dictionary.Dictionary;
import org.apache.arrow.vector.dictionary.DictionaryProvider;
import org.apache.arrow.vector.ipc.message.ArrowBlock;
import org.apache.arrow.vector.ipc.message.ArrowDictionaryBatch;
import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
import org.apache.arrow.vector.ipc.message.IpcOption;
import org.apache.arrow.vector.ipc.message.MessageSerializer;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.DictionaryUtility;
import org.apache.arrow.vector.validate.MetadataV4UnionChecker;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Abstract base class for implementing Arrow writers for IPC over a WriteChannel. */
public abstract class ArrowWriter implements AutoCloseable {

  protected static final Logger LOGGER = LoggerFactory.getLogger(ArrowWriter.class);

  // schema with fields in message format, not memory format
  protected final Schema schema;
  protected final WriteChannel out;

  private final VectorUnloader unloader;
  private final DictionaryProvider dictionaryProvider;
  private final Set<Long> dictionaryIdsUsed = new HashSet<>();

  private final CompressionCodec.Factory compressionFactory;
  private final CompressionUtil.CodecType codecType;
  private final Optional<Integer> compressionLevel;
  private boolean started = false;
  private boolean ended = false;

  private final CompressionCodec codec;

  protected IpcOption option;

  protected ArrowWriter(
      VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out) {
    this(root, provider, out, IpcOption.DEFAULT);
  }

  protected ArrowWriter(
      VectorSchemaRoot root,
      DictionaryProvider provider,
      WritableByteChannel out,
      IpcOption option) {
    this(
        root,
        provider,
        out,
        option,
        NoCompressionCodec.Factory.INSTANCE,
        CompressionUtil.CodecType.NO_COMPRESSION,
        Optional.empty());
  }

  /**
   * Note: fields are not closed when the writer is closed.
   *
   * @param root the vectors to write to the output
   * @param provider where to find the dictionaries
   * @param out the output where to write
   * @param option IPC write options
   * @param compressionFactory Compression codec factory
   * @param codecType Compression codec
   * @param compressionLevel Compression level
   */
  protected ArrowWriter(
      VectorSchemaRoot root,
      DictionaryProvider provider,
      WritableByteChannel out,
      IpcOption option,
      CompressionCodec.Factory compressionFactory,
      CompressionUtil.CodecType codecType,
      Optional<Integer> compressionLevel) {
    this.out = new WriteChannel(out);
    this.option = option;
    this.dictionaryProvider = provider;

    this.compressionFactory = compressionFactory;
    this.codecType = codecType;
    this.compressionLevel = compressionLevel;
    this.codec =
        this.compressionLevel.isPresent()
            ? this.compressionFactory.createCodec(this.codecType, this.compressionLevel.get())
            : this.compressionFactory.createCodec(this.codecType);
    this.unloader =
        new VectorUnloader(root, /*includeNullCount*/ true, codec, /*alignBuffers*/ true);

    List<Field> fields = new ArrayList<>(root.getSchema().getFields().size());

    MetadataV4UnionChecker.checkForUnion(
        root.getSchema().getFields().iterator(), option.metadataVersion);
    // Convert fields with dictionaries to have dictionary type
    for (Field field : root.getSchema().getFields()) {
      fields.add(DictionaryUtility.toMessageFormat(field, provider, dictionaryIdsUsed));
    }

    this.schema = new Schema(fields, root.getSchema().getCustomMetadata());
  }

  public void start() throws IOException {
    ensureStarted();
  }

  /** Writes the record batch currently loaded in this instance's VectorSchemaRoot. */
  public void writeBatch() throws IOException {
    ensureStarted();
    ensureDictionariesWritten(dictionaryProvider, dictionaryIdsUsed);
    try (ArrowRecordBatch batch = unloader.getRecordBatch()) {
      writeRecordBatch(batch);
    }
  }

  protected void writeDictionaryBatch(Dictionary dictionary) throws IOException {
    FieldVector vector = dictionary.getVector();
    long id = dictionary.getEncoding().getId();
    int count = vector.getValueCount();
    VectorSchemaRoot dictRoot =
        new VectorSchemaRoot(
            Collections.singletonList(vector.getField()), Collections.singletonList(vector), count);
    VectorUnloader unloader =
        new VectorUnloader(dictRoot, /*includeNullCount*/ true, this.codec, /*alignBuffers*/ true);
    ArrowRecordBatch batch = unloader.getRecordBatch();
    ArrowDictionaryBatch dictionaryBatch = new ArrowDictionaryBatch(id, batch, false);
    try {
      writeDictionaryBatch(dictionaryBatch);
    } finally {
      try {
        dictionaryBatch.close();
      } catch (Exception e) {
        throw new RuntimeException("Error occurred while closing dictionary.", e);
      }
    }
  }

  protected ArrowBlock writeDictionaryBatch(ArrowDictionaryBatch batch) throws IOException {
    ArrowBlock block = MessageSerializer.serialize(out, batch, option);
    if (LOGGER.isDebugEnabled()) {
      LOGGER.debug(
          "DictionaryRecordBatch at {}, metadata: {}, body: {}",
          block.getOffset(),
          block.getMetadataLength(),
          block.getBodyLength());
    }
    return block;
  }

  protected ArrowBlock writeRecordBatch(ArrowRecordBatch batch) throws IOException {
    ArrowBlock block = MessageSerializer.serialize(out, batch, option);
    if (LOGGER.isDebugEnabled()) {
      LOGGER.debug(
          "RecordBatch at {}, metadata: {}, body: {}",
          block.getOffset(),
          block.getMetadataLength(),
          block.getBodyLength());
    }
    return block;
  }

  public void end() throws IOException {
    ensureStarted();
    ensureEnded();
  }

  public long bytesWritten() {
    return out.getCurrentPosition();
  }

  private void ensureStarted() throws IOException {
    if (!started) {
      started = true;
      startInternal(out);
      // write the schema - for file formats this is duplicated in the footer, but matches
      // the streaming format
      MessageSerializer.serialize(out, schema, option);
    }
  }

  /**
   * Write dictionaries after schema and before recordBatches, dictionaries won't be written if
   * empty stream (only has schema data in IPC).
   */
  protected abstract void ensureDictionariesWritten(
      DictionaryProvider provider, Set<Long> dictionaryIdsUsed) throws IOException;

  private void ensureEnded() throws IOException {
    if (!ended) {
      ended = true;
      endInternal(out);
    }
  }

  protected void startInternal(WriteChannel out) throws IOException {}

  protected void endInternal(WriteChannel out) throws IOException {}

  @Override
  public void close() {
    try {
      end();
      out.close();
    } catch (Exception e) {
      throw new RuntimeException(e);
    }
  }
}