Source:
public static Iterator<byte[]> collectAsPandasDataFrame(Table table, int maxArrowBatchSize) throws Exception { BufferAllocator allocator = getRootAllocator().newChildAllocator("collectAsPandasDataFrame", 0, Long.MAX_VALUE); RowType rowType = (RowType) table.getSchema().toRowDataType().getLogicalType(); VectorSchemaRoot root = VectorSchemaRoot.create(ArrowUtils.toArrowSchema(rowType), allocator); ByteArrayOutputStream baos = new ByteArrayOutputStream(); ArrowStreamWriter arrowStreamWriter = new ArrowStreamWriter(root, null, baos); arrowStreamWriter.start(); ArrowWriter arrowWriter; Iterator<Row> results = table.execute().collect(); Iterator convertedResults; if (isBlinkPlanner(table)) { arrowWriter = createRowDataArrowWriter(root, rowType); convertedResults = new Iterator<RowData>() { @Override public boolean hasNext() { return results.hasNext(); } @Override public RowData next() { TableSchema convertedTableSchema = SelectTableSinkSchemaConverter.changeDefaultConversionClass(table.getSchema()); DataFormatConverters.DataFormatConverter converter = DataFormatConverters.getConverterForDataType(convertedTableSchema.toRowDataType()); return (RowData) converter.toInternal(results.next()); } }; } else { arrowWriter = createRowArrowWriter(root, rowType); convertedResults = results; } return new Iterator<byte[]>() { @Override public boolean hasNext() { return convertedResults.hasNext(); } @Override public byte[] next() { try { int i = 0; while (convertedResults.hasNext() && i < maxArrowBatchSize) { i++; arrowWriter.write(convertedResults.next()); } arrowWriter.finish(); arrowStreamWriter.writeBatch(); return baos.toByteArray(); } catch (Throwable t) { String msg = "Failed to serialize the data of the table"; LOG.error(msg, t); throw new RuntimeException(msg, t); } finally { arrowWriter.reset(); baos.reset(); if (!hasNext()) { root.close(); allocator.close(); } } } }; }

Target:
public static CustomIterator<byte[]> collectAsPandasDataFrame(Table table, int maxArrowBatchSize) throws Exception { BufferAllocator allocator = getRootAllocator().newChildAllocator("collectAsPandasDataFrame", 0, Long.MAX_VALUE); RowType rowType = (RowType) table.getSchema().toRowDataType().getLogicalType(); VectorSchemaRoot root = VectorSchemaRoot.create(ArrowUtils.toArrowSchema(rowType), allocator); ByteArrayOutputStream baos = new ByteArrayOutputStream(); ArrowStreamWriter arrowStreamWriter = new ArrowStreamWriter(root, null, baos); arrowStreamWriter.start(); ArrowWriter arrowWriter; Iterator<Row> results = table.execute().collect(); Iterator convertedResults; if (isBlinkPlanner(table)) { arrowWriter = createRowDataArrowWriter(root, rowType); convertedResults = new Iterator<RowData>() { @Override public boolean hasNext() { return results.hasNext(); } @Override public RowData next() { TableSchema convertedTableSchema = SelectTableSinkSchemaConverter.changeDefaultConversionClass(table.getSchema()); DataFormatConverters.DataFormatConverter converter = DataFormatConverters.getConverterForDataType(convertedTableSchema.toRowDataType()); return (RowData) converter.toInternal(results.next()); } }; } else { arrowWriter = createRowArrowWriter(root, rowType); convertedResults = results; } return new CustomIterator<byte[]>() { @Override public boolean hasNext() { return convertedResults.hasNext(); } @Override public byte[] next() { try { int i = 0; while (convertedResults.hasNext() && i < maxArrowBatchSize) { i++; arrowWriter.write(convertedResults.next()); } arrowWriter.finish(); arrowStreamWriter.writeBatch(); return baos.toByteArray(); } catch (Throwable t) { String msg = "Failed to serialize the data of the table"; LOG.error(msg, t); throw new RuntimeException(msg, t); } finally { arrowWriter.reset(); baos.reset(); if (!hasNext()) { root.close(); allocator.close(); } } } }; }