diff --git a/core/src/main/java/io/substrait/expression/MaskExpression.java b/core/src/main/java/io/substrait/expression/MaskExpression.java new file mode 100644 index 000000000..730e3d222 --- /dev/null +++ b/core/src/main/java/io/substrait/expression/MaskExpression.java @@ -0,0 +1,422 @@ +package io.substrait.expression; + +import io.substrait.util.VisitationContext; +import java.util.List; +import java.util.Optional; +import org.immutables.value.Value; + +/** + * A mask expression that selectively removes fields from complex types (struct, list, map). + * + *

This corresponds to the {@code Expression.MaskExpression} message in the Substrait protobuf + * specification. It is used in {@code ReadRel} to describe column projection — the subset of a + * relation's schema that should actually be read. + * + * @see Substrait Field References + */ +@Value.Enclosing +@Value.Immutable +public interface MaskExpression { + + /** + * The top-level struct selection describing which fields to include. + * + * @return the top-level struct selection + */ + StructSelect getSelect(); + + /** + * When {@code true}, a struct that has only a single selected field will not be + * unwrapped into its child type. + * + * @return {@code true} if singular structs should be maintained, {@code false} otherwise + */ + @Value.Default + default boolean getMaintainSingularStruct() { + return false; + } + + /** + * Creates a new builder for constructing a MaskExpression. + * + * @return a new builder instance + */ + static ImmutableMaskExpression.Builder builder() { + return ImmutableMaskExpression.builder(); + } + + // --------------------------------------------------------------------------- + // Select – a union of StructSelect | ListSelect | MapSelect + // --------------------------------------------------------------------------- + + /** A selection on a complex type – one of StructSelect, ListSelect, or MapSelect. */ + interface Select { + /** + * Accepts a visitor to process this select node. + * + * @param the return type of the visitor + * @param the context type + * @param the exception type that may be thrown + * @param visitor the visitor to accept + * @param context the visitation context + * @return the result of the visitation + * @throws E if an error occurs during visitation + */ + R accept( + MaskExpressionVisitor visitor, C context) throws E; + } + + // --------------------------------------------------------------------------- + // Struct selection + // --------------------------------------------------------------------------- + + /** Selects a subset of fields from a struct type. */ + @Value.Immutable + interface StructSelect extends Select { + /** + * Returns the list of struct items being selected. + * + * @return the list of struct items + */ + List getStructItems(); + + /** + * Creates a new builder for constructing a StructSelect. + * + * @return a new builder instance + */ + static ImmutableMaskExpression.StructSelect.Builder builder() { + return ImmutableMaskExpression.StructSelect.builder(); + } + + @Override + default R accept( + MaskExpressionVisitor visitor, C context) throws E { + return visitor.visit(this, context); + } + } + + /** Selects a single field from a struct, with an optional nested child selection. */ + @Value.Immutable + interface StructItem { + /** + * Returns the zero-based field index within the struct. + * + * @return the field index + */ + int getField(); + + /** + * Returns the optional child selection for nested complex types. + * + * @return the optional child selection + */ + Optional getChild(); + + /** + * Creates a new builder for constructing a ListSelect. + * + * @return a new builder instance + */ + static ImmutableMaskExpression.ListSelect.Builder builder() { + return ImmutableMaskExpression.ListSelect.builder(); + } + + @Override + default R accept( + MaskExpressionVisitor visitor, C context) throws E { + return visitor.visit(this, context); + } + } + + /** A single selection within a list – either an element or a slice. */ + @Value.Immutable + interface ListSelectItem { + /** + * Returns the optional list element selection. + * + * @return the optional list element + */ + Optional getItem(); + + /** + * Returns the optional list slice selection. + * + * @return the optional list slice + */ + Optional getSlice(); + + /** + * Creates a new builder for constructing a ListSelectItem. + * + * @return a new builder instance + */ + static ImmutableMaskExpression.ListSelectItem.Builder builder() { + return ImmutableMaskExpression.ListSelectItem.builder(); + } + + /** + * Creates a ListSelectItem for a single element selection. + * + * @param element the list element to select + * @return a new ListSelectItem instance + */ + static ListSelectItem ofItem(ListElement element) { + return builder().item(element).build(); + } + + /** + * Creates a ListSelectItem for a slice selection. + * + * @param slice the list slice to select + * @return a new ListSelectItem instance + */ + static ListSelectItem ofSlice(ListSlice slice) { + return builder().slice(slice).build(); + } + } + + /** Selects a single element from a list by zero-based index. */ + @Value.Immutable + interface ListElement { + /** + * Returns the zero-based element index within the list. + * + * @return the element index + */ + int getField(); + + /** + * Creates a new builder for constructing a ListElement. + * + * @return a new builder instance + */ + static ImmutableMaskExpression.ListElement.Builder builder() { + return ImmutableMaskExpression.ListElement.builder(); + } + + /** + * Creates a ListElement for a single element selection. + * + * @param field the zero-based element index within the list + * @return a new ListElement instance + */ + static ListElement of(int field) { + return builder().field(field).build(); + } + } + + /** Selects a contiguous range of elements from a list. */ + @Value.Immutable + interface ListSlice { + /** + * Returns the zero-based start index of the slice (inclusive). + * + * @return the start index + */ + int getStart(); + + /** + * Returns the zero-based end index of the slice (exclusive). + * + * @return the end index + */ + int getEnd(); + + /** + * Creates a new builder for constructing a ListSlice. + * + * @return a new builder instance + */ + static ImmutableMaskExpression.ListSlice.Builder builder() { + return ImmutableMaskExpression.ListSlice.builder(); + } + + /** + * Creates a ListSlice for a contiguous range of elements. + * + * @param start the zero-based start index (inclusive) + * @param end the zero-based end index (exclusive) + * @return a new ListSlice instance + */ + static ListSlice of(int start, int end) { + return builder().start(start).end(end).build(); + } + } + + // --------------------------------------------------------------------------- + // Map selection + // --------------------------------------------------------------------------- + + /** Selects entries from a map type by exact key or key expression. */ + @Value.Immutable + interface MapSelect extends Select { + /** + * Returns the optional exact key for map selection. + * + * @return the optional map key + */ + Optional getKey(); + + /** + * Returns the optional key expression for wildcard map selection. + * + * @return the optional map key expression + */ + Optional getExpression(); + + /** + * Returns the optional child selection applied to each selected map value. + * + * @return the optional child selection + */ + Optional