leventov commented on a change in pull request #8157: Enum of ResponseContext keys URL: https://github.com/apache/incubator-druid/pull/8157#discussion_r307368997
########## File path: processing/src/main/java/org/apache/druid/query/context/ResponseContext.java ########## @@ -76,56 +138,128 @@ public static ResponseContext createEmpty() return DefaultResponseContext.createEmpty(); } - protected abstract Map<String, Object> getDelegate(); - - public Object put(String key, Object value) + public static ResponseContext deserialize(String responseContext, ObjectMapper objectMapper) throws IOException { - return getDelegate().put(key, value); + final Map<String, Object> delegate = objectMapper.readValue( + responseContext, + JacksonUtils.TYPE_REFERENCE_MAP_STRING_OBJECT + ); + return new ResponseContext() + { + @Override + protected Map<String, Object> getDelegate() + { + return delegate; + } + }; } - public Object get(String key) + protected abstract Map<String, Object> getDelegate(); + + public Object put(Key key, Object value) { - return getDelegate().get(key); + return getDelegate().put(key.name, value); } - public Object remove(String key) + public Object get(Key key) { - return getDelegate().remove(key); + return getDelegate().get(key.name); } - public void putAll(Map<? extends String, ?> m) + public Object remove(Key key) { - getDelegate().putAll(m); + return getDelegate().remove(key.name); } - public void putAll(ResponseContext responseContext) + /** + * Merges a new value associated with a key with an old value. + */ + public Object merge(Key key, Object value) { - getDelegate().putAll(responseContext.getDelegate()); + return getDelegate().merge(key.name, value, key.mergeFunction); } - public int size() + /** + * Merges a response context into current. + * This method merges only keys from the enum {@link Key}. + */ + public void merge(ResponseContext responseContext) { - return getDelegate().size(); + for (Key key : Key.values()) { + final Object newValue = responseContext.get(key); + if (newValue != null) { + merge(key, newValue); + } + } } - public String serializeWith(ObjectMapper objectMapper) throws JsonProcessingException + /** + * Serializes the context given that the resulting string length is less than the provided limit. + * The method removes max-length fields one by one if the resulting string length is greater than the limit. + * The resulting string might be correctly deserialized as a {@link ResponseContext}. + */ + public SerializationResult serializeWith(ObjectMapper objectMapper, int maxLength) throws JsonProcessingException { - return objectMapper.writeValueAsString(getDelegate()); + final String fullSerializedString = objectMapper.writeValueAsString(getDelegate()); + if (fullSerializedString.length() <= maxLength) { + return new SerializationResult(fullSerializedString, fullSerializedString, false); + } else { + final HashMap<String, Object> copiedMap = new HashMap<>(getDelegate()); + final PriorityQueue<Pair<String, String>> serializedPairs = new PriorityQueue<>((o1, o2) -> { + Preconditions.checkNotNull(o1.rhs); + Preconditions.checkNotNull(o2.rhs); + return o2.rhs.length() - o1.rhs.length(); + }); + for (Map.Entry<String, Object> e : copiedMap.entrySet()) { + serializedPairs.add(new Pair<>(e.getKey(), objectMapper.writeValueAsString(e.getValue()))); + } + while (!copiedMap.isEmpty()) { + final Pair<String, String> maxLengthPair = serializedPairs.poll(); + Preconditions.checkNotNull(maxLengthPair); Review comment: May do `!serializedPairs.isEmpty()` in the loop condition and remove this line ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org With regards, Apache Git Services --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@druid.apache.org For additional commands, e-mail: commits-h...@druid.apache.org