Add generic Set support to streams (#54769) (#55123)

This commit adds support for reading and writing sets as generic values
in stream input and output.

closes #54708
This commit is contained in:
Ryan Ernst 2020-04-16 14:29:38 -07:00 committed by GitHub
parent 22c55180c1
commit 62246aa9c9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 46 additions and 0 deletions

View File

@ -70,6 +70,7 @@ import java.util.EnumSet;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
@ -735,6 +736,10 @@ public abstract class StreamInput extends InputStream {
return readGeoPoint();
case 23:
return readZonedDateTime();
case 24:
return readCollection(StreamInput::readGenericValue, LinkedHashSet::new, Collections.emptySet());
case 25:
return readCollection(StreamInput::readGenericValue, HashSet::new, Collections.emptySet());
default:
throw new IOException("Can't read unknown type [" + type + "]");
}

View File

@ -70,6 +70,7 @@ import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
@ -776,6 +777,14 @@ public abstract class StreamOutput extends OutputStream {
o.writeString(zoneId.equals("Z") ? DateTimeZone.UTC.getID() : zoneId);
o.writeLong(zonedDateTime.toInstant().toEpochMilli());
});
writers.put(Set.class, (o, v) -> {
if (v instanceof LinkedHashSet) {
o.writeByte((byte) 24);
} else {
o.writeByte((byte) 25);
}
o.writeCollection((Set<?>) v, StreamOutput::writeGenericValue);
});
WRITERS = Collections.unmodifiableMap(writers);
}
@ -797,6 +806,8 @@ public abstract class StreamOutput extends OutputStream {
type = Object[].class;
} else if (value instanceof Map) {
type = Map.class;
} else if (value instanceof Set) {
type = Set.class;
} else if (value instanceof ReadableInstant) {
type = ReadableInstant.class;
} else if (value instanceof BytesReference) {

View File

@ -21,6 +21,7 @@ package org.elasticsearch.common.io.stream;
import org.apache.lucene.util.BytesRef;
import org.elasticsearch.common.CheckedBiConsumer;
import org.elasticsearch.common.CheckedConsumer;
import org.elasticsearch.common.CheckedFunction;
import org.elasticsearch.common.bytes.BytesArray;
import org.elasticsearch.common.bytes.BytesReference;
@ -38,6 +39,7 @@ import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
@ -434,4 +436,32 @@ public class StreamTests extends ESTestCase {
}
}
public void testGenericSet() throws IOException {
Set<String> set = new HashSet<>(Arrays.asList("a", "b", "c", "d", "e"));
assertGenericRoundtrip(set);
// reverse order in normal set so linked hashset does not match the order
List<String> list = new ArrayList<>(set);
Collections.reverse(list);
assertGenericRoundtrip(new LinkedHashSet<>(list));
}
private void assertSerialization(CheckedConsumer<StreamOutput, IOException> outputAssertions,
CheckedConsumer<StreamInput, IOException> inputAssertions) throws IOException {
try (BytesStreamOutput output = new BytesStreamOutput()) {
outputAssertions.accept(output);
final BytesReference bytesReference = output.bytes();
final StreamInput input = bytesReference.streamInput();
inputAssertions.accept(input);
}
}
private void assertGenericRoundtrip(Object original) throws IOException {
assertSerialization(output -> {
output.writeGenericValue(original);
}, input -> {
Object read = input.readGenericValue();
assertThat(read, equalTo(original));
});
}
}