org.apache.calcite.rel.type.RelDataTypeFieldImpl Java Examples

@Test
public void testNumericValuesRelRowTypeAdjustment() {
  final int INListLength = 20;

  // Build RowType & Tuples
  RelDataTypeField relDataType = new RelDataTypeFieldImpl("ROW_VALUE", 0, new BasicSqlType(RelDataTypeSystemImpl.REL_DATA_TYPE_SYSTEM, SqlTypeName.ANY));
  RelDataType rowType = new RelRecordType(StructKind.FULLY_QUALIFIED, Arrays.asList(relDataType));
  ImmutableList.Builder<ImmutableList<RexLiteral>> tuples = new ImmutableList.Builder<>();
  for (int i = 0; i < INListLength; i++) {
    tuples.add(new ImmutableList.Builder<RexLiteral>().add(new RexBuilder(typeFactory).makeExactLiteral(new BigDecimal(i))).build());
  }

  // Check original types.
  assertEquals(1, rowType.getFieldCount());
  assertEquals(SqlTypeName.ANY, rowType.getFieldList().get(0).getType().getSqlTypeName());

  // Construct ValuesRel
  final ValuesRel valuesRel = new ValuesRel(cluster, rowType, tuples.build(), traits);

  // Check the adjusted types.
  RelDataType adjustedRowType = valuesRel.getRowType();
  assertEquals(1, adjustedRowType.getFieldCount());
  assertEquals(SqlTypeName.INTEGER, adjustedRowType.getFieldList().get(0).getType().getSqlTypeName());
}
 

/** Test case for
 * <a href="https://issues.apache.org/jira/browse/CALCITE-2464">[CALCITE-2464]
 * Allow to set nullability for columns of structured types</a>. */
@Test void createStructTypeWithNullability() {
  SqlTypeFixture f = new SqlTypeFixture();
  RelDataTypeFactory typeFactory = f.typeFactory;
  List<RelDataTypeField> fields = new ArrayList<>();
  RelDataTypeField field0 = new RelDataTypeFieldImpl(
          "i", 0, typeFactory.createSqlType(SqlTypeName.INTEGER));
  RelDataTypeField field1 = new RelDataTypeFieldImpl(
          "s", 1, typeFactory.createSqlType(SqlTypeName.VARCHAR));
  fields.add(field0);
  fields.add(field1);
  final RelDataType recordType = new RelRecordType(fields); // nullable false by default
  final RelDataType copyRecordType = typeFactory.createTypeWithNullability(recordType, true);
  assertFalse(recordType.isNullable());
  assertTrue(copyRecordType.isNullable());
}
 

/**
 * Gets a list of extended columns with field indices to the underlying table.
 */
public static List<RelDataTypeField> getExtendedColumns(
    SqlValidator validator, SqlValidatorTable table, SqlNodeList extendedColumns) {
  final ImmutableList.Builder<RelDataTypeField> extendedFields =
      ImmutableList.builder();
  final ExtensibleTable extTable = table.unwrap(ExtensibleTable.class);
  int extendedFieldOffset =
      extTable == null
          ? table.getRowType().getFieldCount()
          : extTable.getExtendedColumnOffset();
  for (final Pair<SqlIdentifier, SqlDataTypeSpec> pair : pairs(extendedColumns)) {
    final SqlIdentifier identifier = pair.left;
    final SqlDataTypeSpec type = pair.right;
    extendedFields.add(
        new RelDataTypeFieldImpl(identifier.toString(),
            extendedFieldOffset++,
            type.deriveType(validator)));
  }
  return extendedFields.build();
}
 

private void recursiveCreatePdxType(PdxInstance pdxInstance,
    List<RelDataTypeField> list, String fieldNamePrefix) {

  for (String fieldName : pdxInstance.getFieldNames()) {
    Object field = pdxInstance.getField(fieldName);
    final Type fieldType = field.getClass();
    if (fieldType instanceof PdxInstance) {
      recursiveCreatePdxType(
          (PdxInstance) field, list, fieldNamePrefix + fieldName + ".");
    } else {
      list.add(
          new RelDataTypeFieldImpl(
              fieldNamePrefix + fieldName,
              list.size(),
              createType(fieldType)));
    }
  }
}
 

