Use Kryo for seraliazation

wrangler-core/pom.xml

@@ -170,6 +170,11 @@
       <artifactId>guava-retrying</artifactId>
       <version>${guava.retrying.version}</version>
     </dependency>
+    <dependency>
+      <groupId>com.esotericsoftware</groupId>
+      <artifactId>kryo</artifactId>
+      <version>4.0.2</version>
+    </dependency>
     <dependency>
       <groupId>org.antlr</groupId>
       <artifactId>antlr4</artifactId>

wrangler-core/src/main/java/io/cdap/wrangler/utils/RowSerializer.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright © 2024 Cask Data, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package io.cdap.wrangler.utils;
+
+import com.esotericsoftware.kryo.Kryo;
+import com.esotericsoftware.kryo.Serializer;
+import com.esotericsoftware.kryo.io.Input;
+import com.esotericsoftware.kryo.io.Output;
+import com.esotericsoftware.kryo.serializers.FieldSerializer;
+import com.esotericsoftware.kryo.serializers.JavaSerializer;
+import com.google.gson.Gson;
+import com.google.gson.JsonArray;
+import com.google.gson.JsonElement;
+import com.google.gson.JsonNull;
+import com.google.gson.JsonObject;
+import com.google.gson.JsonPrimitive;
+import io.cdap.wrangler.api.Row;
+import java.sql.Time;
+import java.sql.Timestamp;
+import java.time.LocalDate;
+import java.time.LocalTime;
+import java.time.ZonedDateTime;
+import java.util.ArrayList;
+import java.util.Date;
+import java.util.List;
+import java.util.Map;
+
+public class RowSerializer {
+
+  private final Kryo kryo;
+  private static final Gson GSON = new Gson();
+
+  public RowSerializer(){
+    kryo = new Kryo();
+    // Register all classes from SchemaConverter
+    kryo.register(Row.class);
+    kryo.register(ArrayList.class);
+    kryo.register(LocalDate.class);
+    kryo.register(LocalTime.class);
+    kryo.register(ZonedDateTime.class);
+    kryo.register(Map.class);
+    kryo.register(JsonNull.class);
+    // JsonPrimitive does not have no-arg constructor hence we need a
+    // custom serializer
+    kryo.register(JsonPrimitive.class, new JsonPrimitiveSerializer());
+    kryo.register(JsonArray.class);
+    kryo.register(JsonObject.class);
+    // Support deprecated util.date classes
+    kryo.register(Date.class);
+    kryo.register(java.sql.Date.class);
+    kryo.register(Time.class);
+    kryo.register(Timestamp.class);
+  }
+
+  public byte[] fromRows(List<Row> rows) {
+    Output output = new Output(1024,-1);
+    kryo.writeClassAndObject(output, rows);
+    return output.getBuffer();
+  }
+
+  public List<Row> toRows(byte[] bytes) {
+    Input input = new Input(bytes);
+    List<Row> result = (List<Row>) kryo.readClassAndObject(input);
+    return result;
+  }
+
+  static class JsonPrimitiveSerializer extends Serializer<JsonElement> {
+
+    @Override
+    public void write(Kryo kryo, Output output, JsonElement object) {
+      output.writeString(GSON.toJson(object));
+    }
+
+    @Override
+    public JsonElement read(Kryo kryo, Input input, Class<JsonElement> type) {
+      return GSON.fromJson(input.readString(), type);
+    }
+  }
+}

wrangler-core/src/test/java/io/cdap/wrangler/utils/JsonTestData.java

@@ -194,4 +194,5 @@ public final class JsonTestData {
     + "   }"
     + " }";
   public static final String EMPTY_OBJECT = "{ \"dividesplitdetails\":{\"type0\":[]}}";
+  public static final String NULL_OBJECT = "{ \"dividesplitdetails\":{\"type0\":null, \"type1\":0}}";
 }

wrangler-core/src/test/java/io/cdap/wrangler/utils/RowSerializerTest.java

