content: /*
   * 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.c;
  
  import static org.assertj.core.api.Assertions.assertThat;
  import static org.junit.jupiter.api.Assertions.assertThrows;
  
  import java.io.IOException;
  import java.nio.charset.StandardCharsets;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.function.Function;
  import java.util.stream.Collectors;
  import org.apache.arrow.memory.BufferAllocator;
  import org.apache.arrow.memory.RootAllocator;
  import org.apache.arrow.util.AutoCloseables;
  import org.apache.arrow.vector.FieldVector;
  import org.apache.arrow.vector.IntVector;
  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.ViewVarBinaryVector;
  import org.apache.arrow.vector.ViewVarCharVector;
  import org.apache.arrow.vector.compare.Range;
  import org.apache.arrow.vector.compare.RangeEqualsVisitor;
  import org.apache.arrow.vector.dictionary.Dictionary;
  import org.apache.arrow.vector.dictionary.DictionaryProvider;
  import org.apache.arrow.vector.ipc.ArrowReader;
  import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
  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.junit.jupiter.api.AfterEach;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Test;
  
  final class StreamTest {
    private RootAllocator allocator = null;
  
    @BeforeEach
    public void setUp() {
      allocator = new RootAllocator(Long.MAX_VALUE);
    }
  
    @AfterEach
    public void tearDown() {
      allocator.close();
    }
  
    @Test
    public void testRoundtripInts() throws Exception {
      final Schema schema =
          new Schema(Collections.singletonList(Field.nullable("ints", new ArrowType.Int(32, true))));
      final List<ArrowRecordBatch> batches = new ArrayList<>();
      try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
        final IntVector ints = (IntVector) root.getVector(0);
        VectorUnloader unloader = new VectorUnloader(root);
  
        root.allocateNew();
        ints.setSafe(0, 1);
        ints.setSafe(1, 2);
        ints.setSafe(2, 4);
        ints.setSafe(3, 8);
        root.setRowCount(4);
        batches.add(unloader.getRecordBatch());
  
        root.allocateNew();
        ints.setSafe(0, 1);
        ints.setNull(1);
        ints.setSafe(2, 4);
        ints.setNull(3);
        root.setRowCount(4);
        batches.add(unloader.getRecordBatch());
        roundtrip(schema, batches);
      }
    }
  
    @Test
    public void roundtripStrings() throws Exception {
      final Schema schema =
          new Schema(
              Arrays.asList(
                  Field.nullable("ints", new ArrowType.Int(32, true)),
                  Field.nullable("strs", new ArrowType.Utf8())));
      final List<ArrowRecordBatch> batches = new ArrayList<>();
      try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
        final IntVector ints = (IntVector) root.getVector(0);
        final VarCharVector strs = (VarCharVector) root.getVector(1);
        VectorUnloader unloader = new VectorUnloader(root);
  
        root.allocateNew();
        ints.setSafe(0, 1);
        ints.setSafe(1, 2);
        ints.setSafe(2, 4);
        ints.setSafe(3, 8);
        strs.setSafe(0, "".getBytes(StandardCharsets.UTF_8));
        strs.setSafe(1, "a".getBytes(StandardCharsets.UTF_8));
        strs.setSafe(2, "bc".getBytes(StandardCharsets.UTF_8));
        strs.setSafe(3, "defg".getBytes(StandardCharsets.UTF_8));
        root.setRowCount(4);
        batches.add(unloader.getRecordBatch());
  
        root.allocateNew();
        ints.setSafe(0, 1);
        ints.setNull(1);
        ints.setSafe(2, 4);
        ints.setNull(3);
        strs.setSafe(0, "".getBytes(StandardCharsets.UTF_8));
        strs.setNull(1);
        strs.setSafe(2, "bc".getBytes(StandardCharsets.UTF_8));
        strs.setNull(3);
        root.setRowCount(4);
        batches.add(unloader.getRecordBatch());
        roundtrip(schema, batches);
      }
    }
  
    @Test
    public void roundtripStringViews() throws Exception {
      final Schema schema =
          new Schema(
              Arrays.asList(
                  Field.nullable("ints", new ArrowType.Int(32, true)),
                  Field.nullable("string_views", new ArrowType.Utf8View())));
      final List<ArrowRecordBatch> batches = new ArrayList<>();
      try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
        final IntVector ints = (IntVector) root.getVector(0);
        final ViewVarCharVector strs = (ViewVarCharVector) root.getVector(1);
        VectorUnloader unloader = new VectorUnloader(root);
  
        root.allocateNew();
        ints.setSafe(0, 1);
        ints.setSafe(1, 2);
        ints.setSafe(2, 4);
        ints.setSafe(3, 8);
        strs.setSafe(0, "".getBytes(StandardCharsets.UTF_8));
        strs.setSafe(1, "a".getBytes(StandardCharsets.UTF_8));
        strs.setSafe(2, "bc1234567890bc".getBytes(StandardCharsets.UTF_8));
        strs.setSafe(3, "defg1234567890defg".getBytes(StandardCharsets.UTF_8));
        root.setRowCount(4);
        batches.add(unloader.getRecordBatch());
  
        root.allocateNew();
        ints.setSafe(0, 1);
        ints.setNull(1);
        ints.setSafe(2, 4);
        ints.setNull(3);
        strs.setSafe(0, "".getBytes(StandardCharsets.UTF_8));
        strs.setNull(1);
        strs.setSafe(2, "bc1234567890bc".getBytes(StandardCharsets.UTF_8));
        strs.setNull(3);
        root.setRowCount(4);
        batches.add(unloader.getRecordBatch());
        roundtrip(schema, batches);
      }
    }
  
    @Test
    public void roundtripBinaryViews() throws Exception {
      final Schema schema =
          new Schema(
              Arrays.asList(
                  Field.nullable("ints", new ArrowType.Int(32, true)),
                  Field.nullable("binary_views", new ArrowType.BinaryView())));
      final List<ArrowRecordBatch> batches = new ArrayList<>();
      try (final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
        final IntVector ints = (IntVector) root.getVector(0);
        final ViewVarBinaryVector strs = (ViewVarBinaryVector) root.getVector(1);
        VectorUnloader unloader = new VectorUnloader(root);
  
        root.allocateNew();
        ints.setSafe(0, 1);
        ints.setSafe(1, 2);
        ints.setSafe(2, 4);
        ints.setSafe(3, 8);
        strs.setSafe(0, new byte[0]);
        strs.setSafe(1, new byte[] {97});
        strs.setSafe(2, new byte[] {98, 99, 49, 50, 51, 52, 53, 54, 55, 56, 57, 48, 98, 99});
        strs.setSafe(
            3,
            new byte[] {
              100, 101, 102, 103, 49, 50, 51, 52, 53, 54, 55, 56, 57, 48, 100, 101, 102, 103
            });
        root.setRowCount(4);
        batches.add(unloader.getRecordBatch());
  
        root.allocateNew();
        ints.setSafe(0, 1);
        ints.setNull(1);
        ints.setSafe(2, 4);
        ints.setNull(3);
        strs.setSafe(0, new byte[0]);
        strs.setNull(1);
        strs.setSafe(2, new byte[] {98, 99, 49, 50, 51, 52, 53, 54, 55, 56, 57, 48, 98, 99});
        strs.setNull(3);
        root.setRowCount(4);
        batches.add(unloader.getRecordBatch());
        roundtrip(schema, batches);
      }
    }
  
    @Test
    public void roundtripDictionary() throws Exception {
      final ArrowType.Int indexType = new ArrowType.Int(32, true);
      final DictionaryEncoding encoding = new DictionaryEncoding(0L, false, indexType);
      final Schema schema =
          new Schema(
              Collections.singletonList(
                  new Field(
                      "dict",
                      new FieldType(/* nullable= */ true, indexType, encoding),
                      Collections.emptyList())));
      final List<ArrowRecordBatch> batches = new ArrayList<>();
      try (final CDataDictionaryProvider provider = new CDataDictionaryProvider();
          final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
        final VarCharVector dictionary = new VarCharVector("values", allocator);
        dictionary.allocateNew();
        dictionary.setSafe(0, "foo".getBytes(StandardCharsets.UTF_8));
        dictionary.setSafe(1, "bar".getBytes(StandardCharsets.UTF_8));
        dictionary.setNull(2);
        dictionary.setValueCount(3);
        provider.put(new Dictionary(dictionary, encoding));
        final IntVector encoded = (IntVector) root.getVector(0);
        VectorUnloader unloader = new VectorUnloader(root);
  
        root.allocateNew();
        encoded.setSafe(0, 0);
        encoded.setSafe(1, 1);
        encoded.setSafe(2, 0);
        encoded.setSafe(3, 2);
        root.setRowCount(4);
        batches.add(unloader.getRecordBatch());
  
        root.allocateNew();
        encoded.setSafe(0, 0);
        encoded.setNull(1);
        encoded.setSafe(2, 1);
        encoded.setNull(3);
        root.setRowCount(4);
        batches.add(unloader.getRecordBatch());
        roundtrip(schema, batches, provider);
      }
    }
  
    @Test
    public void importReleasedStream() {
      try (final ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) {
        Exception e =
            assertThrows(
                IllegalStateException.class, () -> Data.importArrayStream(allocator, stream));
        assertThat(e).hasMessageContaining("Cannot import released ArrowArrayStream");
      }
    }
  
    @Test
    public void getNextError() throws Exception {
      final Schema schema =
          new Schema(Collections.singletonList(Field.nullable("ints", new ArrowType.Int(32, true))));
      final List<ArrowRecordBatch> batches = new ArrayList<>();
      try (final ArrowReader source =
              new InMemoryArrowReader(
                  allocator, schema, batches, new DictionaryProvider.MapDictionaryProvider()) {
                @Override
                public boolean loadNextBatch() throws IOException {
                  throw new IOException("Failed to load batch!");
                }
              };
          final ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) {
        Data.exportArrayStream(allocator, source, stream);
        try (final ArrowReader reader = Data.importArrayStream(allocator, stream)) {
          assertThat(reader.getVectorSchemaRoot().getSchema()).isEqualTo(schema);
          final IOException e = assertThrows(IOException.class, reader::loadNextBatch);
          assertThat(e).hasMessageContaining("Failed to load batch!");
          assertThat(e).hasMessageContaining("[errno ");
        }
      }
    }
  
    @Test
    public void getSchemaError() throws Exception {
      final Schema schema =
          new Schema(Collections.singletonList(Field.nullable("ints", new ArrowType.Int(32, true))));
      final List<ArrowRecordBatch> batches = new ArrayList<>();
      try (final ArrowReader source =
              new InMemoryArrowReader(
                  allocator, schema, batches, new DictionaryProvider.MapDictionaryProvider()) {
                @Override
                protected Schema readSchema() {
                  throw new IllegalArgumentException("Failed to read schema!");
                }
              };
          final ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) {
        Data.exportArrayStream(allocator, source, stream);
        try (final ArrowReader reader = Data.importArrayStream(allocator, stream)) {
          final IOException e = assertThrows(IOException.class, reader::getVectorSchemaRoot);
          assertThat(e).hasMessageContaining("Failed to read schema!");
          assertThat(e).hasMessageContaining("[errno ");
        }
      }
    }
  
    void roundtrip(Schema schema, List<ArrowRecordBatch> batches, DictionaryProvider provider)
        throws Exception {
      ArrowReader source = new InMemoryArrowReader(allocator, schema, batches, provider);
  
      try (final ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator);
          final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
        final VectorLoader loader = new VectorLoader(root);
        Data.exportArrayStream(allocator, source, stream);
  
        try (final ArrowReader reader = Data.importArrayStream(allocator, stream)) {
          assertThat(reader.getVectorSchemaRoot().getSchema()).isEqualTo(schema);
  
          for (ArrowRecordBatch batch : batches) {
            assertThat(reader.loadNextBatch()).isTrue();
            loader.load(batch);
  
            assertThat(reader.getVectorSchemaRoot().getRowCount()).isEqualTo(root.getRowCount());
  
            for (int i = 0; i < root.getFieldVectors().size(); i++) {
              final FieldVector expected = root.getVector(i);
              final FieldVector actual = reader.getVectorSchemaRoot().getVector(i);
              assertVectorsEqual(expected, actual);
            }
          }
          assertThat(reader.loadNextBatch()).isFalse();
          assertThat(reader.getDictionaryIds()).isEqualTo(provider.getDictionaryIds());
          for (Map.Entry<Long, Dictionary> entry : reader.getDictionaryVectors().entrySet()) {
            final FieldVector expected = provider.lookup(entry.getKey()).getVector();
            final FieldVector actual = entry.getValue().getVector();
            assertVectorsEqual(expected, actual);
          }
        }
      }
    }
  
    void roundtrip(Schema schema, List<ArrowRecordBatch> batches) throws Exception {
      roundtrip(schema, batches, new CDataDictionaryProvider());
    }
  
    private static void assertVectorsEqual(FieldVector expected, FieldVector actual) {
      assertThat(actual.getField().getType()).isEqualTo(expected.getField().getType());
      assertThat(actual.getValueCount()).isEqualTo(expected.getValueCount());
      final Range range =
          new Range(/* leftStart= */ 0, /* rightStart= */ 0, expected.getValueCount());
      assertThat(new RangeEqualsVisitor(expected, actual).rangeEquals(range))
          .as("Vectors were not equal.\nExpected: %s\nGot: %s", expected, actual)
          .isTrue();
    }
  
    /** An ArrowReader backed by a fixed list of batches. */
    static class InMemoryArrowReader extends ArrowReader {
      private final Schema schema;
      private final List<ArrowRecordBatch> batches;
      private final DictionaryProvider provider;
      private int nextBatch;
  
      InMemoryArrowReader(
          BufferAllocator allocator,
          Schema schema,
          List<ArrowRecordBatch> batches,
          DictionaryProvider provider) {
        super(allocator);
        this.schema = schema;
        this.batches = batches;
        this.provider = provider;
        this.nextBatch = 0;
      }
  
      @Override
      public Dictionary lookup(long id) {
        return provider.lookup(id);
      }
  
      @Override
      public Set<Long> getDictionaryIds() {
        return provider.getDictionaryIds();
      }
  
      @Override
      public Map<Long, Dictionary> getDictionaryVectors() {
        return getDictionaryIds().stream()
            .collect(Collectors.toMap(Function.identity(), this::lookup));
      }
  
      @Override
      public boolean loadNextBatch() throws IOException {
        if (nextBatch < batches.size()) {
          VectorLoader loader = new VectorLoader(getVectorSchemaRoot());
          loader.load(batches.get(nextBatch++));
          return true;
        }
        return false;
      }
  
      @Override
      public long bytesRead() {
        return 0;
      }
  
      @Override
      protected void closeReadSource() throws IOException {
        try {
          AutoCloseables.close(batches);
        } catch (Exception e) {
          throw new IOException(e);
        }
      }
  
      @Override
      protected Schema readSchema() {
        return schema;
      }
    }
  }
  