/**
 * See <a href="http://stackoverflow.com/questions/16966629/what-is-the-difference-between-getfields-and-getdeclaredfields-in-java-reflectio">
 *   the difference between fields and declared fields</a>.
 */
@Override public RelDataType createStructType(Class type) {

  final List<RelDataTypeField> list = new ArrayList<>();
  for (Field field : type.getDeclaredFields()) {
    if (!Modifier.isStatic(field.getModifiers())) {
      // FIXME: watch out for recursion
      final Type fieldType = field.getType();
      list.add(
          new RelDataTypeFieldImpl(
              field.getName(),
              list.size(),
              createType(fieldType)));
    }
  }
  return canonize(new JavaRecordType(list, type));
}
 

private synchronized void initialize(JavaTypeFactory typeFactory) {
	if (this.initialized) {
		return;
	}

	this.methodsForFields = new ArrayList<>();
	List<RelDataTypeField> fields = new ArrayList<>();

	for (Method method : this.referenceInterface.getMethods()) {
		if (method.getParameterCount() == 0) {
			String fieldName = computeFieldName(method.getName());
			if (fieldName != null) {
				this.methodsForFields.add(extractorForMethod(method));
				Class<?> retType = method.getReturnType();
				if (retType.isEnum()) {
					retType = String.class;
				}
				fields.add(new RelDataTypeFieldImpl(fieldName.toUpperCase(), fields.size(), typeFactory.createType(retType)));
			}
		}
	}

	this.rowType = new MyDataType(fields, referenceInterface);
	this.initialized = true;
}
 

private RelDataType getViewRowType() {

    List<RelDataTypeField> relDataTypeFields =
        ImmutableList.<RelDataTypeField>of(
            new RelDataTypeFieldImpl("id", 1, getIntegerJavaType()),
            new RelDataTypeFieldImpl("name", 2, getStringJavaType()),
            new RelDataTypeFieldImpl("description", 3, getStringJavaType()),
            new RelDataTypeFieldImpl("cost", 4, getIntegerJavaType()),
            new RelDataTypeFieldImpl("query", 5, getStringJavaType()),
            new RelDataTypeFieldImpl("destination_id", 6, getIntegerJavaType()),
            new RelDataTypeFieldImpl("schema_name", 7, getStringJavaType()),
            new RelDataTypeFieldImpl("table_name", 8, getStringJavaType()),
            new RelDataTypeFieldImpl("ds_set_id", 9, getIntegerJavaType()));

    return new RelRecordType(relDataTypeFields);
  }
 

private RelDataType getDataSourceRowType() throws SQLException {

    List<RelDataTypeField> relDataTypeFields =
        ImmutableList.<RelDataTypeField>of(
            new RelDataTypeFieldImpl("id", 1, getIntegerJavaType()),
            new RelDataTypeFieldImpl("type", 2, getStringJavaType()),
            new RelDataTypeFieldImpl("url", 3, getStringJavaType()),
            new RelDataTypeFieldImpl("name", 4, getStringJavaType()),
            new RelDataTypeFieldImpl("ds_set_id", 5, getIntegerJavaType()),
            new RelDataTypeFieldImpl("datasource_type", 6, getStringJavaType()),
            new RelDataTypeFieldImpl("auth_token", 7, getStringJavaType()),
            new RelDataTypeFieldImpl("dbtap_id", 8, getIntegerJavaType()),
            new RelDataTypeFieldImpl("username", 9, getStringJavaType()),
            new RelDataTypeFieldImpl("password", 10, getStringJavaType()));

    return new RelRecordType(relDataTypeFields);
  }
 

