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.dataset;
  
  import static org.hamcrest.MatcherAssert.assertThat;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertNotNull;
  import static org.junit.jupiter.api.Assertions.assertThrows;
  
  import com.google.common.collect.ImmutableMap;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.Map;
  import java.util.Optional;
  import org.apache.arrow.c.ArrowSchema;
  import org.apache.arrow.c.CDataDictionaryProvider;
  import org.apache.arrow.c.Data;
  import org.apache.arrow.dataset.file.FileFormat;
  import org.apache.arrow.dataset.file.FileSystemDatasetFactory;
  import org.apache.arrow.dataset.jni.NativeMemoryPool;
  import org.apache.arrow.dataset.scanner.ScanOptions;
  import org.apache.arrow.dataset.scanner.Scanner;
  import org.apache.arrow.dataset.scanner.csv.CsvConvertOptions;
  import org.apache.arrow.dataset.scanner.csv.CsvFragmentScanOptions;
  import org.apache.arrow.dataset.source.Dataset;
  import org.apache.arrow.dataset.source.DatasetFactory;
  import org.apache.arrow.memory.BufferAllocator;
  import org.apache.arrow.memory.RootAllocator;
  import org.apache.arrow.vector.ValueIterableVector;
  import org.apache.arrow.vector.ipc.ArrowReader;
  import org.apache.arrow.vector.types.pojo.ArrowType;
  import org.apache.arrow.vector.types.pojo.Field;
  import org.apache.arrow.vector.types.pojo.Schema;
  import org.apache.arrow.vector.util.Text;
  import org.hamcrest.collection.IsIterableContainingInOrder;
  import org.junit.jupiter.api.Test;
  
  public class TestFragmentScanOptions {
  
    private CsvFragmentScanOptions create(
        ArrowSchema cSchema,
        Map<String, String> convertOptionsMap,
        Map<String, String> readOptions,
        Map<String, String> parseOptions) {
      CsvConvertOptions convertOptions = new CsvConvertOptions(convertOptionsMap);
      convertOptions.setArrowSchema(cSchema);
      return new CsvFragmentScanOptions(convertOptions, readOptions, parseOptions);
    }
  
    @Test
    public void testCsvConvertOptions() throws Exception {
      final Schema schema =
          new Schema(
              Arrays.asList(
                  Field.nullable("Id", new ArrowType.Int(32, true)),
                  Field.nullable("Name", new ArrowType.Utf8()),
                  Field.nullable("Language", new ArrowType.Utf8())),
              null);
      String path = "file://" + getClass().getResource("/").getPath() + "/data/student.csv";
      BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
      try (ArrowSchema cSchema = ArrowSchema.allocateNew(allocator);
          ArrowSchema cSchema2 = ArrowSchema.allocateNew(allocator);
          CDataDictionaryProvider provider = new CDataDictionaryProvider()) {
        Data.exportSchema(allocator, schema, provider, cSchema);
        Data.exportSchema(allocator, schema, provider, cSchema2);
        CsvFragmentScanOptions fragmentScanOptions1 =
            create(cSchema, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of("delimiter", ";"));
        CsvFragmentScanOptions fragmentScanOptions2 =
            create(cSchema2, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of("delimiter", ";"));
        ScanOptions options =
            new ScanOptions.Builder(/*batchSize*/ 32768)
                .columns(Optional.empty())
                .fragmentScanOptions(fragmentScanOptions1)
                .build();
        try (DatasetFactory datasetFactory =
                new FileSystemDatasetFactory(
                    allocator,
                    NativeMemoryPool.getDefault(),
                    FileFormat.CSV,
                    path,
                    Optional.of(fragmentScanOptions2));
            Dataset dataset = datasetFactory.finish();
            Scanner scanner = dataset.newScan(options);
            ArrowReader reader = scanner.scanBatches()) {
          assertEquals(schema.getFields(), reader.getVectorSchemaRoot().getSchema().getFields());
          int rowCount = 0;
          while (reader.loadNextBatch()) {
            final ValueIterableVector<Integer> idVector =
                (ValueIterableVector<Integer>) reader.getVectorSchemaRoot().getVector("Id");
            assertThat(idVector.getValueIterable(), IsIterableContainingInOrder.contains(1, 2, 3));
            rowCount += reader.getVectorSchemaRoot().getRowCount();
          }
          assertEquals(3, rowCount);
        }
      }
    }
  
    @Test
    public void testCsvConvertOptionsDelimiterNotSet() throws Exception {
      final Schema schema =
          new Schema(
              Arrays.asList(
                  Field.nullable("Id", new ArrowType.Int(32, true)),
                  Field.nullable("Name", new ArrowType.Utf8()),
                  Field.nullable("Language", new ArrowType.Utf8())),
              null);
      String path = "file://" + getClass().getResource("/").getPath() + "/data/student.csv";
      BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
      try (ArrowSchema cSchema = ArrowSchema.allocateNew(allocator);
          ArrowSchema cSchema2 = ArrowSchema.allocateNew(allocator);
          CDataDictionaryProvider provider = new CDataDictionaryProvider()) {
        Data.exportSchema(allocator, schema, provider, cSchema);
        Data.exportSchema(allocator, schema, provider, cSchema2);
        CsvFragmentScanOptions fragmentScanOptions1 =
            create(cSchema, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of());
        CsvFragmentScanOptions fragmentScanOptions2 =
            create(cSchema2, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of());
        ScanOptions options =
            new ScanOptions.Builder(/*batchSize*/ 32768)
                .columns(Optional.empty())
                .fragmentScanOptions(fragmentScanOptions1)
                .build();
        try (DatasetFactory datasetFactory =
                new FileSystemDatasetFactory(
                    allocator,
                    NativeMemoryPool.getDefault(),
                    FileFormat.CSV,
                    path,
                    Optional.of(fragmentScanOptions2));
            Dataset dataset = datasetFactory.finish();
            Scanner scanner = dataset.newScan(options);
            ArrowReader reader = scanner.scanBatches()) {
          int rowCount = 0;
          while (reader.loadNextBatch()) {
            final ValueIterableVector<Text> idVector =
                (ValueIterableVector<Text>)
                    reader.getVectorSchemaRoot().getVector("Id;Name;Language");
            assertThat(
                idVector.getValueIterable(),
                IsIterableContainingInOrder.contains(
                    new Text("1;Juno;Java"), new Text("2;Peter;Python"), new Text("3;Celin;C++")));
            rowCount += reader.getVectorSchemaRoot().getRowCount();
          }
          assertEquals(3, rowCount);
        }
      }
    }
  
    @Test
    public void testCsvConvertOptionsNoOption() throws Exception {
      final Schema schema =
          new Schema(
              Collections.singletonList(Field.nullable("Id;Name;Language", new ArrowType.Utf8())),
              null);
      String path = "file://" + getClass().getResource("/").getPath() + "/data/student.csv";
      BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
      ScanOptions options =
          new ScanOptions.Builder(/*batchSize*/ 32768).columns(Optional.empty()).build();
      try (DatasetFactory datasetFactory =
              new FileSystemDatasetFactory(
                  allocator, NativeMemoryPool.getDefault(), FileFormat.CSV, path);
          Dataset dataset = datasetFactory.finish();
          Scanner scanner = dataset.newScan(options);
          ArrowReader reader = scanner.scanBatches()) {
  
        assertEquals(schema.getFields(), reader.getVectorSchemaRoot().getSchema().getFields());
        int rowCount = 0;
        while (reader.loadNextBatch()) {
          final ValueIterableVector<Text> idVector =
              (ValueIterableVector<Text>) reader.getVectorSchemaRoot().getVector("Id;Name;Language");
          assertThat(
              idVector.getValueIterable(),
              IsIterableContainingInOrder.contains(
                  new Text("1;Juno;Java"), new Text("2;Peter;Python"), new Text("3;Celin;C++")));
          rowCount += reader.getVectorSchemaRoot().getRowCount();
        }
        assertEquals(3, rowCount);
      }
    }
  
    @Test
    public void testCsvReadParseAndReadOptions() throws Exception {
      final Schema schema =
          new Schema(
              Arrays.asList(
                  Field.nullable("Id", new ArrowType.Int(64, true)),
                  Field.nullable("Name", new ArrowType.Utf8()),
                  Field.nullable("Language", new ArrowType.Utf8())),
              null);
      String path = "file://" + getClass().getResource("/").getPath() + "/data/student.csv";
      BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
      CsvFragmentScanOptions fragmentScanOptions =
          new CsvFragmentScanOptions(
              new CsvConvertOptions(ImmutableMap.of()),
              ImmutableMap.of("skip_rows_after_names", "1"),
              ImmutableMap.of("delimiter", ";"));
      ScanOptions options =
          new ScanOptions.Builder(/*batchSize*/ 32768)
              .columns(Optional.empty())
              .fragmentScanOptions(fragmentScanOptions)
              .build();
      try (DatasetFactory datasetFactory =
              new FileSystemDatasetFactory(
                  allocator,
                  NativeMemoryPool.getDefault(),
                  FileFormat.CSV,
                  path,
                  Optional.of(fragmentScanOptions));
          Dataset dataset = datasetFactory.finish();
          Scanner scanner = dataset.newScan(options);
          ArrowReader reader = scanner.scanBatches()) {
  
        assertEquals(schema.getFields(), reader.getVectorSchemaRoot().getSchema().getFields());
        int rowCount = 0;
        while (reader.loadNextBatch()) {
          final ValueIterableVector<Long> idVector =
              (ValueIterableVector<Long>) reader.getVectorSchemaRoot().getVector("Id");
          assertThat(idVector.getValueIterable(), IsIterableContainingInOrder.contains(2L, 3L));
          rowCount += reader.getVectorSchemaRoot().getRowCount();
        }
        assertEquals(2, rowCount);
      }
    }
  
    @Test
    public void testCsvReadOtherOptions() throws Exception {
      String path = "file://" + getClass().getResource("/").getPath() + "/data/student.csv";
      BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
      Map<String, String> convertOption =
          ImmutableMap.of(
              "check_utf8",
              "true",
              "null_values",
              "NULL",
              "true_values",
              "True",
              "false_values",
              "False",
              "quoted_strings_can_be_null",
              "true",
              "auto_dict_encode",
              "false",
              "auto_dict_max_cardinality",
              "3456",
              "decimal_point",
              ".",
              "include_missing_columns",
              "false");
      Map<String, String> readOption =
          ImmutableMap.of(
              "use_threads",
              "true",
              "block_size",
              "1024",
              "skip_rows",
              "12",
              "skip_rows_after_names",
              "12",
              "autogenerate_column_names",
              "false");
      Map<String, String> parseOption =
          ImmutableMap.of(
              "delimiter",
              ".",
              "quoting",
              "true",
              "quote_char",
              "'",
              "double_quote",
              "False",
              "escaping",
              "true",
              "escape_char",
              "v",
              "newlines_in_values",
              "false",
              "ignore_empty_lines",
              "true");
      CsvFragmentScanOptions fragmentScanOptions =
          new CsvFragmentScanOptions(new CsvConvertOptions(convertOption), readOption, parseOption);
      ScanOptions options =
          new ScanOptions.Builder(/*batchSize*/ 32768)
              .columns(Optional.empty())
              .fragmentScanOptions(fragmentScanOptions)
              .build();
      try (DatasetFactory datasetFactory =
              new FileSystemDatasetFactory(
                  allocator, NativeMemoryPool.getDefault(), FileFormat.CSV, path);
          Dataset dataset = datasetFactory.finish();
          Scanner scanner = dataset.newScan(options)) {
        assertNotNull(scanner);
      }
    }
  
    @Test
    public void testCsvInvalidOption() throws Exception {
      String path = "file://" + getClass().getResource("/").getPath() + "/data/student.csv";
      BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
      Map<String, String> convertOption = ImmutableMap.of("not_exists_key_check_utf8", "true");
      CsvFragmentScanOptions fragmentScanOptions =
          new CsvFragmentScanOptions(
              new CsvConvertOptions(convertOption), ImmutableMap.of(), ImmutableMap.of());
      ScanOptions options =
          new ScanOptions.Builder(/*batchSize*/ 32768)
              .columns(Optional.empty())
              .fragmentScanOptions(fragmentScanOptions)
              .build();
      try (DatasetFactory datasetFactory =
              new FileSystemDatasetFactory(
                  allocator, NativeMemoryPool.getDefault(), FileFormat.CSV, path);
          Dataset dataset = datasetFactory.finish()) {
        assertThrows(RuntimeException.class, () -> dataset.newScan(options));
      }
  
      CsvFragmentScanOptions fragmentScanOptionsFaultValue =
          new CsvFragmentScanOptions(
              new CsvConvertOptions(ImmutableMap.of()),
              ImmutableMap.of("", ""),
              ImmutableMap.of("escape_char", "vbvb"));
      ScanOptions optionsFault =
          new ScanOptions.Builder(/*batchSize*/ 32768)
              .columns(Optional.empty())
              .fragmentScanOptions(fragmentScanOptionsFaultValue)
              .build();
      try (DatasetFactory datasetFactory =
              new FileSystemDatasetFactory(
                  allocator, NativeMemoryPool.getDefault(), FileFormat.CSV, path);
          Dataset dataset = datasetFactory.finish()) {
        assertThrows(RuntimeException.class, () -> dataset.newScan(optionsFault));
      }
    }
  }
  