Skip to content

Commit 41617d2

Browse files
jievinceNicole00
jievince
authored and
Nicole00
committed
support geo encoder (#380)
* support geo encoder * debug * remove debug log * let WKBWriter use machine byte order * add null test for geo * debug * format code
1 parent f4cba88 commit 41617d2

File tree

9 files changed

+212
-11
lines changed

9 files changed

+212
-11
lines changed

‎.gitignore‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,14 @@ target/
2929
.idea/
3030
.eclipse/
3131
*.iml
32+
.vscode/
33+
.settings
34+
.project
35+
client/.classpath
3236

3337
spark-importer.ipr
3438
spark-importer.iws
3539

3640
.DS_Store
41+
42+
examples/

‎client/src/main/java/com/vesoft/nebula/encoder/NebulaCodecImpl.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import com.vesoft.nebula.meta.ColumnDef;
99
import com.vesoft.nebula.meta.ColumnTypeDef;
1010
import com.vesoft.nebula.meta.EdgeItem;
11+
import com.vesoft.nebula.meta.GeoShape;
1112
import com.vesoft.nebula.meta.Schema;
1213
import com.vesoft.nebula.meta.TagItem;
1314
import java.nio.ByteBuffer;
@@ -201,11 +202,13 @@ private SchemaProviderImpl genSchemaProvider(long ver, Schema schema) {
201202
boolean nullable = col.isSetNullable() && col.isNullable();
202203
boolean hasDefault = col.isSetDefault_value();
203204
int len = type.isSetType_length() ? type.getType_length() : 0;
205+
GeoShape geoShape = type.isSetGeo_shape() ? type.getGeo_shape() : GeoShape.ANY;
204206
schemaProvider.addField(new String(col.getName()),
205207
type.type.getValue(),
206208
len,
207209
nullable,
208-
hasDefault ? col.getDefault_value() : null);
210+
hasDefault ? col.getDefault_value() : null,
211+
geoShape.getValue());
209212
}
210213
return schemaProvider;
211214
}

‎client/src/main/java/com/vesoft/nebula/encoder/RowWriter.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
import com.vesoft.nebula.Date;
99
import com.vesoft.nebula.DateTime;
10+
import com.vesoft.nebula.Geography;
1011
import com.vesoft.nebula.Time;
1112

