refactor: address PR review — share vector builder and use in-memory ArrowStreamReader

Rework the ResultSet unification to address two reviewer requests on #175:

1. Share the vector-building code with the parameter-encoding path instead
   of having a dedicated MetadataArrowBuilder. VectorPopulator now exposes a
   row-indexed primitive (setCell) used by both callers. The existing
   single-row parameter-binding overload and a new many-row metadata
   overload both funnel through it, and all the individual vector setters
   are parameterised by row index.

2. Keep ArrowStreamReaderCursor on its original ArrowStreamReader-only
   interface. The metadata path now serialises a populated VSR to Arrow
   IPC bytes and wraps the result in a ByteArrayInputStream-backed
   ArrowStreamReader, so both streaming and metadata result sets travel
   through exactly the same reader/cursor plumbing.

Supporting changes:

- typeName overrides (e.g. "TEXT" for JDBC-spec metadata columns) now
  round-trip through Arrow via a jdbc:type_name field-metadata key rather
  than a columns-override parameter on StreamingResultSet. HyperTypeToArrow
  stamps the key on write; ArrowToHyperTypeMapper.toColumnMetadata reads it
  back.
- StreamingResultSet drops the ofInMemory(...) factory and the columns
  override; callers construct an ArrowStreamReader + BufferAllocator pair
  and hand them to of(reader, allocator, queryId, zone). The cursor owns
  both and closes reader-then-allocator on close.
- QueryResultArrowStream.toArrowStreamReader returns a simple Result
  holder (reader + allocator) instead of an AutoCloseable bundle.
- MetadataResultSets is the single entry point for Arrow-backed metadata
  result sets; MetadataArrowBuilder is deleted.
- Empty metadata results skip writeBatch() entirely so ArrowStreamReaderCursor
  doesn't interpret a zero-row batch as "at least one row available".
- Tests updated to the new API; StreamingResultSetMethodTest builds its
  in-memory ResultSet the same way as the metadata path (IPC round-trip).

Author: mkaufmann

jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/core/ArrowStreamReaderCursor.java

@@ -8,6 +8,7 @@
 
 import com.salesforce.datacloud.jdbc.core.accessor.QueryJDBCAccessor;
 import com.salesforce.datacloud.jdbc.core.accessor.QueryJDBCAccessorFactory;
+import java.io.IOException;
 import java.sql.SQLException;
 import java.time.ZoneId;
 import java.util.List;
@@ -16,83 +17,47 @@
 import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
 import lombok.val;
