fix: correctly set type for scalar subquery when converting proto to …

…pojo
substrait-io · Nov 21, 2024 · 623ee12 · 623ee12
1 parent c3c4415
commit 623ee12
Showing 1 changed file with 10 additions and 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
@@ -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 -> {