From 41cba16fba072e2101cd4a1041449c1560918449 Mon Sep 17 00:00:00 2001
From: Chaitanya Gohel <gashutos@amazon.com>
Date: Tue, 31 Oct 2023 23:24:09 +0530
Subject: [PATCH] Making indexFieldData provate in FloatValuesComparatorSource

Signed-off-by: Chaitanya Gohel <gashutos@amazon.com>
---
 .../FloatValuesComparatorSource.java          |  4 ++--
 .../HalfFloatValuesComparatorSource.java      | 20 +++++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/server/src/main/java/org/opensearch/index/fielddata/fieldcomparator/FloatValuesComparatorSource.java b/server/src/main/java/org/opensearch/index/fielddata/fieldcomparator/FloatValuesComparatorSource.java
index abf872dd235e5..04a34cd418520 100644
--- a/server/src/main/java/org/opensearch/index/fielddata/fieldcomparator/FloatValuesComparatorSource.java
+++ b/server/src/main/java/org/opensearch/index/fielddata/fieldcomparator/FloatValuesComparatorSource.java
@@ -61,7 +61,7 @@
  */
 public class FloatValuesComparatorSource extends IndexFieldData.XFieldComparatorSource {
 
-    protected final IndexNumericFieldData indexFieldData;
+    private final IndexNumericFieldData indexFieldData;
 
     public FloatValuesComparatorSource(
         IndexNumericFieldData indexFieldData,
@@ -78,7 +78,7 @@ public SortField.Type reducedType() {
         return SortField.Type.FLOAT;
     }
 
-    protected NumericDoubleValues getNumericDocValues(LeafReaderContext context, float missingValue) throws IOException {
+    private NumericDoubleValues getNumericDocValues(LeafReaderContext context, float missingValue) throws IOException {
         final SortedNumericDoubleValues values = indexFieldData.load(context).getDoubleValues();
         if (nested == null) {
             return FieldData.replaceMissing(sortMode.select(values), missingValue);
diff --git a/server/src/main/java/org/opensearch/index/fielddata/fieldcomparator/HalfFloatValuesComparatorSource.java b/server/src/main/java/org/opensearch/index/fielddata/fieldcomparator/HalfFloatValuesComparatorSource.java
index f63c4110bbdda..7e3936be1d8a5 100644
--- a/server/src/main/java/org/opensearch/index/fielddata/fieldcomparator/HalfFloatValuesComparatorSource.java
+++ b/server/src/main/java/org/opensearch/index/fielddata/fieldcomparator/HalfFloatValuesComparatorSource.java
@@ -10,9 +10,14 @@
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.NumericDocValues;
+import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.FieldComparator;
 import org.apache.lucene.search.LeafFieldComparator;
+import org.apache.lucene.util.BitSet;
+import org.opensearch.index.fielddata.FieldData;
 import org.opensearch.index.fielddata.IndexNumericFieldData;
+import org.opensearch.index.fielddata.NumericDoubleValues;
+import org.opensearch.index.fielddata.SortedNumericDoubleValues;
 import org.opensearch.index.search.comparators.HalfFloatComparator;
 import org.opensearch.search.MultiValueMode;
 
@@ -24,6 +29,8 @@
  * @opensearch.internal
  */
 public class HalfFloatValuesComparatorSource extends FloatValuesComparatorSource {
+    private final IndexNumericFieldData indexFieldData;
+
     public HalfFloatValuesComparatorSource(
         IndexNumericFieldData indexFieldData,
         Object missingValue,
@@ -31,6 +38,7 @@ public HalfFloatValuesComparatorSource(
         Nested nested
     ) {
         super(indexFieldData, missingValue, sortMode, nested);
+        this.indexFieldData = indexFieldData;
     }
 
     @Override
@@ -52,4 +60,16 @@ protected NumericDocValues getNumericDocValues(LeafReaderContext context, String
             }
         };
     }
+
+    private NumericDoubleValues getNumericDocValues(LeafReaderContext context, float missingValue) throws IOException {
+        final SortedNumericDoubleValues values = indexFieldData.load(context).getDoubleValues();
+        if (nested == null) {
+            return FieldData.replaceMissing(sortMode.select(values), missingValue);
+        } else {
+            final BitSet rootDocs = nested.rootDocs(context);
+            final DocIdSetIterator innerDocs = nested.innerDocs(context);
+            final int maxChildren = nested.getNestedSort() != null ? nested.getNestedSort().getMaxChildren() : Integer.MAX_VALUE;
+            return sortMode.select(values, missingValue, rootDocs, innerDocs, context.reader().maxDoc(), maxChildren);
+        }
+    }
 }