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.adapter.avro;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertNull;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  import java.io.File;
  import java.io.FileInputStream;
  import java.io.FileOutputStream;
  import java.io.InputStream;
  import java.lang.reflect.Method;
  import java.nio.ByteBuffer;
  import java.util.ArrayList;
  import java.util.List;
  import org.apache.arrow.memory.BufferAllocator;
  import org.apache.arrow.memory.RootAllocator;
  import org.apache.arrow.vector.FieldVector;
  import org.apache.arrow.vector.VectorSchemaRoot;
  import org.apache.arrow.vector.complex.ListVector;
  import org.apache.arrow.vector.complex.StructVector;
  import org.apache.arrow.vector.util.Text;
  import org.apache.avro.Schema;
  import org.apache.avro.generic.GenericDatumWriter;
  import org.apache.avro.generic.GenericRecord;
  import org.apache.avro.io.BinaryDecoder;
  import org.apache.avro.io.BinaryEncoder;
  import org.apache.avro.io.DatumWriter;
  import org.apache.avro.io.DecoderFactory;
  import org.apache.avro.io.EncoderFactory;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.io.TempDir;
  
  public class AvroTestBase {
  
    @TempDir public File TMP;
  
    protected AvroToArrowConfig config;
  
    @BeforeEach
    public void init() {
      BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
      config = new AvroToArrowConfigBuilder(allocator).build();
    }
  
    public static Schema getSchema(String schemaName) throws Exception {
      try {
        // Attempt to use JDK 9 behavior of getting the module then the resource stream from the
        // module.
        // Note that this code is caller-sensitive.
        Method getModuleMethod = Class.class.getMethod("getModule");
        Object module = getModuleMethod.invoke(TestWriteReadAvroRecord.class);
        Method getResourceAsStreamFromModule =
            module.getClass().getMethod("getResourceAsStream", String.class);
        try (InputStream is =
            (InputStream) getResourceAsStreamFromModule.invoke(module, "/schema/" + schemaName)) {
          return new Schema.Parser().parse(is);
        }
      } catch (NoSuchMethodException ex) {
        // Use JDK8 behavior.
        try (InputStream is =
            TestWriteReadAvroRecord.class.getResourceAsStream("/schema/" + schemaName)) {
          return new Schema.Parser().parse(is);
        }
      }
    }
  
    protected VectorSchemaRoot writeAndRead(Schema schema, List data) throws Exception {
      File dataFile = new File(TMP, "test.avro");
  
      try (FileOutputStream fos = new FileOutputStream(dataFile);
          FileInputStream fis = new FileInputStream(dataFile)) {
  
        BinaryEncoder encoder = new EncoderFactory().directBinaryEncoder(fos, null);
        DatumWriter<Object> writer = new GenericDatumWriter<>(schema);
        BinaryDecoder decoder = new DecoderFactory().directBinaryDecoder(fis, null);
  
        for (Object value : data) {
          writer.write(value, encoder);
        }
  
        return AvroToArrow.avroToArrow(schema, decoder, config);
      }
    }
  
    protected void checkArrayResult(List<List<?>> expected, ListVector vector) {
      assertEquals(expected.size(), vector.getValueCount());
      for (int i = 0; i < expected.size(); i++) {
        checkArrayElement(expected.get(i), vector.getObject(i));
      }
    }
  
    protected void checkArrayElement(List expected, List actual) {
      assertEquals(expected.size(), actual.size());
      for (int i = 0; i < expected.size(); i++) {
        Object value1 = expected.get(i);
        Object value2 = actual.get(i);
        if (value1 == null) {
          assertTrue(value2 == null);
          continue;
        }
        if (value2 instanceof byte[]) {
          value2 = ByteBuffer.wrap((byte[]) value2);
        } else if (value2 instanceof Text) {
          value2 = value2.toString();
        }
        assertEquals(value1, value2);
      }
    }
  
    protected void checkPrimitiveResult(List data, FieldVector vector) {
      assertEquals(data.size(), vector.getValueCount());
      for (int i = 0; i < data.size(); i++) {
        Object value1 = data.get(i);
        Object value2 = vector.getObject(i);
        if (value1 == null) {
          assertTrue(value2 == null);
          continue;
        }
        if (value2 instanceof byte[]) {
          value2 = ByteBuffer.wrap((byte[]) value2);
          if (value1 instanceof byte[]) {
            value1 = ByteBuffer.wrap((byte[]) value1);
          }
        } else if (value2 instanceof Text) {
          value2 = value2.toString();
        } else if (value2 instanceof Byte) {
          value2 = ((Byte) value2).intValue();
        }
        assertEquals(value1, value2);
      }
    }
  
    protected void checkRecordResult(Schema schema, List<GenericRecord> data, VectorSchemaRoot root) {
      assertEquals(data.size(), root.getRowCount());
      assertEquals(schema.getFields().size(), root.getFieldVectors().size());
  
      for (int i = 0; i < schema.getFields().size(); i++) {
        ArrayList fieldData = new ArrayList();
        for (GenericRecord record : data) {
          fieldData.add(record.get(i));
        }
  
        checkPrimitiveResult(fieldData, root.getFieldVectors().get(i));
      }
    }
  
    protected void checkNestedRecordResult(
        Schema schema, List<GenericRecord> data, VectorSchemaRoot root) {
      assertEquals(data.size(), root.getRowCount());
      assertTrue(schema.getFields().size() == 1);
  
      final Schema nestedSchema = schema.getFields().get(0).schema();
      final StructVector structVector = (StructVector) root.getFieldVectors().get(0);
  
      for (int i = 0; i < nestedSchema.getFields().size(); i++) {
        ArrayList fieldData = new ArrayList();
        for (GenericRecord record : data) {
          GenericRecord nestedRecord = (GenericRecord) record.get(0);
          fieldData.add(nestedRecord.get(i));
        }
  
        checkPrimitiveResult(fieldData, structVector.getChildrenFromFields().get(i));
      }
    }
  
    // belows are for iterator api
  
    protected void checkArrayResult(List<List<?>> expected, List<ListVector> vectors) {
      int valueCount = vectors.stream().mapToInt(v -> v.getValueCount()).sum();
      assertEquals(expected.size(), valueCount);
  
      int index = 0;
      for (ListVector vector : vectors) {
        for (int i = 0; i < vector.getValueCount(); i++) {
          checkArrayElement(expected.get(index++), vector.getObject(i));
        }
      }
    }
  
    protected void checkRecordResult(
        Schema schema, List<GenericRecord> data, List<VectorSchemaRoot> roots) {
      roots.forEach(
          root -> {
            assertEquals(schema.getFields().size(), root.getFieldVectors().size());
          });
  
      for (int i = 0; i < schema.getFields().size(); i++) {
        List fieldData = new ArrayList();
        List<FieldVector> vectors = new ArrayList<>();
        for (GenericRecord record : data) {
          fieldData.add(record.get(i));
        }
        final int columnIndex = i;
        roots.forEach(root -> vectors.add(root.getFieldVectors().get(columnIndex)));
  
        checkPrimitiveResult(fieldData, vectors);
      }
    }
  
    protected void checkPrimitiveResult(List data, List<FieldVector> vectors) {
      int valueCount = vectors.stream().mapToInt(v -> v.getValueCount()).sum();
      assertEquals(data.size(), valueCount);
  
      int index = 0;
      for (FieldVector vector : vectors) {
        for (int i = 0; i < vector.getValueCount(); i++) {
          Object value1 = data.get(index++);
          Object value2 = vector.getObject(i);
          if (value1 == null) {
            assertNull(value2);
            continue;
          }
          if (value2 instanceof byte[]) {
            value2 = ByteBuffer.wrap((byte[]) value2);
            if (value1 instanceof byte[]) {
              value1 = ByteBuffer.wrap((byte[]) value1);
            }
          } else if (value2 instanceof Text) {
            value2 = value2.toString();
          }
          assertEquals(value1, value2);
        }
      }
    }
  }
  