Skip to content

Commit a61ac61

Browse files
yuancuyuancu
committed
Simplify the overriding of reverse() for IP comparators
Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent 3f7c6a6 commit a61ac61

3 files changed

Lines changed: 52 additions & 106 deletions

File tree

‎core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java‎

Lines changed: 6 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
import java.lang.reflect.InvocationTargetException;
1313
import java.lang.reflect.Method;
1414
import java.util.List;
15-
import java.util.concurrent.atomic.AtomicReference;
1615
import java.util.function.Supplier;
1716
import org.apache.calcite.adapter.enumerable.NullPolicy;
1817
import org.apache.calcite.adapter.enumerable.RexImpTable;
@@ -110,19 +109,12 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
110109

111110
// IP comparing functions
112111
public static final SqlOperator NOT_EQUALS_IP =
113-
CompareIpFunction.notEquals()
114-
.withReverse(lookupOperator("NOT_EQUALS_IP"))
115-
.toUDF("NOT_EQUALS_IP");
116-
public static final SqlOperator EQUALS_IP =
117-
CompareIpFunction.equals().withReverse(lookupOperator("EQUALS_IP")).toUDF("EQUALS_IP");
118-
public static final SqlOperator GREATER_IP =
119-
CompareIpFunction.greater().withReverse(lookupOperator("LESS_IP")).toUDF("GREATER_IP");
120-
public static final SqlOperator GTE_IP =
121-
CompareIpFunction.greaterOrEquals().withReverse(lookupOperator("LTE_IP")).toUDF("GTE_IP");
122-
public static final SqlOperator LESS_IP =
123-
CompareIpFunction.less().withReverse(lookupOperator("GREATER_IP")).toUDF("LESS_IP");
124-
public static final SqlOperator LTE_IP =
125-
CompareIpFunction.lessOrEquals().withReverse(lookupOperator("GTE_IP")).toUDF("LTE_IP");
112+
CompareIpFunction.notEquals().toUDF("NOT_EQUALS_IP");
113+
public static final SqlOperator EQUALS_IP = CompareIpFunction.equals().toUDF("EQUALS_IP");
114+
public static final SqlOperator GREATER_IP = CompareIpFunction.greater().toUDF("GREATER_IP");
115+
public static final SqlOperator GTE_IP = CompareIpFunction.greaterOrEquals().toUDF("GTE_IP");
116+
public static final SqlOperator LESS_IP = CompareIpFunction.less().toUDF("LESS_IP");
117+
public static final SqlOperator LTE_IP = CompareIpFunction.lessOrEquals().toUDF("LTE_IP");
126118

127119
// Condition function
128120
public static final SqlOperator EARLIEST = new EarliestFunction().toUDF("EARLIEST");
@@ -391,25 +383,4 @@ private static Expression invokeCalciteImplementor(
391383
method.setAccessible(true);
392384
return (Expression) method.invoke(rexCallImplementor, translator, call, List.of(field));
393385
}
394-
395-
/**
396-
* Creates a lazy supplier for looking up a SQL operator by name.
397-
*
398-
* <p>This method enables operators to reference each other, even when they have circular
399-
* dependencies. Instead of looking up the operator immediately, it returns a supplier that will
400-
* perform the lookup only when needed.
401-
*
402-
* <p>For example, {@code LESS_IP} needs to reference {@code GREATER_IP} as its reverse, and vice
403-
* versa. Using this lazy approach, both can be defined before either is used.
404-
*
405-
* @param name The name of the operator to look up
406-
* @return A supplier that will look up the operator when called
407-
*/
408-
private static Supplier<SqlOperator> lookupOperator(String name) {
409-
return () -> {
410-
AtomicReference<SqlOperator> ref = new AtomicReference<>();
411-
instance().lookUpOperators(name, false, ref::set);
412-
return ref.get();
413-
};
414-
}
415386
}

‎core/src/main/java/org/opensearch/sql/expression/function/UserDefinedFunctionBuilder.java‎

Lines changed: 1 addition & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -6,17 +6,13 @@
66
package org.opensearch.sql.expression.function;
77