@@ -0,0 +1,117 @@
+/*
+ * Copyright © 2024 Cask Data, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package io.cdap.wrangler.utils;
+
+import com.google.common.collect.Lists;
+import com.google.gson.JsonParser;
+import io.cdap.wrangler.TestingRig;
+import io.cdap.wrangler.api.RecipePipeline;
+import io.cdap.wrangler.api.Row;
+import java.math.BigDecimal;
+import java.math.BigInteger;
+import java.time.LocalDate;
+import java.time.LocalDateTime;
+import java.time.LocalTime;
+import java.time.ZoneId;
+import java.time.ZonedDateTime;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import org.junit.Assert;
+import org.junit.Test;
+
+public class RowSerializerTest {
+
+  private static final String[] TESTS = new String[]{
+      JsonTestData.BASIC,
+      JsonTestData.SIMPLE_JSON_OBJECT,
+      JsonTestData.ARRAY_OF_OBJECTS,
+      JsonTestData.JSON_ARRAY_WITH_OBJECT,
+      JsonTestData.COMPLEX_1,
+      JsonTestData.ARRAY_OF_NUMBERS,
+      JsonTestData.ARRAY_OF_STRING,
+      JsonTestData.COMPLEX_2,
+      JsonTestData.EMPTY_OBJECT,
+      JsonTestData.NULL_OBJECT,
+      JsonTestData.FB_JSON
+  };
+
+  private static final String[] directives = new String[]{
+      "set-column body json:Parse(body)"
+  };
+
+  @Test
+  public void testJsonTypes() throws Exception {
+    SchemaConverter converter = new SchemaConverter();
+    RecordConvertor recordConvertor = new RecordConvertor();
+    JsonParser parser = new JsonParser();
+    RecipePipeline executor = TestingRig.execute(directives);
+    for (String test : TESTS) {
+      Row row = new Row("body", test);
+
+      List<Row> expectedRows = executor.execute(Lists.newArrayList(row));
+      byte[] serializedRows = new RowSerializer().fromRows(expectedRows);
+      List<Row> gotRows = new RowSerializer().toRows(serializedRows);
+      Assert.assertArrayEquals(expectedRows.toArray(), gotRows.toArray());
+    }
+  }
+
+  @Test
+  public void testLogicalTypes() throws Exception {
+    Row testRow = new Row();
+    testRow.add("id", 1);
+    testRow.add("name", "abc");
+    testRow.add("date", LocalDate.of(2018, 11, 11));
+    testRow.add("time", LocalTime.of(11, 11, 11));
+    testRow.add("timestamp", ZonedDateTime.of(2018, 11, 11, 11, 11, 11, 0, ZoneId.of("UTC")));
+    testRow.add("bigdecimal", new BigDecimal(new BigInteger("123456"), 5));
+    testRow.add("datetime", LocalDateTime.now());
+    List<Row> expectedRows = Collections.singletonList(testRow);
+    byte[] serializedRows = new RowSerializer().fromRows(expectedRows);
+    List<Row> gotRows = new RowSerializer().toRows(serializedRows);
+    Assert.assertArrayEquals(expectedRows.toArray(), gotRows.toArray());
+  }
+
+  @Test
+  public void testCollectionTypes() throws Exception {
+    List<Integer> list = new ArrayList<>();
+    list.add(null);
+    list.add(1);
+    list.add(2);
+    Set<Integer> set = new HashSet<>();
+    set.add(null);
+    set.add(1);
+    set.add(2);
+    Map<String, Integer> map = new HashMap<>();
+    map.put("null", null);
+    map.put("1", 1);
+    map.put("2", 2);
+
+    Row testRow = new Row();
+    testRow.add("list", list);
+    testRow.add("set", set);
+    testRow.add("map", map);
+
+    List<Row> expectedRows = Collections.singletonList(testRow);
+    byte[] serializedRows = new RowSerializer().fromRows(expectedRows);
+    List<Row> gotRows = new RowSerializer().toRows(serializedRows);
+    Assert.assertArrayEquals(expectedRows.toArray(), gotRows.toArray());
+  }
+}

wrangler-service/src/main/java/io/cdap/wrangler/service/directive/RemoteExecutionTask.java

@@ -44,6 +44,7 @@
 import io.cdap.wrangler.registry.UserDirectiveRegistry;
 import io.cdap.wrangler.utils.ObjectSerDe;
 
+import io.cdap.wrangler.utils.RowSerializer;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -57,6 +58,7 @@ public class RemoteExecutionTask implements RunnableTask {
 
   private static final Gson GSON = new Gson();
 
+
   @Override
   public void run(RunnableTaskContext runnableTaskContext) throws Exception {
     RemoteDirectiveRequest directiveRequest = GSON.fromJson(runnableTaskContext.getParam(),
@@ -121,7 +123,9 @@ public void run(RunnableTaskContext runnableTaskContext) throws Exception {
       }
 
       runnableTaskContext.setTerminateOnComplete(hasUDD.get() || EL.isUsed());
-      runnableTaskContext.writeResult(objectSerDe.toByteArray(rows));
+
+      // TODO(samik) add feature gate for using RowSerializer
+      runnableTaskContext.writeResult(new RowSerializer().fromRows(rows));
     } catch (DirectiveParseException | ClassNotFoundException | CompileException e) {
       throw new BadRequestException(e.getMessage(), e);
     }

wrangler-service/src/main/java/io/cdap/wrangler/service/directive/WorkspaceHandler.java

@@ -77,8 +77,8 @@
 import io.cdap.wrangler.schema.TransientStoreKeys;
 import io.cdap.wrangler.store.recipe.RecipeStore;
 import io.cdap.wrangler.store.workspace.WorkspaceStore;
-import io.cdap.wrangler.utils.ObjectSerDe;
 import io.cdap.wrangler.utils.RowHelper;
+import io.cdap.wrangler.utils.RowSerializer;
 import io.cdap.wrangler.utils.SchemaConverter;
 import io.cdap.wrangler.utils.StructuredToRowTransformer;
 import org.apache.commons.lang3.StringEscapeUtils;
@@ -624,7 +624,8 @@ private <E extends Exception> List<Row> executeRemotely(String namespace, List<S
       .withNamespace(namespace)
       .build();
     byte[] bytes = getContext().runTask(runnableTaskRequest);
-    return new ObjectSerDe<List<Row>>().toObject(bytes);
+    // TODO(samik) add feature gate for using RowSerializer
+    return new RowSerializer().toRows(bytes);
   }
 
   private List<Row> getSample(SampleResponse sampleResponse) {