push down binary comparisons on date and date nanos union type fields

Author: fang-xing-esql

x-pack/plugin/esql/qa/testFixtures/src/main/resources/union_types.csv-spec

@@ -18,6 +18,7 @@
 import org.elasticsearch.xpack.esql.Column;
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
 import org.elasticsearch.xpack.esql.VerificationException;
+import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
 import org.elasticsearch.xpack.esql.analysis.AnalyzerRules.ParameterizedAnalyzerRule;
 import org.elasticsearch.xpack.esql.common.Failure;
 import org.elasticsearch.xpack.esql.core.capabilities.Resolvables;
@@ -1797,11 +1798,15 @@ private static LogicalPlan planWithoutSyntheticAttributes(LogicalPlan plan) {
     private static class ImplicitCastingForUnionTypedFields extends ParameterizedRule<LogicalPlan, LogicalPlan, AnalyzerContext> {
         @Override
         public LogicalPlan apply(LogicalPlan plan, AnalyzerContext context) {
+            if (EsqlCapabilities.Cap.IMPLICIT_CASTING_DATE_AND_DATE_NANOS.isEnabled() == false) {
+                return plan;
+            }
             // This rule should be applied after ResolveUnionTypes, so that the InvalidMappedFields with explicit casting are converted into
-            // MultiTypeEsField, and don't be double cast here.
+            // MultiTypeEsField, and don't get double cast here.
             Map<FieldAttribute, Alias> invalidMappedFieldCasted = new HashMap<>();
             LogicalPlan transformedPlan = plan.transformUp(LogicalPlan.class, p -> {
-                if (p instanceof UnaryPlan == false && p instanceof LookupJoin == false) {
+                // exclude LookupJoin for now, as it doesn't support date_nanos as join key yet
+                if (p instanceof UnaryPlan == false) {
                     return p;
                 }
                 Set<FieldAttribute> invalidMappedFields = invalidMappedFieldsInLogicalPlan(p);
@@ -1829,11 +1834,10 @@ public LogicalPlan apply(LogicalPlan plan, AnalyzerContext context) {
                     // If there are new aliases created, create a new eval child with new aliases for the current plan。
                     // How many children does a LogicalPlan have? Only deal with UnaryPlan and LookupJoin for now.
                     if (newAliases.isEmpty() == false) { // create a new eval child plan
-                        if (p instanceof UnaryPlan u) { // unary plan
-                            Eval eval = new Eval(u.source(), u.child(), newAliases.values().stream().toList());
-                            p = u.replaceChild(eval);
-                        }
-                        // TODO Lookup join does not work on date_nanos field today, joining on a date_nanos field does not find a match.
+                        UnaryPlan u = (UnaryPlan) p; // this must be a unary plan, as it is checked at the beginning of plan loop
+                        Eval eval = new Eval(u.source(), u.child(), newAliases.values().stream().toList());
+                        p = u.replaceChild(eval);
+                        // TODO Lookup join does not work on date_nanos field yet, joining on a date_nanos field does not find a match.
                         // And lookup up join is a special case as a lookup join has two children, after date_nanos is supported as a join
                         // key, the transformation needs to take it into account.
                     }
@@ -1899,7 +1903,10 @@ private static Expression castInvalidMappedField(DataType targetType, FieldAttri
          */
         private static Set<FieldAttribute> invalidMappedFieldsInLogicalPlan(LogicalPlan plan) {
             Set<FieldAttribute> fas = new HashSet<>();
-            if (plan instanceof EsRelation == false) { // not all the union typed fields from indices need to be cast implicitly
+            // Invalid mapped fields are legal at EsRelation level, as long as they are not used elsewhere. In the final output, if they
+            // have not been dropped, implicit cast will be added for them, so that we can return not null values, the implicit casting is
+            // deferred to when the fields are used or returned.
+            if (plan instanceof EsRelation == false) {
                 plan.forEachExpression(FieldAttribute.class, fa -> {
                     if (fa.field() instanceof InvalidMappedField) {
                         fas.add(fa);
@@ -1939,7 +1946,8 @@ private static LogicalPlan castInvalidMappedFieldInFinalOutput(LogicalPlan logic
                         logicalPlan = unary.replaceChild(eval);
                     }
                     // TODO LookupJoin is a binary plan, it does not create a new field, ideally adding an Eval on top of it should be fine,
-                    // however because the output of a LookupJoin does not include InvalidMappedFields, it is a bug need to be addressed
+                    // however because the output of a LookupJoin does not include InvalidMappedFields even the LHS output has
+                    // InvalidMappedFields, it is a bug need to be addressed
                 }
             }
             return logicalPlan;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
+import org.elasticsearch.xpack.esql.core.expression.TypedAttribute;
 import org.elasticsearch.xpack.esql.core.expression.function.scalar.ScalarFunction;
 import org.elasticsearch.xpack.esql.core.expression.predicate.operator.comparison.BinaryComparison;
 import org.elasticsearch.xpack.esql.core.querydsl.query.Query;
@@ -279,7 +280,13 @@ private RangeQuery translate(TranslatorHandler handler) {
             }
         }
         logger.trace("Building range query with format string [{}]", format);
-        return new RangeQuery(source(), handler.nameOf(value), l, includeLower(), u, includeUpper(), format, zoneId);
+        // This is a similar check as in EsqlBinaryComparison
+        // Extract the real field name from MultiTypeEsField, and use it in the push down query if it is found
+        TypedAttribute attribute = LucenePushdownPredicates.checkIsPushableAttribute(value);
+        String name = handler.nameOf(attribute);
+        String fieldNameFromMultiTypeEsField = LucenePushdownPredicates.extractFieldNameFromMultiTypeEsField(attribute);
+        name = fieldNameFromMultiTypeEsField != null ? fieldNameFromMultiTypeEsField : name;
+        return new RangeQuery(source(), name, l, includeLower(), u, includeUpper(), format, zoneId);
     }
 
     @Override

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/predicate/Range.java

@@ -378,6 +378,9 @@ public Expression singleValueField() {
     private Query translate(TranslatorHandler handler) {
         TypedAttribute attribute = LucenePushdownPredicates.checkIsPushableAttribute(left());
         String name = handler.nameOf(attribute);
+        // Extract the real field name from MultiTypeEsField, and use it in the push down query if it is found
+        String fieldNameFromMultiTypeEsField = LucenePushdownPredicates.extractFieldNameFromMultiTypeEsField(attribute);
+        name = fieldNameFromMultiTypeEsField != null ? fieldNameFromMultiTypeEsField : name;
         Object value = valueOf(FoldContext.small() /* TODO remove me */, right());
         String format = null;
         boolean isDateLiteralComparison = false;
@@ -452,7 +455,10 @@ private Query translate(TranslatorHandler handler) {
             return new RangeQuery(source(), name, null, false, value, true, format, zoneId);
         }
         if (this instanceof Equals || this instanceof NotEquals) {
-            name = LucenePushdownPredicates.pushableAttributeName(attribute);
+            // Extract the real field name from MultiTypeEsField, and use it in the push down query if it is found
+            name = fieldNameFromMultiTypeEsField != null
+                ? fieldNameFromMultiTypeEsField
+                : LucenePushdownPredicates.pushableAttributeName(attribute);
 
             Query query;
             if (isDateLiteralComparison) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/predicate/operator/comparison/EsqlBinaryComparison.java

@@ -7,14 +7,22 @@
 
 package org.elasticsearch.xpack.esql.optimizer.rules.physical.local;
 
+import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
 import org.elasticsearch.xpack.esql.core.expression.MetadataAttribute;
 import org.elasticsearch.xpack.esql.core.expression.TypedAttribute;
 import org.elasticsearch.xpack.esql.core.type.DataType;
+import org.elasticsearch.xpack.esql.core.type.MultiTypeEsField;
 import org.elasticsearch.xpack.esql.core.util.Check;
+import org.elasticsearch.xpack.esql.expression.function.scalar.convert.AbstractConvertFunction;
 import org.elasticsearch.xpack.esql.stats.SearchStats;
 
+import java.util.Map;
+import java.util.function.Predicate;
+
+import static org.elasticsearch.xpack.esql.core.type.DataType.DATE_NANOS;
+
 /**
  * When deciding if a filter or topN can be pushed down to Lucene, we need to check a few things on the field.
  * Exactly what is checked depends on the type of field and the query. For example, we have the following possible combinations:
@@ -94,6 +102,46 @@ static String pushableAttributeName(TypedAttribute attribute) {
             : attribute.name();
     }
 
+    /**
+     * Extract the real field name from a MultiTypeEsField, limit to MultiTypeEsField that has date_nanos type only.
+     *
+     * For example, the name of a MultiTypeEsField can be $$myfield$converted_to$date_nanos, and the real field name extract from the
+     * MultiTypeEsField is myfield, this method return myfield given a MultiTypeEsField.
+     *
+     * If the real field name is found, and the original field data types contain only date and date_nanos types, return the real field
+     * name, so that the real field name will be used to check for eligibility of being pushed down, and the real field name will be used
+     * in the push down query, instead of the name of the MultiTypeEsField, which should not match any field in an index.
+     *
+     * This method can be extended to support the other data types in the future if there is a need.
+     */
+    static String extractFieldNameFromMultiTypeEsField(TypedAttribute attribute) {
+        if (EsqlCapabilities.Cap.IMPLICIT_CASTING_DATE_AND_DATE_NANOS.isEnabled()
+            && attribute instanceof FieldAttribute fa
+            && fa.field() instanceof MultiTypeEsField multiTypeEsField
+            && fa.dataType() == DATE_NANOS
+            &&  // limit to casting to date_nanos only
+            mixedDateAndDateNanosOnly(multiTypeEsField, DataType::isMillisOrNanos) // limit to mixed date and date_nanos only
+        ) {
+            return fa.fieldName();
+        }
+        return null;
+    }
+
+    /**
+     * Check if the original field types in a MultiTypeEsField satisfy the required data types defined in the predicate.
+     */
+    private static boolean mixedDateAndDateNanosOnly(MultiTypeEsField multiTypeEsField, Predicate<DataType> predicate) {
+        Map<String, Expression> indexToConversionExpressions = multiTypeEsField.getIndexToConversionExpressions();
+        for (Map.Entry<String, Expression> entry : indexToConversionExpressions.entrySet()) {
+            Expression conversionFunction = entry.getValue();
+            if (conversionFunction instanceof AbstractConvertFunction abstractConvertFunction
+                && predicate.test(abstractConvertFunction.field().dataType()) == false) {
+                return false;
+            }
+        }
+        return true;
+    }
+
     /**
      * The default implementation of this has no access to SearchStats, so it can only make decisions based on the FieldAttribute itself.
      * In particular, it assumes TEXT fields have no exact subfields (underlying keyword field),
@@ -131,10 +179,14 @@ public boolean hasExactSubfield(FieldAttribute attr) {
 
             @Override
             public boolean isIndexedAndHasDocValues(FieldAttribute attr) {
+                // If this is a MultiTypeEsField cast to date_nanos, make it eligible for being pushed down by checking the real
+                // field name against SearchStats
+                String fieldNameFromMultiTypeEsField = LucenePushdownPredicates.extractFieldNameFromMultiTypeEsField(attr);
+                String name = fieldNameFromMultiTypeEsField != null ? fieldNameFromMultiTypeEsField : attr.name();
                 // We still consider the value of isAggregatable here, because some fields like ScriptFieldTypes are always aggregatable
                 // But this could hide issues with fields that are not indexed but are aggregatable
                 // This is the original behaviour for ES|QL, but is it correct?
-                return attr.field().isAggregatable() || stats.isIndexed(attr.name()) && stats.hasDocValues(attr.name());
+                return attr.field().isAggregatable() || stats.isIndexed(name) && stats.hasDocValues(name);
             }
 
             @Override

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/LucenePushdownPredicates.java

@@ -32,7 +32,6 @@
 import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
-import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.index.EsIndex;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
 import org.elasticsearch.xpack.esql.plan.GeneratingPlan;
@@ -209,7 +208,6 @@ public boolean expressionsResolved() {
         return policyName.resolved()
             && matchField instanceof EmptyAttribute == false // matchField not defined in the query, needs to be resolved from the policy
             && matchField.resolved()
-            && matchField.dataType() != DataType.UNSUPPORTED
             && Resolvables.resolved(enrichFields());
     }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Enrich.java

@@ -13,6 +13,7 @@
 import org.elasticsearch.xpack.esql.capabilities.PostAnalysisVerificationAware;
 import org.elasticsearch.xpack.esql.capabilities.TelemetryAware;
 import org.elasticsearch.xpack.esql.common.Failures;
+import org.elasticsearch.xpack.esql.core.capabilities.Resolvables;
 import org.elasticsearch.xpack.esql.core.expression.Alias;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
@@ -131,12 +132,7 @@ private List<Alias> renameAliases(List<Alias> originalAttributes, List<String> n
 
     @Override
     public boolean expressionsResolved() {
-        for (Alias a : fields) {
-            if (a.resolved() == false || a.dataType() == DataType.UNSUPPORTED) {
-                return false;
-            }
-        }
-        return true;
+        return Resolvables.resolved(fields);
     }
 
     @Override

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Eval.java

@@ -12,7 +12,6 @@
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.xpack.esql.core.capabilities.Resolvables;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
-import org.elasticsearch.xpack.esql.core.type.DataType;
 
 import java.io.IOException;
 import java.util.List;
@@ -51,7 +50,6 @@ public boolean expressionsResolved() {
         return type.resolved()
             && Resolvables.resolved(matchFields)
             && Resolvables.resolved(leftFields)
-            && Resolvables.resolved(rightFields)
-            && leftFields.stream().noneMatch(e -> e.dataType() == DataType.UNSUPPORTED);
+            && Resolvables.resolved(rightFields);
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/join/JoinConfig.java

@@ -15,6 +15,7 @@
 import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
 import org.elasticsearch.xpack.esql.core.expression.MetadataAttribute;
 import org.elasticsearch.xpack.esql.core.querydsl.query.Query;
+import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.LucenePushdownPredicates;
 import org.elasticsearch.xpack.esql.querydsl.query.SingleValueQuery;
 
 /**
@@ -41,7 +42,10 @@ public Query asQuery(Expression e) {
     private static Query wrapFunctionQuery(Expression field, Query query) {
         if (field instanceof FieldAttribute fa) {
             fa = fa.getExactInfo().hasExact() ? fa.exactAttribute() : fa;
-            return new SingleValueQuery(query, fa.name());
+            // Extract the real field name from MultiTypeEsField, and use it in the push down query if it is found
+            String fieldNameFromMultiTypeEsField = LucenePushdownPredicates.extractFieldNameFromMultiTypeEsField(fa);
+            String fieldName = fieldNameFromMultiTypeEsField != null ? fieldNameFromMultiTypeEsField : fa.name();
+            return new SingleValueQuery(query, fieldName);
         }
         if (field instanceof MetadataAttribute) {
             return query; // MetadataAttributes are always single valued

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/TranslatorHandler.java

@@ -11,6 +11,8 @@
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xpack.core.enrich.EnrichPolicy;
 import org.elasticsearch.xpack.esql.EsqlTestUtils;
+import org.elasticsearch.xpack.esql.core.type.EsField;
+import org.elasticsearch.xpack.esql.core.type.InvalidMappedField;
 import org.elasticsearch.xpack.esql.enrich.ResolvedEnrichPolicy;
 import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
 import org.elasticsearch.xpack.esql.index.EsIndex;
@@ -24,8 +26,10 @@
 import org.elasticsearch.xpack.esql.session.Configuration;
 
 import java.util.ArrayList;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import static org.elasticsearch.xpack.core.enrich.EnrichPolicy.GEO_MATCH_TYPE;
 import static org.elasticsearch.xpack.core.enrich.EnrichPolicy.MATCH_TYPE;
@@ -209,4 +213,26 @@ public static void loadEnrichPolicyResolution(EnrichResolution enrich, String po
     public static IndexResolution tsdbIndexResolution() {
         return loadMapping("tsdb-mapping.json", "test");
     }
+
+    public static IndexResolution indexWithDateDateNanosUnionType() {
+        // this method is shared by AnalyzerTest, QueryTranslatorTests and LocalPhysicalPlanOptimizerTests
+        String dateDateNanos = "date_and_date_nanos"; // mixed date and date_nanos
+        String dateDateNanosLong = "date_and_date_nanos_and_long"; // mixed date, date_nanos and long
+        LinkedHashMap<String, Set<String>> typesToIndices1 = new LinkedHashMap<>();
+        typesToIndices1.put("date", Set.of("index1", "index2"));
+        typesToIndices1.put("date_nanos", Set.of("index3"));
+        LinkedHashMap<String, Set<String>> typesToIndices2 = new LinkedHashMap<>();
+        typesToIndices2.put("date", Set.of("index1"));
+        typesToIndices2.put("date_nanos", Set.of("index2"));
+        typesToIndices2.put("long", Set.of("index3"));
+        EsField dateDateNanosField = new InvalidMappedField(dateDateNanos, typesToIndices1);
+        EsField dateDateNanosLongField = new InvalidMappedField(dateDateNanosLong, typesToIndices2);
+        EsIndex index = new EsIndex(
+            "test*",
+            Map.of(dateDateNanos, dateDateNanosField, dateDateNanosLong, dateDateNanosLongField),
+            Map.of("index1", IndexMode.STANDARD, "index2", IndexMode.STANDARD, "index3", IndexMode.STANDARD)
+        );
+        return IndexResolution.valid(index);
+
+    }
 }