org.apache.spark.sql.catalyst.expressions.Expression Java Examples

public static UnsafeProjection projection(Schema expectedSchema, Schema actualSchema) {
  StructType struct = SparkSchemaUtil.convert(actualSchema);

  List<AttributeReference> refs = JavaConverters.seqAsJavaListConverter(struct.toAttributes()).asJava();
  List<Attribute> attrs = Lists.newArrayListWithExpectedSize(struct.fields().length);
  List<Expression> exprs = Lists.newArrayListWithExpectedSize(struct.fields().length);

  for (AttributeReference ref : refs) {
    attrs.add(ref.toAttribute());
  }

  for (Types.NestedField field : expectedSchema.columns()) {
    int indexInIterSchema = struct.fieldIndex(field.name());
    exprs.add(refs.get(indexInIterSchema));
  }

  return UnsafeProjection.create(
      JavaConverters.asScalaBufferConverter(exprs).asScala().toSeq(),
      JavaConverters.asScalaBufferConverter(attrs).asScala().toSeq());
}
 

static
private Expression translateInternal(String sqlStatement){
	StructType schema = new StructType()
		.add("flag", DataTypes.BooleanType)
		.add("x1", DataTypes.DoubleType)
		.add("x2", DataTypes.DoubleType)
		.add("status", DataTypes.IntegerType);

	LogicalPlan logicalPlan = DatasetUtil.createAnalyzedLogicalPlan(ExpressionTranslatorTest.sparkSession, schema, sqlStatement);

	List<Expression> expressions = JavaConversions.seqAsJavaList(logicalPlan.expressions());
	if(expressions.size() != 1){
		throw new IllegalArgumentException();
	}

	return expressions.get(0);
}
 

private static com.netflix.iceberg.expressions.Expression convertIn(Expression expr,
                                                                    Collection<Object> values) {
  if (expr instanceof Attribute) {
    Attribute attr = (Attribute) expr;
    com.netflix.iceberg.expressions.Expression converted = alwaysFalse();
    for (Object item : values) {
      converted = or(converted, equal(attr.name(), item));
    }
    return converted;
  }

  return null;
}
 

private static List<Object> convertLiterals(List<Expression> values) {
  List<Object> converted = Lists.newArrayListWithExpectedSize(values.size());

  for (Expression value : values) {
    if (value instanceof Literal) {
      Literal lit = (Literal) value;
      converted.add(valueFromSpark(lit));
    } else {
      return null;
    }
  }

  return converted;
}
 

private static com.netflix.iceberg.expressions.Expression convert(Operation op,
                                                                  Expression left,
                                                                  Expression right) {
  Pair<Transform, String> attrPair = null;
  Operation leftOperation = null;
  Literal lit = null;

  if (right instanceof Literal) {
    lit = (Literal) right;
    attrPair = convertAttr(left);
    leftOperation = op;
  } else if (left instanceof Literal) {
    lit = (Literal) left;
    attrPair = convertAttr(right);
    leftOperation = op.flipLR();
  }

  if (attrPair != null) {
    switch (attrPair.first()) {
      case IDENTITY:
        return predicate(leftOperation, attrPair.second(), valueFromSpark(lit));
      case YEAR:
        return filter(leftOperation, attrPair.second(), (int) lit.value(),
            SparkExpressions::yearToTimestampMicros);
      case DAY:
        return filter(leftOperation, attrPair.second(), (int) lit.value(),
            SparkExpressions::dayToTimestampMicros);
      default:
    }
  }

  return null;
}
 

private static com.netflix.iceberg.expressions.Expression filter(
    Operation op, String name, int value, Function<Integer, Long> startTsMicros) {
  switch (op) {
    case LT:
      return predicate(Operation.LT, name, tsLiteral(startTsMicros.apply(value)));
    case LT_EQ:
      return predicate(Operation.LT, name, tsLiteral(startTsMicros.apply(value + 1)));
    case GT:
      return predicate(Operation.GT_EQ, name, tsLiteral(startTsMicros.apply(value + 1)));
    case GT_EQ:
      return predicate(Operation.GT_EQ, name, tsLiteral(startTsMicros.apply(value)));
    case EQ:
      return and(
          predicate(Operation.GT_EQ, name, tsLiteral(startTsMicros.apply(value))),
          predicate(Operation.LT, name, tsLiteral(startTsMicros.apply(value + 1)))
      );
    case NOT_EQ:
      return or(
          predicate(Operation.GT_EQ, name, tsLiteral(startTsMicros.apply(value + 1))),
          predicate(Operation.LT, name, tsLiteral(startTsMicros.apply(value)))
      );
    case IN:
    case NOT_IN:
    default:
      throw new IllegalArgumentException("Cannot convert operation to year filter: " + op);
  }
}
 

/**
 * Wrap a Beam coder into a Spark Encoder using Catalyst Expression Encoders (which uses java code
 * generation).
 */
public static <T> Encoder<T> fromBeamCoder(Coder<T> coder) {
  Class<? super T> clazz = coder.getEncodedTypeDescriptor().getRawType();
  ClassTag<T> classTag = ClassTag$.MODULE$.apply(clazz);
  List<Expression> serializers =
      Collections.singletonList(
          new EncodeUsingBeamCoder<>(new BoundReference(0, new ObjectType(clazz), true), coder));

  return new ExpressionEncoder<>(
      SchemaHelpers.binarySchema(),
      false,
      JavaConversions.collectionAsScalaIterable(serializers).toSeq(),
      new DecodeUsingBeamCoder<>(
          new Cast(new GetColumnByOrdinal(0, BinaryType), BinaryType), classTag, coder),
      classTag);
}
 

