Skip to content

Commit

Permalink
ESQL: Move randomLiteral (elastic#111647)
Browse files Browse the repository at this point in the history
This moves `randomLiteral` from the function testing infrastructure into
`EsqlTestUtils` so it can be used a little more widely. That'll make
other refactoring a little easier later.
  • Loading branch information
nik9000 authored Aug 6, 2024
1 parent a78a21b commit 10b5117
Show file tree
Hide file tree
Showing 17 changed files with 91 additions and 65 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,12 @@

package org.elasticsearch.xpack.esql;

import org.apache.lucene.document.InetAddressPoint;
import org.apache.lucene.sandbox.document.HalfFloatPoint;
import org.apache.lucene.util.BytesRef;
import org.elasticsearch.common.breaker.CircuitBreaker;
import org.elasticsearch.common.breaker.NoopCircuitBreaker;
import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.regex.Regex;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.util.BigArrays;
Expand All @@ -22,8 +25,11 @@
import org.elasticsearch.core.PathUtils;
import org.elasticsearch.core.SuppressForbidden;
import org.elasticsearch.core.Tuple;
import org.elasticsearch.geo.GeometryTestUtils;
import org.elasticsearch.geo.ShapeTestUtils;
import org.elasticsearch.index.IndexMode;
import org.elasticsearch.test.ESTestCase;
import org.elasticsearch.xcontent.json.JsonXContent;
import org.elasticsearch.xpack.esql.action.EsqlQueryResponse;
import org.elasticsearch.xpack.esql.analysis.EnrichResolution;
import org.elasticsearch.xpack.esql.analysis.Verifier;
Expand Down Expand Up @@ -56,12 +62,14 @@
import org.elasticsearch.xpack.esql.session.EsqlConfiguration;
import org.elasticsearch.xpack.esql.stats.Metrics;
import org.elasticsearch.xpack.esql.stats.SearchStats;
import org.elasticsearch.xpack.versionfield.Version;
import org.junit.Assert;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.UncheckedIOException;
import java.net.URL;
import java.net.URLConnection;
import java.nio.charset.StandardCharsets;
Expand All @@ -71,6 +79,8 @@
import java.nio.file.Path;
import java.nio.file.SimpleFileVisitor;
import java.nio.file.attribute.BasicFileAttributes;
import java.time.Duration;
import java.time.Period;
import java.util.ArrayList;
import java.util.EnumSet;
import java.util.Iterator;
Expand All @@ -87,14 +97,27 @@
import static java.util.Collections.emptyList;
import static java.util.Collections.emptyMap;
import static java.util.Collections.unmodifiableMap;
import static org.elasticsearch.test.ESTestCase.between;
import static org.elasticsearch.test.ESTestCase.randomAlphaOfLength;
import static org.elasticsearch.test.ESTestCase.randomBoolean;
import static org.elasticsearch.test.ESTestCase.randomByte;
import static org.elasticsearch.test.ESTestCase.randomDouble;
import static org.elasticsearch.test.ESTestCase.randomFloat;
import static org.elasticsearch.test.ESTestCase.randomFrom;
import static org.elasticsearch.test.ESTestCase.randomInt;
import static org.elasticsearch.test.ESTestCase.randomIntBetween;
import static org.elasticsearch.test.ESTestCase.randomIp;
import static org.elasticsearch.test.ESTestCase.randomLong;
import static org.elasticsearch.test.ESTestCase.randomLongBetween;
import static org.elasticsearch.test.ESTestCase.randomMillisUpToYear9999;
import static org.elasticsearch.test.ESTestCase.randomShort;
import static org.elasticsearch.test.ESTestCase.randomZone;
import static org.elasticsearch.test.ListMatcher.matchesList;
import static org.elasticsearch.test.MapMatcher.assertMap;
import static org.elasticsearch.xpack.esql.core.tree.Source.EMPTY;
import static org.elasticsearch.xpack.esql.core.type.DataType.INTEGER;
import static org.elasticsearch.xpack.esql.core.util.SpatialCoordinateTypes.CARTESIAN;
import static org.elasticsearch.xpack.esql.core.util.SpatialCoordinateTypes.GEO;
import static org.hamcrest.Matchers.instanceOf;
import static org.junit.Assert.assertTrue;

Expand Down Expand Up @@ -600,4 +623,53 @@ public static Tuple<String, String> pathAndName(String string) {
return new Tuple<>(folder, file);
}

/**
* Generate a random value of the appropriate type to fit into blocks of {@code e}.
*/
public static Literal randomLiteral(DataType type) {
return new Literal(Source.EMPTY, switch (type) {
case BOOLEAN -> randomBoolean();
case BYTE -> randomByte();
case SHORT -> randomShort();
case INTEGER, COUNTER_INTEGER -> randomInt();
case UNSIGNED_LONG, LONG, COUNTER_LONG -> randomLong();
case DATE_PERIOD -> Period.of(randomIntBetween(-1000, 1000), randomIntBetween(-13, 13), randomIntBetween(-32, 32));
case DATETIME -> randomMillisUpToYear9999();
case DOUBLE, SCALED_FLOAT, COUNTER_DOUBLE -> randomDouble();
case FLOAT -> randomFloat();
case HALF_FLOAT -> HalfFloatPoint.sortableShortToHalfFloat(HalfFloatPoint.halfFloatToSortableShort(randomFloat()));
case KEYWORD -> new BytesRef(randomAlphaOfLength(5));
case IP -> new BytesRef(InetAddressPoint.encode(randomIp(randomBoolean())));
case TIME_DURATION -> Duration.ofMillis(randomLongBetween(-604800000L, 604800000L)); // plus/minus 7 days
case TEXT -> new BytesRef(randomAlphaOfLength(50));
case VERSION -> randomVersion().toBytesRef();
case GEO_POINT -> GEO.asWkb(GeometryTestUtils.randomPoint());
case CARTESIAN_POINT -> CARTESIAN.asWkb(ShapeTestUtils.randomPoint());
case GEO_SHAPE -> GEO.asWkb(GeometryTestUtils.randomGeometry(randomBoolean()));
case CARTESIAN_SHAPE -> CARTESIAN.asWkb(ShapeTestUtils.randomGeometry(randomBoolean()));
case NULL -> null;
case SOURCE -> {
try {
yield BytesReference.bytes(
JsonXContent.contentBuilder().startObject().field(randomAlphaOfLength(3), randomAlphaOfLength(10)).endObject()
).toBytesRef();
} catch (IOException e) {
throw new UncheckedIOException(e);
}
}
case UNSUPPORTED, OBJECT, DOC_DATA_TYPE, TSID_DATA_TYPE, PARTIAL_AGG -> throw new IllegalArgumentException(
"can't make random values for [" + type.typeName() + "]"
);
}, type);
}

static Version randomVersion() {
// TODO degenerate versions and stuff
return switch (between(0, 2)) {
case 0 -> new Version(Integer.toString(between(0, 100)));
case 1 -> new Version(between(0, 100) + "." + between(0, 100));
case 2 -> new Version(between(0, 100) + "." + between(0, 100) + "." + between(0, 100));
default -> throw new IllegalArgumentException();
};
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,9 @@
import com.carrotsearch.randomizedtesting.ClassModel;
import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;

import org.apache.lucene.document.InetAddressPoint;
import org.apache.lucene.sandbox.document.HalfFloatPoint;
import org.apache.lucene.util.BytesRef;
import org.elasticsearch.common.Strings;
import org.elasticsearch.common.breaker.CircuitBreaker;
import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.common.util.BigArrays;
import org.elasticsearch.common.util.MockBigArrays;
Expand All @@ -27,8 +24,6 @@
import org.elasticsearch.compute.operator.DriverContext;
import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
import org.elasticsearch.core.PathUtils;
import org.elasticsearch.geo.GeometryTestUtils;
import org.elasticsearch.geo.ShapeTestUtils;
import org.elasticsearch.indices.CrankyCircuitBreakerService;
import org.elasticsearch.logging.LogManager;
import org.elasticsearch.logging.Logger;
Expand Down Expand Up @@ -71,21 +66,17 @@
import org.elasticsearch.xpack.esql.parser.ExpressionBuilder;
import org.elasticsearch.xpack.esql.planner.Layout;
import org.elasticsearch.xpack.esql.planner.PlannerUtils;
import org.elasticsearch.xpack.versionfield.Version;
import org.hamcrest.Matcher;
import org.junit.After;
import org.junit.AfterClass;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.math.BigInteger;
import java.nio.file.Files;
import java.nio.file.Path;
import java.time.Duration;
import java.time.Period;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
Expand All @@ -104,9 +95,8 @@

import static java.util.Map.entry;
import static org.elasticsearch.compute.data.BlockUtils.toJavaObject;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
import static org.elasticsearch.xpack.esql.SerializationTestUtils.assertSerialization;
import static org.elasticsearch.xpack.esql.core.util.SpatialCoordinateTypes.CARTESIAN;
import static org.elasticsearch.xpack.esql.core.util.SpatialCoordinateTypes.GEO;
import static org.hamcrest.Matchers.either;
import static org.hamcrest.Matchers.endsWith;
import static org.hamcrest.Matchers.equalTo;
Expand Down Expand Up @@ -141,46 +131,6 @@ public abstract class AbstractFunctionTestCase extends ESTestCase {
entry("is_not_null", IsNotNull.class)
);

/**
* Generate a random value of the appropriate type to fit into blocks of {@code e}.
*/
public static Literal randomLiteral(DataType type) {
return new Literal(Source.EMPTY, switch (type) {
case BOOLEAN -> randomBoolean();
case BYTE -> randomByte();
case SHORT -> randomShort();
case INTEGER, COUNTER_INTEGER -> randomInt();
case UNSIGNED_LONG, LONG, COUNTER_LONG -> randomLong();
case DATE_PERIOD -> Period.of(randomIntBetween(-1000, 1000), randomIntBetween(-13, 13), randomIntBetween(-32, 32));
case DATETIME -> randomMillisUpToYear9999();
case DOUBLE, SCALED_FLOAT, COUNTER_DOUBLE -> randomDouble();
case FLOAT -> randomFloat();
case HALF_FLOAT -> HalfFloatPoint.sortableShortToHalfFloat(HalfFloatPoint.halfFloatToSortableShort(randomFloat()));
case KEYWORD -> new BytesRef(randomAlphaOfLength(5));
case IP -> new BytesRef(InetAddressPoint.encode(randomIp(randomBoolean())));
case TIME_DURATION -> Duration.ofMillis(randomLongBetween(-604800000L, 604800000L)); // plus/minus 7 days
case TEXT -> new BytesRef(randomAlphaOfLength(50));
case VERSION -> randomVersion().toBytesRef();
case GEO_POINT -> GEO.asWkb(GeometryTestUtils.randomPoint());
case CARTESIAN_POINT -> CARTESIAN.asWkb(ShapeTestUtils.randomPoint());
case GEO_SHAPE -> GEO.asWkb(GeometryTestUtils.randomGeometry(randomBoolean()));
case CARTESIAN_SHAPE -> CARTESIAN.asWkb(ShapeTestUtils.randomGeometry(randomBoolean()));
case NULL -> null;
case SOURCE -> {
try {
yield BytesReference.bytes(
JsonXContent.contentBuilder().startObject().field(randomAlphaOfLength(3), randomAlphaOfLength(10)).endObject()
).toBytesRef();
} catch (IOException e) {
throw new UncheckedIOException(e);
}
}
case UNSUPPORTED, OBJECT, DOC_DATA_TYPE, TSID_DATA_TYPE, PARTIAL_AGG -> throw new IllegalArgumentException(
"can't make random values for [" + type.typeName() + "]"
);
}, type);
}

protected TestCaseSupplier.TestCase testCase;

/**
Expand Down Expand Up @@ -1294,16 +1244,6 @@ public void allMemoryReleased() {
}
}

static Version randomVersion() {
// TODO degenerate versions and stuff
return switch (between(0, 2)) {
case 0 -> new Version(Integer.toString(between(0, 100)));
case 1 -> new Version(between(0, 100) + "." + between(0, 100));
case 2 -> new Version(between(0, 100) + "." + between(0, 100) + "." + between(0, 100));
default -> throw new IllegalArgumentException();
};
}

/**
* All string types (keyword, text, match_only_text, etc).
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import java.util.stream.Stream;

import static org.elasticsearch.compute.data.BlockUtils.toJavaObject;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.nullValue;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import java.util.List;
import java.util.function.Supplier;

import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
import static org.hamcrest.Matchers.equalTo;

@FunctionName("to_base64")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
import java.util.stream.LongStream;
import java.util.stream.Stream;

import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
import static org.elasticsearch.xpack.esql.core.util.SpatialCoordinateTypes.CARTESIAN;
import static org.elasticsearch.xpack.esql.core.util.SpatialCoordinateTypes.GEO;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import java.util.List;
import java.util.function.Supplier;

import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
import static org.elasticsearch.xpack.esql.core.util.SpatialCoordinateTypes.CARTESIAN;
import static org.elasticsearch.xpack.esql.core.util.SpatialCoordinateTypes.GEO;
import static org.hamcrest.Matchers.equalTo;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import java.util.List;
import java.util.function.Supplier;

import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
import static org.elasticsearch.xpack.esql.core.util.SpatialCoordinateTypes.CARTESIAN;
import static org.elasticsearch.xpack.esql.core.util.SpatialCoordinateTypes.GEO;
import static org.hamcrest.Matchers.equalTo;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import java.util.List;
import java.util.function.Supplier;

import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.nullValue;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
import java.util.function.Supplier;

import static org.elasticsearch.compute.data.BlockUtils.toJavaObject;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
import static org.hamcrest.Matchers.equalTo;

public class CoalesceTests extends AbstractScalarFunctionTestCase {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import java.util.List;
import java.util.function.Supplier;

import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
import static org.hamcrest.Matchers.equalTo;

public class IsNotNullTests extends AbstractScalarFunctionTestCase {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import java.util.List;
import java.util.function.Supplier;

import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
import static org.hamcrest.Matchers.equalTo;

public class IsNullTests extends AbstractScalarFunctionTestCase {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import java.util.function.Function;
import java.util.function.Supplier;

import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.nullValue;
import static org.hamcrest.Matchers.startsWith;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import java.util.function.Supplier;

import static org.elasticsearch.compute.data.BlockUtils.toJavaObject;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
import static org.hamcrest.Matchers.equalTo;

public class NegTests extends AbstractScalarFunctionTestCase {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import java.util.List;
import java.util.function.Supplier;

import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
import static org.elasticsearch.xpack.esql.core.type.DateUtils.asDateTime;
import static org.elasticsearch.xpack.esql.core.type.DateUtils.asMillis;
import static org.elasticsearch.xpack.esql.core.util.NumericUtils.ZERO_AS_UNSIGNED_LONG;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import java.util.function.Supplier;

import static org.elasticsearch.xpack.esql.EsqlTestUtils.of;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
import static org.elasticsearch.xpack.esql.core.expression.Literal.NULL;
import static org.elasticsearch.xpack.esql.core.tree.Source.EMPTY;
import static org.elasticsearch.xpack.esql.core.type.DataType.CARTESIAN_POINT;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,13 @@
import org.elasticsearch.xpack.esql.TestBlockFactory;
import org.elasticsearch.xpack.esql.core.expression.Attribute;
import org.elasticsearch.xpack.esql.core.tree.Source;
import org.elasticsearch.xpack.esql.expression.function.AbstractFunctionTestCase;
import org.elasticsearch.xpack.esql.plan.logical.AbstractLogicalPlanSerializationTests;

import java.io.IOException;
import java.util.List;

import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;

public class LocalRelationSerializationTests extends AbstractLogicalPlanSerializationTests<LocalRelation> {
public static LocalRelation randomLocalRelation() {
Source source = randomSource();
Expand All @@ -31,7 +32,7 @@ private static LocalSupplier randomLocalSupplier(List<Attribute> attributes) {
for (int b = 0; b < blocks.length; b++) {
blocks[b] = BlockUtils.constantBlock(
TestBlockFactory.getNonBreakingInstance(),
AbstractFunctionTestCase.randomLiteral(attributes.get(b).dataType()).value(),
randomLiteral(attributes.get(b).dataType()).value(),
1
);
}
Expand Down
Loading

0 comments on commit 10b5117

Please sign in to comment.