bk

star-whale · Oct 11, 2023 · dcc7ed6 · dcc7ed6
1 parent 4065aab
commit dcc7ed6
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 9 deletions.
diff --git a/server/controller/src/main/java/ai/starwhale/mlops/datastore/ColumnSchema.java b/server/controller/src/main/java/ai/starwhale/mlops/datastore/ColumnSchema.java
@@ -32,6 +32,8 @@ public class ColumnSchema {
     private static Set<String> pythonTypeConstants = new ConcurrentHashSet<>();
     private String name;
     private int index;
+    // offset is used to indicate the order of the elements in a list or tuple
+    private Integer offset;
     private ColumnType type;
     private String pythonType;
     private ColumnSchema elementSchema;
@@ -50,6 +52,7 @@ public ColumnSchema(@NonNull String name, int index) {
     public ColumnSchema(@NonNull ColumnSchema schema) {
         this.name = schema.name;
         this.index = schema.index;
+        this.offset = schema.offset;
         this.type = schema.type;
         this.pythonType = schema.pythonType;
         if (schema.elementSchema != null) {
@@ -73,6 +76,7 @@ public ColumnSchema(@NonNull ColumnSchema schema) {
     }
 
     public ColumnSchema(@NonNull ColumnSchemaDesc schema, int index) {
+        // ensure the name is not null (wal do not allow null name)
         this.name = schema.getName() == null ? "" : schema.getName();
         this.index = index;
         this.type = ColumnType.valueOf(schema.getType());
@@ -84,7 +88,9 @@ public ColumnSchema(@NonNull ColumnSchemaDesc schema, int index) {
                 if (schema.getAttributes() != null && !schema.getAttributes().isEmpty()) {
                     this.sparseElementSchema = new HashMap<>();
                     for (var attr : schema.getAttributes()) {
-                        this.sparseElementSchema.put(attr.getIndex(), new ColumnSchema(attr, attr.getIndex()));
+                        var it = new ColumnSchema(attr, 0);
+                        it.setOffset(attr.getIndex());
+                        this.sparseElementSchema.put(attr.getIndex(), it);
                     }
                 }
                 break;
@@ -120,7 +126,9 @@ public ColumnSchema(@NonNull Wal.ColumnSchema schema) {
                 if (!schema.getAttributesList().isEmpty()) {
                     this.sparseElementSchema = new HashMap<>();
                     for (var attr : schema.getAttributesList()) {
-                        this.sparseElementSchema.put(attr.getColumnIndex(), new ColumnSchema(attr));
+                        var it = new ColumnSchema(attr);
+                        it.setOffset(attr.getColumnIndex());
+                        this.sparseElementSchema.put(attr.getColumnIndex(), it);
                     }
                 }
                 break;
@@ -147,7 +155,7 @@ public ColumnSchema(@NonNull Wal.ColumnSchema schema) {
     public ColumnSchemaDesc toColumnSchemaDesc() {
         var builder = ColumnSchemaDesc.builder()
                 .name(this.name)
-                .index(this.index)
+                .index(this.offset)
                 .type(this.type.name());
         switch (this.type) {
             case LIST:
@@ -178,7 +186,7 @@ public ColumnSchemaDesc toColumnSchemaDesc() {
     public Wal.ColumnSchema.Builder toWal() {
         var builder = Wal.ColumnSchema.newBuilder()
                 .setColumnName(this.name)
-                .setColumnIndex(this.index)
+                .setColumnIndex(this.offset != null ? this.offset : this.index)
                 .setColumnType(this.type.name());
         switch (this.type) {
             case LIST:
@@ -324,7 +332,8 @@ public void update(@NonNull Wal.ColumnSchema schema) {
                     for (var attr : schema.getAttributesList()) {
                         var attrSchema = this.sparseElementSchema.get(attr.getColumnIndex());
                         if (attrSchema == null) {
-                            attrSchema = new ColumnSchema(attr.getColumnName(), attr.getColumnIndex());
+                            attrSchema = new ColumnSchema(attr.getColumnName(), 0);
+                            attrSchema.setOffset(attr.getColumnIndex());
                             this.sparseElementSchema.put(attr.getColumnIndex(), attrSchema);
                         }
                         attrSchema.update(attr);
@@ -394,8 +403,8 @@ public boolean isSameType(@NonNull ColumnSchema other) {
                 if (this.sparseElementSchema.size() != other.sparseElementSchema.size()) {
                     return false;
                 }
-                for (var attrName : this.sparseElementSchema.keySet()) {
-                    if (!this.sparseElementSchema.get(attrName).isSameType(other.sparseElementSchema.get(attrName))) {
+                for (var offset : this.sparseElementSchema.keySet()) {
+                    if (!this.sparseElementSchema.get(offset).isSameType(other.sparseElementSchema.get(offset))) {
                         return false;
                     }
                 }

diff --git a/server/controller/src/main/java/ai/starwhale/mlops/datastore/impl/WalRecordDecoder.java b/server/controller/src/main/java/ai/starwhale/mlops/datastore/impl/WalRecordDecoder.java
@@ -110,7 +110,7 @@ public static BaseValue decodeValue(ColumnSchema columnSchema, Wal.Column col) {
     private static BaseValue decodeList(ColumnSchema columnSchema, @NonNull Wal.Column col) {
         var ret = new ListValue();
         var values = col.getListValueList();
-        for (var i = 0; i < values.size(); i++ ) {
+        for (var i = 0; i < values.size(); i++) {
             ColumnSchema schema = null;
             if (columnSchema != null) {
                 var sparse = columnSchema.getSparseElementSchema();
@@ -128,7 +128,7 @@ private static BaseValue decodeList(ColumnSchema columnSchema, @NonNull Wal.Colu
 
     private static BaseValue decodeTuple(ColumnSchema columnSchema, @NonNull Wal.Column col) {
         var ret = new TupleValue();
-        ret.addAll((ListValue)WalRecordDecoder.decodeList(columnSchema, col));
+        ret.addAll((ListValue) WalRecordDecoder.decodeList(columnSchema, col));
         return ret;
     }