/*
 * 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 static java.nio.channels.Channels.newChannel;
import static java.util.Arrays.asList;
import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt;
import static org.apache.arrow.vector.TestUtils.newVarCharVector;
import static org.apache.arrow.vector.TestUtils.newVector;
import static org.apache.arrow.vector.testing.ValueVectorDataPopulator.setVector;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.channels.Channels;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.function.BiFunction;
import java.util.stream.Collectors;
import org.apache.arrow.flatbuf.FieldNode;
import org.apache.arrow.flatbuf.Message;
import org.apache.arrow.flatbuf.RecordBatch;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.util.Collections2;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.NullVector;
import org.apache.arrow.vector.TestUtils;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorLoader;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.VectorUnloader;
import org.apache.arrow.vector.compare.Range;
import org.apache.arrow.vector.compare.RangeEqualsVisitor;
import org.apache.arrow.vector.compare.TypeEqualsVisitor;
import org.apache.arrow.vector.compare.VectorEqualsVisitor;
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.dictionary.Dictionary;
import org.apache.arrow.vector.dictionary.DictionaryEncoder;
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.ArrowFieldNode;
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.MetadataVersion;
import org.apache.arrow.vector.types.Types.MinorType;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.DictionaryEncoding;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel;
import org.apache.arrow.vector.util.DictionaryUtility;
import org.apache.arrow.vector.util.TransferPair;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

public class TestArrowReaderWriter {

  private BufferAllocator allocator;

  private VarCharVector dictionaryVector1;
  private VarCharVector dictionaryVector2;
  private VarCharVector dictionaryVector3;
  private StructVector dictionaryVector4;

  private Dictionary dictionary1;
  private Dictionary dictionary2;
  private Dictionary dictionary3;
  private Dictionary dictionary4;

  private Schema schema;
  private Schema encodedSchema;

  @BeforeEach
  public void init() {
    allocator = new RootAllocator(Long.MAX_VALUE);

    dictionaryVector1 = newVarCharVector("D1", allocator);
    setVector(
        dictionaryVector1,
        "foo".getBytes(StandardCharsets.UTF_8),
        "bar".getBytes(StandardCharsets.UTF_8),
        "baz".getBytes(StandardCharsets.UTF_8));

    dictionaryVector2 = newVarCharVector("D2", allocator);
    setVector(
        dictionaryVector2,
        "aa".getBytes(StandardCharsets.UTF_8),
        "bb".getBytes(StandardCharsets.UTF_8),
        "cc".getBytes(StandardCharsets.UTF_8));

    dictionaryVector3 = newVarCharVector("D3", allocator);
    setVector(
        dictionaryVector3,
        "foo".getBytes(StandardCharsets.UTF_8),
        "bar".getBytes(StandardCharsets.UTF_8),
        "baz".getBytes(StandardCharsets.UTF_8),
        "aa".getBytes(StandardCharsets.UTF_8),
        "bb".getBytes(StandardCharsets.UTF_8),
        "cc".getBytes(StandardCharsets.UTF_8));

    dictionaryVector4 = newVector(StructVector.class, "D4", MinorType.STRUCT, allocator);
    final Map<String, List<Integer>> dictionaryValues4 = new HashMap<>();
    dictionaryValues4.put("a", Arrays.asList(1, 2, 3));
    dictionaryValues4.put("b", Arrays.asList(4, 5, 6));
    setVector(dictionaryVector4, dictionaryValues4);

    dictionary1 =
        new Dictionary(
            dictionaryVector1,
            new DictionaryEncoding(/*id=*/ 1L, /*ordered=*/ false, /*indexType=*/ null));
    dictionary2 =
        new Dictionary(
            dictionaryVector2,
            new DictionaryEncoding(/*id=*/ 2L, /*ordered=*/ false, /*indexType=*/ null));
    dictionary3 =
        new Dictionary(
            dictionaryVector3,
            new DictionaryEncoding(/*id=*/ 1L, /*ordered=*/ false, /*indexType=*/ null));
    dictionary4 =
        new Dictionary(
            dictionaryVector4,
            new DictionaryEncoding(/*id=*/ 3L, /*ordered=*/ false, /*indexType=*/ null));
  }

  @AfterEach
  public void terminate() throws Exception {
    dictionaryVector1.close();
    dictionaryVector2.close();
    dictionaryVector3.close();
    dictionaryVector4.close();
    allocator.close();
  }

  ArrowBuf buf(byte[] bytes) {
    ArrowBuf buffer = allocator.buffer(bytes.length);
    buffer.writeBytes(bytes);
    return buffer;
  }

  byte[] array(ArrowBuf buf) {
    byte[] bytes = new byte[checkedCastToInt(buf.readableBytes())];
    buf.readBytes(bytes);
    return bytes;
  }

  @Test
  public void test() throws IOException {
    Schema schema =
        new Schema(
            asList(
                new Field(
                    "testField",
                    FieldType.nullable(new ArrowType.Int(8, true)),
                    Collections.<Field>emptyList())));
    ArrowType type = schema.getFields().get(0).getType();
    FieldVector vector = TestUtils.newVector(FieldVector.class, "testField", type, allocator);
    vector.initializeChildrenFromFields(schema.getFields().get(0).getChildren());

    byte[] validity = new byte[] {(byte) 255, 0};
    // second half is "undefined"
    byte[] values = new byte[] {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16};

    ByteArrayOutputStream out = new ByteArrayOutputStream();
    try (VectorSchemaRoot root = new VectorSchemaRoot(schema.getFields(), asList(vector), 16);
        ArrowFileWriter writer = new ArrowFileWriter(root, null, newChannel(out))) {
      ArrowBuf validityb = buf(validity);
      ArrowBuf valuesb = buf(values);
      ArrowRecordBatch batch =
          new ArrowRecordBatch(16, asList(new ArrowFieldNode(16, 8)), asList(validityb, valuesb));
      VectorLoader loader = new VectorLoader(root);
      loader.load(batch);
      writer.writeBatch();

      validityb.close();
      valuesb.close();
      batch.close();
    }

    byte[] byteArray = out.toByteArray();

    try (SeekableReadChannel channel =
            new SeekableReadChannel(new ByteArrayReadableSeekableByteChannel(byteArray));
        ArrowFileReader reader = new ArrowFileReader(channel, allocator)) {
      Schema readSchema = reader.getVectorSchemaRoot().getSchema();
      assertEquals(schema, readSchema);
      // TODO: dictionaries
      List<ArrowBlock> recordBatches = reader.getRecordBlocks();
      assertEquals(1, recordBatches.size());
      reader.loadNextBatch();
      VectorUnloader unloader = new VectorUnloader(reader.getVectorSchemaRoot());
      ArrowRecordBatch recordBatch = unloader.getRecordBatch();
      List<ArrowFieldNode> nodes = recordBatch.getNodes();
      assertEquals(1, nodes.size());
      ArrowFieldNode node = nodes.get(0);
      assertEquals(16, node.getLength());
      assertEquals(8, node.getNullCount());
      List<ArrowBuf> buffers = recordBatch.getBuffers();
      assertEquals(2, buffers.size());
      assertArrayEquals(validity, array(buffers.get(0)));
      assertArrayEquals(values, array(buffers.get(1)));

      // Read just the header. This demonstrates being able to read without need to
      // deserialize the buffer.
      ByteBuffer headerBuffer = ByteBuffer.allocate(recordBatches.get(0).getMetadataLength());
      headerBuffer.put(byteArray, (int) recordBatches.get(0).getOffset(), headerBuffer.capacity());
      // new format prefix_size ==8
      headerBuffer.position(8);
      Message messageFB = Message.getRootAsMessage(headerBuffer);
      RecordBatch recordBatchFB = (RecordBatch) messageFB.header(new RecordBatch());
      assertEquals(2, recordBatchFB.buffersLength());
      assertEquals(1, recordBatchFB.nodesLength());
      FieldNode nodeFB = recordBatchFB.nodes(0);
      assertEquals(16, nodeFB.length());
      assertEquals(8, nodeFB.nullCount());

      recordBatch.close();
    }
  }

  @Test
  public void testWriteReadNullVector() throws IOException {

    int valueCount = 3;

    NullVector nullVector = new NullVector("vector");
    nullVector.setValueCount(valueCount);

    Schema schema = new Schema(asList(nullVector.getField()));

    ByteArrayOutputStream out = new ByteArrayOutputStream();
    try (VectorSchemaRoot root =
            new VectorSchemaRoot(schema.getFields(), asList(nullVector), valueCount);
        ArrowFileWriter writer = new ArrowFileWriter(root, null, newChannel(out))) {
      ArrowRecordBatch batch =
          new ArrowRecordBatch(
              valueCount, asList(new ArrowFieldNode(valueCount, 0)), Collections.emptyList());
      VectorLoader loader = new VectorLoader(root);
      loader.load(batch);
      writer.writeBatch();
    }

    byte[] byteArray = out.toByteArray();

    try (SeekableReadChannel channel =
            new SeekableReadChannel(new ByteArrayReadableSeekableByteChannel(byteArray));
        ArrowFileReader reader = new ArrowFileReader(channel, allocator)) {
      Schema readSchema = reader.getVectorSchemaRoot().getSchema();
      assertEquals(schema, readSchema);
      List<ArrowBlock> recordBatches = reader.getRecordBlocks();
      assertEquals(1, recordBatches.size());

      assertTrue(reader.loadNextBatch());
      assertEquals(1, reader.getVectorSchemaRoot().getFieldVectors().size());

      NullVector readNullVector =
          (NullVector) reader.getVectorSchemaRoot().getFieldVectors().get(0);
      assertEquals(valueCount, readNullVector.getValueCount());
    }
  }

  @Test
  public void testWriteReadWithDictionaries() throws IOException {
    DictionaryProvider.MapDictionaryProvider provider =
        new DictionaryProvider.MapDictionaryProvider();
    provider.put(dictionary1);

    VarCharVector vector1 = newVarCharVector("varchar1", allocator);
    vector1.allocateNewSafe();
    vector1.set(0, "foo".getBytes(StandardCharsets.UTF_8));
    vector1.set(1, "bar".getBytes(StandardCharsets.UTF_8));
    vector1.set(3, "baz".getBytes(StandardCharsets.UTF_8));
    vector1.set(4, "bar".getBytes(StandardCharsets.UTF_8));
    vector1.set(5, "baz".getBytes(StandardCharsets.UTF_8));
    vector1.setValueCount(6);
    FieldVector encodedVector1 = (FieldVector) DictionaryEncoder.encode(vector1, dictionary1);
    vector1.close();

    VarCharVector vector2 = newVarCharVector("varchar2", allocator);
    vector2.allocateNewSafe();
    vector2.set(0, "bar".getBytes(StandardCharsets.UTF_8));
    vector2.set(1, "baz".getBytes(StandardCharsets.UTF_8));
    vector2.set(2, "foo".getBytes(StandardCharsets.UTF_8));
    vector2.set(3, "foo".getBytes(StandardCharsets.UTF_8));
    vector2.set(4, "foo".getBytes(StandardCharsets.UTF_8));
    vector2.set(5, "bar".getBytes(StandardCharsets.UTF_8));
    vector2.setValueCount(6);
    FieldVector encodedVector2 = (FieldVector) DictionaryEncoder.encode(vector2, dictionary1);
    vector2.close();

    List<Field> fields = Arrays.asList(encodedVector1.getField(), encodedVector2.getField());
    List<FieldVector> vectors = Collections2.asImmutableList(encodedVector1, encodedVector2);
    try (VectorSchemaRoot root =
            new VectorSchemaRoot(fields, vectors, encodedVector1.getValueCount());
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        ArrowFileWriter writer = new ArrowFileWriter(root, provider, newChannel(out)); ) {

      writer.start();
      writer.writeBatch();
      writer.end();

      try (SeekableReadChannel channel =
              new SeekableReadChannel(new ByteArrayReadableSeekableByteChannel(out.toByteArray()));
          ArrowFileReader reader = new ArrowFileReader(channel, allocator)) {
        Schema readSchema = reader.getVectorSchemaRoot().getSchema();
        assertEquals(root.getSchema(), readSchema);
        assertEquals(1, reader.getDictionaryBlocks().size());
        assertEquals(1, reader.getRecordBlocks().size());

        reader.loadNextBatch();
        assertEquals(2, reader.getVectorSchemaRoot().getFieldVectors().size());
      }
    }
  }

  @Test
  public void testWriteReadWithStructDictionaries() throws IOException {
    DictionaryProvider.MapDictionaryProvider provider =
        new DictionaryProvider.MapDictionaryProvider();
    provider.put(dictionary4);

    try (final StructVector vector =
        newVector(StructVector.class, "D4", MinorType.STRUCT, allocator)) {
      final Map<String, List<Integer>> values = new HashMap<>();
      // Index: 0, 2, 1, 2, 1, 0, 0
      values.put("a", Arrays.asList(1, 3, 2, 3, 2, 1, 1));
      values.put("b", Arrays.asList(4, 6, 5, 6, 5, 4, 4));
      setVector(vector, values);
      FieldVector encodedVector = (FieldVector) DictionaryEncoder.encode(vector, dictionary4);

      List<Field> fields = Arrays.asList(encodedVector.getField());
      List<FieldVector> vectors = Collections2.asImmutableList(encodedVector);
      try (VectorSchemaRoot root =
              new VectorSchemaRoot(fields, vectors, encodedVector.getValueCount());
          ByteArrayOutputStream out = new ByteArrayOutputStream();
          ArrowFileWriter writer = new ArrowFileWriter(root, provider, newChannel(out)); ) {

        writer.start();
        writer.writeBatch();
        writer.end();

        try (SeekableReadChannel channel =
                new SeekableReadChannel(
                    new ByteArrayReadableSeekableByteChannel(out.toByteArray()));
            ArrowFileReader reader = new ArrowFileReader(channel, allocator)) {
          final VectorSchemaRoot readRoot = reader.getVectorSchemaRoot();
          final Schema readSchema = readRoot.getSchema();
          assertEquals(root.getSchema(), readSchema);
          assertEquals(1, reader.getDictionaryBlocks().size());
          assertEquals(1, reader.getRecordBlocks().size());

          reader.loadNextBatch();
          assertEquals(1, readRoot.getFieldVectors().size());
          assertEquals(1, reader.getDictionaryVectors().size());

          // Read the encoded vector and check it
          final FieldVector readEncoded = readRoot.getVector(0);
          assertEquals(encodedVector.getValueCount(), readEncoded.getValueCount());
          assertTrue(
              new RangeEqualsVisitor(encodedVector, readEncoded)
                  .rangeEquals(new Range(0, 0, encodedVector.getValueCount())));

          // Read the dictionary
          final Map<Long, Dictionary> readDictionaryMap = reader.getDictionaryVectors();
          final Dictionary readDictionary =
              readDictionaryMap.get(readEncoded.getField().getDictionary().getId());
          assertNotNull(readDictionary);

          // Assert the dictionary vector is correct
          final FieldVector readDictionaryVector = readDictionary.getVector();
          assertEquals(dictionaryVector4.getValueCount(), readDictionaryVector.getValueCount());
          final BiFunction<ValueVector, ValueVector, Boolean> typeComparatorIgnoreName =
              (v1, v2) -> new TypeEqualsVisitor(v1, false, true).equals(v2);
          assertTrue(
              new RangeEqualsVisitor(
                      dictionaryVector4, readDictionaryVector, typeComparatorIgnoreName)
                  .rangeEquals(new Range(0, 0, dictionaryVector4.getValueCount())),
              "Dictionary vectors are not equal");

          // Assert the decoded vector is correct
          try (final ValueVector readVector =
              DictionaryEncoder.decode(readEncoded, readDictionary)) {
            assertEquals(vector.getValueCount(), readVector.getValueCount());
            assertTrue(
                new RangeEqualsVisitor(vector, readVector, typeComparatorIgnoreName)
                    .rangeEquals(new Range(0, 0, vector.getValueCount())),
                "Decoded vectors are not equal");
          }
        }
      }
    }
  }

  @Test
  public void testEmptyStreamInFileIPC() throws IOException {

    DictionaryProvider.MapDictionaryProvider provider =
        new DictionaryProvider.MapDictionaryProvider();
    provider.put(dictionary1);

    VarCharVector vector = newVarCharVector("varchar", allocator);
    vector.allocateNewSafe();
    vector.set(0, "foo".getBytes(StandardCharsets.UTF_8));
    vector.set(1, "bar".getBytes(StandardCharsets.UTF_8));
    vector.set(3, "baz".getBytes(StandardCharsets.UTF_8));
    vector.set(4, "bar".getBytes(StandardCharsets.UTF_8));
    vector.set(5, "baz".getBytes(StandardCharsets.UTF_8));
    vector.setValueCount(6);

    FieldVector encodedVector1A = (FieldVector) DictionaryEncoder.encode(vector, dictionary1);
    vector.close();

    List<Field> fields = Arrays.asList(encodedVector1A.getField());
    List<FieldVector> vectors = Collections2.asImmutableList(encodedVector1A);

    try (VectorSchemaRoot root =
            new VectorSchemaRoot(fields, vectors, encodedVector1A.getValueCount());
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        ArrowFileWriter writer = new ArrowFileWriter(root, provider, newChannel(out))) {

      writer.start();
      writer.end();

      try (SeekableReadChannel channel =
              new SeekableReadChannel(new ByteArrayReadableSeekableByteChannel(out.toByteArray()));
          ArrowFileReader reader = new ArrowFileReader(channel, allocator)) {
        Schema readSchema = reader.getVectorSchemaRoot().getSchema();
        assertEquals(root.getSchema(), readSchema);
        assertEquals(1, reader.getDictionaryVectors().size());
        assertEquals(0, reader.getDictionaryBlocks().size());
        assertEquals(0, reader.getRecordBlocks().size());
      }
    }
  }

  @Test
  public void testEmptyStreamInStreamingIPC() throws IOException {

    DictionaryProvider.MapDictionaryProvider provider =
        new DictionaryProvider.MapDictionaryProvider();
    provider.put(dictionary1);

    VarCharVector vector = newVarCharVector("varchar", allocator);
    vector.allocateNewSafe();
    vector.set(0, "foo".getBytes(StandardCharsets.UTF_8));
    vector.set(1, "bar".getBytes(StandardCharsets.UTF_8));
    vector.set(3, "baz".getBytes(StandardCharsets.UTF_8));
    vector.set(4, "bar".getBytes(StandardCharsets.UTF_8));
    vector.set(5, "baz".getBytes(StandardCharsets.UTF_8));
    vector.setValueCount(6);

    FieldVector encodedVector = (FieldVector) DictionaryEncoder.encode(vector, dictionary1);
    vector.close();

    List<Field> fields = Arrays.asList(encodedVector.getField());
    try (VectorSchemaRoot root =
            new VectorSchemaRoot(
                fields, Arrays.asList(encodedVector), encodedVector.getValueCount());
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        ArrowStreamWriter writer = new ArrowStreamWriter(root, provider, newChannel(out))) {

      writer.start();
      writer.end();

      try (ArrowStreamReader reader =
          new ArrowStreamReader(
              new ByteArrayReadableSeekableByteChannel(out.toByteArray()), allocator)) {
        Schema readSchema = reader.getVectorSchemaRoot().getSchema();
        assertEquals(root.getSchema(), readSchema);
        assertEquals(1, reader.getDictionaryVectors().size());
        assertFalse(reader.loadNextBatch());
      }
    }
  }

  @Test
  public void testDictionaryReplacement() throws Exception {
    VarCharVector vector1 = newVarCharVector("varchar1", allocator);
    setVector(
        vector1,
        "foo".getBytes(StandardCharsets.UTF_8),
        "bar".getBytes(StandardCharsets.UTF_8),
        "baz".getBytes(StandardCharsets.UTF_8),
        "bar".getBytes(StandardCharsets.UTF_8));

    FieldVector encodedVector1 = (FieldVector) DictionaryEncoder.encode(vector1, dictionary1);

    VarCharVector vector2 = newVarCharVector("varchar2", allocator);
    setVector(
        vector2,
        "foo".getBytes(StandardCharsets.UTF_8),
        "foo".getBytes(StandardCharsets.UTF_8),
        "foo".getBytes(StandardCharsets.UTF_8),
        "foo".getBytes(StandardCharsets.UTF_8));

    FieldVector encodedVector2 = (FieldVector) DictionaryEncoder.encode(vector2, dictionary1);

    DictionaryProvider.MapDictionaryProvider provider =
        new DictionaryProvider.MapDictionaryProvider();
    provider.put(dictionary1);
    List<Field> schemaFields = new ArrayList<>();
    schemaFields.add(
        DictionaryUtility.toMessageFormat(encodedVector1.getField(), provider, new HashSet<>()));
    schemaFields.add(
        DictionaryUtility.toMessageFormat(encodedVector2.getField(), provider, new HashSet<>()));
    Schema schema = new Schema(schemaFields);

    ByteArrayOutputStream outStream = new ByteArrayOutputStream();
    WriteChannel out = new WriteChannel(newChannel(outStream));

    // write schema
    MessageSerializer.serialize(out, schema);

    List<AutoCloseable> closeableList = new ArrayList<>();

    // write non-delta dictionary with id=1
    serializeDictionaryBatch(out, dictionary3, false, closeableList);

    // write non-delta dictionary with id=1
    serializeDictionaryBatch(out, dictionary1, false, closeableList);

    // write recordBatch2
    serializeRecordBatch(out, Arrays.asList(encodedVector1, encodedVector2), closeableList);

    // write eos
    out.writeIntLittleEndian(0);

    try (ArrowStreamReader reader =
        new ArrowStreamReader(
            new ByteArrayReadableSeekableByteChannel(outStream.toByteArray()), allocator)) {
      assertEquals(1, reader.getDictionaryVectors().size());
      assertTrue(reader.loadNextBatch());
      FieldVector dictionaryVector = reader.getDictionaryVectors().get(1L).getVector();
      // make sure the delta dictionary is concatenated.
      assertTrue(VectorEqualsVisitor.vectorEquals(dictionaryVector, dictionaryVector1, null));
      assertFalse(reader.loadNextBatch());
    }

    vector1.close();
    vector2.close();
    AutoCloseables.close(closeableList);
  }

  @Test
  public void testDeltaDictionary() throws Exception {
    VarCharVector vector1 = newVarCharVector("varchar1", allocator);
    setVector(
        vector1,
        "foo".getBytes(StandardCharsets.UTF_8),
        "bar".getBytes(StandardCharsets.UTF_8),
        "baz".getBytes(StandardCharsets.UTF_8),
        "bar".getBytes(StandardCharsets.UTF_8));

    FieldVector encodedVector1 = (FieldVector) DictionaryEncoder.encode(vector1, dictionary1);

    VarCharVector vector2 = newVarCharVector("varchar2", allocator);
    setVector(
        vector2,
        "foo".getBytes(StandardCharsets.UTF_8),
        "aa".getBytes(StandardCharsets.UTF_8),
        "bb".getBytes(StandardCharsets.UTF_8),
        "cc".getBytes(StandardCharsets.UTF_8));

    FieldVector encodedVector2 = (FieldVector) DictionaryEncoder.encode(vector2, dictionary3);

    DictionaryProvider.MapDictionaryProvider provider =
        new DictionaryProvider.MapDictionaryProvider();
    provider.put(dictionary1);
    provider.put(dictionary3);
    List<Field> schemaFields = new ArrayList<>();
    schemaFields.add(
        DictionaryUtility.toMessageFormat(encodedVector1.getField(), provider, new HashSet<>()));
    schemaFields.add(
        DictionaryUtility.toMessageFormat(encodedVector2.getField(), provider, new HashSet<>()));
    Schema schema = new Schema(schemaFields);

    ByteArrayOutputStream outStream = new ByteArrayOutputStream();
    WriteChannel out = new WriteChannel(newChannel(outStream));

    // write schema
    MessageSerializer.serialize(out, schema);

    List<AutoCloseable> closeableList = new ArrayList<>();

    // write non-delta dictionary with id=1
    serializeDictionaryBatch(out, dictionary1, false, closeableList);

    // write delta dictionary with id=1
    Dictionary deltaDictionary =
        new Dictionary(dictionaryVector2, new DictionaryEncoding(1L, false, null));
    serializeDictionaryBatch(out, deltaDictionary, true, closeableList);
    deltaDictionary.getVector().close();

    // write recordBatch2
    serializeRecordBatch(out, Arrays.asList(encodedVector1, encodedVector2), closeableList);

    // write eos
    out.writeIntLittleEndian(0);

    try (ArrowStreamReader reader =
        new ArrowStreamReader(
            new ByteArrayReadableSeekableByteChannel(outStream.toByteArray()), allocator)) {
      assertEquals(1, reader.getDictionaryVectors().size());
      assertTrue(reader.loadNextBatch());
      FieldVector dictionaryVector = reader.getDictionaryVectors().get(1L).getVector();
      // make sure the delta dictionary is concatenated.
      assertTrue(VectorEqualsVisitor.vectorEquals(dictionaryVector, dictionaryVector3, null));
      assertFalse(reader.loadNextBatch());
    }

    vector1.close();
    vector2.close();
    AutoCloseables.close(closeableList);
  }

  // Tests that the ArrowStreamWriter re-emits dictionaries when they change
  @Test
  public void testWriteReadStreamWithDictionaryReplacement() throws Exception {
    DictionaryProvider.MapDictionaryProvider provider =
        new DictionaryProvider.MapDictionaryProvider();
    provider.put(dictionary1);

    String[] batch0 = {"foo", "bar", "baz", "bar", "baz"};
    String[] batch1 = {"foo", "aa", "bar", "bb", "baz", "cc"};

    VarCharVector vector = newVarCharVector("varchar", allocator);
    vector.allocateNewSafe();
    for (int i = 0; i < batch0.length; ++i) {
      vector.set(i, batch0[i].getBytes(StandardCharsets.UTF_8));
    }
    vector.setValueCount(batch0.length);
    FieldVector encodedVector1 = (FieldVector) DictionaryEncoder.encode(vector, dictionary1);

    List<Field> fields = Arrays.asList(encodedVector1.getField());
    try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
      try (VectorSchemaRoot root =
              new VectorSchemaRoot(
                  fields, Arrays.asList(encodedVector1), encodedVector1.getValueCount());
          ArrowStreamWriter writer = new ArrowStreamWriter(root, provider, newChannel(out))) {
        writer.start();

        // Write batch with initial data and dictionary
        writer.writeBatch();

        // Create data for the next batch, using an extended dictionary with the same id
        vector.reset();
        for (int i = 0; i < batch1.length; ++i) {
          vector.set(i, batch1[i].getBytes(StandardCharsets.UTF_8));
        }
        vector.setValueCount(batch1.length);

        // Re-encode and move encoded data into the vector schema root
        provider.put(dictionary3);
        FieldVector encodedVector2 = (FieldVector) DictionaryEncoder.encode(vector, dictionary3);
        TransferPair transferPair = encodedVector2.makeTransferPair(root.getVector(0));
        transferPair.transfer();

        // Write second batch
        root.setRowCount(batch1.length);
        writer.writeBatch();

        writer.end();
      }

      try (ArrowStreamReader reader =
          new ArrowStreamReader(
              new ByteArrayReadableSeekableByteChannel(out.toByteArray()), allocator)) {
        VectorSchemaRoot root = reader.getVectorSchemaRoot();

        // Read and verify first batch
        assertTrue(reader.loadNextBatch());
        assertEquals(batch0.length, root.getRowCount());
        FieldVector readEncoded1 = root.getVector(0);
        long dictionaryId = readEncoded1.getField().getDictionary().getId();
        try (VarCharVector decodedValues =
            (VarCharVector) DictionaryEncoder.decode(readEncoded1, reader.lookup(dictionaryId))) {
          for (int i = 0; i < batch0.length; ++i) {
            assertEquals(batch0[i], new String(decodedValues.get(i), StandardCharsets.UTF_8));
          }
        }

        // Read and verify second batch
        assertTrue(reader.loadNextBatch());
        assertEquals(batch1.length, root.getRowCount());
        FieldVector readEncoded2 = root.getVector(0);
        dictionaryId = readEncoded2.getField().getDictionary().getId();
        try (VarCharVector decodedValues =
            (VarCharVector) DictionaryEncoder.decode(readEncoded2, reader.lookup(dictionaryId))) {
          for (int i = 0; i < batch1.length; ++i) {
            assertEquals(batch1[i], new String(decodedValues.get(i), StandardCharsets.UTF_8));
          }
        }

        assertFalse(reader.loadNextBatch());
      }
    }

    vector.close();
  }

  private void serializeDictionaryBatch(
      WriteChannel out, Dictionary dictionary, boolean isDelta, List<AutoCloseable> closeables)
      throws IOException {

    FieldVector dictVector = dictionary.getVector();
    VectorSchemaRoot root =
        new VectorSchemaRoot(
            Collections.singletonList(dictVector.getField()),
            Collections.singletonList(dictVector),
            dictVector.getValueCount());
    ArrowDictionaryBatch batch =
        new ArrowDictionaryBatch(
            dictionary.getEncoding().getId(), new VectorUnloader(root).getRecordBatch(), isDelta);
    MessageSerializer.serialize(out, batch);
    closeables.add(batch);
    closeables.add(root);
  }

  private void serializeRecordBatch(
      WriteChannel out, List<FieldVector> vectors, List<AutoCloseable> closeables)
      throws IOException {

    List<Field> fields = vectors.stream().map(v -> v.getField()).collect(Collectors.toList());
    VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, vectors.get(0).getValueCount());
    VectorUnloader unloader = new VectorUnloader(root);
    ArrowRecordBatch batch = unloader.getRecordBatch();
    MessageSerializer.serialize(out, batch);
    closeables.add(batch);
    closeables.add(root);
  }

  @Test
  public void testReadInterleavedData() throws IOException {
    List<ArrowRecordBatch> batches = createRecordBatches();

    ByteArrayOutputStream outStream = new ByteArrayOutputStream();
    WriteChannel out = new WriteChannel(newChannel(outStream));

    // write schema
    MessageSerializer.serialize(out, schema);

    // write dictionary1
    FieldVector dictVector1 = dictionary1.getVector();
    VectorSchemaRoot dictRoot1 =
        new VectorSchemaRoot(
            Collections.singletonList(dictVector1.getField()),
            Collections.singletonList(dictVector1),
            dictVector1.getValueCount());
    ArrowDictionaryBatch dictionaryBatch1 =
        new ArrowDictionaryBatch(1, new VectorUnloader(dictRoot1).getRecordBatch());
    MessageSerializer.serialize(out, dictionaryBatch1);
    dictionaryBatch1.close();
    dictRoot1.close();

    // write recordBatch1
    MessageSerializer.serialize(out, batches.get(0));

    // write dictionary2
    FieldVector dictVector2 = dictionary2.getVector();
    VectorSchemaRoot dictRoot2 =
        new VectorSchemaRoot(
            Collections.singletonList(dictVector2.getField()),
            Collections.singletonList(dictVector2),
            dictVector2.getValueCount());
    ArrowDictionaryBatch dictionaryBatch2 =
        new ArrowDictionaryBatch(2, new VectorUnloader(dictRoot2).getRecordBatch());
    MessageSerializer.serialize(out, dictionaryBatch2);
    dictionaryBatch2.close();
    dictRoot2.close();

    // write recordBatch1
    MessageSerializer.serialize(out, batches.get(1));

    // write eos
    out.writeIntLittleEndian(0);

    try (ArrowStreamReader reader =
        new ArrowStreamReader(
            new ByteArrayReadableSeekableByteChannel(outStream.toByteArray()), allocator)) {
      Schema readSchema = reader.getVectorSchemaRoot().getSchema();
      assertEquals(encodedSchema, readSchema);
      assertEquals(2, reader.getDictionaryVectors().size());
      assertTrue(reader.loadNextBatch());
      assertTrue(reader.loadNextBatch());
      assertFalse(reader.loadNextBatch());
    }

    batches.forEach(batch -> batch.close());
  }

  private List<ArrowRecordBatch> createRecordBatches() {
    List<ArrowRecordBatch> batches = new ArrayList<>();

    DictionaryProvider.MapDictionaryProvider provider =
        new DictionaryProvider.MapDictionaryProvider();
    provider.put(dictionary1);
    provider.put(dictionary2);

    VarCharVector vectorA1 = newVarCharVector("varcharA1", allocator);
    vectorA1.allocateNewSafe();
    vectorA1.set(0, "foo".getBytes(StandardCharsets.UTF_8));
    vectorA1.set(1, "bar".getBytes(StandardCharsets.UTF_8));
    vectorA1.set(3, "baz".getBytes(StandardCharsets.UTF_8));
    vectorA1.set(4, "bar".getBytes(StandardCharsets.UTF_8));
    vectorA1.set(5, "baz".getBytes(StandardCharsets.UTF_8));
    vectorA1.setValueCount(6);

    VarCharVector vectorA2 = newVarCharVector("varcharA2", allocator);
    vectorA2.setValueCount(6);
    FieldVector encodedVectorA1 = (FieldVector) DictionaryEncoder.encode(vectorA1, dictionary1);
    vectorA1.close();
    FieldVector encodedVectorA2 = (FieldVector) DictionaryEncoder.encode(vectorA1, dictionary2);
    vectorA2.close();

    List<Field> fields = Arrays.asList(encodedVectorA1.getField(), encodedVectorA2.getField());
    List<FieldVector> vectors = Collections2.asImmutableList(encodedVectorA1, encodedVectorA2);
    VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, encodedVectorA1.getValueCount());
    VectorUnloader unloader = new VectorUnloader(root);
    batches.add(unloader.getRecordBatch());
    root.close();

    VarCharVector vectorB1 = newVarCharVector("varcharB1", allocator);
    vectorB1.setValueCount(6);

    VarCharVector vectorB2 = newVarCharVector("varcharB2", allocator);
    vectorB2.allocateNew();
    vectorB2.setValueCount(6);
    vectorB2.set(0, "aa".getBytes(StandardCharsets.UTF_8));
    vectorB2.set(1, "aa".getBytes(StandardCharsets.UTF_8));
    vectorB2.set(3, "bb".getBytes(StandardCharsets.UTF_8));
    vectorB2.set(4, "bb".getBytes(StandardCharsets.UTF_8));
    vectorB2.set(5, "cc".getBytes(StandardCharsets.UTF_8));
    vectorB2.setValueCount(6);
    FieldVector encodedVectorB1 = (FieldVector) DictionaryEncoder.encode(vectorB1, dictionary1);
    vectorB1.close();
    FieldVector encodedVectorB2 = (FieldVector) DictionaryEncoder.encode(vectorB2, dictionary2);
    vectorB2.close();

    List<Field> fieldsB = Arrays.asList(encodedVectorB1.getField(), encodedVectorB2.getField());
    List<FieldVector> vectorsB = Collections2.asImmutableList(encodedVectorB1, encodedVectorB2);
    VectorSchemaRoot rootB = new VectorSchemaRoot(fieldsB, vectorsB, 6);
    VectorUnloader unloaderB = new VectorUnloader(rootB);
    batches.add(unloaderB.getRecordBatch());
    rootB.close();

    List<Field> schemaFields = new ArrayList<>();
    schemaFields.add(
        DictionaryUtility.toMessageFormat(encodedVectorA1.getField(), provider, new HashSet<>()));
    schemaFields.add(
        DictionaryUtility.toMessageFormat(encodedVectorA2.getField(), provider, new HashSet<>()));
    schema = new Schema(schemaFields);

    encodedSchema =
        new Schema(Arrays.asList(encodedVectorA1.getField(), encodedVectorA2.getField()));

    return batches;
  }

  @Test
  public void testLegacyIpcBackwardsCompatibility() throws Exception {
    Schema schema = new Schema(asList(Field.nullable("field", new ArrowType.Int(32, true))));
    IntVector vector = new IntVector("vector", allocator);
    final int valueCount = 2;
    vector.setValueCount(valueCount);
    vector.setSafe(0, 1);
    vector.setSafe(1, 2);
    ArrowRecordBatch batch =
        new ArrowRecordBatch(
            valueCount,
            asList(new ArrowFieldNode(valueCount, 0)),
            asList(vector.getValidityBuffer(), vector.getDataBuffer()));

    ByteArrayOutputStream outStream = new ByteArrayOutputStream();
    WriteChannel out = new WriteChannel(newChannel(outStream));

    // write legacy ipc format
    IpcOption option = new IpcOption(true, MetadataVersion.DEFAULT);
    MessageSerializer.serialize(out, schema, option);
    MessageSerializer.serialize(out, batch);

    ReadChannel in = new ReadChannel(newChannel(new ByteArrayInputStream(outStream.toByteArray())));
    Schema readSchema = MessageSerializer.deserializeSchema(in);
    assertEquals(schema, readSchema);
    ArrowRecordBatch readBatch = MessageSerializer.deserializeRecordBatch(in, allocator);
    assertEquals(batch.getLength(), readBatch.getLength());
    assertEquals(batch.computeBodyLength(), readBatch.computeBodyLength());
    readBatch.close();

    // write ipc format with continuation
    option = IpcOption.DEFAULT;
    MessageSerializer.serialize(out, schema, option);
    MessageSerializer.serialize(out, batch);

    ReadChannel in2 =
        new ReadChannel(newChannel(new ByteArrayInputStream(outStream.toByteArray())));
    Schema readSchema2 = MessageSerializer.deserializeSchema(in2);
    assertEquals(schema, readSchema2);
    ArrowRecordBatch readBatch2 = MessageSerializer.deserializeRecordBatch(in2, allocator);
    assertEquals(batch.getLength(), readBatch2.getLength());
    assertEquals(batch.computeBodyLength(), readBatch2.computeBodyLength());
    readBatch2.close();

    batch.close();
    vector.close();
  }

  @Test
  public void testChannelReadFully() throws IOException {
    final ByteBuffer buf = ByteBuffer.allocate(4).order(ByteOrder.nativeOrder());
    buf.putInt(200);
    buf.rewind();

    try (ReadChannel channel =
            new ReadChannel(Channels.newChannel(new ByteArrayInputStream(buf.array())));
        ArrowBuf arrBuf = allocator.buffer(8)) {
      arrBuf.setInt(0, 100);
      arrBuf.writerIndex(4);
      assertEquals(4, arrBuf.writerIndex());

      long n = channel.readFully(arrBuf, 4);
      assertEquals(4, n);
      assertEquals(8, arrBuf.writerIndex());

      assertEquals(100, arrBuf.getInt(0));
      assertEquals(200, arrBuf.getInt(4));
    }
  }

  @Test
  public void testChannelReadFullyEos() throws IOException {
    final ByteBuffer buf = ByteBuffer.allocate(4).order(ByteOrder.nativeOrder());
    buf.putInt(10);
    buf.rewind();

    try (ReadChannel channel =
            new ReadChannel(Channels.newChannel(new ByteArrayInputStream(buf.array())));
        ArrowBuf arrBuf = allocator.buffer(8)) {
      int n = channel.readFully(arrBuf.nioBuffer(0, 8));
      assertEquals(4, n);

      // the input has only 4 bytes, so the number of bytes read should be 4
      assertEquals(4, channel.bytesRead());

      // the first 4 bytes have been read successfully.
      assertEquals(10, arrBuf.getInt(0));
    }
  }

  @Test
  public void testCustomMetaData() throws IOException {

    VarCharVector vector = newVarCharVector("varchar1", allocator);

    List<Field> fields = Arrays.asList(vector.getField());
    List<FieldVector> vectors = Collections2.asImmutableList(vector);
    Map<String, String> metadata = new HashMap<>();
    metadata.put("key1", "value1");
    metadata.put("key2", "value2");
    try (VectorSchemaRoot root = new VectorSchemaRoot(fields, vectors, vector.getValueCount());
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        ArrowFileWriter writer = new ArrowFileWriter(root, null, newChannel(out), metadata); ) {

      writer.start();
      writer.end();

      try (SeekableReadChannel channel =
              new SeekableReadChannel(new ByteArrayReadableSeekableByteChannel(out.toByteArray()));
          ArrowFileReader reader = new ArrowFileReader(channel, allocator)) {
        reader.getVectorSchemaRoot();

        Map<String, String> readMeta = reader.getMetaData();
        assertEquals(2, readMeta.size());
        assertEquals("value1", readMeta.get("key1"));
        assertEquals("value2", readMeta.get("key2"));
      }
    }
  }

  /**
   * This test case covers the case for which the footer size is extremely large (much larger than
   * the file size). Due to integer overflow, our implementation fails detect the problem, which
   * leads to extremely large memory allocation and eventually causing an OutOfMemoryError.
   */
  @Test
  public void testFileFooterSizeOverflow() {
    // copy of org.apache.arrow.vector.ipc.ArrowMagic#MAGIC
    final byte[] magicBytes = "ARROW1".getBytes(StandardCharsets.UTF_8);

    // prepare input data
    byte[] data = new byte[30];
    System.arraycopy(magicBytes, 0, data, 0, ArrowMagic.MAGIC_LENGTH);
    int footerLength = Integer.MAX_VALUE;
    byte[] footerLengthBytes =
        ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN).putInt(footerLength).array();
    int footerOffset = data.length - ArrowMagic.MAGIC_LENGTH - 4;
    System.arraycopy(footerLengthBytes, 0, data, footerOffset, 4);
    System.arraycopy(magicBytes, 0, data, footerOffset + 4, ArrowMagic.MAGIC_LENGTH);

    // test file reader
    InvalidArrowFileException e =
        assertThrows(
            InvalidArrowFileException.class,
            () -> {
              try (SeekableReadChannel channel =
                      new SeekableReadChannel(new ByteArrayReadableSeekableByteChannel(data));
                  ArrowFileReader reader = new ArrowFileReader(channel, allocator)) {
                reader.getVectorSchemaRoot().getSchema();
              }
            });

    assertEquals("invalid footer length: " + footerLength, e.getMessage());
  }
}