88
import java.util.Collections;
9-
import java.util.function.Supplier;
109
import org.apache.calcite.schema.ImplementableFunction;
1110
import org.apache.calcite.sql.SqlIdentifier;
1211
import org.apache.calcite.sql.SqlKind;
13-
import org.apache.calcite.sql.SqlOperator;
14-
import org.apache.calcite.sql.SqlSyntax;
1512
import org.apache.calcite.sql.parser.SqlParserPos;
1613
import org.apache.calcite.sql.type.InferTypes;
1714
import org.apache.calcite.sql.type.SqlReturnTypeInference;
1815
import org.apache.calcite.sql.validate.SqlUserDefinedFunction;
19-
import org.checkerframework.checker.nullness.qual.Nullable;
2016

2117
/**
2218
* The interface helps to construct a SqlUserDefinedFunction
@@ -36,18 +32,6 @@ public interface UserDefinedFunctionBuilder {
3632

3733
UDFOperandMetadata getOperandMetadata();
3834

39-
default SqlKind getKind() {
40-
return SqlKind.OTHER_FUNCTION;
41-
}
42-
43-
default SqlSyntax getSqlSyntax() {
44-
return SqlSyntax.FUNCTION;
45-
}
46-
47-
default Supplier<SqlOperator> getReverse() {
48-
return null;
49-
}
50-
5135
default SqlUserDefinedFunction toUDF(String functionName) {
5236
return toUDF(functionName, true);
5337
}
@@ -66,7 +50,7 @@ default SqlUserDefinedFunction toUDF(String functionName, boolean isDeterministi
6650
new SqlIdentifier(Collections.singletonList(functionName), null, SqlParserPos.ZERO, null);
6751
return new SqlUserDefinedFunction(
6852
udfLtrimIdentifier,
69-
getKind(),
53+
SqlKind.OTHER_FUNCTION,
7054
getReturnTypeInference(),
7155
InferTypes.ANY_NULLABLE,
7256
getOperandMetadata(),
@@ -75,19 +59,6 @@ default SqlUserDefinedFunction toUDF(String functionName, boolean isDeterministi
7559
public boolean isDeterministic() {
7660
return isDeterministic;
7761
}
78-
79-
@Override
80-
public @Nullable SqlOperator reverse() {
81-
if (getReverse() == null) {
82-
return null;
83-
}
84-
return getReverse().get();
85-
}
86-
87-
@Override
88-
public SqlSyntax getSyntax() {
89-
return getSqlSyntax();
90-
}
9162
};
9263
}
9364
}

‎core/src/main/java/org/opensearch/sql/expression/function/udf/ip/CompareIpFunction.java‎

Lines changed: 45 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -5,24 +5,30 @@
55

66
package org.opensearch.sql.expression.function.udf.ip;
77

8+
import java.util.Collections;
89
import java.util.List;
910
import java.util.Locale;
10-
import java.util.function.Supplier;
1111
import org.apache.calcite.adapter.enumerable.NotNullImplementor;
1212
import org.apache.calcite.adapter.enumerable.NullPolicy;
1313
import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
1414
import org.apache.calcite.linq4j.tree.ConstantExpression;
1515
import org.apache.calcite.linq4j.tree.Expression;
1616
import org.apache.calcite.linq4j.tree.Expressions;
1717
import org.apache.calcite.rex.RexCall;
18+
import org.apache.calcite.sql.SqlIdentifier;
1819
import org.apache.calcite.sql.SqlKind;
1920
import org.apache.calcite.sql.SqlOperator;
2021
import org.apache.calcite.sql.SqlSyntax;
22+
import org.apache.calcite.sql.parser.SqlParserPos;
23+
import org.apache.calcite.sql.type.InferTypes;
2124
import org.apache.calcite.sql.type.ReturnTypes;
2225
import org.apache.calcite.sql.type.SqlReturnTypeInference;
26+
import org.apache.calcite.sql.validate.SqlUserDefinedFunction;
27+
import org.checkerframework.checker.nullness.qual.Nullable;
2328
import org.opensearch.sql.data.model.ExprIpValue;
2429
import org.opensearch.sql.data.type.ExprCoreType;
2530
import org.opensearch.sql.expression.function.ImplementorUDF;
31+
import org.opensearch.sql.expression.function.PPLBuiltinOperators;
2632
import org.opensearch.sql.expression.function.UDFOperandMetadata;
2733

2834
/**
@@ -38,50 +44,10 @@
3844
*/
3945
public class CompareIpFunction extends ImplementorUDF {
4046
private final SqlKind kind;
41-
private Supplier<SqlOperator> reverse;
4247

4348
private CompareIpFunction(SqlKind kind) {
4449
super(new CompareImplementor(kind), NullPolicy.ANY);
4550
this.kind = kind;
46-
// Will be set later
47-
this.reverse = null;
48-
}
49-
50-
@Override
51-
public SqlSyntax getSqlSyntax() {
52-
return SqlSyntax.BINARY;
53-
}
54-
55-
@Override
56-
public SqlKind getKind() {
57-
return kind;
58-
}
59-
60-
@Override
61-
public Supplier<SqlOperator> getReverse() {
62-
return reverse;
63-
}
64-
65-
/**
66-
* Sets the reverse operator supplier for this comparison function.
67-
*
68-
* <p>This method is used to establish the reversed relationship between comparison operators
69-
* (e.g., "less than" and "greater than"). When the query optimizer normalizes expressions, it may
70-
* need to transform "b > a" to "a < b".
71-
*
72-
* <p>E.g. in the {@code hashCode} method of {@link org.apache.calcite.rex.RexNormalize}#L115, it
73-
* always converts <i>B [comparator] A</i> to <i>A [reverse_comparator] B</i> if the ordinal of
74-
* the reverse of the comparator is smaller.
75-
*
76-
* <p>IP comparison functions use this to inform the optimizer that ip_less_than is the reverse of
77-
* ip_greater_than, allowing for proper query normalization.
78-
*
79-
* @param supplier The supplier that provides the reverse SQL operator
80-
* @return This CompareIpFunction instance for method chaining
81-
*/
82-
public CompareIpFunction withReverse(Supplier<SqlOperator> supplier) {
83-
this.reverse = supplier;
84-
return this;
8551
}
8652

8753
public static CompareIpFunction less() {
@@ -108,6 +74,44 @@ public static CompareIpFunction notEquals() {
10874
return new CompareIpFunction(SqlKind.NOT_EQUALS);
10975
}
11076

77+
@Override
78+
public SqlUserDefinedFunction toUDF(String functionName, boolean isDeterministic) {
79+
SqlIdentifier udfIdentifier =
80+
new SqlIdentifier(Collections.singletonList(functionName), null, SqlParserPos.ZERO, null);
81+
return new SqlUserDefinedFunction(
82+
udfIdentifier,
83+
kind,
84+
getReturnTypeInference(),
85+
InferTypes.ANY_NULLABLE,
86+
getOperandMetadata(),
87+
getFunction()) {
88+
@Override
89+
public boolean isDeterministic() {
90+
return isDeterministic;
91+
}
92+
93+
@Override
94+
public @Nullable SqlOperator reverse() {
95+
return switch (kind) {
96+
case LESS_THAN -> PPLBuiltinOperators.GREATER_IP;
97+
case GREATER_THAN -> PPLBuiltinOperators.LESS_IP;
98+
case LESS_THAN_OR_EQUAL -> PPLBuiltinOperators.GTE_IP;
99+
case GREATER_THAN_OR_EQUAL -> PPLBuiltinOperators.LTE_IP;
100+
case EQUALS -> PPLBuiltinOperators.EQUALS_IP;
101+
case NOT_EQUALS -> PPLBuiltinOperators.NOT_EQUALS_IP;
102+
default -> throw new IllegalArgumentException(
103+
String.format(
104+
Locale.ROOT, "CompareIpFunction is not supposed to be of kind: %s", kind));
105+
};
106+
}
107+
108+
@Override
109+
public SqlSyntax getSyntax() {
110+
return SqlSyntax.BINARY;
111+
}
112+
};
113+
}
114+
111115
@Override
112116
public SqlReturnTypeInference getReturnTypeInference() {
113117
return ReturnTypes.BOOLEAN_FORCE_NULLABLE;

0 commit comments

Comments
 (0)