|
8 | 8 |
|
9 | 9 | package org.opensearch.plugin.insights.rules.model;
|
10 | 10 |
|
| 11 | +import org.apache.lucene.util.ArrayUtil; |
11 | 12 | import org.opensearch.core.common.io.stream.StreamInput;
|
12 | 13 | import org.opensearch.core.common.io.stream.StreamOutput;
|
| 14 | +import org.opensearch.core.common.io.stream.Writeable; |
| 15 | +import org.opensearch.core.tasks.resourcetracker.TaskResourceInfo; |
13 | 16 |
|
14 | 17 | import java.io.IOException;
|
| 18 | +import java.util.Collections; |
| 19 | +import java.util.HashMap; |
| 20 | +import java.util.List; |
15 | 21 | import java.util.Locale;
|
| 22 | +import java.util.Map; |
16 | 23 |
|
17 | 24 | /**
|
18 | 25 | * Valid attributes for a search query record
|
@@ -65,14 +72,78 @@ static Attribute readFromStream(final StreamInput in) throws IOException {
|
65 | 72 | /**
|
66 | 73 | * Write Attribute to a StreamOutput
|
67 | 74 | *
|
68 |
| - * @param out the StreamOutput to write |
| 75 | + * @param out the StreamOutput to write |
69 | 76 | * @param attribute the Attribute to write
|
70 | 77 | * @throws IOException IOException
|
71 | 78 | */
|
72 | 79 | static void writeTo(final StreamOutput out, final Attribute attribute) throws IOException {
|
73 | 80 | out.writeString(attribute.toString());
|
74 | 81 | }
|
75 | 82 |
|
| 83 | + /** |
| 84 | + * Write Attribute value to a StreamOutput |
| 85 | + * |
| 86 | + * @param out the StreamOutput to write |
| 87 | + * @param attributeValue the Attribute value to write |
| 88 | + */ |
| 89 | + @SuppressWarnings("unchecked") |
| 90 | + public static void writeValueTo(StreamOutput out, Object attributeValue) throws IOException { |
| 91 | + if (attributeValue instanceof List) { |
| 92 | + out.writeList((List<? extends Writeable>) attributeValue); |
| 93 | + } else { |
| 94 | + out.writeGenericValue(attributeValue); |
| 95 | + } |
| 96 | + } |
| 97 | + |
| 98 | + /** |
| 99 | + * Read attribute value from the input stream given the Attribute type |
| 100 | + * |
| 101 | + * @param in the {@link StreamInput} input to read |
| 102 | + * @param attribute attribute type to differentiate between Source and others |
| 103 | + * @return parse value |
| 104 | + * @throws IOException IOException |
| 105 | + */ |
| 106 | + public static Object readAttributeValue(StreamInput in, Attribute attribute) throws IOException { |
| 107 | + if (attribute == Attribute.TASK_RESOURCE_USAGES) { |
| 108 | + return in.readList(TaskResourceInfo::readFromStream); |
| 109 | + } else { |
| 110 | + return in.readGenericValue(); |
| 111 | + } |
| 112 | + } |
| 113 | + |
| 114 | + /** |
| 115 | + * Read attribute map from the input stream |
| 116 | + * |
| 117 | + * @param in the {@link StreamInput} to read |
| 118 | + * @return parsed attribute map |
| 119 | + * @throws IOException IOException |
| 120 | + */ |
| 121 | + public static Map<Attribute, Object> readAttributeMap(StreamInput in) throws IOException { |
| 122 | + int size = readArraySize(in); |
| 123 | + if (size == 0) { |
| 124 | + return Collections.emptyMap(); |
| 125 | + } |
| 126 | + Map<Attribute, Object> map = new HashMap<>(size); |
| 127 | + |
| 128 | + for (int i = 0; i < size; i++) { |
| 129 | + Attribute key = readFromStream(in); |
| 130 | + Object value = readAttributeValue(in, key); |
| 131 | + map.put(key, value); |
| 132 | + } |
| 133 | + return map; |
| 134 | + } |
| 135 | + |
| 136 | + private static int readArraySize(StreamInput in) throws IOException { |
| 137 | + final int arraySize = in.readVInt(); |
| 138 | + if (arraySize > ArrayUtil.MAX_ARRAY_LENGTH) { |
| 139 | + throw new IllegalStateException("array length must be <= to " + ArrayUtil.MAX_ARRAY_LENGTH + " but was: " + arraySize); |
| 140 | + } |
| 141 | + if (arraySize < 0) { |
| 142 | + throw new NegativeArraySizeException("array size must be positive but was: " + arraySize); |
| 143 | + } |
| 144 | + return arraySize; |
| 145 | + } |
| 146 | + |
76 | 147 | @Override
|
77 | 148 | public String toString() {
|
78 | 149 | return this.name().toLowerCase(Locale.ROOT);
|
|
0 commit comments