@Test
public void testCharValuesRelRowTypeAdjustment() {
  final int INListLength = 20;

  // Build RowType & Tuples
  RelDataTypeField relDataType = new RelDataTypeFieldImpl("ROW_VALUE", 0, new BasicSqlType(RelDataTypeSystemImpl.REL_DATA_TYPE_SYSTEM, SqlTypeName.ANY));
  RelDataType rowType = new RelRecordType(StructKind.FULLY_QUALIFIED, Arrays.asList(relDataType));
  ImmutableList.Builder<ImmutableList<RexLiteral>> tuples = new ImmutableList.Builder<>();
  for (int i = 0; i < INListLength; ++i) {
    tuples.add(new ImmutableList.Builder<RexLiteral>().add(new RexBuilder(typeFactory).makeLiteral(charLiteralBuilder(i))).build());
  }

  // Check original types.
  assertEquals(1, rowType.getFieldCount());
  assertEquals(SqlTypeName.ANY, rowType.getFieldList().get(0).getType().getSqlTypeName());

  // Construct ValuesRel
  final ValuesRel valuesRel = new ValuesRel(cluster, rowType, tuples.build(), traits);

  // Check the adjusted types.
  RelDataType adjustedRowType = valuesRel.getRowType();
  assertEquals(1, adjustedRowType.getFieldCount());
  assertEquals(SqlTypeName.CHAR, adjustedRowType.getFieldList().get(0).getType().getSqlTypeName());
  assertEquals(INListLength - 1, adjustedRowType.getFieldList().get(0).getType().getPrecision());
}
 

/**
 * Decode all fields that are global dictionary encoded.
 * It adds a dictionary lookup operations on top of current operation.
 * @return dictionary lookup operator who input is current prel.
 */
DictionaryLookupPrel decodeAllFields() {
  final List<RelDataTypeField> newFieldList = Lists.newArrayList();
  final List<RelDataTypeField> oldFieldList = prel.getRowType().getFieldList();
  final List<GlobalDictionaryFieldInfo> fieldsToDecode = Lists.newArrayList();
  for (int i = 0; i < fields.length; ++i) {
    if (fields[i] != null) {
      final RelDataTypeField oldField = oldFieldList.get(i);
      fieldsToDecode.add(fields[i]);
      newFieldList.add(new RelDataTypeFieldImpl(oldField.getName(), oldField.getIndex(), fields[i].getRelDataTypeField().getType()));
    } else {
      newFieldList.add(oldFieldList.get(i));
    }
  }

  return new DictionaryLookupPrel(
    prel.getCluster(),
    prel.getTraitSet(),
    prel,
    toRowDataType(newFieldList, prel.getCluster().getTypeFactory()),
    fieldsToDecode);
}
 

/**
 * Gets a list of extended columns with field indices to the underlying table.
 */
public static List<RelDataTypeField> getExtendedColumns(
    RelDataTypeFactory typeFactory, SqlValidatorTable table, SqlNodeList extendedColumns) {
  final ImmutableList.Builder<RelDataTypeField> extendedFields =
      ImmutableList.builder();
  final ExtensibleTable extTable = table.unwrap(ExtensibleTable.class);
  int extendedFieldOffset =
      extTable == null
          ? table.getRowType().getFieldCount()
          : extTable.getExtendedColumnOffset();
  for (final Pair<SqlIdentifier, SqlDataTypeSpec> pair : pairs(extendedColumns)) {
    final SqlIdentifier identifier = pair.left;
    final SqlDataTypeSpec type = pair.right;
    extendedFields.add(
        new RelDataTypeFieldImpl(identifier.toString(),
            extendedFieldOffset++,
            type.deriveType(typeFactory)));
  }
  return extendedFields.build();
}
 

public RelDataTypeField getField(RelDataTypeFactory typeFactory, String fieldName) {

    /* First check if this field name exists in our field list */
    for (RelDataTypeField f : fields) {
      if (fieldName.equalsIgnoreCase(f.getName())) {
        return f;
      }
    }

    /* This field does not exist in our field list add it */
    final SqlTypeName typeName = DynamicRecordType.isDynamicStarColName(fieldName)
        ? SqlTypeName.DYNAMIC_STAR : SqlTypeName.ANY;

    // This field does not exist in our field list add it
    RelDataTypeField newField = new RelDataTypeFieldImpl(
        fieldName,
        fields.size(),
        typeFactory.createTypeWithNullability(typeFactory.createSqlType(typeName), true));

    /* Add the name to our list of field names */
    fields.add(newField);

    return newField;
  }
 

