fix parquet batching

Author: jakedorne

src/main/java/com/github/mmolimar/kafka/connect/fs/file/reader/AbstractFileReader.java

@@ -76,8 +76,9 @@ protected void setOffset(long offset) {
     public final boolean hasNext() {
         checkClosed();
         try {
-            return (batchSize <= 0 || offset == 0 || offset % batchSize != 0 || (offset % batchSize == 0 && seeked)) &&
+            boolean hasNext = (batchSize <= 0 || offset == 0 || offset % batchSize != 0 || (offset % batchSize == 0 && seeked)) &&
                     hasNextRecord();
+            return hasNext;
         } catch (ConnectException ce) {
             throw ce;
         } catch (Exception e) {

src/main/java/com/github/mmolimar/kafka/connect/fs/file/reader/ParquetFileReader.java

@@ -90,7 +90,8 @@ public void seekFile(long offset) throws IOException {
             this.closed = false;
             setOffset(0);
         }
-        while (hasNext() && currentOffset() < offset) {
+
+        while (hasNextRecord() && currentOffset() < offset) {
             nextRecord();
         }
     }

src/test/java/com/github/mmolimar/kafka/connect/fs/InMemoryFsContext.java

@@ -0,0 +1,27 @@
+package com.github.mmolimar.kafka.connect.fs;
+
+import org.apache.kafka.connect.source.SourceTaskContext;
+import org.apache.kafka.connect.storage.OffsetStorageReader;
+
+import java.util.Map;
+
+public class InMemoryFsContext implements SourceTaskContext {
+
+    public OffsetStorageReader reader;
+    public FsSourceTaskConfig configs;
+
+    public InMemoryFsContext() {
+        reader = new InMemoryFsOffsetStorageReader();
+    }
+
+    @Override
+    public Map<String, String> configs() {
+        System.out.println("TURNS OUT YOU PROBABLY NEED TO IMPLEMENT THIS.");
+        return null;
+    }
+
+    @Override
+    public OffsetStorageReader offsetStorageReader() {
+        return reader;
+    }
+}

src/test/java/com/github/mmolimar/kafka/connect/fs/InMemoryFsOffsetStorageReader.java

@@ -0,0 +1,26 @@
+package com.github.mmolimar.kafka.connect.fs;
+
+import org.apache.kafka.connect.storage.OffsetStorageReader;
+
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
+
+public class InMemoryFsOffsetStorageReader implements OffsetStorageReader {
+    @Override
+    public <T> Map<String, Object> offset(Map<String, T> partition) {
+        return InMemoryOffsetStore.Offsets.get(partition);
+    }
+
+    @Override
+    public <T> Map<Map<String, T>, Map<String, Object>> offsets(Collection<Map<String, T>> partitions) {
+        Map<Map<String, T>, Map<String, Object>> offsets = new HashMap<>();
+        for (Map.Entry<Map<String, Object>, Map<String, Object>> entry : InMemoryOffsetStore.Offsets.entrySet()) {
+            if (partitions.contains(entry.getKey())) {
+
+                offsets.put((Map<String, T>) entry.getKey(), entry.getValue());
+            }
+        }
+        return offsets;
+    }
+}

src/test/java/com/github/mmolimar/kafka/connect/fs/InMemoryFsSourceTask.java

@@ -0,0 +1,32 @@
+package com.github.mmolimar.kafka.connect.fs;
+
+import org.apache.kafka.connect.source.SourceRecord;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+public class InMemoryFsSourceTask extends FsSourceTask {
+
+    private List<SourceRecord> batchRecords = new ArrayList<>();
+
+    @Override
+    public void commit() throws InterruptedException {
+        if (batchRecords == null) return;
+        batchRecords.forEach(sourceRecord -> InMemoryOffsetStore.Offsets.put(
+                (Map<String, Object>) sourceRecord.sourcePartition() ,
+                (Map<String, Object>) sourceRecord.sourceOffset()
+        ));
+    }
+
+    @Override
+    public List<SourceRecord> poll() {
+        List<SourceRecord> records = super.poll();
+        batchRecords = records;
+        try {
+            this.commit();
+        } catch (InterruptedException exception) {}
+
+        return records;
+    }
+}

src/test/java/com/github/mmolimar/kafka/connect/fs/InMemoryOffsetStore.java

@@ -0,0 +1,8 @@
+package com.github.mmolimar.kafka.connect.fs;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class InMemoryOffsetStore {
+    public static Map<Map<String, Object>, Map<String, Object>> Offsets = new HashMap<>();
+}

src/test/java/com/github/mmolimar/kafka/connect/fs/file/reader/ParquetFileReaderTest.java

@@ -1,5 +1,6 @@
 package com.github.mmolimar.kafka.connect.fs.file.reader;
 
+import com.github.mmolimar.kafka.connect.fs.FsSourceTaskConfig;
 import org.apache.avro.AvroRuntimeException;
 import org.apache.avro.Schema;
 import org.apache.avro.SchemaBuilder;
@@ -26,6 +27,7 @@
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.NoSuchElementException;
 import java.util.UUID;
 import java.util.stream.IntStream;
 
@@ -180,6 +182,36 @@ public void readerWithUnparseableSchema(ReaderFsTestConfig fsConfig) {
         });
     }
 
+    @ParameterizedTest
+    @MethodSource("fileSystemConfigProvider")
+    public void readAllDataInBatches(ReaderFsTestConfig fsConfig) {
+        Map<String, Object> readerConfig = getReaderConfig();
+        int batchSize = 10;
+        readerConfig.put(ParquetFileReader.FILE_READER_PARQUET_SCHEMA, readerSchema.toString());
+        readerConfig.put(AgnosticFileReader.FILE_READER_AGNOSTIC_EXTENSIONS_PARQUET, getFileExtension());
+        readerConfig.put(FsSourceTaskConfig.FILE_READER_BATCH_SIZE, batchSize);
+
+        AbstractFileReader<?> reader = (AbstractFileReader<?>) getReader(fsConfig.getFs(), fsConfig.getDataFile(), readerConfig);
+        assertTrue(reader.hasNext());
+
+        int recordCount = 0;
+        int batchCount = 0;
+
+        while (reader.hasNextBatch()) {
+            reader.nextBatch();
+            while (reader.hasNext()) {
+                Struct record = reader.next();
+                checkData(record, recordCount);
+                recordCount++;
+            }
+            batchCount++;
+        }
+
+        assertThrows(NoSuchElementException.class, reader::nextBatch);
+        assertEquals(NUM_RECORDS / batchSize, batchCount, "The number of batches processed does not match");
+        assertEquals(NUM_RECORDS, recordCount, "The number of records in the file does not match");
+    }
+
     @Override
     protected Map<String, Object> getReaderConfig() {
         return new HashMap<>();

src/test/java/com/github/mmolimar/kafka/connect/fs/file/reader/ReaderFsTestConfig.java

@@ -8,7 +8,7 @@
 import java.util.HashMap;
 import java.util.Map;
 
-interface ReaderFsTestConfig extends FsTestConfig {
+public interface ReaderFsTestConfig extends FsTestConfig {
 
     void setDataFile(Path dataFile);