+import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.ipc.ArrowStreamReader;
 
 /**
- * Row cursor over a {@link VectorSchemaRoot} that drives the {@link StreamingResultSet}.
+ * Row cursor over an {@link ArrowStreamReader} that drives the {@link StreamingResultSet}.
  *
- * <p>The cursor is source-agnostic: a {@link BatchLoader} loads the next batch into the vector
- * schema root, whether that comes from an Arrow IPC stream or a pre-populated in-memory batch.
- * This is the single codepath that the driver exposes to JDBC callers — streaming query results
- * and materialised metadata results both funnel through here.
- *
- * <p>The cursor owns an optional {@link AutoCloseable} holding the resources that back the
- * vector schema root (allocator, underlying reader, etc.). Closing the cursor closes that holder,
- * guaranteeing root-allocator hygiene without requiring each call site to manage the allocator
- * separately.
+ * <p>The cursor owns the supplied {@link BufferAllocator} alongside the reader: closing the
+ * cursor closes the reader (which releases ArrowBuf accounting) and then the allocator (which
+ * returns its budget). This is the single place that guarantees root-allocator hygiene for the
+ * driver; callers of {@link StreamingResultSet#of} hand ownership over and do not close the
+ * allocator themselves.
  */
 @Slf4j
 class ArrowStreamReaderCursor implements AutoCloseable {
 
     private static final int INIT_ROW_NUMBER = -1;
 
-    private final VectorSchemaRoot root;
-    private final BatchLoader batchLoader;
-    private final AutoCloseable ownedResources;
+    private final ArrowStreamReader reader;
+    private final BufferAllocator allocator;
     private final ZoneId sessionZone;
 
     @lombok.Getter
     private int rowsSeen = 0;
 
     private final AtomicInteger currentIndex = new AtomicInteger(INIT_ROW_NUMBER);
 
-    /**
-     * Loads the next batch of rows into the vector schema root.
-     *
-     * <p>Implementations should return {@code true} if the vector schema root now holds rows from
-     * a newly-loaded batch, and {@code false} if the source has no more data.
-     */
-    @FunctionalInterface
-    interface BatchLoader {
-        boolean loadNextBatch() throws Exception;
+    ArrowStreamReaderCursor(ArrowStreamReader reader, BufferAllocator allocator, ZoneId sessionZone) {
+        this.reader = reader;
+        this.allocator = allocator;
+        this.sessionZone = sessionZone;
     }
 
-    /**
-     * Create a cursor that pulls batches from an {@link ArrowStreamReader}. The reader (and the
-     * allocator it was constructed with) are owned by the cursor — closing the cursor closes the
-     * supplied {@code ownedResources}.
-     */
     @SneakyThrows
-    static ArrowStreamReaderCursor streaming(
-            ArrowStreamReader reader, AutoCloseable ownedResources, ZoneId sessionZone) {
-        val root = reader.getVectorSchemaRoot();
-        BatchLoader loader = reader::loadNextBatch;
-        return new ArrowStreamReaderCursor(root, loader, ownedResources, sessionZone);
-    }
-
-    /**
-     * Create a cursor over a single pre-populated {@link VectorSchemaRoot}. The root (and any
-     * backing allocator wrapped in {@code ownedResources}) are owned by the cursor — closing the
-     * cursor closes the supplied {@code ownedResources}.
-     */
-    static ArrowStreamReaderCursor inMemory(VectorSchemaRoot root, AutoCloseable ownedResources, ZoneId sessionZone) {
-        // The VSR is already populated, so there is nothing more to load — the cursor walks the
-        // row count until exhausted and then reports end-of-stream.
-        return new ArrowStreamReaderCursor(root, () -> false, ownedResources, sessionZone);
-    }
-
-    private ArrowStreamReaderCursor(
-            VectorSchemaRoot root, BatchLoader batchLoader, AutoCloseable ownedResources, ZoneId sessionZone) {
-        this.root = root;
-        this.batchLoader = batchLoader;
-        this.ownedResources = ownedResources;
-        this.sessionZone = sessionZone;
+    private VectorSchemaRoot getSchemaRoot() {
+        return reader.getVectorSchemaRoot();
     }
 
     List<QueryJDBCAccessor> createAccessors() {
-        return root.getFieldVectors().stream()
+        return getSchemaRoot().getFieldVectors().stream()
                 .map(rethrowFunction(this::createAccessor))
                 .collect(Collectors.toList());
     }
@@ -101,8 +66,8 @@ private QueryJDBCAccessor createAccessor(FieldVector vector) throws SQLException
         return QueryJDBCAccessorFactory.createAccessor(vector, currentIndex::get, sessionZone);
     }
 
-    private boolean loadNextBatch() throws Exception {
-        if (batchLoader.loadNextBatch()) {
+    private boolean loadNextBatch() throws IOException {
+        if (reader.loadNextBatch()) {
             currentIndex.set(0);
             return true;
         }
@@ -112,7 +77,7 @@ private boolean loadNextBatch() throws Exception {
     @SneakyThrows
     public boolean next() {
         val current = currentIndex.incrementAndGet();
-        val total = root.getRowCount();
+        val total = getSchemaRoot().getRowCount();
 
         try {
             val next = current < total || loadNextBatch();
@@ -133,8 +98,12 @@ public boolean next() {
     @SneakyThrows
     @Override
     public void close() {
-        if (ownedResources != null) {
-            ownedResources.close();
+        // Close the reader first: it releases the buffers accounted against the allocator, so the
+        // allocator's closing budget check passes. Reversing the order trips a leak detector.
+        try {
+            reader.close();
+        } finally {
+            allocator.close();
         }
     }
 }

jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/core/DataCloudConnection.java

@@ -220,7 +220,7 @@ public DataCloudResultSet getRowBasedResultSet(String queryId, long offset, long
                     QueryResultArrowStream.OUTPUT_FORMAT);
             val arrowStream = SQLExceptionQueryResultIterator.createSqlExceptionArrowStreamReader(
                     iterator, connectionProperties.isIncludeCustomerDetailInReason(), queryId, null);
-            return StreamingResultSet.of(arrowStream, queryId);
+            return StreamingResultSet.of(arrowStream.getReader(), arrowStream.getAllocator(), queryId);
         } catch (StatusRuntimeException ex) {
             throw QueryExceptionHandler.createException(
                     connectionProperties.isIncludeCustomerDetailInReason(), null, queryId, ex);
@@ -263,7 +263,7 @@ public DataCloudResultSet getChunkBasedResultSet(String queryId, long chunkId, l
                     QueryResultArrowStream.OUTPUT_FORMAT);
             val arrowStream = SQLExceptionQueryResultIterator.createSqlExceptionArrowStreamReader(
                     iterator, connectionProperties.isIncludeCustomerDetailInReason(), queryId, null);
-            return StreamingResultSet.of(arrowStream, queryId);
+            return StreamingResultSet.of(arrowStream.getReader(), arrowStream.getAllocator(), queryId);
         } catch (StatusRuntimeException ex) {
             throw QueryExceptionHandler.createException(
                     connectionProperties.isIncludeCustomerDetailInReason(), null, queryId, ex);

jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/core/DataCloudStatement.java

@@ -196,8 +196,11 @@ public ResultSet executeQuery(String sql) throws SQLException {
             val iterator = executeAdaptiveQuery(sql);
             val arrowStream = SQLExceptionQueryResultIterator.createSqlExceptionArrowStreamReader(
                     iterator, includeCustomerDetail, iterator.getQueryStatus().getQueryId(), sql);
-            resultSet =
-                    StreamingResultSet.of(arrowStream, iterator.getQueryStatus().getQueryId(), sessionZone);
+            resultSet = StreamingResultSet.of(
+                    arrowStream.getReader(),
+                    arrowStream.getAllocator(),
+                    iterator.getQueryStatus().getQueryId(),
+                    sessionZone);
             log.info(
                     "executeAdaptiveQuery completed. queryId={}, sessionZone={}",
                     queryHandle.getQueryStatus().getQueryId(),
@@ -437,7 +440,8 @@ public ResultSet getResultSet() throws SQLException {
                                     adaptiveIterator.getQueryStatus().getQueryId(),
                                     null);
                             resultSet = StreamingResultSet.of(
-                                    arrowStream,
+                                    arrowStream.getReader(),
+                                    arrowStream.getAllocator(),
                                     adaptiveIterator.getQueryStatus().getQueryId(),
                                     sessionZone);
                         } else if (resultSet == null) {

jdbc-core/src/main/java/com/salesforce/datacloud/jdbc/core/StreamingResultSet.java

@@ -9,9 +9,7 @@
 import com.salesforce.datacloud.jdbc.core.resultset.ForwardOnlyResultSet;
 import com.salesforce.datacloud.jdbc.core.resultset.ReadOnlyResultSet;
 import com.salesforce.datacloud.jdbc.core.resultset.ResultSetWithPositionalGetters;
-import com.salesforce.datacloud.jdbc.protocol.QueryResultArrowStream;
 import com.salesforce.datacloud.jdbc.protocol.data.ArrowToHyperTypeMapper;
-import com.salesforce.datacloud.jdbc.protocol.data.ColumnMetadata;
 import com.salesforce.datacloud.jdbc.util.ThrowingJdbcSupplier;
 import com.salesforce.datacloud.query.v3.QueryStatus;
 import java.io.IOException;
@@ -38,13 +36,13 @@
 import java.sql.Types;
 import java.time.ZoneId;
 import java.util.Calendar;
-import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
 import lombok.val;
-import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.ipc.ArrowStreamReader;
 
 @Slf4j
 public class StreamingResultSet
@@ -75,28 +73,41 @@ private StreamingResultSet(
         this.closed = false;
     }
 
-    public static StreamingResultSet of(QueryResultArrowStream.Result arrowStream, String queryId) throws SQLException {
-        return of(arrowStream, queryId, ZoneId.systemDefault());
+    public static StreamingResultSet of(ArrowStreamReader reader, BufferAllocator allocator, String queryId)
+            throws SQLException {
+        return of(reader, allocator, queryId, ZoneId.systemDefault());
     }
 
     /**
-     * Creates a StreamingResultSet with a specified session timezone.
+     * Creates a StreamingResultSet from an {@link ArrowStreamReader} and its backing allocator.
+     *
+     * <p>Ownership of both the reader and the allocator transfers to the returned result set —
+     * closing the result set closes the reader and then the allocator, in that order, so Arrow's
+     * buffer accounting clears before the allocator's budget check. Callers must not close
+     * either separately.
      *
-     * <p>Ownership of {@code arrowStream} (both the reader and its backing allocator) transfers
-     * to the returned result set — callers must not close it separately.
+     * <p>The column metadata (including any {@link ColumnMetadata#getTypeName()} override
+     * stamped under {@link com.salesforce.datacloud.jdbc.protocol.data.HyperTypeToArrow#JDBC_TYPE_NAME_METADATA_KEY})
+     * is derived from the Arrow schema via {@link ArrowToHyperTypeMapper#toColumnMetadata(org.apache.arrow.vector.types.pojo.Field)}.
      *
-     * @param arrowStream The Arrow stream containing query results, owned by the result set
-     * @param queryId The query identifier
-     * @param sessionZone The session timezone to use for timestamp conversions
-     * @return A new StreamingResultSet
-     * @throws SQLException If an error occurs during ResultSet creation
+     * @param reader The Arrow stream, owned by the result set.
+     * @param allocator The allocator backing the reader, owned by the result set.
+     * @param queryId The query identifier.
+     * @param sessionZone The session timezone used for timestamp conversions.
      */
-    public static StreamingResultSet of(QueryResultArrowStream.Result arrowStream, String queryId, ZoneId sessionZone)
+    public static StreamingResultSet of(
+            ArrowStreamReader reader, BufferAllocator allocator, String queryId, ZoneId sessionZone)
             throws SQLException {
         try {
-            val schemaRoot = arrowStream.getReader().getVectorSchemaRoot();
-            val cursor = ArrowStreamReaderCursor.streaming(arrowStream.getReader(), arrowStream, sessionZone);
-            return build(cursor, schemaRoot, queryId, null);
+            val schemaRoot = reader.getVectorSchemaRoot();
+            val columns = schemaRoot.getSchema().getFields().stream()
+                    .map(ArrowToHyperTypeMapper::toColumnMetadata)
+                    .collect(Collectors.toList());
+            val metadata = new DataCloudResultSetMetaData(columns);
+            val cursor = new ArrowStreamReaderCursor(reader, allocator, sessionZone);
+            val accessors = cursor.createAccessors().toArray(new QueryJDBCAccessor[0]);
+            val columnNameResolver = new ColumnNameResolver(columns);
+            return new StreamingResultSet(cursor, queryId, metadata, accessors, columnNameResolver);
         } catch (IOException ex) {
             throw new SQLException("Unexpected error during ResultSet creation", "XX000", ex);
         } catch (IllegalArgumentException ex) {
@@ -105,66 +116,6 @@ public static StreamingResultSet of(QueryResultArrowStream.Result arrowStream, S
         }
     }
 
-    /**
-     * Creates a StreamingResultSet over a pre-populated in-memory {@link VectorSchemaRoot}.
-     *
-     * <p>Used by the metadata path (e.g. {@code DatabaseMetaData.getTables}), where the rows are
-     * materialised into Arrow up front rather than streamed from the server. Ownership of
-     * {@code ownedResources} — typically the {@link org.apache.arrow.memory.BufferAllocator} and
-     * {@link VectorSchemaRoot} — transfers to the returned result set, which closes them when
-     * it is closed.
-     *
-     * @param columns optional column-metadata override. When non-null it takes precedence over
-     *     what would be derived from the Arrow schema so that {@link ColumnMetadata#getTypeName()}
-     *     overrides (e.g. {@code "TEXT"} for {@code getTables} rather than the derived
-     *     {@code "VARCHAR"}) are preserved on the way through.
-     */
-    public static StreamingResultSet ofInMemory(
-            VectorSchemaRoot schemaRoot,
-            AutoCloseable ownedResources,
-            String queryId,
-            ZoneId sessionZone,
-            List<ColumnMetadata> columns)
-            throws SQLException {
-        try {
-            val cursor = ArrowStreamReaderCursor.inMemory(schemaRoot, ownedResources, sessionZone);
-            return build(cursor, schemaRoot, queryId, columns);
-        } catch (IllegalArgumentException ex) {
-            throw new SQLException("Unsupported column type in result set: " + ex.getMessage(), "0A000", ex);
-        }
-    }
-
-    /**
-     * Overload that derives the column metadata from the Arrow schema. Prefer the overload that
-     * takes an explicit {@code columns} list if callers need to preserve JDBC-spec type names
-     * (e.g. {@code "TEXT"}).
-     */
-    public static StreamingResultSet ofInMemory(
-            VectorSchemaRoot schemaRoot, AutoCloseable ownedResources, String queryId, ZoneId sessionZone)
-            throws SQLException {
-        return ofInMemory(schemaRoot, ownedResources, queryId, sessionZone, null);
-    }
-
-    private static StreamingResultSet build(
-            ArrowStreamReaderCursor cursor,
-            VectorSchemaRoot schemaRoot,
-            String queryId,
-            List<ColumnMetadata> columnOverride)
-            throws SQLException {
-        final List<ColumnMetadata> columns;
-        if (columnOverride != null) {
-            columns = columnOverride;
-        } else {
-            columns = schemaRoot.getSchema().getFields().stream()
-                    .map(field -> new ColumnMetadata(field.getName(), ArrowToHyperTypeMapper.toHyperType(field)))
-                    .collect(Collectors.toList());
-        }
-        val metadata = new DataCloudResultSetMetaData(columns);
-        val accessors = cursor.createAccessors().toArray(new QueryJDBCAccessor[0]);
-        val columnNameResolver = new ColumnNameResolver(columns);
-        return new StreamingResultSet(cursor, queryId, metadata, accessors, columnNameResolver);
-    }
-
     // --- Core ResultSet navigation ---
 
     @Override