Skip to content

Commit

Permalink
support geo encoder
Browse files Browse the repository at this point in the history
  • Loading branch information
@jievince
jievince committed Nov 9, 2021
1 parent 3d56b9c commit 8225a5b
Show file tree
Hide file tree
Showing 9 changed files with 199 additions and 7 deletions.
6 changes: 6 additions & 0 deletions .gitignore
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,14 @@ target/
.idea/
.eclipse/
*.iml
.vscode/
.settings
.project
client/.classpath

spark-importer.ipr
spark-importer.iws

.DS_Store

examples/
5 changes: 4 additions & 1 deletion client/src/main/java/com/vesoft/nebula/encoder/NebulaCodecImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import com.vesoft.nebula.meta.ColumnDef;
import com.vesoft.nebula.meta.ColumnTypeDef;
import com.vesoft.nebula.meta.EdgeItem;
import com.vesoft.nebula.meta.GeoShape;
import com.vesoft.nebula.meta.Schema;
import com.vesoft.nebula.meta.TagItem;
import java.nio.ByteBuffer;
Expand Down Expand Up @@ -202,11 +203,13 @@ private SchemaProviderImpl genSchemaProvider(long ver, Schema schema) {
boolean nullable = col.isSetNullable() && col.isNullable();
boolean hasDefault = col.isSetDefault_value();
int len = type.isSetType_length() ? type.getType_length() : 0;
GeoShape geoShape = type.isSetGeo_shape() ? type.getGeo_shape() : GeoShape.ANY;
schemaProvider.addField(new String(col.getName()),
type.type.getValue(),
len,
nullable,
hasDefault ? col.getDefault_value() : null);
hasDefault ? col.getDefault_value() : null,
geoShape.getValue());
}
return schemaProvider;
}
Expand Down
3 changes: 3 additions & 0 deletions client/src/main/java/com/vesoft/nebula/encoder/RowWriter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

import com.vesoft.nebula.Date;
import com.vesoft.nebula.DateTime;
import com.vesoft.nebula.Geography;
import com.vesoft.nebula.Time;

