java/vector/src/main/java/org/apache/arrow/vector/ipc/ArrowFileWriter.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.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;

import org.apache.arrow.util.VisibleForTesting;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.compression.CompressionCodec;
import org.apache.arrow.vector.compression.CompressionUtil;
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.ArrowFooter;
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.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * {@link ArrowWriter} that writes out a Arrow files (https://arrow.apache.org/docs/format/IPC.html#file-format).
 */
public class ArrowFileWriter extends ArrowWriter {

  private static final Logger LOGGER = LoggerFactory.getLogger(ArrowFileWriter.class);

  // All ArrowBlocks written are saved in these lists to be passed to ArrowFooter in endInternal.
  private final List<ArrowBlock> dictionaryBlocks = new ArrayList<>();
  private final List<ArrowBlock> recordBlocks = new ArrayList<>();

  private Map<String, String> metaData;
  private boolean dictionariesWritten = false;

  public ArrowFileWriter(VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out) {
    super(root, provider, out);
  }

  public ArrowFileWriter(VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out,
      Map<String, String> metaData) {
    super(root, provider, out);
    this.metaData = metaData;
  }

  public ArrowFileWriter(VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out,
      IpcOption option) {
    super(root, provider, out, option);
  }

  public ArrowFileWriter(VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out,
                         Map<String, String> metaData, IpcOption option) {
    super(root, provider, out, option);
    this.metaData = metaData;
  }

  public ArrowFileWriter(VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out,
                         Map<String, String> metaData, IpcOption option, CompressionCodec.Factory compressionFactory,
                         CompressionUtil.CodecType codecType) {
    this(root, provider, out, metaData, option, compressionFactory, codecType, Optional.empty());
  }

  public ArrowFileWriter(VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out,
                         Map<String, String> metaData, IpcOption option, CompressionCodec.Factory compressionFactory,
                         CompressionUtil.CodecType codecType, Optional<Integer> compressionLevel) {
    super(root, provider, out, option, compressionFactory, codecType, compressionLevel);
    this.metaData = metaData;
  }

  @Override
  protected void startInternal(WriteChannel out) throws IOException {
    ArrowMagic.writeMagic(out, true);
  }

  @Override
  protected ArrowBlock writeDictionaryBatch(ArrowDictionaryBatch batch) throws IOException {
    ArrowBlock block = super.writeDictionaryBatch(batch);
    dictionaryBlocks.add(block);
    return block;
  }

  @Override
  protected ArrowBlock writeRecordBatch(ArrowRecordBatch batch) throws IOException {
    ArrowBlock block = super.writeRecordBatch(batch);
    recordBlocks.add(block);
    return block;
  }

  @Override
  protected void endInternal(WriteChannel out) throws IOException {
    if (!option.write_legacy_ipc_format) {
      out.writeIntLittleEndian(MessageSerializer.IPC_CONTINUATION_TOKEN);
    }
    out.writeIntLittleEndian(0);

    long footerStart = out.getCurrentPosition();
    out.write(new ArrowFooter(schema, dictionaryBlocks, recordBlocks, metaData, option.metadataVersion), false);
    int footerLength = (int) (out.getCurrentPosition() - footerStart);
    if (footerLength <= 0) {
      throw new InvalidArrowFileException("invalid footer");
    }
    out.writeIntLittleEndian(footerLength);
    LOGGER.debug("Footer starts at {}, length: {}", footerStart, footerLength);
    ArrowMagic.writeMagic(out, false);
    LOGGER.debug("magic written, now at {}", out.getCurrentPosition());
  }

  @Override
  protected void ensureDictionariesWritten(DictionaryProvider provider, Set<Long> dictionaryIdsUsed)
      throws IOException {
    if (dictionariesWritten) {
      return;
    }
    dictionariesWritten = true;
    // Write out all dictionaries required.
    // Replacement dictionaries are not supported in the IPC file format.
    for (long id : dictionaryIdsUsed) {
      Dictionary dictionary = provider.lookup(id);
      writeDictionaryBatch(dictionary);
    }
  }

  @VisibleForTesting
  public List<ArrowBlock> getRecordBlocks() {
    return recordBlocks;
  }

  @VisibleForTesting
  public List<ArrowBlock> getDictionaryBlocks() {
    return dictionaryBlocks;
  }
}