Decouple QueryData on the wire from decoding

client/trino-cli/src/main/java/io/trino/cli/OutputHandler.java

@@ -65,11 +65,8 @@ public void processRows(StatementClient client)
         BlockingQueue<List<?>> rowQueue = new ArrayBlockingQueue<>(MAX_QUEUED_ROWS);
         CompletableFuture<Void> readerFuture = CompletableFuture.runAsync(() -> {
             while (client.isRunning()) {
-                Iterable<List<Object>> data = client.currentData().getData();
-                if (data != null) {
-                    for (List<Object> row : data) {
-                        putOrThrow(rowQueue, row);
-                    }
+                for (List<Object> row : client.currentRows()) {
+                    putOrThrow(rowQueue, row);
                 }
                 client.advance();
             }

client/trino-cli/src/main/java/io/trino/cli/Query.java

@@ -220,7 +220,7 @@ private boolean isInteractive(Optional<String> pager)
 
     private void processInitialStatusUpdates(WarningsPrinter warningsPrinter)
     {
-        while (client.isRunning() && (client.currentData().getData() == null)) {
+        while (client.isRunning() && client.currentRows().isNull()) {
             warningsPrinter.print(client.currentStatusInfo().getWarnings(), true, false);
             try {
                 client.advance();

client/trino-cli/src/main/java/io/trino/cli/StatusPrinter.java

@@ -105,7 +105,7 @@ public void printInitialStatusUpdates(Terminal terminal)
             while (client.isRunning()) {
                 try {
                     // exit status loop if there is pending output
-                    if (client.currentData().getData() != null) {
+                    if (!client.currentRows().isNull()) {
                         return;
                     }
 

client/trino-cli/src/main/java/io/trino/cli/TableNameCompleter.java

@@ -18,7 +18,6 @@
 import com.google.common.cache.LoadingCache;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
-import io.trino.client.QueryData;
 import io.trino.client.StatementClient;
 import org.gaul.modernizer_maven_annotations.SuppressModernizer;
 import org.jline.reader.Candidate;
@@ -86,11 +85,8 @@ private List<String> queryMetadata(String query)
         ImmutableList.Builder<String> cache = ImmutableList.builder();
         try (StatementClient client = queryRunner.startInternalQuery(query)) {
             while (client.isRunning() && !Thread.currentThread().isInterrupted()) {
-                QueryData results = client.currentData();
-                if (results.getData() != null) {
-                    for (List<Object> row : results.getData()) {
-                        cache.add((String) row.get(0));
-                    }
+                for (List<Object> row : client.currentRows()) {
+                    cache.add((String) row.get(0));
                 }
                 client.advance();
             }

client/trino-client/src/main/java/io/trino/client/QueryData.java

@@ -17,8 +17,17 @@
 
 import java.util.List;
 
+/**
+ * Used for representing both raw JSON values and spooled metadata.
+ */
 public interface QueryData
 {
+    @Deprecated
     @Nullable
-    Iterable<List<Object>> getData();
+    default Iterable<List<Object>> getData()
+    {
+        throw new UnsupportedOperationException("getData() is deprecated for removal, use StatementClient.currentRows() instead");
+    }
+
+    boolean isNull();
 }

client/trino-client/src/main/java/io/trino/client/QueryDataDecoder.java

@@ -14,7 +14,6 @@
 package io.trino.client;
 
 import io.trino.client.spooling.DataAttributes;
-import io.trino.client.spooling.encoding.QueryDataAccess;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -24,13 +23,13 @@ public interface QueryDataDecoder
 {
     interface Factory
     {
-        QueryDataDecoder create(List<Column> columns, DataAttributes queryAttributes);
+        QueryDataDecoder create(List<Column> columns, DataAttributes attributes);
 
         String encoding();
     }
 
     /**
-     * Decodes the input stream into a QueryDataAccess object.
+     * Decodes the input stream into a lazy ResultRows.
      * <p>
      * Decoder is responsible for closing input stream when
      * all values are decoded or exception was thrown.
@@ -40,7 +39,7 @@ interface Factory
      *
      * @throws IOException if an I/O error occurs
      */
-    QueryDataAccess decode(InputStream input, DataAttributes segmentAttributes)
+    ResultRows decode(InputStream input, DataAttributes segmentAttributes)
             throws IOException;
 
     String encoding();

client/trino-client/src/main/java/io/trino/client/QueryResults.java

@@ -183,12 +183,6 @@ public String toString()
 
     private static boolean hasData(QueryData data)
     {
-        if (data == null) {
-            return false;
-        }
-        if (data instanceof RawQueryData) {
-            return data.getData() != null;
-        }
-        return true;
+        return data != null && !data.isNull();
     }
 }

client/trino-client/src/main/java/io/trino/client/RawQueryData.java

@@ -13,12 +13,13 @@
  */
 package io.trino.client;
 
+import jakarta.annotation.Nonnull;
 import jakarta.annotation.Nullable;
 
 import java.util.List;
 
+import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.Iterables.unmodifiableIterable;
-import static io.trino.client.FixJsonDataUtils.fixData;
 
 /**
  * Class represents QueryData serialized to JSON array of arrays of objects.
@@ -34,9 +35,9 @@ private RawQueryData(Iterable<List<Object>> values)
         this.iterable = values == null ? null : unmodifiableIterable(values);
     }
 
-    @Override
-    public Iterable<List<Object>> getData()
+    public @Nonnull Iterable<List<Object>> getIterable()
     {
+        checkState(iterable != null, "cannot return a null iterable");
         return iterable;
     }
 
@@ -45,9 +46,9 @@ public static QueryData of(@Nullable Iterable<List<Object>> values)
         return new RawQueryData(values);
     }
 
-    // JSON encoding loses type information. In order for it to be usable, we need to fix types.
-    public QueryData fixTypes(List<Column> columns)
+    @Override
+    public boolean isNull()
     {
-        return RawQueryData.of(fixData(columns, iterable));
+        return iterable == null;
     }
 }

client/trino-client/src/main/java/io/trino/client/ResultRows.java

@@ -0,0 +1,54 @@
+/*
+ * 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 io.trino.client;
+
+import jakarta.annotation.Nonnull;
+
+import java.util.Iterator;
+import java.util.List;
+
+import static java.util.Collections.emptyIterator;
+
+/**
+ * Allows iterating over decoded result data in row-wise manner.
+ */
+public interface ResultRows
+        extends Iterable<List<Object>>
+{
+    ResultRows NULL_ROWS = new ResultRows() {
+        @Override
+        public boolean isNull()
+        {
+            // This should be the only instance of this method returning true,
+            // as this means "no rows yet" which is different from "empty rows".
+            return true;
+        }
+
+        @Override
+        public @Nonnull Iterator<List<Object>> iterator()
+        {
+            return emptyIterator();
+        }
+    };
+
+    static ResultRows fromIterableRows(Iterable<List<Object>> values)
+    {
+        return values::iterator;
+    }
+
+    default boolean isNull()
+    {
+        return false;
+    }
+}

client/trino-client/src/main/java/io/trino/client/ResultRowsDecoder.java

@@ -0,0 +1,167 @@
+/*
+ * 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 io.trino.client;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterators;
+import io.trino.client.spooling.DataAttributes;
+import io.trino.client.spooling.EncodedQueryData;
+import io.trino.client.spooling.InlineSegment;
+import io.trino.client.spooling.Segment;
+import io.trino.client.spooling.SegmentLoader;
+import io.trino.client.spooling.SpooledSegment;
+import io.trino.client.spooling.encoding.QueryDataDecoders;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.UncheckedIOException;
+import java.util.List;
+import java.util.Optional;
+
+import static com.google.common.base.Verify.verify;
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static io.trino.client.FixJsonDataUtils.fixData;
+import static io.trino.client.ResultRows.NULL_ROWS;
+import static java.util.Collections.emptyList;
+import static java.util.Objects.requireNonNull;
+
+/**
+ * Class responsible for decoding any QueryData type.
+ */
+public class ResultRowsDecoder
+        implements AutoCloseable
+{
+    private final SegmentLoader loader;
+    private QueryDataDecoder.Factory decoderFactory;
+    private List<Column> columns = emptyList();
+
+    public ResultRowsDecoder()
+    {
+        this(new SegmentLoader());
+    }
+
+    public ResultRowsDecoder(SegmentLoader loader)
+    {
+        this.loader = requireNonNull(loader, "loader is null");
+    }
+
+    public ResultRowsDecoder withEncoding(String encoding)
+    {
+        if (decoderFactory != null) {
+            if (!encoding.equals(decoderFactory.encoding())) {
+                throw new IllegalStateException("Already set encoding " + encoding + " is not equal to " + decoderFactory.encoding());
+            }
+        }
+        else {
+            this.decoderFactory = QueryDataDecoders.get(encoding);
+        }
+        return this;
+    }
+
+    public ResultRowsDecoder withColumns(List<Column> columns)
+    {
+        if (this.columns.isEmpty()) {
+            this.columns = ImmutableList.copyOf(columns);
+        }
+        else if (!columns.equals(this.columns)) {
+            throw new IllegalStateException("Already set columns " + columns + " are not equal to " + this.columns);
+        }
+
+        return this;
+    }
+
+    public ResultRows toRows(QueryData data)
+    {
+        if (data == null) {
+            return NULL_ROWS; // for backward compatibility instead of null
+        }
+
+        verify(!columns.isEmpty(), "columns must be set");
+        if (data instanceof RawQueryData) {
+            RawQueryData rawData = (RawQueryData) data;
+            if (rawData.isNull()) {
+                return NULL_ROWS; // for backward compatibility instead of null
+            }
+            return () -> fixData(columns, rawData.getIterable()).iterator();
+        }
+
+        if (data instanceof EncodedQueryData) {
+            verify(decoderFactory != null, "decoderFactory must be set");
+            // we don't need query-level attributes for now
+            QueryDataDecoder decoder = decoderFactory.create(columns, DataAttributes.empty());
+            EncodedQueryData encodedData = (EncodedQueryData) data;
+            verify(decoder.encoding().equals(encodedData.getEncoding()), "encoding %s is not equal to %s", encodedData.getEncoding(), decoder.encoding());
+
+            List<ResultRows> resultRows = encodedData.getSegments()
+                    .stream()
+                    .map(segment -> segmentToRows(decoder, segment))
+                    .collect(toImmutableList());
+
+            return concat(resultRows);
+        }
+
+        throw new UnsupportedOperationException("Unsupported data type: " + data.getClass().getName());
+    }
+
+    private ResultRows segmentToRows(QueryDataDecoder decoder, Segment segment)
+    {
+        if (segment instanceof InlineSegment) {
+            InlineSegment inlineSegment = (InlineSegment) segment;
+            try {
+                return decoder.decode(new ByteArrayInputStream(inlineSegment.getData()), inlineSegment.getMetadata());
+            }
+            catch (IOException e) {
+                throw new UncheckedIOException(e);
+            }
+        }
+
+        if (segment instanceof SpooledSegment) {
+            SpooledSegment spooledSegment = (SpooledSegment) segment;
+
+            try {
+                InputStream stream = loader.load(spooledSegment);
+                return decoder.decode(stream, spooledSegment.getMetadata());
+            }
+            catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+        throw new UnsupportedOperationException("Unsupported segment type: " + segment.getClass().getName());
+    }
+
+    public Optional<String> getEncoding()
+    {
+        return Optional.ofNullable(decoderFactory)
+                .map(QueryDataDecoder.Factory::encoding);
+    }
+
+    @Override
+    public void close()
+            throws Exception
+    {
+        loader.close();
+    }
+
+    private static ResultRows concat(List<ResultRows> resultRows)
+    {
+        return () -> Iterators.concat(resultRows
+                .stream()
+                .filter(rows -> !rows.isNull())
+                .map(ResultRows::iterator)
+                .collect(toImmutableList())
+                .iterator());
+    }
+}