Add support for source fallback with the boolean field type (#89052)

This change adds a SourceValueFetcherSortedBooleanIndexFieldData to support boolean doc values 
for source fallback.

docs/changelog/89052.yaml

@@ -0,0 +1,5 @@
+pr: 89052
+summary: Add support for source fallback with the boolean field type
+area: Mapping
+type: enhancement
+issues: []

modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/50_script_doc_values.yml

@@ -9,6 +9,9 @@ setup:
                     properties:
                         boolean:
                             type: boolean
+                        boolean_no_doc_values:
+                            type: boolean
+                            doc_values: false
                         date:
                             type: date
                         nanos:
@@ -79,6 +82,7 @@ setup:
             body:
                 rank: 1
                 boolean: true
+                boolean_no_doc_values: true
                 date: 2017-01-01T12:11:12
                 nanos: 2015-01-01T12:10:30.123456789Z
                 geo_point: 41.12,-71.34
@@ -117,6 +121,7 @@ setup:
           body:
               rank: 3
               boolean: [true, false, true]
+              boolean_no_doc_values: [true, false, true]
               ip: ["10.1.2.3", "2001:db8::2:1"]
               date: [2017-01-01T12:11:12, 2018-01-01T12:11:12]
               nanos: [2015-01-01T12:10:30.123456789Z, 2015-01-01T12:10:30.987654321Z]
@@ -256,6 +261,120 @@ setup:
                             source: "field('boolean').size()"
     - match: { hits.hits.0.fields.field.0: 0 }
 
+---
+"boolean_no_doc_values":
+  - do:
+      catch: bad_request
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query: { term: { _id: "1" } }
+          script_fields:
+            field:
+              script:
+                source: "doc['boolean_no_doc_values'].get(0)"
+  - match: { error.failed_shards.0.reason.caused_by.type: "illegal_argument_exception" }
+
+  - do:
+      catch: bad_request
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query: { term: { _id: "1" } }
+          script_fields:
+            field:
+              script:
+                source: "doc['boolean_no_doc_values'].value"
+  - match: { error.failed_shards.0.reason.caused_by.type: "illegal_argument_exception" }
+
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query: { term: { _id: "1" } }
+          script_fields:
+            field:
+              script:
+                source: "field('boolean_no_doc_values').get(false)"
+  - match: { hits.hits.0.fields.field.0: true }
+
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query: { term: { _id: "1" } }
+          script_fields:
+            field:
+              script:
+                source: "/* avoid yaml stash for '$' */ $('boolean_no_doc_values', false)"
+  - match: { hits.hits.0.fields.field.0: true }
+
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query: { term: { _id: "2" } }
+          script_fields:
+            field:
+              script:
+                source: "field('boolean_no_doc_values').get(false)"
+  - match: { hits.hits.0.fields.field.0: false }
+
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query: { term: { _id: "2" } }
+          script_fields:
+            field:
+              script:
+                source: "/* avoid yaml stash for '$' */ $('boolean_no_doc_values', false)"
+  - match: { hits.hits.0.fields.field.0: false }
+
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query: { term: { _id: "1" } }
+          script_fields:
+            field:
+              script:
+                source: "field('boolean_no_doc_values').get(1, false)"
+  - match: { hits.hits.0.fields.field.0: false }
+
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query: { term: { _id: "1" } }
+          script_fields:
+            field:
+              script:
+                source: "int total = 0; for (boolean b : field('boolean_no_doc_values')) { total += b ? 1 : 0; } total;"
+  - match: { hits.hits.0.fields.field.0: 1 }
+
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query: { term: { _id: "3" } }
+          script_fields:
+            field:
+              script:
+                source: "int total = 0; for (boolean b : field('boolean_no_doc_values')) { total += b ? 1 : 0; } total + field('boolean').size();"
+  - match: { hits.hits.0.fields.field.0: 5 }
+
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        body:
+          query: { term: { _id: "2" } }
+          script_fields:
+            field:
+              script:
+                source: "field('boolean_no_doc_values').size()"
+  - match: { hits.hits.0.fields.field.0: 0 }
+
 ---
 "date":
     - skip:

server/src/main/java/org/elasticsearch/index/fielddata/SourceValueFetcherSortedBooleanIndexFieldData.java

@@ -0,0 +1,154 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.index.fielddata;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SortedNumericDocValues;
+import org.elasticsearch.index.mapper.ValueFetcher;
+import org.elasticsearch.indices.breaker.CircuitBreakerService;
+import org.elasticsearch.script.field.DocValuesScriptFieldFactory;
+import org.elasticsearch.script.field.ToScriptFieldFactory;
+import org.elasticsearch.search.aggregations.support.ValuesSourceType;
+import org.elasticsearch.search.lookup.SourceLookup;
+
+import java.io.IOException;
+import java.util.Collections;
+
+public class SourceValueFetcherSortedBooleanIndexFieldData extends SourceValueFetcherIndexFieldData<SortedNumericDocValues> {
+
+    public static class Builder extends SourceValueFetcherIndexFieldData.Builder<SortedNumericDocValues> {
+
+        public Builder(
+            String fieldName,
+            ValuesSourceType valuesSourceType,
+            ValueFetcher valueFetcher,
+            SourceLookup sourceLookup,
+            ToScriptFieldFactory<SortedNumericDocValues> toScriptFieldFactory
+        ) {
+            super(fieldName, valuesSourceType, valueFetcher, sourceLookup, toScriptFieldFactory);
+        }
+
+        @Override
+        public SourceValueFetcherSortedBooleanIndexFieldData build(IndexFieldDataCache cache, CircuitBreakerService breakerService) {
+            return new SourceValueFetcherSortedBooleanIndexFieldData(
+                fieldName,
+                valuesSourceType,
+                valueFetcher,
+                sourceLookup,
+                toScriptFieldFactory
+            );
+        }
+    }
+
+    protected SourceValueFetcherSortedBooleanIndexFieldData(
+        String fieldName,
+        ValuesSourceType valuesSourceType,
+        ValueFetcher valueFetcher,
+        SourceLookup sourceLookup,
+        ToScriptFieldFactory<SortedNumericDocValues> toScriptFieldFactory
+    ) {
+        super(fieldName, valuesSourceType, valueFetcher, sourceLookup, toScriptFieldFactory);
+    }
+
+    @Override
+    public SourceValueFetcherLeafFieldData<SortedNumericDocValues> loadDirect(LeafReaderContext context) throws Exception {
+        return new SourceValueFetcherSortedBooleanLeafFieldData(toScriptFieldFactory, context, valueFetcher, sourceLookup);
+    }
+
+    private static class SourceValueFetcherSortedBooleanLeafFieldData extends SourceValueFetcherLeafFieldData<SortedNumericDocValues> {
+
+        private SourceValueFetcherSortedBooleanLeafFieldData(
+            ToScriptFieldFactory<SortedNumericDocValues> toScriptFieldFactory,
+            LeafReaderContext leafReaderContext,
+            ValueFetcher valueFetcher,
+            SourceLookup sourceLookup
+        ) {
+            super(toScriptFieldFactory, leafReaderContext, valueFetcher, sourceLookup);
+        }
+
+        @Override
+        public DocValuesScriptFieldFactory getScriptFieldFactory(String name) {
+            return toScriptFieldFactory.getScriptFieldFactory(
+                new SourceValueFetcherSortedBooleanDocValues(leafReaderContext, valueFetcher, sourceLookup),
+                name
+            );
+        }
+    }
+
+    private static class SourceValueFetcherSortedBooleanDocValues extends SortedNumericDocValues implements ValueFetcherDocValues {
+
+        private final LeafReaderContext leafReaderContext;
+
+        private final ValueFetcher valueFetcher;
+        private final SourceLookup sourceLookup;
+
+        private int trueCount;
+        private int falseCount;
+        private int iteratorIndex;
+
+        private SourceValueFetcherSortedBooleanDocValues(
+            LeafReaderContext leafReaderContext,
+            ValueFetcher valueFetcher,
+            SourceLookup sourceLookup
+        ) {
+            this.leafReaderContext = leafReaderContext;
+            this.valueFetcher = valueFetcher;
+            this.sourceLookup = sourceLookup;
+        }
+
+        @Override
+        public boolean advanceExact(int doc) throws IOException {
+            sourceLookup.setSegmentAndDocument(leafReaderContext, doc);
+
+            for (Object value : valueFetcher.fetchValues(sourceLookup, Collections.emptyList())) {
+                assert value instanceof Boolean;
+                if ((Boolean) value) {
+                    ++trueCount;
+                } else {
+                    ++falseCount;
+                }
+            }
+
+            iteratorIndex = 0;
+
+            return true;
+        }
+
+        @Override
+        public int docValueCount() {
+            return trueCount + falseCount;
+        }
+
+        @Override
+        public long nextValue() throws IOException {
+            assert iteratorIndex < trueCount + falseCount;
+            return iteratorIndex++ < falseCount ? 0L : 1L;
+        }
+
+        @Override
+        public int docID() {
+            throw new UnsupportedOperationException("not supported for source fallback");
+        }
+
+        @Override
+        public int nextDoc() throws IOException {
+            throw new UnsupportedOperationException("not supported for source fallback");
+        }
+
+        @Override
+        public int advance(int target) throws IOException {
+            throw new UnsupportedOperationException("not supported for source fallback");
+        }
+
+        @Override
+        public long cost() {
+            throw new UnsupportedOperationException("not supported for source fallback");
+        }
+    }
+}

server/src/main/java/org/elasticsearch/index/mapper/BooleanFieldMapper.java

@@ -30,6 +30,7 @@ import org.elasticsearch.index.analysis.NamedAnalyzer;
 import org.elasticsearch.index.fielddata.FieldDataContext;
 import org.elasticsearch.index.fielddata.IndexFieldData;
 import org.elasticsearch.index.fielddata.IndexNumericFieldData.NumericType;
+import org.elasticsearch.index.fielddata.SourceValueFetcherSortedBooleanIndexFieldData;
 import org.elasticsearch.index.fielddata.plain.SortedNumericIndexFieldData;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.script.BooleanFieldScript;
@@ -37,6 +38,7 @@ import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptCompiler;
 import org.elasticsearch.script.field.BooleanDocValuesField;
 import org.elasticsearch.search.DocValueFormat;
+import org.elasticsearch.search.aggregations.support.CoreValuesSourceType;
 import org.elasticsearch.search.lookup.FieldValues;
 import org.elasticsearch.search.lookup.SearchLookup;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -48,6 +50,7 @@ import java.util.Collection;
 import java.util.Collections;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 
 /**
  * A field mapper for boolean fields.
@@ -198,7 +201,11 @@ public class BooleanFieldMapper extends FieldMapper {
             if (this.scriptValues != null) {
                 return FieldValues.valueFetcher(this.scriptValues, context);
             }
-            return new SourceValueFetcher(name(), context, nullValue) {
+            return sourceValueFetcher(context.isSourceEnabled() ? context.sourcePath(name()) : Collections.emptySet());
+        }
+
+        private SourceValueFetcher sourceValueFetcher(Set<String> sourcePaths) {
+            return new SourceValueFetcher(sourcePaths, nullValue) {
                 @Override
                 protected Boolean parseSourceValue(Object value) {
                     if (value instanceof Boolean) {
@@ -255,8 +262,30 @@ public class BooleanFieldMapper extends FieldMapper {
 
         @Override
         public IndexFieldData.Builder fielddataBuilder(FieldDataContext fieldDataContext) {
-            failIfNoDocValues();
-            return new SortedNumericIndexFieldData.Builder(name(), NumericType.BOOLEAN, BooleanDocValuesField::new);
+            FielddataOperation operation = fieldDataContext.fielddataOperation();
+
+            if (operation == FielddataOperation.SEARCH) {
+                failIfNoDocValues();
+            }
+
+            if ((operation == FielddataOperation.SEARCH || operation == FielddataOperation.SCRIPT) && hasDocValues()) {
+                return new SortedNumericIndexFieldData.Builder(name(), NumericType.BOOLEAN, BooleanDocValuesField::new);
+            }
+
+            if (operation == FielddataOperation.SCRIPT) {
+                SearchLookup searchLookup = fieldDataContext.lookupSupplier().get();
+                Set<String> sourcePaths = fieldDataContext.sourcePathsLookup().apply(name());
+
+                return new SourceValueFetcherSortedBooleanIndexFieldData.Builder(
+                    name(),
+                    CoreValuesSourceType.BOOLEAN,
+                    sourceValueFetcher(sourcePaths),
+                    searchLookup.source(),
+                    BooleanDocValuesField::new
+                );
+            }
+
+            throw new IllegalStateException("unknown field data type [" + operation.name() + "]");
         }
 
         @Override

server/src/main/java/org/elasticsearch/script/field/BooleanDocValuesField.java

@@ -43,7 +43,7 @@ public class BooleanDocValuesField extends AbstractScriptFieldFactory<Boolean>
         if (input.advanceExact(docId)) {
             resize(input.docValueCount());
             for (int i = 0; i < count; i++) {
-                values[i] = input.nextValue() == 1;
+                values[i] = input.nextValue() == 1L;
             }
         } else {
             resize(0);