static
public org.dmg.pmml.Expression translate(Expression expression, boolean compact){
	org.dmg.pmml.Expression pmmlExpression = translateInternal(expression);

	if(compact){
		ExpressionCompactor expressionCompactor = new ExpressionCompactor();

		expressionCompactor.applyTo(pmmlExpression);
	}

	return pmmlExpression;
}
 

static
public void checkValue(Object expectedValue, String sqlExpression){
	Expression expression = translateInternal("SELECT (" + sqlExpression + ") FROM __THIS__");

	Object sparkValue = expression.eval(InternalRow.empty());

	if(expectedValue instanceof String){
		assertEquals(expectedValue, sparkValue.toString());
	} else

	if(expectedValue instanceof Integer){
		assertEquals(expectedValue, ((Number)sparkValue).intValue());
	} else

	if(expectedValue instanceof Float){
		assertEquals(expectedValue, ((Number)sparkValue).floatValue());
	} else

	if(expectedValue instanceof Double){
		assertEquals(expectedValue, ((Number)sparkValue).doubleValue());
	} else

	{
		assertEquals(expectedValue, sparkValue);
	}

	org.dmg.pmml.Expression pmmlExpression = ExpressionTranslator.translate(expression);

	EvaluationContext context = new VirtualEvaluationContext();
	context.declareAll(Collections.emptyMap());

	FieldValue value = ExpressionUtil.evaluate(pmmlExpression, context);

	Object pmmlValue = FieldValueUtil.getValue(value);
	assertEquals(expectedValue, pmmlValue);
}
 

static
private String formatMessage(Expression expression){

	if(expression == null){
		return null;
	}

	return "Spark SQL function \'" + expression + "\' (class " + (expression.getClass()).getName() + ") is not supported";
}
 

@Override
public <T> Expression notIn(BoundReference<T> ref,
                            com.netflix.iceberg.expressions.Literal<T> lit) {
  throw new UnsupportedOperationException("Not implemented: notIn");
}
 

@Override
public Expression child() {
  return child;
}
 

private void pushFilters(DataSourceReader reader,
                         Expression... expressions) {
  Assert.assertTrue(reader instanceof SupportsPushDownCatalystFilters);
  SupportsPushDownCatalystFilters filterable = (SupportsPushDownCatalystFilters) reader;
  filterable.pushCatalystFilters(expressions);
}
 

static
public org.dmg.pmml.Expression translate(Expression expression){
	return translate(expression, true);
}
 

static
private Object evaluate(String sqlExpression){
	Expression expression = translateInternal("SELECT (" + sqlExpression + ") FROM __THIS__");

	return expression.eval(InternalRow.empty());
}
 

static
private org.dmg.pmml.Expression translate(String sqlExpression){
	Expression expression = translateInternal("SELECT (" + sqlExpression + ") FROM __THIS__");

	return ExpressionTranslator.translate(expression);
}
 

static
private void checkExpression(org.dmg.pmml.Expression expected, String string){
	org.dmg.pmml.Expression actual = translate(string);

	assertTrue(ReflectionUtil.equals(expected, actual));
}
 

public EncodeUsingBeamCoder(Expression child, Coder<T> coder) {
  this.child = child;
  this.coder = coder;
}
 

@Override
public Expression child() {
  return child;
}
 

public DecodeUsingBeamCoder(Expression child, ClassTag<T> classTag, Coder<T> coder) {
  this.child = child;
  this.classTag = classTag;
  this.coder = coder;
}
 

@Override
public <T> Expression in(BoundReference<T> ref,
                         com.netflix.iceberg.expressions.Literal<T> lit) {
  throw new UnsupportedOperationException("Not implemented: in");
}
 

@Override
public <T> Expression notEq(BoundReference<T> ref,
                            com.netflix.iceberg.expressions.Literal<T> lit) {
  return column(ref).notEqual(lit.value()).expr();
}
 

@Override
public <T> Expression eq(BoundReference<T> ref,
                         com.netflix.iceberg.expressions.Literal<T> lit) {
  return column(ref).equalTo(lit.value()).expr();
}
 

@Override
public <T> Expression gtEq(BoundReference<T> ref,
                           com.netflix.iceberg.expressions.Literal<T> lit) {
  return column(ref).geq(lit.value()).expr();
}
 

@Override
public <T> Expression gt(BoundReference<T> ref,
                         com.netflix.iceberg.expressions.Literal<T> lit) {
  return column(ref).gt(lit.value()).expr();
}
 

@Override
public <T> Expression ltEq(BoundReference<T> ref,
                           com.netflix.iceberg.expressions.Literal<T> lit) {
  return column(ref).leq(lit.value()).expr();
}
 

@Override
public <T> Expression lt(BoundReference<T> ref,
                         com.netflix.iceberg.expressions.Literal<T> lit) {
  return column(ref).lt(lit.value()).expr();
}
 

@Override
public <T> Expression notNull(BoundReference<T> ref) {
  return column(ref).isNotNull().expr();
}
 

@Override
public <T> Expression isNull(BoundReference<T> ref) {
  return column(ref).isNull().expr();
}
 

@Override
public Expression or(Expression left, Expression right) {
  return Or$.MODULE$.apply(left, right);
}