/*
 * 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.parquet.hadoop;

import static org.assertj.core.api.Assertions.assertThat;

import java.io.IOException;
import org.apache.hadoop.fs.Path;
import org.apache.parquet.example.data.Group;
import org.apache.parquet.example.data.simple.SimpleGroup;
import org.apache.parquet.hadoop.example.GroupReadSupport;
import org.junit.jupiter.api.Test;

public class TestInterOpReadByteStreamSplit {
  private static final String FLOATS_FILE = "byte_stream_split.zstd.parquet";
  private static final String EXTENDED_FILE = "byte_stream_split_extended.gzip.parquet";
  private static final String CHANGESET = "74278bc";
  private final InterOpTester interop = new InterOpTester();

  @Test
  public void testReadFloats() throws IOException {
    Path floatsFile = interop.GetInterOpFile(FLOATS_FILE, CHANGESET);
    final int expectRows = 300;

    try (ParquetReader<Group> reader =
        ParquetReader.builder(new GroupReadSupport(), floatsFile).build()) {
      for (int i = 0; i < expectRows; ++i) {
        Group group = reader.read();
        assertThat(group).isNotNull();
        float fval = group.getFloat("f32", 0);
        double dval = group.getDouble("f64", 0);
        // Values are from the normal distribution
        assertThat(Math.abs(fval)).isLessThan(4.0f);
        assertThat(Math.abs(dval)).isLessThan(4.0);
        switch (i) {
          case 0:
            assertThat(fval).isEqualTo(1.7640524f);
            assertThat(dval).isEqualTo(-1.3065268517353166);
            break;
          case 1:
            assertThat(fval).isEqualTo(0.4001572f);
            assertThat(dval).isEqualTo(1.658130679618188);
            break;
          case expectRows - 2:
            assertThat(fval).isEqualTo(-0.39944902f);
            assertThat(dval).isEqualTo(-0.9301565025243212);
            break;
          case expectRows - 1:
            assertThat(fval).isEqualTo(0.37005588f);
            assertThat(dval).isEqualTo(-0.17858909208732915);
            break;
        }
      }
      assertThat(reader.read()).isNull();
    }
  }

  private void compareColumnValues(Path path, int expectRows, String leftCol, String rightCol) throws IOException {
    try (ParquetReader<Group> reader =
        ParquetReader.builder(new GroupReadSupport(), path).build()) {
      for (int i = 0; i < expectRows; ++i) {
        SimpleGroup group = (SimpleGroup) reader.read();
        assertThat(group).isNotNull();
        Object left = group.getObject(leftCol, 0);
        Object right = group.getObject(rightCol, 0);
        assertThat(right).isEqualTo(left);
      }
      assertThat(reader.read()).isNull();
    }
  }

  @Test
  public void testReadAllSupportedTypes() throws IOException {
    Path extendedFile = interop.GetInterOpFile(EXTENDED_FILE, CHANGESET);
    final int expectRows = 200;
    compareColumnValues(extendedFile, expectRows, "float_plain", "float_byte_stream_split");
    compareColumnValues(extendedFile, expectRows, "double_plain", "double_byte_stream_split");
    compareColumnValues(extendedFile, expectRows, "int32_plain", "int32_byte_stream_split");
    compareColumnValues(extendedFile, expectRows, "int64_plain", "int64_byte_stream_split");
    compareColumnValues(extendedFile, expectRows, "float16_plain", "float16_byte_stream_split");
    compareColumnValues(extendedFile, expectRows, "flba5_plain", "flba5_byte_stream_split");
    compareColumnValues(extendedFile, expectRows, "decimal_plain", "decimal_byte_stream_split");
  }
}
