Update length of avro, rc, and line input files if memory input files…

… are created
trinodb · Oct 25, 2024 · 19995ec · 19995ec
1 parent f8d6289
commit 19995ec
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 7 deletions.
diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/avro/AvroPageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/avro/AvroPageSourceFactory.java
@@ -122,13 +122,13 @@ public Optional<ReaderPageSource> createPageSource(
         }
 
         try {
-            length = min(inputFile.length() - start, length);
             if (estimatedFileSize < BUFFER_SIZE.toBytes()) {
                 try (TrinoInputStream input = inputFile.newStream()) {
                     byte[] data = input.readAllBytes();
                     inputFile = new MemoryInputFile(path, Slices.wrappedBuffer(data));
                 }
             }
+            length = min(inputFile.length() - start, length);
         }
         catch (TrinoException e) {
             throw e;

diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/LinePageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/line/LinePageSourceFactory.java
@@ -53,6 +53,7 @@
 import static io.trino.plugin.hive.util.HiveUtil.getFooterCount;
 import static io.trino.plugin.hive.util.HiveUtil.getHeaderCount;
 import static io.trino.plugin.hive.util.HiveUtil.splitError;
+import static java.lang.Math.min;
 import static java.util.Objects.requireNonNull;
 
 public abstract class LinePageSourceFactory
@@ -126,11 +127,6 @@ public Optional<ReaderPageSource> createPageSource(
                     schema.serdeProperties());
         }
 
-        // Skip empty inputs
-        if (length == 0) {
-            return Optional.of(noProjectionAdaptation(new EmptyPageSource()));
-        }
-
         TrinoFileSystem trinoFileSystem = fileSystemFactory.create(session);
         TrinoInputFile inputFile = trinoFileSystem.newInputFile(path);
         try {
@@ -141,6 +137,13 @@ public Optional<ReaderPageSource> createPageSource(
                     inputFile = new MemoryInputFile(path, Slices.wrappedBuffer(data));
                 }
             }
+            length = min(inputFile.length() - start, length);
+
+            // Skip empty inputs
+            if (length <= 0) {
+                return Optional.of(noProjectionAdaptation(new EmptyPageSource()));
+            }
+
             LineReader lineReader = lineReaderFactory.createLineReader(inputFile, start, length, headerCount, footerCount);
             // Split may be empty after discovering the real file size and skipping headers
             if (lineReader.isClosed()) {

diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/rcfile/RcFilePageSourceFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/rcfile/RcFilePageSourceFactory.java
@@ -124,7 +124,6 @@ else if (serializationLibraryName.equals(COLUMNAR_SERDE_CLASS)) {
         TrinoFileSystem trinoFileSystem = fileSystemFactory.create(session);
         TrinoInputFile inputFile = trinoFileSystem.newInputFile(path);
         try {
-            length = min(inputFile.length() - start, length);
             if (!inputFile.exists()) {
                 throw new TrinoException(HIVE_CANNOT_OPEN_SPLIT, "File does not exist");
             }
@@ -134,6 +133,7 @@ else if (serializationLibraryName.equals(COLUMNAR_SERDE_CLASS)) {
                     inputFile = new MemoryInputFile(path, Slices.wrappedBuffer(data));
                 }
             }
+            length = min(inputFile.length() - start, length);
         }
         catch (TrinoException e) {
             throw e;