public interface RowWriter {
Expand All @@ -34,5 +35,7 @@ public interface RowWriter {

void write(int index, DateTime v);

void write(int index, Geography v);

byte[] encodeStr();
}
99 changes: 99 additions & 0 deletions client/src/main/java/com/vesoft/nebula/encoder/RowWriterImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,13 @@

package com.vesoft.nebula.encoder;

import com.vesoft.nebula.Coordinate;
import com.vesoft.nebula.Date;
import com.vesoft.nebula.DateTime;
import com.vesoft.nebula.Geography;
import com.vesoft.nebula.LineString;
import com.vesoft.nebula.Point;
import com.vesoft.nebula.Polygon;
import com.vesoft.nebula.Time;
import com.vesoft.nebula.Value;
import com.vesoft.nebula.meta.PropertyType;
Expand All @@ -16,6 +21,8 @@
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import org.locationtech.jts.geom.GeometryFactory;
import org.locationtech.jts.io.WKBWriter;

public class RowWriterImpl implements RowWriter {
private final SchemaProviderImpl schema;
Expand Down Expand Up @@ -510,6 +517,7 @@ public void write(int index, byte[] v) {
}
int offset = headerLen + numNullBytes + field.offset();
switch (typeEnum) {
case GEOGRAPHY:
case STRING: {
strList.add(v);
outOfSpaceStr = true;
Expand Down Expand Up @@ -626,6 +634,27 @@ public void write(int index, DateTime v) {
isSet.set(index, true);
}

@Override
public void write(int index, Geography v) {
SchemaProvider.Field field = schema.field(index);
PropertyType typeEnum = PropertyType.findByValue(field.type());
if (typeEnum == null) {
throw new RuntimeException("Incorrect field type " + field.type());
}
if (typeEnum == PropertyType.GEOGRAPHY) {
if (field.geoShape() != 0 && field.geoShape() != v.getSetField()) {
throw new RuntimeException("Incorrect geo shape, expect "
+ field.geoShape() + ", got "
+ v.getSetField());
}
} else {
throw new RuntimeException("Value: " + v + "'s type is unexpected");
}
org.locationtech.jts.geom.Geometry jtsGeom = convertGeographyToJTSGeometry(v);
byte[] wkb = new org.locationtech.jts.io.WKBWriter().write(jtsGeom);
write(index, wkb);
}

@Override
public byte[] encodeStr() {
return buf.array();
Expand Down Expand Up @@ -667,6 +696,8 @@ public void setValue(int index, Object value) {
write(index, (Date)value);
} else if (value instanceof DateTime) {
write(index, (DateTime)value);
} else if (value instanceof Geography) {
write(index, (Geography)value);
} else {
throw new RuntimeException("Unsupported value object `" + value.getClass() + "\"");
}
Expand Down Expand Up @@ -709,6 +740,9 @@ public void setValue(int index, Value value) {
case Value.DTVAL:
write(index, value.getDtVal());
break;
case Value.GGVAL:
write(index, value.getGgVal());
break;
default:
throw new RuntimeException(
"Unknown value: " + value.getFieldValue().getClass()
Expand Down Expand Up @@ -791,6 +825,9 @@ public void checkUnsetFields() {
// case Value.DTVAL:
// write(i, defVal.getDtVal());
// break;
// case Value.GGVAL:
// write(i, defVal.getGgVal());
// break;
// default:
// throw new RuntimeException("Unsupported default value type");
// }
Expand Down Expand Up @@ -868,4 +905,66 @@ private long getTimestamp() {
long nanoTime = System.nanoTime();
return curTime + (nanoTime - nanoTime / 1000000 * 1000000) / 1000;
}

public org.locationtech.jts.geom.Geometry
convertGeographyToJTSGeometry(Geography geog) {
GeometryFactory geomFactory = new GeometryFactory();
switch (geog.getSetField()) {
case Geography.PTVAL: {
Point point = geog.getPtVal();
Coordinate coord = point.getCoord();
return geomFactory.createPoint(
new org.locationtech.jts.geom.Coordinate(coord.x, coord.y));
}
case Geography.LSVAL: {
LineString line = geog.getLsVal();
List<Coordinate> coordList = line.getCoordList();

List<org.locationtech.jts.geom.Coordinate> jtsCoordList =
new ArrayList<>();
for (int i = 0; i < coordList.size(); ++i) {
jtsCoordList.add(new org.locationtech.jts.geom.Coordinate(
coordList.get(i).x, coordList.get(i).y));
}
org.locationtech.jts.geom.Coordinate[] jtsCoordArray =
new org.locationtech.jts.geom.Coordinate[jtsCoordList.size()];
return geomFactory.createLineString(
jtsCoordList.toArray(jtsCoordArray));
}
case Geography.PGVAL: {
Polygon polygon = geog.getPgVal();
List<List<Coordinate>> coordListList = polygon.getCoordListList();
if (coordListList.isEmpty()) {
throw new RuntimeException("Polygon must at least contain one loop");
}

List<org.locationtech.jts.geom.LinearRing> rings = new ArrayList<>();
for (int i = 0; i < coordListList.size(); ++i) {
List<Coordinate> coordList = coordListList.get(i);
List<org.locationtech.jts.geom.Coordinate> jtsCoordList =
new ArrayList<>();
for (int j = 0; j < coordList.size(); ++j) {
jtsCoordList.add(new org.locationtech.jts.geom.Coordinate(
coordList.get(j).x, coordList.get(j).y));
}
org.locationtech.jts.geom.Coordinate[] jtsCoordArray =
new org.locationtech.jts.geom.Coordinate[jtsCoordList.size()];
rings.add(geomFactory.createLinearRing(
jtsCoordList.toArray(jtsCoordArray)));
}
org.locationtech.jts.geom.LinearRing shell = rings.get(0);
if (rings.size() == 1) {
return geomFactory.createPolygon(shell);
} else {
rings.remove(0);
org.locationtech.jts.geom.LinearRing[] holesArrary =
new org.locationtech.jts.geom.LinearRing[rings.size() - 1];
return geomFactory.createPolygon(shell, rings.toArray(holesArrary));
}
}
default:
throw new RuntimeException("Unknown geography: "
+ geog.getFieldValue().getClass());
}
}
}
2 changes: 2 additions & 0 deletions client/src/main/java/com/vesoft/nebula/encoder/SchemaProvider.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@ public interface Field {
public int offset();

public int nullFlagPos();

public int geoShape();
}

public long getVersion();
Expand Down
18 changes: 15 additions & 3 deletions client/src/main/java/com/vesoft/nebula/encoder/SchemaProviderImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ static class SchemaField implements Field {
private final int size;
private final int offset;
private final int nullFlagPos;
private final int geoShape;

public SchemaField(String name,
int type,
Expand All @@ -35,7 +36,8 @@ public SchemaField(String name,
byte[] defaultValue,
int size,
int offset,
int nullFlagPos) {
int nullFlagPos,
int geoShape) {
this.name = name;
this.type = type;
this.nullable = nullable;
Expand All @@ -44,6 +46,7 @@ public SchemaField(String name,
this.size = size;
this.offset = offset;
this.nullFlagPos = nullFlagPos;
this.geoShape = geoShape;
}

@Override
Expand Down Expand Up @@ -85,6 +88,11 @@ public int offset() {
public int nullFlagPos() {
return nullFlagPos;
}

@Override
public int geoShape() {
return geoShape;
}
}

public SchemaProviderImpl(long ver) {
Expand Down Expand Up @@ -170,7 +178,8 @@ public void addField(String name,
int type,
int fixedStrLen,
boolean nullable,
byte[] defaultValue) {
byte[] defaultValue,
int geoShape) {
int size = fieldSize(type, fixedStrLen);

int offset = 0;
Expand All @@ -191,7 +200,8 @@ public void addField(String name,
defaultValue,
size,
offset,
nullFlagPos));
nullFlagPos,
geoShape));
fieldNameIndex.put(name, fields.size() - 1);
}

Expand Down Expand Up @@ -242,6 +252,8 @@ public int fieldSize(int type, int fixedStrLimit) {
+ Byte.BYTES // minute
+ Byte.BYTES // sec
+ Integer.BYTES; // microsec
case GEOGRAPHY:
return 8; // wkb offset + wkb length
default:
throw new RuntimeException("Incorrect field type " + type);
}
Expand Down
9 changes: 9 additions & 0 deletions client/src/test/java/com/vesoft/nebula/encoder/MetaCacheImplTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,15 @@ private Schema genNoDefaultVal() {
new ColumnTypeDef(PropertyType.INT32));
columnDef.setNullable(true);
columns.add(columnDef);
columnDef = new ColumnDef(("Col16").getBytes(),
new ColumnTypeDef(PropertyType.GEOGRAPHY));
columns.add(columnDef);
columnDef = new ColumnDef(("Col17").getBytes(),
new ColumnTypeDef(PropertyType.GEOGRAPHY));
columns.add(columnDef);
columnDef = new ColumnDef(("Col18").getBytes(),
new ColumnTypeDef(PropertyType.GEOGRAPHY));
columns.add(columnDef);
return new Schema(columns, null);
}

Expand Down
54 changes: 51 additions & 3 deletions client/src/test/java/com/vesoft/nebula/encoder/TestEncoder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,15 @@

package test.java.com.vesoft.nebula.encoder;

import com.vesoft.nebula.Coordinate;
import com.vesoft.nebula.Date;
import com.vesoft.nebula.DateTime;
import com.vesoft.nebula.Geography;
import com.vesoft.nebula.HostAddr;
import com.vesoft.nebula.LineString;
import com.vesoft.nebula.NullType;
import com.vesoft.nebula.Point;
import com.vesoft.nebula.Polygon;
import com.vesoft.nebula.Time;
import com.vesoft.nebula.Value;
import com.vesoft.nebula.encoder.MetaCacheImplTest;
Expand All @@ -35,11 +40,19 @@ public class TestEncoder {
final String allTypeValueExpectResult = "090cc001081000200000004000000000000000db0f494069"
+ "57148b0abf05405d0000000c0000004e6562756c61204772617068bb334e5e000000"
+ "00e40702140a1e2d00000000e40702140a1e2d00000000000000000000000000000000"
+ "48656c6c6f20776f726c6421";
+ "48656c6c6f20776f726c6421010100000000000000006066409a999999997956400102"
+ "000000030000000000000000000000000000000000f03f000000000000f03f00000000"
+ "0000004000000000000008400000000000001c4001030000000200000004000000cdcc"
+ "cccccc2c5bc0000000000080414000000000000059c00000000000404740cdcccccccc"
+ "ac56c03333333333734140cdcccccccc2c5bc000000000008041400400000066666666"
+ "660659c03333333333b344409a99999999b959c0cdcccccccccc424033333333333358"
+ "c00000000000c0424066666666660659c03333333333b34440fe148c1853d00500";

private List<String> getCols() {
return Arrays.asList("Col01","Col02", "Col03", "Col04", "Col05", "Col06",
"Col07","Col08", "Col09", "Col10", "Col11", "Col12", "Col13", "Col14");
"Col07","Col08", "Col09", "Col10", "Col11", "Col12", "Col13", "Col14",
// Purposely skip the col15
"Col16", "Col17", "Col18");
}

private List<Object> getValues() {
Expand All @@ -62,8 +75,39 @@ private List<Object> getValues() {
dateValue.setDVal(new Date((short)2020, (byte)2, (byte)20));
final Value nullVal = new Value();
nullVal.setNVal(NullType.__NULL__);
// geograph point
// POINT(179.0 89.9)
final Value geogPointVal = new Value();
geogPointVal.setGgVal(Geography.ptVal(new Point(new Coordinate(179.0, 89.9))));
// geography linestring
// LINESTRING(0 1, 1 2, 3 7)
final Value geogLineStringVal = new Value();
List<Coordinate> line = new ArrayList<>();
line.add(new Coordinate(0, 1));
line.add(new Coordinate(1, 2));
line.add(new Coordinate(3, 7));
geogLineStringVal.setGgVal(Geography.lsVal(new LineString(line)));
// geography polygon
// POLYGON((-108.7 35.0, -100.0 46.5, -90.7 34.9, -108.7 35.0),
// (-100.1 41.4, -102.9 37.6, -96.8 37.5, -100.1 41.4))
final Value geogPolygonVal = new Value();
List<Coordinate> shell = new ArrayList<>();
shell.add(new Coordinate(-108.7, 35.0));
shell.add(new Coordinate(-100.0, 46.5));
shell.add(new Coordinate(-90.7, 34.9));
shell.add(new Coordinate(-108.7, 35.0));
List<Coordinate> hole = new ArrayList<>();
hole.add(new Coordinate(-100.1, 41.4));
hole.add(new Coordinate(-102.9, 37.6));
hole.add(new Coordinate(-96.8, 37.5));
hole.add(new Coordinate(-100.1, 41.4));
List<List<Coordinate>> rings = new ArrayList<List<Coordinate>>();
rings.add(shell);
rings.add(hole);
geogPolygonVal.setGgVal(Geography.pgVal(new Polygon(rings)));
return Arrays.asList(true, 8, 16, 32, intVal, pi, e, strVal, fixVal,
timestampVal, dateValue, timeVal, datetimeValue, nullVal);
timestampVal, dateValue, timeVal, datetimeValue, nullVal,
geogPointVal, geogLineStringVal, geogPolygonVal);
}

public int getSpaceVidLen(String spaceName) throws RuntimeException {
Expand Down Expand Up @@ -228,6 +272,10 @@ public void testEncodeVertexValue() {
// File file = new File("encode_java.txt");
// FileOutputStream fileOutputStream = new FileOutputStream(file);
// fileOutputStream.write(encodeStr);
System.out.println("hex vs expect");
System.out.println(hexStr);
System.out.println("***********");
System.out.println(allTypeValueExpectResult);
Assert.assertArrayEquals(allTypeValueExpectResult.getBytes(),
hexStr.substring(0, hexStr.length() - 16).getBytes());
} catch (Exception exception) {
Expand Down
Loading

0 comments on commit 8225a5b

Please sign in to comment.