protected RelDataType convertRowType(RelDataType origRowType, RelDataTypeFactory typeFactory) {
  if ( getRowKeyIndex(origRowType, origScan)>=0 ) { // row key already present
    return origRowType;
  }
  List<RelDataTypeField> fields = new ArrayList<>();

  fields.addAll(origRowType.getFieldList());
  fields.add(new RelDataTypeFieldImpl(
      ((DbGroupScan)IndexPlanUtils.getGroupScan(origScan)).getRowKeyName(), fields.size(),
          typeFactory.createSqlType(SqlTypeName.ANY)));
  return new RelRecordType(fields);
}
 

@Override public Void visitCall(RexCall call) {
  if (call.getOperator() == RexBuilder.GET_OPERATOR) {
    RexLiteral literal = (RexLiteral) call.getOperands().get(1);
    extraFields.add(
        new RelDataTypeFieldImpl(
            (String) literal.getValue2(),
            -1,
            call.getType()));
  }
  return super.visitCall(call);
}
 

@Override
public Void visitCall(RexCall call) {
    if (call.getOperator() == RexBuilder.GET_OPERATOR) {
        RexLiteral literal = (RexLiteral) call.getOperands().get(1);
        extraFields.add(new RelDataTypeFieldImpl((String) literal.getValue2(), -1, call.getType()));
    }
    return super.visitCall(call);
}
 

public RelDataType createStructType(Class type) {
  final List<RelDataTypeField> list = new ArrayList<>();
  for (Field field : type.getFields()) {
    if (!Modifier.isStatic(field.getModifiers())) {
      // FIXME: watch out for recursion
      final Type fieldType = fieldType(field);
      list.add(
          new RelDataTypeFieldImpl(
              field.getName(),
              list.size(),
              createType(fieldType)));
    }
  }
  return canonize(new JavaRecordType(list, type));
}
 

public RelDataType createPdxType(PdxInstance pdxInstance) {
  final List<RelDataTypeField> list = new ArrayList<>();
  for (String fieldName : pdxInstance.getFieldNames()) {
    Object field = pdxInstance.getField(fieldName);

    Type fieldType;

    if (field == null) {
      fieldType = String.class;
    } else if (field instanceof PdxInstance) {
      // Map Nested PDX structures as String. This relates with
      // GeodeUtils.convert case when clazz is Null.
      fieldType = Map.class;
      // RelDataType boza = createPdxType((PdxInstance) field);
    } else {
      fieldType = field.getClass();
    }

    list.add(
        new RelDataTypeFieldImpl(
            fieldName,
            list.size(),
            createType(fieldType)));
  }

  return canonize(new RelRecordType(list));
}
 

public static StructuredRelDataType create(FlinkTypeFactory factory, StructuredType structuredType) {
	final List<RelDataTypeField> fields = new ArrayList<>();
	for (int i = 0; i < structuredType.getAttributes().size(); i++) {
		final StructuredAttribute attribute = structuredType.getAttributes().get(i);
		final RelDataTypeField field = new RelDataTypeFieldImpl(
			attribute.getName(),
			i,
			factory.createFieldTypeFromLogicalType(attribute.getType()));
		fields.add(field);
	}
	return new StructuredRelDataType(structuredType, fields);
}
 