1213
public interface RowWriter {
@@ -33,5 +34,7 @@ public interface RowWriter {
3334

3435
void write(int index, DateTime v);
3536

37+
void write(int index, Geography v);
38+
3639
byte[] encodeStr();
3740
}

‎client/src/main/java/com/vesoft/nebula/encoder/RowWriterImpl.java‎

Lines changed: 107 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,13 @@
55

66
package com.vesoft.nebula.encoder;
77

8+
import com.vesoft.nebula.Coordinate;
89
import com.vesoft.nebula.Date;
910
import com.vesoft.nebula.DateTime;
11+
import com.vesoft.nebula.Geography;
12+
import com.vesoft.nebula.LineString;
13+
import com.vesoft.nebula.Point;
14+
import com.vesoft.nebula.Polygon;
1015
import com.vesoft.nebula.Time;
1116
import com.vesoft.nebula.Value;
1217
import com.vesoft.nebula.meta.PropertyType;
@@ -15,6 +20,10 @@
1520
import java.util.ArrayList;
1621
import java.util.Collections;
1722
import java.util.List;
23+
import org.apache.commons.codec.binary.Hex;
24+
import org.locationtech.jts.geom.GeometryFactory;
25+
import org.locationtech.jts.io.ByteOrderValues;
26+
import org.locationtech.jts.io.WKBWriter;
1827

1928
public class RowWriterImpl implements RowWriter {
2029
private final SchemaProviderImpl schema;
@@ -509,6 +518,7 @@ public void write(int index, byte[] v) {
509518
}
510519
int offset = headerLen + numNullBytes + field.offset();
511520
switch (typeEnum) {
521+
case GEOGRAPHY:
512522
case STRING: {
513523
strList.add(v);
514524
outOfSpaceStr = true;
@@ -625,6 +635,31 @@ public void write(int index, DateTime v) {
625635
isSet.set(index, true);
626636
}
627637

638+
@Override
639+
public void write(int index, Geography v) {
640+
SchemaProvider.Field field = schema.field(index);
641+
PropertyType typeEnum = PropertyType.findByValue(field.type());
642+
if (typeEnum == null) {
643+
throw new RuntimeException("Incorrect field type " + field.type());
644+
}
645+
if (typeEnum == PropertyType.GEOGRAPHY) {
646+
if (field.geoShape() != 0 && field.geoShape() != v.getSetField()) {
647+
throw new RuntimeException("Incorrect geo shape, expect "
648+
+ field.geoShape() + ", got "
649+
+ v.getSetField());
650+
}
651+
} else {
652+
throw new RuntimeException("Value: " + v + "'s type is unexpected");
653+
}
654+
org.locationtech.jts.geom.Geometry jtsGeom = convertGeographyToJTSGeometry(v);
655+
byte[] wkb = new org.locationtech.jts.io
656+
.WKBWriter(2, this.byteOrder == ByteOrder.BIG_ENDIAN
657+
? ByteOrderValues.BIG_ENDIAN
658+
: ByteOrderValues.LITTLE_ENDIAN)
659+
.write(jtsGeom);
660+
write(index, wkb);
661+
}
662+
628663
@Override
629664
public byte[] encodeStr() {
630665
return buf.array();
@@ -666,6 +701,8 @@ public void setValue(int index, Object value) {
666701
write(index, (Date)value);
667702
} else if (value instanceof DateTime) {
668703
write(index, (DateTime)value);
704+
} else if (value instanceof Geography) {
705+
write(index, (Geography)value);
669706
} else {
670707
throw new RuntimeException("Unsupported value object `" + value.getClass() + "\"");
671708
}
@@ -708,6 +745,9 @@ public void setValue(int index, Value value) {
708745
case Value.DTVAL:
709746
write(index, value.getDtVal());
710747
break;
748+
case Value.GGVAL:
749+
write(index, value.getGgVal());
750+
break;
711751
default:
712752
throw new RuntimeException(
713753
"Unknown value: " + value.getFieldValue().getClass()
@@ -790,6 +830,9 @@ public void checkUnsetFields() {
790830
// case Value.DTVAL:
791831
// write(i, defVal.getDtVal());
792832
// break;
833+
// case Value.GGVAL:
834+
// write(i, defVal.getGgVal());
835+
// break;
793836
// default:
794837
// throw new RuntimeException("Unsupported default value type");
795838
// }
@@ -821,7 +864,8 @@ public ByteBuffer processOutOfSpace() {
821864
if (typeEnum == null) {
822865
throw new RuntimeException("Incorrect field type " + field.type());
823866
}
824-
if (typeEnum != PropertyType.STRING) {
867+
if (typeEnum != PropertyType.STRING
868+
&& typeEnum != PropertyType.GEOGRAPHY) {
825869
continue;
826870
}
827871
int offset = headerLen + numNullBytes + field.offset();
@@ -867,4 +911,66 @@ private long getTimestamp() {
867911
long nanoTime = System.nanoTime();
868912
return curTime + (nanoTime - nanoTime / 1000000 * 1000000) / 1000;
869913
}
914+
915+
public org.locationtech.jts.geom.Geometry
916+
convertGeographyToJTSGeometry(Geography geog) {
917+
GeometryFactory geomFactory = new GeometryFactory();
918+
switch (geog.getSetField()) {
919+
case Geography.PTVAL: {
920+
Point point = geog.getPtVal();
921+
Coordinate coord = point.getCoord();
922+
return geomFactory.createPoint(
923+
new org.locationtech.jts.geom.Coordinate(coord.x, coord.y));
924+
}
925+
case Geography.LSVAL: {
926+
LineString line = geog.getLsVal();
927+
List<Coordinate> coordList = line.getCoordList();
928+
929+
List<org.locationtech.jts.geom.Coordinate> jtsCoordList =
930+
new ArrayList<>();
931+
for (int i = 0; i < coordList.size(); ++i) {
932+
jtsCoordList.add(new org.locationtech.jts.geom.Coordinate(
933+
coordList.get(i).x, coordList.get(i).y));
934+
}
935+
org.locationtech.jts.geom.Coordinate[] jtsCoordArray =
936+
new org.locationtech.jts.geom.Coordinate[jtsCoordList.size()];
937+
return geomFactory.createLineString(
938+
jtsCoordList.toArray(jtsCoordArray));
939+
}
940+
case Geography.PGVAL: {
941+
Polygon polygon = geog.getPgVal();
942+
List<List<Coordinate>> coordListList = polygon.getCoordListList();
943+
if (coordListList.isEmpty()) {
944+
throw new RuntimeException("Polygon must at least contain one loop");
945+
}
946+
947+
List<org.locationtech.jts.geom.LinearRing> rings = new ArrayList<>();
948+
for (int i = 0; i < coordListList.size(); ++i) {
949+
List<Coordinate> coordList = coordListList.get(i);
950+
List<org.locationtech.jts.geom.Coordinate> jtsCoordList =
951+
new ArrayList<>();
952+
for (int j = 0; j < coordList.size(); ++j) {
953+
jtsCoordList.add(new org.locationtech.jts.geom.Coordinate(
954+
coordList.get(j).x, coordList.get(j).y));
955+
}
956+
org.locationtech.jts.geom.Coordinate[] jtsCoordArray =
957+
new org.locationtech.jts.geom.Coordinate[jtsCoordList.size()];
958+
rings.add(geomFactory.createLinearRing(
959+
jtsCoordList.toArray(jtsCoordArray)));
960+
}
961+
org.locationtech.jts.geom.LinearRing shell = rings.get(0);
962+
if (rings.size() == 1) {
963+
return geomFactory.createPolygon(shell);
964+
} else {
965+
rings.remove(0);
966+
org.locationtech.jts.geom.LinearRing[] holesArrary =
967+
new org.locationtech.jts.geom.LinearRing[rings.size() - 1];
968+
return geomFactory.createPolygon(shell, rings.toArray(holesArrary));
969+
}
970+
}
971+
default:
972+
throw new RuntimeException("Unknown geography: "
973+
+ geog.getFieldValue().getClass());
974+
}
975+
}
870976
}

‎client/src/main/java/com/vesoft/nebula/encoder/SchemaProvider.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ public interface Field {
2222
public int offset();
2323

2424
public int nullFlagPos();
25+
26+
public int geoShape();
2527
}
2628

2729
public long getVersion();

‎client/src/main/java/com/vesoft/nebula/encoder/SchemaProviderImpl.java‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ static class SchemaField implements Field {
2626
private final int size;
2727
private final int offset;
2828
private final int nullFlagPos;
29+
private final int geoShape;
2930

3031
public SchemaField(String name,
3132
int type,
@@ -34,7 +35,8 @@ public SchemaField(String name,
3435
byte[] defaultValue,
3536
int size,
3637
int offset,
37-
int nullFlagPos) {
38+
int nullFlagPos,
39+
int geoShape) {
3840
this.name = name;
3941
this.type = type;
4042
this.nullable = nullable;
@@ -43,6 +45,7 @@ public SchemaField(String name,
4345
this.size = size;
4446
this.offset = offset;
4547
this.nullFlagPos = nullFlagPos;
48+
this.geoShape = geoShape;
4649
}
4750

4851
@Override
@@ -84,6 +87,11 @@ public int offset() {
8487
public int nullFlagPos() {
8588
return nullFlagPos;
8689
}
90+
91+
@Override
92+
public int geoShape() {
93+
return geoShape;
94+
}
8795
}
8896

8997
public SchemaProviderImpl(long ver) {
@@ -169,7 +177,8 @@ public void addField(String name,
169177
int type,
170178
int fixedStrLen,
171179
boolean nullable,
172-
byte[] defaultValue) {
180+
byte[] defaultValue,
181+
int geoShape) {
173182
int size = fieldSize(type, fixedStrLen);
174183

175184
int offset = 0;
@@ -190,7 +199,8 @@ public void addField(String name,
190199
defaultValue,
191200
size,
192201
offset,
193-
nullFlagPos));
202+
nullFlagPos,
203+
geoShape));
194204
fieldNameIndex.put(name, fields.size() - 1);
195205
}
196206

@@ -241,6 +251,8 @@ public int fieldSize(int type, int fixedStrLimit) {
241251
+ Byte.BYTES // minute
242252
+ Byte.BYTES // sec
243253
+ Integer.BYTES; // microsec
254+
case GEOGRAPHY:
255+
return 8; // wkb offset + wkb length
244256
default:
245257
throw new RuntimeException("Incorrect field type " + type);
246258
}

‎client/src/test/java/com/vesoft/nebula/encoder/MetaCacheImplTest.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import com.vesoft.nebula.meta.ColumnDef;
1212
import com.vesoft.nebula.meta.ColumnTypeDef;
1313
import com.vesoft.nebula.meta.EdgeItem;
14+
import com.vesoft.nebula.meta.GeoShape;
1415
import com.vesoft.nebula.meta.PropertyType;
1516
import com.vesoft.nebula.meta.Schema;
1617
import com.vesoft.nebula.meta.SpaceDesc;
@@ -82,6 +83,19 @@ private Schema genNoDefaultVal() {
8283
new ColumnTypeDef(PropertyType.INT32));
8384
columnDef.setNullable(true);
8485
columns.add(columnDef);
86+
columnDef = new ColumnDef(("Col16").getBytes(),
87+
new ColumnTypeDef(PropertyType.GEOGRAPHY, (short)0, GeoShape.POINT));
88+
columns.add(columnDef);
89+
columnDef = new ColumnDef(("Col17").getBytes(),
90+
new ColumnTypeDef(PropertyType.GEOGRAPHY, (short)0, GeoShape.LINESTRING));
91+
columns.add(columnDef);
92+
columnDef = new ColumnDef(("Col18").getBytes(),
93+
new ColumnTypeDef(PropertyType.GEOGRAPHY, (short)0, GeoShape.POLYGON));
94+
columns.add(columnDef);
95+
columnDef = new ColumnDef(("Col19").getBytes(),
96+
new ColumnTypeDef(PropertyType.GEOGRAPHY, (short)0, GeoShape.ANY));
97+
columnDef.setNullable(true);
98+
columns.add(columnDef);
8599
return new Schema(columns, null);
86100
}
87101

0 commit comments

Comments
 (0)