hvanhovell commented on code in PR #42011:
URL: https://github.com/apache/spark/pull/42011#discussion_r1267313704
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/client/SparkResult.scala:
##########
@@ -76,61 +71,63 @@ private[sql] class SparkResult[T](
}
}
- private def processResponses(stopOnFirstNonEmptyResponse: Boolean): Boolean
= {
- while (responses.hasNext) {
+ private def processResponses(
+ stopOnSchema: Boolean = false,
+ stopOnArrowSchema: Boolean = false,
+ stopOnFirstNonEmptyResponse: Boolean = false): Boolean = {
+ var nonEmpty = false
+ var stop = false
+ while (!stop && responses.hasNext) {
val response = responses.next()
if (response.hasSchema) {
// The original schema should arrive before ArrowBatches.
structType =
DataTypeProtoConverter.toCatalystType(response.getSchema).asInstanceOf[StructType]
- } else if (response.hasArrowBatch) {
+ stop |= stopOnSchema
+ }
+ if (response.hasArrowBatch) {
val ipcStreamBytes = response.getArrowBatch.getData
- val reader = new ArrowStreamReader(ipcStreamBytes.newInput(),
allocator)
- try {
- val root = reader.getVectorSchemaRoot
- if (structType == null) {
- // If the schema is not available yet, fallback to the schema from
Arrow.
- structType = ArrowUtils.fromArrowSchema(root.getSchema)
- }
- // TODO: create encoders that directly operate on arrow vectors.
- if (boundEncoder == null) {
- boundEncoder = createEncoder(structType)
- .resolveAndBind(DataTypeUtils.toAttributes(structType))
- }
- while (reader.loadNextBatch()) {
- val rowCount = root.getRowCount
- if (rowCount > 0) {
- val vectors = root.getFieldVectors.asScala
- .map(v => new ArrowColumnVector(transferToNewVector(v)))
- .toArray[ColumnVector]
- idxToBatches.put(nextBatchIndex, new ColumnarBatch(vectors,
rowCount))
- nextBatchIndex += 1
- numRecords += rowCount
- if (stopOnFirstNonEmptyResponse) {
- return true
- }
- }
+ val reader = new MessageIterator(ipcStreamBytes.newInput(), allocator)
+ if (arrowSchema == null) {
+ arrowSchema = reader.schema
+ stop |= stopOnArrowSchema
+ } else if (arrowSchema != reader.schema) {
+ // Uh oh...
Review Comment:
Done. The only issue is that this incredibly difficult to test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