public static QueryOperationCatalogViewTable createCalciteTable(
		QueryOperationCatalogView catalogView,
		TableSchema resolvedSchema) {
	return new QueryOperationCatalogViewTable(catalogView, typeFactory -> {
		final FlinkTypeFactory flinkTypeFactory = (FlinkTypeFactory) typeFactory;
		final RelDataType relType = flinkTypeFactory.buildLogicalRowType(resolvedSchema);
		Boolean[] nullables = resolvedSchema
			.getTableColumns()
			.stream()
			.map(c -> c.getType().getLogicalType().isNullable())
			.toArray(Boolean[]::new);
		final List<RelDataTypeField> fields = relType
			.getFieldList()
			.stream()
			.map(f -> {
				boolean nullable = nullables[f.getIndex()];
				if (nullable != f.getType().isNullable()
					&& !FlinkTypeFactory.isTimeIndicatorType(f.getType())) {
					return new RelDataTypeFieldImpl(
						f.getName(),
						f.getIndex(),
						flinkTypeFactory.createTypeWithNullability(f.getType(), nullable));
				} else {
					return f;
				}
			})
			.collect(Collectors.toList());
		return flinkTypeFactory.createStructType(fields);
	});
}
 

@Override
public void implementRewrite(RewriteImplementor implementor) {
    implementor.visitChild(this, this.left);
    implementor.visitChild(this, this.right);

    this.rowType = this.deriveRowType();
    if (this.isTopJoin) {

        // add dynamic field
        Map<TblColRef, RelDataType> dynFields = this.context.dynamicFields;
        if (!dynFields.isEmpty()) {
            List<TblColRef> newCols = Lists.newArrayList(this.columnRowType.getAllColumns());
            List<RelDataTypeField> newFieldList = Lists.newArrayList();
            int paramIndex = this.rowType.getFieldList().size();
            for (TblColRef fieldCol : dynFields.keySet()) {
                RelDataType fieldType = dynFields.get(fieldCol);

                RelDataTypeField newField = new RelDataTypeFieldImpl(fieldCol.getName(), paramIndex++, fieldType);
                newFieldList.add(newField);

                newCols.add(fieldCol);
            }

            // rebuild row type
            RelDataTypeFactory.FieldInfoBuilder fieldInfo = getCluster().getTypeFactory().builder();
            fieldInfo.addAll(this.rowType.getFieldList());
            fieldInfo.addAll(newFieldList);
            this.rowType = getCluster().getTypeFactory().createStructType(fieldInfo);

            this.columnRowType = new ColumnRowType(newCols);
        }
    }
}
 

public static DrillProjectRel convert(Project project, ConversionContext context) throws InvalidRelException{
  RelNode input = context.toRel(project.getInput());
  List<RelDataTypeField> fields = Lists.newArrayList();
  List<RexNode> exps = Lists.newArrayList();
  for(NamedExpression expr : project.getSelections()){
    fields.add(new RelDataTypeFieldImpl(expr.getRef().getRootSegment().getPath(), fields.size(), context.getTypeFactory().createSqlType(SqlTypeName.ANY) ));
    exps.add(context.toRex(expr.getExpr()));
  }
  return new DrillProjectRel(context.getCluster(), context.getLogicalTraits(), input, exps, new RelRecordType(fields));
}
 

private List<RelDataTypeField> getRelFields(List<SqlSchema.SqlField> fields) {
  List<RelDataTypeField> relFields = new ArrayList<>();

  for (SqlSchema.SqlField field : fields) {
    String fieldName = field.getFieldName();
    int fieldPos = field.getPosition();
    RelDataType dataType = getRelDataType(field.getFieldSchema());
    relFields.add(new RelDataTypeFieldImpl(fieldName, fieldPos, dataType));
  }

  return relFields;
}
 

/**
 * For each aggregate call creates field based on its name with bigint type.
 * Constructs record type for created fields.
 *
 * @param aggregateRel aggregate relation expression
 * @param fieldNames field names
 * @return record type
 */
public static RelDataType constructDataType(Aggregate aggregateRel, Collection<String> fieldNames) {
  List<RelDataTypeField> fields = new ArrayList<>();
  int fieldIndex = 0;
  for (String name : fieldNames) {
    RelDataTypeField field = new RelDataTypeFieldImpl(
        name,
        fieldIndex++,
        aggregateRel.getCluster().getTypeFactory().createSqlType(SqlTypeName.BIGINT));
    fields.add(field);
  }
  return new RelRecordType(fields);
}
 

