/*
 * Licensed 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.lance;

import org.lance.index.IndexOptions;
import org.lance.index.IndexParams;
import org.lance.index.IndexType;
import org.lance.index.scalar.ScalarIndexParams;
import org.lance.ipc.ColumnOrdering;
import org.lance.ipc.LanceScanner;
import org.lance.ipc.MaterializationStyle;
import org.lance.ipc.ScanOptions;
import org.lance.ipc.ScanStats;

import org.apache.arrow.dataset.scanner.Scanner;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
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.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.MethodSource;

import java.io.IOException;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

public class ScannerTest {
  private static Dataset dataset;

  @BeforeAll
  static void setup() {}

  @AfterAll
  static void tearDown() {
    // Cleanup resources used by the tests
    if (dataset != null) {
      dataset.close();
    }
  }

  @Test
  void testScannerPerformanceOptions() {
    MaterializationStyle materialization =
        MaterializationStyle.allEarlyExcept(Collections.singletonList("name"));
    ScanOptions options =
        new ScanOptions.Builder()
            .batchSize(1024)
            .batchSizeBytes(64 * 1024)
            .ioBufferSize(8 * 1024 * 1024)
            .batchReadahead(4)
            .fragmentReadahead(2)
            .scanInOrder(false)
            .lateMaterialization(materialization)
            .build();

    assertEquals(1024L, options.getBatchSize().orElseThrow());
    assertEquals(64 * 1024L, options.getBatchSizeBytes().orElseThrow());
    assertEquals(8 * 1024 * 1024L, options.getIoBufferSize().orElseThrow());
    assertEquals(4, options.getBatchReadahead());
    assertEquals(2, options.getFragmentReadahead().orElseThrow());
    assertFalse(options.isScanInOrder());
    assertEquals(materialization, options.getLateMaterialization().orElseThrow());
    assertEquals("heuristic", MaterializationStyle.heuristic().toRustString());
    assertEquals("all_late", MaterializationStyle.allLate().toRustString());
    assertEquals("all_early", MaterializationStyle.allEarly().toRustString());
    assertEquals("all_early_except", materialization.toRustString());
    assertEquals(Collections.singletonList("name"), materialization.getColumns());

    ScanOptions defaults = new ScanOptions.Builder().build();
    assertTrue(defaults.getBatchSizeBytes().isEmpty());
    assertTrue(defaults.getIoBufferSize().isEmpty());
    assertTrue(defaults.getFragmentReadahead().isEmpty());
    assertTrue(defaults.isScanInOrder());
    assertTrue(defaults.getLateMaterialization().isEmpty());

    assertThrows(
        IllegalArgumentException.class, () -> new ScanOptions.Builder().batchSizeBytes(0).build());
    assertThrows(
        IllegalArgumentException.class, () -> new ScanOptions.Builder().ioBufferSize(0).build());
    assertThrows(
        IllegalArgumentException.class,
        () -> new ScanOptions.Builder().fragmentReadahead(0).build());
    assertThrows(
        IllegalArgumentException.class,
        () -> new ScanOptions.Builder().batchSizeBytes(1024).strictBatchSize(true).build());
    assertTrue(MaterializationStyle.allEarlyExcept(Collections.emptyList()).getColumns().isEmpty());
  }

  static Stream<MaterializationStyle> materializationStyles() {
    return Stream.of(
        MaterializationStyle.heuristic(),
        MaterializationStyle.allLate(),
        MaterializationStyle.allEarly(),
        MaterializationStyle.allEarlyExcept(Collections.emptyList()));
  }

  @ParameterizedTest
  @MethodSource("materializationStyles")
  void testMaterializationStylesAcrossJni(
      MaterializationStyle materializationStyle, @TempDir Path tempDir) throws Exception {
    String datasetPath =
        tempDir.resolve("materialization_" + materializationStyle.getMode().name()).toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();

      try (Dataset dataset = testDataset.write(1, 40);
          LanceScanner scanner =
              dataset.newScan(
                  new ScanOptions.Builder()
                      .filter("id < 20")
                      .lateMaterialization(materializationStyle)
                      .build());
          ArrowReader reader = scanner.scanBatches()) {
        int rowCount = 0;
        while (reader.loadNextBatch()) {
          rowCount += reader.getVectorSchemaRoot().getRowCount();
        }
        assertEquals(20, rowCount);
      }
    }
  }

  @Test
  void testDatasetScanner(@TempDir Path tempDir) throws IOException {
    String datasetPath = tempDir.resolve("dataset_scanner").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      int totalRows = 40;
      int batchRows = 20;
      try (Dataset dataset = testDataset.write(1, totalRows)) {
        Scanner scanner = dataset.newScan(batchRows);
        testDataset.validateScanResults(dataset, scanner, totalRows, batchRows);
      }
    }
  }

  @Test
  void testDatasetScannerFilter(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("dataset_scanner_filter").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      // write id with value from 0 to 39
      try (Dataset dataset = testDataset.write(1, 40)) {
        try (Scanner scanner =
            dataset.newScan(new ScanOptions.Builder().filter("id < 20").build())) {
          testDataset.validateScanResults(dataset, scanner, 20, 20);
        }
      }
    }
  }

  @Test
  void testDatasetScannerColumns(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("dataset_scanner_columns").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      int totalRows = 40;
      int batchRows = 20;
      try (Dataset dataset = testDataset.write(1, totalRows)) {
        try (Scanner scanner =
            dataset.newScan(
                new ScanOptions.Builder()
                    .batchSize(batchRows)
                    .columns(Arrays.asList("id"))
                    .build())) {
          try (ArrowReader reader = scanner.scanBatches()) {
            VectorSchemaRoot root = reader.getVectorSchemaRoot();
            int index = 0;
            while (reader.loadNextBatch()) {
              List<FieldVector> fieldVectors = root.getFieldVectors();
              assertEquals(1, fieldVectors.size());
              FieldVector fieldVector = fieldVectors.get(0);
              assertEquals(ArrowType.ArrowTypeID.Int, fieldVector.getField().getType().getTypeID());
              assertEquals(batchRows, fieldVector.getValueCount());
              IntVector vector = (IntVector) fieldVector;
              for (int i = 0; i < batchRows; i++) {
                assertEquals(index, vector.get(i));
                index++;
              }
            }
          }
        }
      }
    }
  }

  @Test
  void testDatasetScannerSchema(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("dataset_scanner_schema").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      int totalRows = 40;
      try (Dataset dataset = testDataset.write(1, totalRows)) {
        try (Scanner scanner =
            dataset.newScan(
                new ScanOptions.Builder()
                    .batchSize(totalRows)
                    .columns(Arrays.asList("id"))
                    .build())) {
          Schema expectedSchema =
              new Schema(Arrays.asList(Field.nullable("id", new ArrowType.Int(32, true))));
          assertEquals(expectedSchema, scanner.schema());
        }
      }
    }
  }

  @Test
  void testDatasetScannerCountRows(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("dataset_scanner_count").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      // write id with value from 0 to 39
      try (Dataset dataset = testDataset.write(1, 40)) {
        try (LanceScanner scanner =
            dataset.newScan(
                new ScanOptions.Builder()
                    .columns(Arrays.asList())
                    .withRowId(true)
                    .filter("id < 20")
                    .build())) {
          assertEquals(20, scanner.countRows());
        }
      }
    }
  }

  @Test
  void testDatasetScannerStats(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("dataset_scanner_stats").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      try (Dataset dataset = testDataset.write(1, 40)) {
        try (LanceScanner scanner =
            dataset.newScan(new ScanOptions.Builder().batchSize(20).build())) {
          assertTrue(scanner.getStats().isEmpty());
          try (ArrowReader reader = scanner.scanBatches()) {
            while (reader.loadNextBatch()) {
              // Consume all batches.
            }
          }
          assertTrue(scanner.getStats().isEmpty());
        }

        try (LanceScanner scanner =
            dataset.newScan(new ScanOptions.Builder().batchSize(20).collectStats(true).build())) {
          assertTrue(scanner.getStats().isEmpty());
          try (ArrowReader reader = scanner.scanBatches()) {
            while (reader.loadNextBatch()) {
              // Consume all batches.
            }
          }

          Optional<ScanStats> statsOpt = scanner.getStats();
          assertTrue(statsOpt.isPresent());
          ScanStats stats = statsOpt.get();
          assertTrue(stats.getBytesRead() > 0 || !stats.getAllCounts().isEmpty());
          // Even without an index on this dataset, the two new counters must
          // still marshal through JNI and default to zero rather than throwing.
          assertTrue(stats.getIndexCacheHits() >= 0);
          assertTrue(stats.getIndexCacheMisses() >= 0);
        }
      }
    }
  }

  @Test
  void testFragmentScanner(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("fragment_scanner").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      int totalRows = 40;
      int batchRows = 20;
      try (Dataset dataset = testDataset.write(1, totalRows)) {
        Fragment fragment = dataset.getFragments().get(0);
        try (Scanner scanner = fragment.newScan(batchRows)) {
          testDataset.validateScanResults(dataset, scanner, totalRows, batchRows);
        }
      }
    }
  }

  @Test
  void testFragmentScannerFilter(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("fragment_scanner_filter").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      // write id with value from 0 to 39
      try (Dataset dataset = testDataset.write(1, 40)) {
        Fragment fragment = dataset.getFragments().get(0);
        try (Scanner scanner =
            fragment.newScan(new ScanOptions.Builder().filter("id < 20").build())) {
          testDataset.validateScanResults(dataset, scanner, 20, 20);
        }
      }
    }
  }

  @Test
  void testFragmentScannerColumns(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("fragment_scanner_columns").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      int totalRows = 40;
      int batchRows = 20;
      try (Dataset dataset = testDataset.write(1, totalRows)) {
        Fragment fragment = dataset.getFragments().get(0);
        try (Scanner scanner =
            fragment.newScan(
                new ScanOptions.Builder()
                    .batchSize(batchRows)
                    .columns(Arrays.asList("id"))
                    .build())) {
          try (ArrowReader reader = scanner.scanBatches()) {
            VectorSchemaRoot root = reader.getVectorSchemaRoot();
            int index = 0;
            while (reader.loadNextBatch()) {
              List<FieldVector> fieldVectors = root.getFieldVectors();
              assertEquals(1, fieldVectors.size());
              FieldVector fieldVector = fieldVectors.get(0);
              assertEquals(ArrowType.ArrowTypeID.Int, fieldVector.getField().getType().getTypeID());
              assertEquals(batchRows, fieldVector.getValueCount());
              IntVector vector = (IntVector) fieldVector;
              for (int i = 0; i < batchRows; i++) {
                assertEquals(index, vector.get(i));
                index++;
              }
            }
          }
        }
      }
    }
  }

  @Test
  void testScanFragment(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("fragment_scanner_single_fragment").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      FragmentMetadata metadata0 = testDataset.createNewFragment(3);
      FragmentMetadata metadata1 = testDataset.createNewFragment(5);
      FragmentMetadata metadata2 = testDataset.createNewFragment(7);
      FragmentOperation.Append appendOp =
          new FragmentOperation.Append(Arrays.asList(metadata0, metadata1, metadata2));
      try (Dataset dataset = Dataset.commit(allocator, datasetPath, appendOp, Optional.of(1L))) {
        List<Fragment> frags = dataset.getFragments();
        assertEquals(3, frags.size());
        validScanResult(dataset, frags.get(0).getId(), 3);
        validScanResult(dataset, frags.get(1).getId(), 5);
        validScanResult(dataset, frags.get(2).getId(), 7);
      }
    }
  }

  @Test
  void testScanFragments(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("fragments_scanner").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      FragmentMetadata metadata0 = testDataset.createNewFragment(3);
      FragmentMetadata metadata1 = testDataset.createNewFragment(5);
      FragmentMetadata metadata2 = testDataset.createNewFragment(7);
      FragmentOperation.Append appendOp =
          new FragmentOperation.Append(Arrays.asList(metadata0, metadata1, metadata2));
      try (Dataset dataset = Dataset.commit(allocator, datasetPath, appendOp, Optional.of(1L))) {
        List<Fragment> frags = dataset.getFragments();
        assertEquals(3, frags.size());
        try (Scanner scanner =
            dataset.newScan(
                new ScanOptions.Builder()
                    .batchSize(1024)
                    .fragmentIds(Arrays.asList(frags.get(1).getId(), frags.get(2).getId()))
                    .build())) {
          try (ArrowReader reader = scanner.scanBatches()) {
            assertEquals(
                dataset.getSchema().getFields(),
                reader.getVectorSchemaRoot().getSchema().getFields());
            int rowcount = 0;
            reader.loadNextBatch();
            int currentRowCount = reader.getVectorSchemaRoot().getRowCount();
            assertEquals(5, currentRowCount);
            rowcount += currentRowCount;
            reader.loadNextBatch();
            currentRowCount = reader.getVectorSchemaRoot().getRowCount();
            assertEquals(7, currentRowCount);
            rowcount += currentRowCount;
            assertEquals(12, rowcount);
          }
        }
      }
    }
  }

  @Test
  void testDatasetScannerLimit(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("dataset_scanner_limit").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      int totalRows = 100;
      int limit = 50;
      try (Dataset dataset = testDataset.write(1, totalRows)) {
        try (Scanner scanner = dataset.newScan(new ScanOptions.Builder().limit(limit).build())) {
          testDataset.validateScanResults(dataset, scanner, limit, limit);
        }
      }
    }
  }

  @Test
  void testDatasetScannerOffset(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("dataset_scanner_offset").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      int totalRows = 100;
      int offset = 50;
      try (Dataset dataset = testDataset.write(1, totalRows)) {
        try (Scanner scanner = dataset.newScan(new ScanOptions.Builder().offset(offset).build())) {
          testDataset.validateScanResults(
              dataset, scanner, totalRows - offset, totalRows - offset, offset);
        }
      }
    }
  }

  @Test
  void testDatasetScannerWithRowId(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("dataset_scanner_with_row_id").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      int totalRows = 50;
      try (Dataset dataset = testDataset.write(1, totalRows)) {
        try (Scanner scanner = dataset.newScan(new ScanOptions.Builder().withRowId(true).build())) {
          try (ArrowReader reader = scanner.scanBatches()) {
            VectorSchemaRoot root = reader.getVectorSchemaRoot();
            assertTrue(
                root.getSchema().getFields().stream()
                    .anyMatch(field -> field.getName().equals("_rowid")));
            while (reader.loadNextBatch()) {
              List<FieldVector> fieldVectors = root.getFieldVectors();
              assertTrue(
                  fieldVectors.stream().anyMatch(vector -> vector.getName().equals("_rowid")));
            }
          }
        }
      }
    }
  }

  @Test
  void testDatasetScannerBatchReadahead(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("dataset_scanner_batch_readahead").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();

      int totalRows = 2000;
      int maxRowsPerFile = 100; // ~20 fragments
      List<FragmentMetadata> fragments = testDataset.createNewFragment(totalRows, maxRowsPerFile);
      assertTrue(fragments.size() > 1, "expected multiple fragments, got " + fragments.size());

      FragmentOperation.Append append = new FragmentOperation.Append(fragments);
      try (Dataset dataset = Dataset.commit(allocator, datasetPath, append, Optional.of(1L))) {
        int batchReadahead = 2; // far below the default (num compute CPUs)
        try (LanceScanner scanner =
            dataset.newScan(
                new ScanOptions.Builder().batchSize(50).batchReadahead(batchReadahead).build())) {
          try (ArrowReader reader = scanner.scanBatches()) {
            int rowCount = 0;
            long idSum = 0;
            while (reader.loadNextBatch()) {
              VectorSchemaRoot root = reader.getVectorSchemaRoot();
              IntVector ids = (IntVector) root.getVector("id");
              for (int i = 0; i < root.getRowCount(); i++) {
                idSum += ids.get(i);
              }
              rowCount += root.getRowCount();
            }
            assertEquals(totalRows, rowCount);
            // ids are the contiguous range [0, totalRows)
            assertEquals((long) totalRows * (totalRows - 1) / 2, idSum);
          }
        }
      }
    }
  }

  @Test
  void testDatasetScannerPerformanceOptions(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("dataset_scanner_performance_options").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();

      int totalRows = 2000;
      int matchingRows = 1000;
      WriteParams writeParams =
          new WriteParams.Builder()
              .withMaxRowsPerFile(100)
              .withDataStorageVersion(LanceConstants.FILE_FORMAT_VERSION_STABLE)
              .build();
      List<FragmentMetadata> fragments = testDataset.createNewFragment(totalRows, writeParams);
      FragmentOperation.Append append = new FragmentOperation.Append(fragments);
      try (Dataset dataset = Dataset.commit(allocator, datasetPath, append, Optional.of(1L))) {
        try (LanceScanner scanner =
                dataset.newScan(
                    new ScanOptions.Builder()
                        .filter("id < " + matchingRows)
                        .batchSize(200)
                        .batchSizeBytes(256)
                        .ioBufferSize(1024 * 1024)
                        .batchReadahead(2)
                        .fragmentReadahead(2)
                        .scanInOrder(false)
                        .lateMaterialization(MaterializationStyle.allEarly())
                        .build());
            ArrowReader reader = scanner.scanBatches()) {
          int rowCount = 0;
          long idSum = 0;
          while (reader.loadNextBatch()) {
            VectorSchemaRoot root = reader.getVectorSchemaRoot();
            assertTrue(
                root.getRowCount() < 200,
                "batchSizeBytes should split batches before the row limit");
            IntVector ids = (IntVector) root.getVector("id");
            for (int i = 0; i < root.getRowCount(); i++) {
              idSum += ids.get(i);
            }
            rowCount += root.getRowCount();
          }
          assertEquals(matchingRows, rowCount);
          assertEquals((long) matchingRows * (matchingRows - 1) / 2, idSum);
        }

        try (LanceScanner scanner =
            dataset.newScan(
                new ScanOptions.Builder()
                    .filter("id < " + matchingRows)
                    .lateMaterialization(
                        MaterializationStyle.allEarlyExcept(Collections.singletonList("name")))
                    .build())) {
          assertEquals(matchingRows, scanner.countRows());
        }
      }
    }
  }

  @Test
  void testDatasetScannerSortBy(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("testDatasetScannerSortBy").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      try (Dataset dataset = testDataset.writeSortByDataset(1)) {
        ColumnOrdering.Builder nameBuilder = new ColumnOrdering.Builder();
        nameBuilder.setColumnName("name");
        nameBuilder.setAscending(true);
        nameBuilder.setNullFirst(false);

        ColumnOrdering.Builder idBuilder = new ColumnOrdering.Builder();
        idBuilder.setColumnName("id");
        idBuilder.setAscending(false);
        idBuilder.setNullFirst(true);

        List<ColumnOrdering> columnOrderings =
            Arrays.asList(nameBuilder.build(), idBuilder.build());
        ScanOptions.Builder scanOptionBuilder = new ScanOptions.Builder();
        scanOptionBuilder
            .columns(Arrays.asList("name", "id"))
            .limit(10)
            .setColumnOrderings(columnOrderings);
        ScanOptions scanOptions = scanOptionBuilder.build();
        try (Scanner scanner = dataset.newScan(scanOptions)) {
          try (ArrowReader reader = scanner.scanBatches()) {
            while (reader.loadNextBatch()) {
              List<FieldVector> fieldVectors = reader.getVectorSchemaRoot().getFieldVectors();
              VarCharVector nameVector = (VarCharVector) fieldVectors.get(0);
              /* dataset context
               * i: |  id   | name | :i
               * 1: |  1    |  P0  | :0
               * 2: | null  |  P1  | :1
               * 3: |  2    |  P2  | :2
               * 5: | null  |  P3  | :3
               * 4: |  2    |  P3  | :4
               * 7: |  4    |  P4  | :5
               * 9: |  5    |  P5  | :6
               * 8: |  4    |  P5  | :7
               * 6: |  3    | null | :8
               * 0: |  0    | null | :9
               */
              assertEquals("P0", new String(nameVector.get(0)));
              assertEquals("P1", new String(nameVector.get(1)));
              assertEquals("P2", new String(nameVector.get(2)));
              assertEquals("P3", new String(nameVector.get(3)));
              assertEquals("P3", new String(nameVector.get(4)));
              assertEquals("P4", new String(nameVector.get(5)));
              assertEquals("P5", new String(nameVector.get(6)));
              assertEquals("P5", new String(nameVector.get(7)));
              assertTrue(nameVector.isNull(8));
              assertTrue(nameVector.isNull(9));

              IntVector idVector = (IntVector) fieldVectors.get(1);
              assertEquals(1, idVector.get(0));
              assertTrue(idVector.isNull(1));
              assertEquals(2, idVector.get(2));
              assertTrue(idVector.isNull(3));
              assertEquals(2, idVector.get(4));
              assertEquals(4, idVector.get(5));
              assertEquals(5, idVector.get(6));
              assertEquals(4, idVector.get(7));
              assertEquals(3, idVector.get(8));
              assertEquals(0, idVector.get(9));
            }
          }
        }
      }
    }
  }

  @Test
  void testDatasetScannerCombinedParams(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("dataset_scanner_combined_params").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      int totalRows = 600;
      int limit = 200;
      int offset = 300;
      int batchSize = 50;
      try (Dataset dataset = testDataset.write(1, totalRows)) {
        try (Scanner scanner =
            dataset.newScan(
                new ScanOptions.Builder()
                    .limit(limit)
                    .offset(offset)
                    .withRowId(true)
                    .batchSize(batchSize)
                    .batchReadahead(3)
                    .build())) {
          try (ArrowReader reader = scanner.scanBatches()) {
            VectorSchemaRoot root = reader.getVectorSchemaRoot();
            List<String> fieldNames =
                root.getSchema().getFields().stream()
                    .map(Field::getName)
                    .collect(Collectors.toList());
            assertTrue(fieldNames.contains("_rowid"), "Schema should contain _rowid column");
            assertTrue(fieldNames.contains("id"), "Schema should contain id column");

            int rowCount = 0;
            int expectedIdStart = offset;
            while (reader.loadNextBatch()) {
              List<FieldVector> fieldVectors = root.getFieldVectors();
              assertTrue(
                  fieldVectors.stream().anyMatch(vector -> vector.getName().equals("_rowid")));
              IntVector idVector = (IntVector) root.getVector("id");
              int batchRowCount = root.getRowCount();
              rowCount += batchRowCount;
              assertTrue(batchRowCount <= batchSize, "Batch size should not exceed " + batchSize);

              for (int i = 0; i < batchRowCount; i++) {
                int expectedId = expectedIdStart + i;
                assertEquals(
                    expectedId,
                    idVector.get(i),
                    "Mismatch at row "
                        + (rowCount - batchRowCount + i)
                        + ". Expected: "
                        + expectedId
                        + ", Actual: "
                        + idVector.get(i));
              }
              expectedIdStart += batchRowCount;
            }
            assertEquals(limit, rowCount, "Total rows should match the limit");
          }
        }
      }
    }
  }

  @Test
  void testUseScalarIndex(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("dataset_scanner_use_scalar_index").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      int totalRows = 100;
      try (Dataset dataset = testDataset.write(1, totalRows)) {
        // Create a scalar index on the 'id' column
        ScalarIndexParams scalarParams = ScalarIndexParams.create("btree", "{}");
        IndexParams indexParams = IndexParams.builder().setScalarIndexParams(scalarParams).build();
        IndexOptions options =
            IndexOptions.builder(Collections.singletonList("id"), IndexType.BTREE, indexParams)
                .withIndexName("id_btree_index")
                .replace(true)
                .build();
        dataset.createIndex(options);

        // Verify index was created
        assertTrue(
            dataset.listIndexes().contains("id_btree_index"),
            "Expected 'id_btree_index' to be in the list of indexes: " + dataset.listIndexes());

        // Test with useScalarIndex = true (default)
        List<Integer> resultsWithIndex = new ArrayList<>();
        try (Scanner scanner =
            dataset.newScan(
                new ScanOptions.Builder()
                    .filter("id < 50")
                    .useScalarIndex(true)
                    .columns(Collections.singletonList("id"))
                    .build())) {
          try (ArrowReader reader = scanner.scanBatches()) {
            while (reader.loadNextBatch()) {
              VectorSchemaRoot root = reader.getVectorSchemaRoot();
              IntVector idVector = (IntVector) root.getVector("id");
              for (int i = 0; i < root.getRowCount(); i++) {
                resultsWithIndex.add(idVector.get(i));
              }
            }
          }
        }

        // Test with useScalarIndex = false
        List<Integer> resultsWithoutIndex = new ArrayList<>();
        try (Scanner scanner =
            dataset.newScan(
                new ScanOptions.Builder()
                    .filter("id < 50")
                    .useScalarIndex(false)
                    .columns(Collections.singletonList("id"))
                    .build())) {
          try (ArrowReader reader = scanner.scanBatches()) {
            while (reader.loadNextBatch()) {
              VectorSchemaRoot root = reader.getVectorSchemaRoot();
              IntVector idVector = (IntVector) root.getVector("id");
              for (int i = 0; i < root.getRowCount(); i++) {
                resultsWithoutIndex.add(idVector.get(i));
              }
            }
          }
        }

        // Results should be the same regardless of whether scalar index is used
        assertEquals(
            resultsWithIndex.size(),
            resultsWithoutIndex.size(),
            "Result count should be the same with or without scalar index");
        assertEquals(50, resultsWithIndex.size(), "Should return 50 rows (id < 50)");
        assertEquals(
            resultsWithIndex,
            resultsWithoutIndex,
            "Results should be identical with or without scalar index");
      }
    }
  }

  @Test
  void testFastSearchSkipsUnindexedFragments(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("dataset_scanner_fast_search_scalar_index").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      try (Dataset dataset = testDataset.write(1, 100)) {
        ScalarIndexParams scalarParams = ScalarIndexParams.create("btree", "{}");
        IndexParams indexParams = IndexParams.builder().setScalarIndexParams(scalarParams).build();
        IndexOptions options =
            IndexOptions.builder(Collections.singletonList("id"), IndexType.BTREE, indexParams)
                .withIndexName("id_btree_index")
                .replace(true)
                .build();
        dataset.createIndex(options);

        FragmentMetadata metadata = testDataset.createNewFragment(10);
        FragmentOperation.Append appendOp =
            new FragmentOperation.Append(Collections.singletonList(metadata));
        try (Dataset appended =
            Dataset.commit(allocator, datasetPath, appendOp, Optional.of(dataset.version()))) {
          try (LanceScanner scanner =
              appended.newScan(new ScanOptions.Builder().filter("id < 5").build())) {
            assertEquals(10, scanner.countRows());
          }

          try (LanceScanner scanner =
              appended.newScan(
                  new ScanOptions.Builder().filter("id < 5").fastSearch(true).build())) {
            assertEquals(5, scanner.countRows());
          }
        }
      }
    }
  }

  @Test
  void testIncludeDeletedRows(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("include_deleted_rows").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      try (Dataset dataset = testDataset.write(1, 10)) {
        assertEquals(10, dataset.countRows());

        // Delete rows where id >= 5
        dataset.delete("id >= 5");
        assertEquals(5, dataset.countRows());

        // Default scan should exclude deleted rows
        try (LanceScanner scanner =
            dataset.newScan(new ScanOptions.Builder().batchSize(20).build())) {
          assertEquals(5, scanner.countRows(), "default scan: should exclude deleted rows");
        }

        // includeDeletedRows=true should surface deleted rows
        // NOTE: includeDeletedRows requires withRowId=true
        try (LanceScanner scanner =
            dataset.newScan(
                new ScanOptions.Builder()
                    .batchSize(20)
                    .withRowId(true)
                    .includeDeletedRows(true)
                    .build())) {
          assertEquals(10, scanner.countRows(), "includeDeletedRows: should include deleted rows");
        }
      }
    }
  }

  @Test
  void testStrictBatchSize(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("strict_batch_size").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      try (Dataset dataset = testDataset.write(1, 25)) {
        int batchSize = 10;

        // With strictBatchSize=true, no batch should exceed batchSize
        try (Scanner scanner =
            dataset.newScan(
                new ScanOptions.Builder().batchSize(batchSize).strictBatchSize(true).build())) {
          try (ArrowReader reader = scanner.scanBatches()) {
            int totalRows = 0;
            while (reader.loadNextBatch()) {
              int rows = reader.getVectorSchemaRoot().getRowCount();
              assertTrue(rows <= batchSize, "strict: batch " + rows + " should be <= " + batchSize);
              totalRows += rows;
            }
            assertEquals(25, totalRows);
          }
        }

        // strictBatchSize=false (default) — batch size may vary
        try (Scanner scanner =
            dataset.newScan(new ScanOptions.Builder().batchSize(batchSize).build())) {
          try (ArrowReader reader = scanner.scanBatches()) {
            int totalRows = 0;
            while (reader.loadNextBatch()) {
              totalRows += reader.getVectorSchemaRoot().getRowCount();
            }
            assertEquals(25, totalRows);
          }
        }
      }
    }
  }

  @Test
  void testDisableScoringAutoprojection(@TempDir Path tempDir) throws Exception {
    String datasetPath = tempDir.resolve("disable_scoring_autoprojection").toString();
    try (BufferAllocator allocator = new RootAllocator()) {
      TestUtils.SimpleTestDataset testDataset =
          new TestUtils.SimpleTestDataset(allocator, datasetPath);
      testDataset.createEmptyDataset().close();
      try (Dataset dataset = testDataset.write(1, 10)) {
        // Smoke test: verify the option is accepted and scan still works
        ScanOptions options =
            new ScanOptions.Builder().batchSize(20).disableScoringAutoprojection(true).build();

        try (LanceScanner scanner = dataset.newScan(options)) {
          assertEquals(
              10,
              scanner.countRows(),
              "scan with disableScoringAutoprojection should return all rows");
        }

        // Also verify it doesn't break when combined with other options
        ScanOptions combinedOptions =
            new ScanOptions.Builder()
                .batchSize(20)
                .filter("id < 5")
                .disableScoringAutoprojection(true)
                .includeDeletedRows(false)
                .strictBatchSize(false)
                .build();

        try (LanceScanner scanner = dataset.newScan(combinedOptions)) {
          assertEquals(
              5,
              scanner.countRows(),
              "scan with disableScoringAutoprojection + filter should work");
        }
      }
    }
  }

  private void validScanResult(Dataset dataset, int fragmentId, int rowCount) throws Exception {
    try (Scanner scanner =
        dataset.newScan(
            new ScanOptions.Builder()
                .batchSize(1024)
                .fragmentIds(Arrays.asList(fragmentId))
                .build())) {
      try (ArrowReader reader = scanner.scanBatches()) {
        assertEquals(
            dataset.getSchema().getFields(), reader.getVectorSchemaRoot().getSchema().getFields());
        reader.loadNextBatch();
        assertEquals(rowCount, reader.getVectorSchemaRoot().getRowCount());
        assertFalse(reader.loadNextBatch());
      }
    }
  }
}
