From 623ee121aeb6998018b955d2a7dc2791cc4f517a Mon Sep 17 00:00:00 2001
From: Arttu Voutilainen <avoutilainen@palantir.com>
Date: Sun, 27 Oct 2024 21:23:51 -0400
Subject: [PATCH] fix: correctly set type for scalar subquery when converting
 proto to pojo

---
 .../expression/proto/ProtoExpressionConverter.java    | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java b/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java
index d2b95d74f..8857614f8 100644
--- a/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java
+++ b/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java
@@ -15,6 +15,7 @@
 import io.substrait.relation.ConsistentPartitionWindow;
 import io.substrait.relation.ProtoRelConverter;
 import io.substrait.type.Type;
+import io.substrait.type.TypeVisitor;
 import io.substrait.type.proto.ProtoTypeConverter;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -196,7 +197,15 @@ public Expression from(io.substrait.proto.Expression expr) {
             var rel = protoRelConverter.from(expr.getSubquery().getScalar().getInput());
             yield ImmutableExpression.ScalarSubquery.builder()
                 .input(rel)
-                .type(rel.getRecordType())
+                .type(rel.getRecordType().accept(new TypeVisitor.TypeThrowsVisitor<Type, RuntimeException>("Expected struct field") {
+                    @Override
+                    public Type visit(Type.Struct type) throws RuntimeException {
+                        if (type.fields().size() != 1) {
+                            throw new UnsupportedOperationException("Scalar subquery must have exactly one field");
+                        }
+                        return type.fields().get(0);
+                    }
+                }))
                 .build();
           }
           case IN_PREDICATE -> {