@Override
public RelDataType createStructType(Class<?> type) {
    final List<RelDataTypeField> list = new ArrayList<>();
    for (Field field : type.getFields()) {
        if (!Modifier.isStatic(field.getModifiers())) {
            // FIXME: watch out for recursion
            final Type fieldType = fieldType(field);
            list.add(new RelDataTypeFieldImpl(field.getName(), list.size(), createType(fieldType)));
        }
    }
    return canonize(new JavaRecordType(list, type));
}
 

/**
 * Decode specified fields that are global dictionary encoded.
 * It adds a dictionary lookup operations on top of current operation.
 * @param decodeFieldIndices field indices to decode
 * @return If at least one field is being decoded then return dictionary lookup operator who input is current prel
 *         otherwise return current prel (this)
 */
PrelWithDictionaryInfo decodeFields(Collection<Integer> decodeFieldIndices) {
  if (decodeFieldIndices.isEmpty()) {
    return this;
  }
  final List<RelDataTypeField> oldFieldList = prel.getRowType().getFieldList();
  final List<RelDataTypeField> newFieldList = Lists.newArrayList();
  final List<GlobalDictionaryFieldInfo> fieldsToDecode = Lists.newArrayList();
  final GlobalDictionaryFieldInfo[] newGlobalDictionaryFieldInfos = new GlobalDictionaryFieldInfo[fields.length];

  boolean decoded = false;
  for (int i = 0; i < fields.length; ++i) {
    if (fields[i] != null && decodeFieldIndices.contains(i)) {
      final RelDataTypeField oldField = oldFieldList.get(i);
      fieldsToDecode.add(fields[i]);
      newFieldList.add(new RelDataTypeFieldImpl(oldField.getName(), oldField.getIndex(), fields[i].getRelDataTypeField().getType()));
      newGlobalDictionaryFieldInfos[i] = null; // don't need to decode anymore
      decoded = true;
    } else {
      // copy through
      newGlobalDictionaryFieldInfos[i] = fields[i];
      newFieldList.add(oldFieldList.get(i));
    }
  }
  // None of the fields have global dictionary or they are already decoded, return current prel
  if (!decoded) {
    return this;
  }
  return new PrelWithDictionaryInfo(new DictionaryLookupPrel(
    prel.getCluster(),
    prel.getTraitSet(),
    prel,
    toRowDataType(newFieldList, prel.getCluster().getTypeFactory()),
    fieldsToDecode), newGlobalDictionaryFieldInfos);
}
 

public RelDataType[] getRowTypeFromEs(RelDataTypeFactory typeFactory) throws IOException {
    RelDataType[] relDataTypes = new RelDataType[2];
    RelDataTypeFactory.FieldInfoBuilder builder = typeFactory.builder();
    RelDataTypeFactory.FieldInfoBuilder builder1 = typeFactory.builder();
    GetMappingsResponse getMappingsResponse = client.admin().indices().prepareGetMappings(esIndexName).get();
    MappingMetaData typeMapping = getMappingsResponse.getMappings().get(esIndexName).get(esTypeName);
    //{"person":{"properties":{"age":{"type":"integer"},"name":{"type":"text","fields":{"raw":{"type":"keyword"}}}}}}
    String json = typeMapping.source().string();
    Map<String, Object> map = new ObjectMapper().readValue(json, new TypeReference<Map<String, Object>>() {
    });
    Map<String, Object> properties = ((Map<String, Map<String, Object>>) map.get(esTypeName)).get("properties");
    int index = 0;
    int index1 = 0;
    //(base-field-name, field-map)
    Stack<Pair<String, Map<String, Object>>> mapStack = new Stack<>();
    mapStack.push(Pair.of(null, properties));
    while (!mapStack.isEmpty()) {
        Pair<String, Map<String, Object>> pair = mapStack.pop();
        String baseFieldName = pair.left;
        for (Map.Entry<String, Object> entry : pair.right.entrySet()) {
            String name = entry.getKey().toUpperCase();
            if (baseFieldName != null) name = baseFieldName + "." + name;

            Map<String, Object> fieldMap = (Map<String, Object>) entry.getValue();
            String type = fieldMap.get("type") != null ? fieldMap.get("type").toString() : null;
            if (type == null)
                throw new IllegalStateException(String.format("type of elasticsearch field '%s' is null", name));
            builder.add(new RelDataTypeFieldImpl(name, index++, typeFactory.createJavaType(ElasticSearchFieldType.of(type).getClazz())));

            if(!name.endsWith(".KEYWORD")){
                builder1.add(new RelDataTypeFieldImpl(name, index1++, typeFactory.createJavaType(ElasticSearchFieldType.of(type).getClazz())));
            }
            //multi-field, that means containing 'fields' attribute
            Map<String, Object> moreFields = fieldMap.get("fields") != null ? (Map<String, Object>) fieldMap.get("fields") : null;
            if (moreFields != null) mapStack.push(Pair.of(name, moreFields));
        }
    }
    relDataTypes[0] = builder.build();
    relDataTypes[1] = builder1.build();
    return relDataTypes;
}
 

