IllegalStateException reading data with struct type (apache#43)

java/dataset/src/main/java/org/apache/arrow/dataset/jni/NativeDataset.java

@@ -42,7 +42,7 @@ public synchronized NativeScanner newScan(ScanOptions options) {
     }
     long scannerId = JniWrapper.get().createScanner(datasetId, options.getColumns().orElse(null),
         options.getBatchSize(), context.getMemoryPool().getNativeInstanceId());
-    return new NativeScanner(context, scannerId);
+    return new NativeScanner(context, options, scannerId);
   }
 
   @Override

java/dataset/src/main/java/org/apache/arrow/dataset/jni/NativeScanner.java

@@ -18,23 +18,17 @@
 package org.apache.arrow.dataset.jni;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.NoSuchElementException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
-import java.util.stream.Collectors;
 
+import org.apache.arrow.dataset.scanner.ScanOptions;
 import org.apache.arrow.dataset.scanner.ScanTask;
 import org.apache.arrow.dataset.scanner.Scanner;
-import org.apache.arrow.memory.ArrowBuf;
-import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.memory.BufferLedger;
-import org.apache.arrow.memory.NativeUnderlyingMemory;
-import org.apache.arrow.memory.util.LargeMemoryUtil;
-import org.apache.arrow.vector.ipc.message.ArrowFieldNode;
+import org.apache.arrow.util.Preconditions;
 import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.apache.arrow.vector.util.SchemaUtility;
@@ -48,15 +42,19 @@ public class NativeScanner implements Scanner {
 
   private final AtomicBoolean executed = new AtomicBoolean(false);
   private final NativeContext context;
+  private final ScanOptions options;
   private final long scannerId;
 
   private final ReadWriteLock lock = new ReentrantReadWriteLock();
   private final Lock writeLock = lock.writeLock();
   private final Lock readLock = lock.readLock();
+
+  private Schema schema = null;
   private boolean closed = false;
 
-  public NativeScanner(NativeContext context, long scannerId) {
+  public NativeScanner(NativeContext context, ScanOptions options, long scannerId) {
     this.context = context;
+    this.options = options;
     this.scannerId = scannerId;
   }
 
@@ -95,6 +93,9 @@ public boolean hasNext() {
           return false;
         }
         peek = UnsafeRecordBatchSerializer.deserializeUnsafe(context.getAllocator(), bytes);
+        if (options.getColumns().isPresent()) {
+          Preconditions.checkState(schema().getFields().size() == options.getColumns().get().length);
+        }
         return true;
       }
 
@@ -123,12 +124,19 @@ public Iterable<? extends NativeScanTask> scan() {
 
   @Override
   public Schema schema() {
+    if (schema != null) {
+      return schema;
+    }
     readLock.lock();
     try {
+      if (schema != null) {
+        return schema;
+      }
       if (closed) {
         throw new NativeInstanceReleasedException();
       }
-      return SchemaUtility.deserialize(JniWrapper.get().getSchemaFromScanner(scannerId), context.getAllocator());
+      schema = SchemaUtility.deserialize(JniWrapper.get().getSchemaFromScanner(scannerId), context.getAllocator());
+      return schema;
     } catch (IOException e) {
       throw new RuntimeException(e);
     } finally {

java/dataset/src/test/java/org/apache/arrow/dataset/file/TestFileSystemDataset.java

@@ -163,6 +163,68 @@ public void testCsvReadTab() throws Exception {
     AutoCloseables.close(vsr, allocator);
   }
 
+
+  @Test
+  public void testStructTypeRead() throws Exception {
+    RootAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+    FileSystemDatasetFactory factory = new FileSystemDatasetFactory(allocator,
+        NativeMemoryPool.getDefault(), ParquetFileFormat.createDefault(), "file://" + resourcePath("data/struct_example.parquet"));
+    ScanOptions options = new ScanOptions(new String[] {"_1"}, Filter.EMPTY, 100);
+    Schema schema = factory.inspect();
+    NativeDataset dataset = factory.finish(schema);
+    NativeScanner nativeScanner = dataset.newScan(options);
+    List<? extends ScanTask> scanTasks = collect(nativeScanner.scan());
+    Assert.assertEquals(1, scanTasks.size());
+    ScanTask scanTask = scanTasks.get(0);
+    ScanTask.BatchIterator itr = scanTask.execute();
+
+    VectorSchemaRoot vsr = VectorSchemaRoot.create(schema, allocator);
+    VectorLoader loader = new VectorLoader(vsr);
+    int rowCount = 0;
+    while (itr.hasNext()) {
+      try (ArrowRecordBatch next = itr.next()) {
+        loader.load(next);
+      }
+      rowCount += vsr.getRowCount();
+
+    }
+    Assert.assertEquals(50, rowCount);
+    assertEquals(1, schema.getFields().size());
+    assertEquals("_1", schema.getFields().get(0).getName());
+    AutoCloseables.close(vsr, allocator);
+  }
+
+  @Test
+  public void testStructTypeReadWithEmptyProjector() throws Exception {
+    RootAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+    FileSystemDatasetFactory factory = new FileSystemDatasetFactory(allocator,
+        NativeMemoryPool.getDefault(), ParquetFileFormat.createDefault(), "file://" + resourcePath("data/struct_example.parquet"));
+    ScanOptions options = new ScanOptions(new String[0], Filter.EMPTY, 100);
+    Schema schema = factory.inspect();
+    NativeDataset dataset = factory.finish(schema);
+    NativeScanner nativeScanner = dataset.newScan(options);
+    List<? extends ScanTask> scanTasks = collect(nativeScanner.scan());
+    Assert.assertEquals(1, scanTasks.size());
+    ScanTask scanTask = scanTasks.get(0);
+    ScanTask.BatchIterator itr = scanTask.execute();
+    Schema scannerSchema = nativeScanner.schema();
+    VectorSchemaRoot vsr = VectorSchemaRoot.create(scannerSchema, allocator);
+    VectorLoader loader = new VectorLoader(vsr);
+    int rowCount = 0;
+    while (itr.hasNext()) {
+      try (ArrowRecordBatch next = itr.next()) {
+        loader.load(next);
+      }
+      rowCount += vsr.getRowCount();
+
+    }
+    Assert.assertEquals(50, rowCount);
+    assertEquals(1, schema.getFields().size());
+    assertEquals("_1", schema.getFields().get(0).getName());
+    assertEquals(0, scannerSchema.getFields().size());
+    AutoCloseables.close(vsr, allocator);
+  }
+
   @Test
   public void testReadPartialFile() throws Exception {
     ParquetWriteSupport writeSupport = ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP.newFolder(), 1, "a");