[FLINK-18286] Implement type inference for GET/FLATTEN

flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/BuiltInFunctionDefinitions.java

@@ -1082,13 +1082,28 @@ public final class BuiltInFunctionDefinitions {
  new BuiltInFunctionDefinition.Builder()
  .name("flatten")
  .kind(OTHER)
- .outputTypeStrategy(TypeStrategies.MISSING)
+ .inputTypeStrategy(sequence(InputTypeStrategies.COMPOSITE))
+ .outputTypeStrategy(callContext -> {
+ throw new UnsupportedOperationException("FLATTEN should be resolved to GET expressions");
+ })
  .build();
  public static final BuiltInFunctionDefinition GET =
  new BuiltInFunctionDefinition.Builder()
  .name("get")
  .kind(OTHER)
- .outputTypeStrategy(TypeStrategies.MISSING)
+ .inputTypeStrategy(
+ sequence(
+ InputTypeStrategies.COMPOSITE,
+ and(
+ InputTypeStrategies.LITERAL,
+ or(
+ logical(LogicalTypeRoot.INTEGER),
+ logical(LogicalTypeFamily.CHARACTER_STRING)
+ )
+ )
+ )
+ )
+ .outputTypeStrategy(TypeStrategies.GET)
  .build();
 
  // window properties

flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/InputTypeStrategies.java

@@ -26,6 +26,7 @@
 import org.apache.flink.table.types.inference.strategies.CastInputTypeStrategy;
 import org.apache.flink.table.types.inference.strategies.CommonInputTypeStrategy;
 import org.apache.flink.table.types.inference.strategies.ComparableTypeStrategy;
+import org.apache.flink.table.types.inference.strategies.CompositeArgumentTypeStrategy;
 import org.apache.flink.table.types.inference.strategies.ConstraintArgumentTypeStrategy;
 import org.apache.flink.table.types.inference.strategies.ExplicitArgumentTypeStrategy;
 import org.apache.flink.table.types.inference.strategies.FamilyArgumentTypeStrategy;
@@ -198,6 +199,11 @@ public static InputTypeStrategy comparable(
  */
  public static final LiteralArgumentTypeStrategy LITERAL_OR_NULL = new LiteralArgumentTypeStrategy(true);
 
+ /**
+ * Strategy that checks that the argument has a composite type.
+ */
+ public static final ArgumentTypeStrategy COMPOSITE = new CompositeArgumentTypeStrategy();
+
  /**
  * Strategy for an argument that corresponds to an explicitly defined type casting.
  * Implicit casts will be inserted if possible.

flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/TypeStrategies.java

@@ -38,6 +38,7 @@
 import org.apache.flink.table.types.logical.LogicalTypeFamily;
 import org.apache.flink.table.types.logical.LogicalTypeRoot;
 import org.apache.flink.table.types.logical.utils.LogicalTypeMerging;
+import org.apache.flink.table.types.utils.DataTypeUtils;
 import org.apache.flink.table.types.utils.TypeConversions;
 
 import java.math.BigDecimal;
@@ -357,6 +358,35 @@ public static TypeStrategy varyingString(TypeStrategy initialStrategy) {
  .map(TypeConversions::fromLogicalToDataType);
  };
 
+ /**
+ * Type strategy that returns a type of a field nested inside a composite type that is described by the second argument.
+ * The second argument must be a literal that describes either the nested field name or index.
+ */
+ public static final TypeStrategy GET = callContext -> {
+ List<DataType> argumentDataTypes = callContext.getArgumentDataTypes();
+ DataType rowDataType = argumentDataTypes.get(0);
+
+ Optional<DataType> result = Optional.empty();
+
+ Optional<String> fieldName = callContext.getArgumentValue(1, String.class);
+ if (fieldName.isPresent()) {
+ result = DataTypeUtils.getField(rowDataType, fieldName.get());
+ }
+
+ Optional<Integer> fieldIndex = callContext.getArgumentValue(1, Integer.class);
+ if (fieldIndex.isPresent()) {
+ result = DataTypeUtils.getField(rowDataType, fieldIndex.get());
+ }
+
+ return result.map(type -> {
+ if (rowDataType.getLogicalType().isNullable()) {
+ return type.nullable();
+ } else {
+ return type;
+ }
+ });
+ };
+
  // --------------------------------------------------------------------------------------------
 
  @SuppressWarnings("BooleanMethodIsAlwaysInverted")

flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/strategies/CompositeArgumentTypeStrategy.java

@@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http:https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.table.types.inference.strategies;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.table.functions.FunctionDefinition;
+import org.apache.flink.table.types.DataType;
+import org.apache.flink.table.types.inference.ArgumentTypeStrategy;
+import org.apache.flink.table.types.inference.CallContext;
+import org.apache.flink.table.types.inference.Signature;
+import org.apache.flink.table.types.logical.utils.LogicalTypeChecks;
+
+import java.util.Optional;
+
+/**
+ * Strategy that checks that the argument has a composite type.
+ */
+@Internal
+public class CompositeArgumentTypeStrategy implements ArgumentTypeStrategy {
+ @Override
+ public Optional<DataType> inferArgumentType(
+ CallContext callContext,
+ int argumentPos,
+ boolean throwOnFailure) {
+ DataType dataType = callContext.getArgumentDataTypes().get(argumentPos);
+ if (!LogicalTypeChecks.isCompositeType(dataType.getLogicalType())) {
+ if (throwOnFailure) {
+ throw callContext.newValidationError(
+ "A composite type expected. Got: %s",
+ dataType);
+ }
+ return Optional.empty();
+ }
+
+ return Optional.of(dataType);
+ }
+
+ @Override
+ public Signature.Argument getExpectedArgument(
+ FunctionDefinition functionDefinition,
+ int argumentPos) {
+ return Signature.Argument.of("<COMPOSITE>");
+ }
+}

flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/utils/DataTypeUtils.java

@@ -127,13 +127,43 @@ public static TableSchema expandCompositeTypeToSchema(DataType dataType) {
  if (dataType instanceof FieldsDataType) {
  return expandCompositeType((FieldsDataType) dataType);
  } else if (dataType.getLogicalType() instanceof LegacyTypeInformationType &&
- dataType.getLogicalType().getTypeRoot() == LogicalTypeRoot.STRUCTURED_TYPE) {
+  dataType.getLogicalType().getTypeRoot() == LogicalTypeRoot.STRUCTURED_TYPE) {
  return expandLegacyCompositeType(dataType);
  }
 
  throw new IllegalArgumentException("Expected a composite type");
  }
 
+ /**
+ * Retrieves a nested field from a composite type at given position.
+ *
+ * <p>Throws an exception for a non composite type. You can use
+ * {@link LogicalTypeChecks#isCompositeType(LogicalType)} to check that.
+ *
+ * @param compositeType Data type to expand. Must be a composite type.
+ * @param index Index of the field to retrieve.
+ * @return The field at the given position.
+ */
+ public static Optional<DataType> getField(DataType compositeType, int index) {
+ TableSchema tableSchema = expandCompositeTypeToSchema(compositeType);
+ return tableSchema.getFieldDataType(index);
+ }
+
+ /**
+ * Retrieves a nested field from a composite type with given name.
+ *
+ * <p>Throws an exception for a non composite type. You can use
+ * {@link LogicalTypeChecks#isCompositeType(LogicalType)} to check that.
+ *
+ * @param compositeType Data type to expand. Must be a composite type.
+ * @param name Name of the field to retrieve.
+ * @return The field with the given name.
+ */
+ public static Optional<DataType> getField(DataType compositeType, String name) {
+ TableSchema tableSchema = expandCompositeTypeToSchema(compositeType);
+ return tableSchema.getFieldDataType(name);
+ }
+
  /**
  * The {@link DataType} class can only partially verify the conversion class. This method can perform
  * the final check when we know if the data type should be used for input.

flink-table/flink-table-common/src/test/java/org/apache/flink/table/types/inference/InputTypeStrategiesTest.java

@@ -582,7 +582,32 @@ public static List<TestSpec> testData() {
  "My constraint says %s must be nullable.",
  args -> args.get(0).getLogicalType().isNullable()))))
  .calledWithArgumentTypes(DataTypes.BOOLEAN().notNull())
- .expectErrorMessage("My constraint says BOOLEAN NOT NULL must be nullable.")
+ .expectErrorMessage("My constraint says BOOLEAN NOT NULL must be nullable."),
+
+ TestSpec
+ .forStrategy(
+ "Composite type strategy with ROW",
+ sequence(InputTypeStrategies.COMPOSITE)
+ )
+ .calledWithArgumentTypes(DataTypes.ROW(DataTypes.FIELD("f0", DataTypes.BIGINT())))
+ .expectSignature("f(<COMPOSITE>)")
+ .expectArgumentTypes(DataTypes.ROW(DataTypes.FIELD("f0", DataTypes.BIGINT()))),
+
+ TestSpec
+ .forStrategy(
+ "Composite type strategy with STRUCTURED type",
+ sequence(InputTypeStrategies.COMPOSITE)
+ )
+ .calledWithArgumentTypes(DataTypes.of(SimpleStructuredType.class).notNull())
+ .expectSignature("f(<COMPOSITE>)")
+ .expectArgumentTypes(DataTypes.of(SimpleStructuredType.class).notNull())
  );
  }
+
+ /**
+ * Simple pojo that should be converted to a Structured type.
+ */
+ public static class SimpleStructuredType {
+ public long f0;
+ }
 }

flink-table/flink-table-common/src/test/java/org/apache/flink/table/types/inference/InputTypeStrategiesTestBase.java

@@ -21,6 +21,7 @@
 import org.apache.flink.table.api.DataTypes;
 import org.apache.flink.table.api.ValidationException;
 import org.apache.flink.table.functions.FunctionKind;
+import org.apache.flink.table.types.AbstractDataType;
 import org.apache.flink.table.types.DataType;
 import org.apache.flink.table.types.inference.utils.CallContextMock;
 import org.apache.flink.table.types.inference.utils.FunctionDefinitionMock;
@@ -198,8 +199,8 @@ TestSpec surroundingStrategy(InputTypeStrategy surroundingStrategy) {
  return this;
  }
 
- TestSpec calledWithArgumentTypes(DataType... dataTypes) {
- this.actualArgumentTypes.add(Arrays.asList(dataTypes));
+ TestSpec calledWithArgumentTypes(AbstractDataType<?>... dataTypes) {
+ this.actualArgumentTypes.add(resolveDataTypes(dataTypes));
  return this;
  }
 
@@ -219,8 +220,8 @@ TestSpec expectSignature(String signature) {
  return this;
  }
 
- TestSpec expectArgumentTypes(DataType... dataTypes) {
- this.expectedArgumentTypes = Arrays.asList(dataTypes);
+ TestSpec expectArgumentTypes(AbstractDataType<?>... dataTypes) {
+ this.expectedArgumentTypes = resolveDataTypes(dataTypes);
  return this;
  }
 
@@ -229,6 +230,13 @@ TestSpec expectErrorMessage(String expectedErrorMessage) {
  return this;
  }
 
+ private List<DataType> resolveDataTypes(AbstractDataType<?>[] dataTypes) {
+ final DataTypeFactoryMock factoryMock = new DataTypeFactoryMock();
+ return Arrays.stream(dataTypes)
+ .map(factoryMock::createDataType)
+ .collect(Collectors.toList());
+ }
+
  @Override
  public String toString() {
  return description != null ? description : strategy.getClass().getSimpleName();