@Test
public void testUpdateSql() throws Exception {
  final String sql = "select foo, bar as b from tbl";
  SqlParser parser = SqlParser.create(sql, new ParserConfig(Quoting.DOUBLE_QUOTE, 128));
  final SqlNode sqlNode = parser.parseStmt();


  final JavaTypeFactory typeFactory = JavaTypeFactoryImpl.INSTANCE;
  final RelDataType rowType = new RelRecordType(Arrays.<RelDataTypeField>asList(
      new RelDataTypeFieldImpl("foo", 0, typeFactory.createSqlType(SqlTypeName.INTEGER)),
      new RelDataTypeFieldImpl("b", 0, typeFactory.createSqlType(SqlTypeName.INTEGER))
  ));

  final QueryMetadata metadata = Mockito.mock(QueryMetadata.class);
  Mockito.when(metadata.getSqlNode()).thenReturn(Optional.of(sqlNode));
  Mockito.when(metadata.getRowType()).thenReturn(rowType);
  Mockito.when(metadata.getQuerySql()).thenReturn(sql);
  Mockito.when(metadata.getReferredTables()).thenReturn(Optional.absent());

  TransformActor actor = new TransformActor(state, false, "test_user", null) {
    @Override
    protected com.dremio.service.jobs.metadata.proto.QueryMetadata getMetadata(SqlQuery query) {
      return com.dremio.service.jobs.metadata.proto.QueryMetadata.newBuilder()
        .setState(QuerySemantics.extract(metadata).get())
        .build();
    }

    @Override
    protected boolean hasMetadata() {
      return true;
    }

    @Override
    protected com.dremio.service.jobs.metadata.proto.QueryMetadata getMetadata() {
      throw new IllegalStateException("not implemented");
    }

    @Override
    protected Optional<BatchSchema> getBatchSchema() {
      throw new IllegalStateException("not implemented");
    }

    @Override
    protected Optional<List<ParentDatasetInfo>> getParents() {
      throw new IllegalStateException("not implemented");
    }

    @Override
    protected Optional<List<FieldOrigin>> getFieldOrigins() {
      throw new IllegalStateException("not implemented");
    }

    @Override
    protected Optional<List<ParentDataset>> getGrandParents() {
      throw new IllegalStateException("not implemented");
    }
  };

  TransformResult result = new TransformUpdateSQL(sql).accept(actor);
  VirtualDatasetState newState = result.getNewState();

  assertEquals(2, newState.getColumnsList().size());
  assertEquals("foo", newState.getColumnsList().get(0).getName());
  assertEquals(new ExpColumnReference("foo").wrap(), newState.getColumnsList().get(0).getValue());
  assertEquals("b", newState.getColumnsList().get(1).getName());
  assertEquals(new ExpColumnReference("bar").wrap(), newState.getColumnsList().get(1).getValue());
}
 

/**
 * if a field in rowType serves only the to-be-replaced column(s), we should replace it with new name "$1",
 * otherwise we should keep this dataTypeField and add a new one for "$1"
 * @param origScan  the original scan whose rowtype is to be rewritten
 * @param indexContext the index plan context
 * @param functionInfo functional index information that may impact rewrite
 * @return
 */
