Merge pull request #73 from nineinchnick/stream-record-set

Iterable for record set
nineinchnick · Jul 8, 2021 · 4d5cf4b · 4d5cf4b
2 parents 4ede777 + 7ac02fe
commit 4d5cf4b
Show file tree

Hide file tree

Showing 21 changed files with 167 additions and 137 deletions.
diff --git a/trino-rest-base/src/main/java/pl/net/was/rest/Rest.java b/trino-rest-base/src/main/java/pl/net/was/rest/Rest.java
@@ -36,7 +36,6 @@
 import io.trino.spi.statistics.TableStatistics;
 import io.trino.spi.type.Type;
 
-import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -61,7 +60,7 @@ default List<SchemaTableName> listTables()
 
     Map<SchemaTableName, List<ColumnMetadata>> listTableColumns(SchemaTablePrefix schemaTablePrefix);
 
-    Collection<? extends List<?>> getRows(RestTableHandle table);
+    Iterable<List<?>> getRows(RestTableHandle table);
 
     Consumer<List> createRowSink(SchemaTableName schemaTableName);
 

diff --git a/trino-rest-base/src/main/java/pl/net/was/rest/RestRecordSetProvider.java b/trino-rest-base/src/main/java/pl/net/was/rest/RestRecordSetProvider.java
@@ -14,6 +14,7 @@
 
 package pl.net.was.rest;
 
+import com.google.common.collect.Iterables;
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ColumnMetadata;
 import io.trino.spi.connector.ConnectorRecordSetProvider;
@@ -28,7 +29,6 @@
 
 import javax.inject.Inject;
 
-import java.util.Collection;
 import java.util.List;
 
 import static java.util.stream.Collectors.toList;
@@ -60,7 +60,7 @@ public RecordSet getRecordSet(
 
         // use the split's table handle, since
         RestTableHandle restTable = split.getTableHandle();
-        Collection<? extends List<?>> rows = rest.getRows(restTable);
+        Iterable<List<?>> rows = rest.getRows(restTable);
         ConnectorTableMetadata tableMetadata = rest.getTableMetadata(restTable.getSchemaTableName());
 
         List<Integer> columnIndexes = restColumnHandles
@@ -77,13 +77,11 @@ public RecordSet getRecordSet(
                 })
                 .collect(toList());
 
-        Collection<? extends List<?>> mappedRows = rows
+        //noinspection StaticPseudoFunctionalStyleMethod
+        Iterable<List<?>> mappedRows = Iterables.transform(rows, row -> columnIndexes
                 .stream()
-                .map(row -> columnIndexes
-                        .stream()
-                        .map(row::get)
-                        .collect(toList()))
-                .collect(toList());
+                .map(row::get)
+                .collect(toList()));
 
         List<Type> mappedTypes = restColumnHandles
                 .stream()