Skip to content

Commit

Permalink
support geo encoder (#380)
Browse files Browse the repository at this point in the history 
* support geo encoder

* debug

* remove debug log

* let WKBWriter use machine byte order

* add null test for geo

* debug

* format code
  • Loading branch information
@jievince
jievince authored Nov 18, 2021
1 parent 4b55c93 commit f29e978
Show file tree
Hide file tree
Showing 9 changed files with 212 additions and 11 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 @@ -8,6 +8,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 @@ -201,11 +202,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 @@ -7,6 +7,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 @@ -33,5 +34,7 @@ public interface RowWriter {

void write(int index, DateTime v);

void write(int index, Geography v);

byte[] encodeStr();
}
108 changes: 107 additions & 1 deletion client/src/main/java/com/vesoft/nebula/encoder/RowWriterImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,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 @@ -15,6 +20,10 @@
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import org.apache.commons.codec.binary.Hex;
import org.locationtech.jts.geom.GeometryFactory;
import org.locationtech.jts.io.ByteOrderValues;
import org.locationtech.jts.io.WKBWriter;

public class RowWriterImpl implements RowWriter {
private final SchemaProviderImpl schema;
Expand Down Expand Up @@ -509,6 +518,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 @@ -625,6 +635,31 @@ 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(2, this.byteOrder == ByteOrder.BIG_ENDIAN
? ByteOrderValues.BIG_ENDIAN
: ByteOrderValues.LITTLE_ENDIAN)
.write(jtsGeom);
write(index, wkb);
}

@Override
public byte[] encodeStr() {
return buf.array();
Expand Down Expand Up @@ -666,6 +701,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 @@ -708,6 +745,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 @@ -790,6 +830,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 @@ -821,7 +864,8 @@ public ByteBuffer processOutOfSpace() {
if (typeEnum == null) {
throw new RuntimeException("Incorrect field type " + field.type());
}
if (typeEnum != PropertyType.STRING) {
if (typeEnum != PropertyType.STRING
&& typeEnum != PropertyType.GEOGRAPHY) {
continue;
}
int offset = headerLen + numNullBytes + field.offset();
Expand Down Expand Up @@ -867,4 +911,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 @@ -22,6 +22,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 @@ -26,6 +26,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 @@ -34,7 +35,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 @@ -43,6 +45,7 @@ public SchemaField(String name,
this.size = size;
this.offset = offset;
this.nullFlagPos = nullFlagPos;
this.geoShape = geoShape;
}

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

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

public SchemaProviderImpl(long ver) {
Expand Down Expand Up @@ -169,7 +177,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 @@ -190,7 +199,8 @@ public void addField(String name,
defaultValue,
size,
offset,
nullFlagPos));
nullFlagPos,
geoShape));
fieldNameIndex.put(name, fields.size() - 1);
}

Expand Down Expand Up @@ -241,6 +251,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
14 changes: 14 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 @@ -11,6 +11,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.PropertyType;
import com.vesoft.nebula.meta.Schema;
import com.vesoft.nebula.meta.SpaceDesc;
Expand Down Expand Up @@ -82,6 +83,19 @@ private Schema genNoDefaultVal() {
new ColumnTypeDef(PropertyType.INT32));
columnDef.setNullable(true);
columns.add(columnDef);
columnDef = new ColumnDef(("Col16").getBytes(),
new ColumnTypeDef(PropertyType.GEOGRAPHY, (short)0, GeoShape.POINT));
columns.add(columnDef);
columnDef = new ColumnDef(("Col17").getBytes(),
new ColumnTypeDef(PropertyType.GEOGRAPHY, (short)0, GeoShape.LINESTRING));
columns.add(columnDef);
columnDef = new ColumnDef(("Col18").getBytes(),
new ColumnTypeDef(PropertyType.GEOGRAPHY, (short)0, GeoShape.POLYGON));
columns.add(columnDef);
columnDef = new ColumnDef(("Col19").getBytes(),
new ColumnTypeDef(PropertyType.GEOGRAPHY, (short)0, GeoShape.ANY));
columnDef.setNullable(true);
columns.add(columnDef);
return new Schema(columns, null);
}

Expand Down
Loading

0 comments on commit f29e978

Please sign in to comment.