public static RelDataType rewriteFunctionalRowType(RelNode origScan, IndexCallContext indexContext,
    FunctionalIndexInfo functionInfo, Collection<SchemaPath> addedPaths) {
  RelDataType origRowType = origScan.getRowType();
  if (!functionInfo.hasFunctional()) {
    return origRowType;
  }

  List<RelDataTypeField> fields = Lists.newArrayList();

  Set<String> leftOutFieldNames  = Sets.newHashSet();
  if (indexContext.getLeftOutPathsInFunctions() != null) {
    for (LogicalExpression expr : indexContext.getLeftOutPathsInFunctions()) {
      leftOutFieldNames.add(((SchemaPath) expr).getRootSegmentPath());
    }
  }

  Set<String> fieldInFunctions  = Sets.newHashSet();
  for (SchemaPath path: functionInfo.allPathsInFunction()) {
    fieldInFunctions.add(path.getRootSegmentPath());
  }

  RelDataTypeFactory typeFactory = origScan.getCluster().getTypeFactory();

  for ( RelDataTypeField field: origRowType.getFieldList()) {
    final String fieldName = field.getName();
    if (fieldInFunctions.contains(fieldName)) {
      if (!leftOutFieldNames.contains(fieldName)) {
        continue;
      }
    }

    fields.add(new RelDataTypeFieldImpl(
        SchemaPath.parseFromString(fieldName).getRootSegmentPath(), fields.size(),
        typeFactory.createSqlType(SqlTypeName.ANY)));
  }

  final Collection<SchemaPath> toAddToRowType = (addedPaths == null)? functionInfo.allNewSchemaPaths() : addedPaths;

  for (SchemaPath dollarPath : toAddToRowType) {
    fields.add(
        new RelDataTypeFieldImpl(dollarPath.getRootSegmentPath(), fields.size(),
            origScan.getCluster().getTypeFactory().createSqlType(SqlTypeName.ANY)));
  }

  return new RelRecordType(fields);
}
 

private RelDataType getCountRowType(RelDataTypeFactory typeFactory) {
  List<RelDataTypeField> fields = Lists.newArrayList();
  fields.add(new RelDataTypeFieldImpl("count", 0, typeFactory.createSqlType(SqlTypeName.BIGINT)));
  return new RelRecordType(fields);
}
 

@Override
public Boolean visitCall(RexCall call) {
  if (call.getOperator().getName().equalsIgnoreCase("contains")) {
    if (!checkOrigin) {
      return true;
    }
    // Check origin
    final Set<RelColumnOrigin> origins;
    if (index >= 0) {
      origins = mq.getColumnOrigins(node, index);
    } else {
      List<RelDataTypeField> fields = ImmutableList.<RelDataTypeField>of(new RelDataTypeFieldImpl("ContainsTemp", 0, call.getType()));
      Project temporary =
          new LogicalProject(node.getCluster(), node.getTraitSet().plus(Convention.NONE), node.getInput(0), ImmutableList.of(call), new RelRecordType(fields));
      origins = mq.getColumnOrigins(temporary, 0);
    }

    boolean supportContains = true;
    if (origins == null) {
      supportContains = false;
    } else {
      for (RelColumnOrigin column : origins) {
        if (column.getOriginTable() == null) {
          supportContains = false;
        } else {
          NamespaceTable namespaceTable2 = column.getOriginTable().unwrap(NamespaceTable.class);
          if (namespaceTable2 != null) {
            if(!namespaceTable2.getStoragePluginId().getCapabilities().getCapability(SourceCapabilities.SUPPORTS_CONTAINS)){
              supportContains = false;
            }
          } else {
            supportContains = false;
          }
        }
      }
    }
    if (!supportContains) {
      throw UserException.unsupportedError().message("Contains operator is not supported for the table, %s", call).build(logger);
    }
    return true;
  }

  for (RexNode op : call.getOperands()) {
    Boolean opResult = op.accept(this);
    if (opResult != null && opResult.booleanValue()) {
      return true;
    }
  }
  return false;
}