From 769a13759f2152541bc69c8f8ecfebbc721824c5 Mon Sep 17 00:00:00 2001 From: Jack Mazanec Date: Wed, 27 Oct 2021 16:46:34 -0700 Subject: [PATCH] Add validation to check max int limit (#159) Signed-off-by: John Mazanec --- .../opensearch/knn/plugin/rest/RestTrainModelHandler.java | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/main/java/org/opensearch/knn/plugin/rest/RestTrainModelHandler.java b/src/main/java/org/opensearch/knn/plugin/rest/RestTrainModelHandler.java index 7ed47270d..bc9283621 100644 --- a/src/main/java/org/opensearch/knn/plugin/rest/RestTrainModelHandler.java +++ b/src/main/java/org/opensearch/knn/plugin/rest/RestTrainModelHandler.java @@ -14,6 +14,7 @@ import com.google.common.collect.ImmutableList; import org.opensearch.client.node.NodeClient; import org.opensearch.common.xcontent.XContentParser; +import org.opensearch.index.mapper.NumberFieldMapper; import org.opensearch.knn.index.KNNMethodContext; import org.opensearch.knn.plugin.KNNPlugin; import org.opensearch.knn.plugin.transport.TrainingJobRouterAction; @@ -105,11 +106,11 @@ private TrainingModelRequest createTransportRequest(RestRequest restRequest) thr } else if (KNN_METHOD.equals(fieldName) && ensureNotSet(fieldName, knnMethodContext)) { knnMethodContext = KNNMethodContext.parse(parser.map()); } else if (DIMENSION.equals(fieldName) && ensureNotSet(fieldName, dimension)) { - dimension = parser.intValue(); + dimension = (Integer) NumberFieldMapper.NumberType.INTEGER.parse(parser.objectBytes(), false); } else if (MAX_VECTOR_COUNT_PARAMETER.equals(fieldName) && ensureNotSet(fieldName, maximumVectorCount)) { - maximumVectorCount = parser.intValue(); + maximumVectorCount = (Integer) NumberFieldMapper.NumberType.INTEGER.parse(parser.objectBytes(), false); } else if (SEARCH_SIZE_PARAMETER.equals(fieldName) && ensureNotSet(fieldName, searchSize)) { - searchSize = parser.intValue(); + searchSize = (Integer) NumberFieldMapper.NumberType.INTEGER.parse(parser.objectBytes(), false); } else if (MODEL_DESCRIPTION.equals(fieldName) && ensureNotSet(fieldName, description)) { description = parser.text(); } else {