VespaDocumentOperation extends EvalFunc<String> { @Override public String exec(Tuple tuple) throws IOException { if (tuple == null || tuple.size() == 0) { if (statusReporter != null) { statusReporter.incrCounter("Vespa Document Operation Counters", "Document operation failed", 1); } return null; } if (template == null || template.length() == 0) { if (statusReporter != null) { statusReporter.incrCounter("Vespa Document Operation Counters", "Document operation failed", 1); } warnLog("No valid document id template found. Skipping.", PigWarning.UDF_WARNING_1); return null; } if (operation == null) { warnLog("No valid operation found. Skipping.", PigWarning.UDF_WARNING_2); return null; } String json = null; try { if (reporter != null) { reporter.progress(); } Schema inputSchema = getInputSchema(); Map<String, Object> fields = TupleTools.tupleMap(inputSchema, tuple); String docId = TupleTools.toString(fields, template); if (verbose) { System.out.println("Processing docId: "+ docId); } json = create(operation, docId, fields, properties, inputSchema); if (json == null || json.length() == 0) { warnLog("No valid document operation could be created.", PigWarning.UDF_WARNING_3); return null; } } catch (Exception e) { if (statusReporter != null) { statusReporter.incrCounter("Vespa Document Operation Counters", "Document operation failed", 1); } StringBuilder sb = new StringBuilder(); sb.append("Caught exception processing input row: \n"); sb.append(tuple.toString()); sb.append("\nException: "); sb.append(getStackTraceAsString(e)); warnLog(sb.toString(), PigWarning.UDF_WARNING_4); return null; } if (statusReporter != null) { statusReporter.incrCounter("Vespa Document Operation Counters", "Document operation ok", 1); } return json; } VespaDocumentOperation(String... params); @Override String exec(Tuple tuple); static String create(Operation op, String docId, Map<String, Object> fields, Properties properties,
                                Schema schema);  }
@Test public void requireThatUDFCorrectlyGeneratesRemoveBagAsMapOperation() throws Exception { DataBag bag = BagFactory.getInstance().newDefaultBag(); Schema innerObjectSchema = new Schema(); Tuple innerObjectTuple = TupleFactory.getInstance().newTuple(); addToTuple("year", DataType.CHARARRAY, "2020", innerObjectSchema, innerObjectTuple); addToTuple("month", DataType.INTEGER, 3, innerObjectSchema, innerObjectTuple); Schema objectSchema = new Schema(); Tuple objectTuple = TupleFactory.getInstance().newTuple(); addToTuple("key", DataType.CHARARRAY, "234566", objectSchema, objectTuple); addToTuple("value", DataType.TUPLE, innerObjectTuple,innerObjectSchema,objectSchema, objectTuple); Schema bagSchema = new Schema(); addToBagWithSchema("firstLayerTuple",DataType.TUPLE,objectTuple,objectSchema,bagSchema,bag); innerObjectSchema = new Schema(); innerObjectTuple = TupleFactory.getInstance().newTuple(); addToTuple("year", DataType.CHARARRAY, "2020", innerObjectSchema, innerObjectTuple); addToTuple("month", DataType.INTEGER, 3, innerObjectSchema, innerObjectTuple); objectSchema = new Schema(); objectTuple = TupleFactory.getInstance().newTuple(); addToTuple("key", DataType.CHARARRAY, "123456", objectSchema, objectTuple); addToTuple("value", DataType.TUPLE, innerObjectTuple,innerObjectSchema,objectSchema, objectTuple); addToBagWithSchema("firstLayerTuple",DataType.TUPLE,objectTuple,objectSchema,bagSchema,bag); Schema schema = new Schema(); Tuple tuple = TupleFactory.getInstance().newTuple(); addToTuple("bag", DataType.BAG, bag, bagSchema, schema, tuple); addToTuple("id", DataType.CHARARRAY, "123", schema, tuple); VespaDocumentOperation docOp = new VespaDocumentOperation("docid=id", "remove-map-fields=bag","operation=update"); docOp.setInputSchema(schema); String json = docOp.exec(tuple); ObjectMapper m = new ObjectMapper(); JsonNode root = m.readTree(json); JsonNode fields = root.get("fields"); assertEquals("{\"remove\":0}", fields.get("bag{123456}").toString()); assertEquals("{\"remove\":0}", fields.get("bag{234566}").toString()); }
@Test public void requireThatUDFCorrectlyGeneratesAddBagAsMapOperation() throws Exception { DataBag bag = BagFactory.getInstance().newDefaultBag(); Schema innerObjectSchema = new Schema(); Tuple innerObjectTuple = TupleFactory.getInstance().newTuple(); addToTuple("year", DataType.CHARARRAY, "2020", innerObjectSchema, innerObjectTuple); addToTuple("month", DataType.INTEGER, 3, innerObjectSchema, innerObjectTuple); Schema objectSchema = new Schema(); Tuple objectTuple = TupleFactory.getInstance().newTuple(); addToTuple("key", DataType.CHARARRAY, "123456", objectSchema, objectTuple); addToTuple("value", DataType.TUPLE, innerObjectTuple,innerObjectSchema,objectSchema, objectTuple); Schema bagSchema = new Schema(); addToBagWithSchema("firstLayerTuple",DataType.TUPLE,objectTuple,objectSchema,bagSchema,bag); Schema schema = new Schema(); Tuple tuple = TupleFactory.getInstance().newTuple(); addToTuple("bag", DataType.BAG, bag, bagSchema, schema, tuple); addToTuple("id", DataType.CHARARRAY, "123", schema, tuple); VespaDocumentOperation docOp = new VespaDocumentOperation("docid=id", "update-map-fields=bag","operation=update"); docOp.setInputSchema(schema); String json = docOp.exec(tuple); ObjectMapper m = new ObjectMapper(); JsonNode root = m.readTree(json); JsonNode fields = root.get("fields"); JsonNode value = fields.get("bag{123456}"); JsonNode assign = value.get("assign"); assertEquals("2020", assign.get("year").getTextValue()); assertEquals(3, assign.get("month").getIntValue()); }
@Test public void requireThatUDFCorrectlyGeneratesAddTensorOperation() throws Exception { Schema schema = new Schema(); Tuple tuple = TupleFactory.getInstance().newTuple(); Map<String, Double> tensor = new HashMap<String, Double>() {{ put("x:label1,y:label2,z:label4", 2.0); put("x:label3", 3.0); }}; addToTuple("id", DataType.CHARARRAY, "123", schema, tuple); addToTuple("tensor", DataType.MAP, tensor, schema, tuple); VespaDocumentOperation docOp = new VespaDocumentOperation("docid=empty", "update-tensor-fields=tensor","operation=update"); docOp.setInputSchema(schema); String json = docOp.exec(tuple); ObjectMapper m = new ObjectMapper(); JsonNode root = m.readTree(json); JsonNode fields = root.get("fields"); JsonNode tensorValue = fields.get("tensor"); JsonNode add = tensorValue.get("add"); JsonNode cells = add.get("cells"); Iterator<JsonNode> cellsIterator = cells.getElements(); JsonNode element = cellsIterator.next(); assertEquals("label1", element.get("address").get("x").getTextValue()); assertEquals("label2", element.get("address").get("y").getTextValue()); assertEquals("label4", element.get("address").get("z").getTextValue()); assertEquals("2.0", element.get("value").toString()); element = cellsIterator.next(); assertEquals("label3", element.get("address").get("x").getTextValue()); assertEquals("3.0", element.get("value").toString()); }
@Test public void requireThatUDFCorrectlyGeneratesRemoveTensorOperation() throws Exception { Schema schema = new Schema(); Tuple tuple = TupleFactory.getInstance().newTuple(); Map<String, Double> tensor = new HashMap<String, Double>() {{ put("x:label1,y:label2,z:label4", 2.0); put("x:label3", 3.0); }}; addToTuple("id", DataType.CHARARRAY, "123", schema, tuple); addToTuple("tensor", DataType.MAP, tensor, schema, tuple); VespaDocumentOperation docOp = new VespaDocumentOperation("docid=empty", "remove-tensor-fields=tensor","operation=update"); docOp.setInputSchema(schema); String json = docOp.exec(tuple); ObjectMapper m = new ObjectMapper(); JsonNode root = m.readTree(json); JsonNode fields = root.get("fields"); JsonNode tensorValue = fields.get("tensor"); JsonNode remove = tensorValue.get("remove"); JsonNode address = remove.get("addresses"); Iterator<JsonNode> addressIterator = address.getElements(); JsonNode element = addressIterator.next(); assertEquals("label1", element.get("x").getTextValue()); element = addressIterator.next(); assertEquals("label2", element.get("y").getTextValue()); element = addressIterator.next(); assertEquals("label4", element.get("z").getTextValue()); element = addressIterator.next(); assertEquals("label3", element.get("x").getTextValue()); }
@Test public void requireThatUDFReturnsNullWhenExceptionHappens() throws IOException { Schema schema = new Schema(); Tuple tuple = TupleFactory.getInstance().newTuple(); Map<String, Double> tensor = new HashMap<String, Double>() {{ put("xlabel1", 2.0); }}; addToTuple("id", DataType.CHARARRAY, "123", schema, tuple); addToTuple("tensor", DataType.MAP, tensor, schema, tuple); VespaDocumentOperation docOp = new VespaDocumentOperation("docid=empty", "create-tensor-fields=tensor"); docOp.setInputSchema(schema); String json = docOp.exec(tuple); assertNull(json); }
@Test public void requireThatUDFGeneratesComplexDataTypes() throws Exception { Schema schema = new Schema(); Tuple tuple = TupleFactory.getInstance().newTuple(); Tuple intTuple = TupleFactory.getInstance().newTuple(); int[] intArray = {1, 2, 3}; for (int i : intArray) { intTuple.append(i); } Tuple stringTuple = TupleFactory.getInstance().newTuple(); String[] stringArray = {"a", "b", "c"}; for (String s : stringArray) { stringTuple.append(s); } DataBag bag = new SortedDataBag(null); bag.add(intTuple); bag.add(stringTuple); Map<String, Object> innerMap = new HashMap<String, Object>() {{ put("a", "string"); put("tuple", intTuple); }}; DataByteArray bytes = new DataByteArray("testdata".getBytes()); Map<String, Object> outerMap = new HashMap<String, Object>() {{ put("string", "value"); put("int", 3); put("float", 3.145); put("bool", true); put("byte", bytes); put("map", innerMap); put("bag", bag); }}; addToTuple("map", DataType.MAP, outerMap, schema, tuple); VespaDocumentOperation docOp = new VespaDocumentOperation("docid=empty"); docOp.setInputSchema(schema); String json = docOp.exec(tuple); ObjectMapper m = new ObjectMapper(); JsonNode root = m.readTree(json); JsonNode fields = root.get("fields"); JsonNode map = fields.get("map"); assertEquals("value", map.get("string").getTextValue()); assertEquals(3, map.get("int").getIntValue()); assertEquals(3.145, map.get("float").getDoubleValue(), 1e-6); assertTrue(map.get("bool").getBooleanValue()); assertEquals("dGVzdGRhdGE=", map.get("byte").getTextValue()); assertEquals("string", map.get("map").get("a").getTextValue()); for (int i = 0; i < intArray.length; ++i) { assertEquals(intArray[i], map.get("map").get("tuple").get(i).asInt()); } JsonNode bagField = map.get("bag"); for (int i = 0; i < intArray.length; ++i) { assertEquals(intArray[i], bagField.get(0).get(i).asInt()); } for (int i = 0; i < stringArray.length; ++i) { assertEquals(stringArray[i], bagField.get(1).get(i).asText()); } }
@Test public void requireThatUDFSupportsTensors() throws IOException { Schema schema = new Schema(); Tuple tuple = TupleFactory.getInstance().newTuple(); Map<String, Double> tensor = new HashMap<String, Double>() {{ put("x:label1,y:label2,z:label4", 2.0); put("x:label3", 3.0); }}; addToTuple("id", DataType.CHARARRAY, "123", schema, tuple); addToTuple("tensor", DataType.MAP, tensor, schema, tuple); VespaDocumentOperation docOp = new VespaDocumentOperation("docid=empty", "create-tensor-fields=tensor"); docOp.setInputSchema(schema); String json = docOp.exec(tuple); ObjectMapper m = new ObjectMapper(); JsonNode root = m.readTree(json); JsonNode fields = root.get("fields"); JsonNode tensorNode = fields.get("tensor"); JsonNode cells = tensorNode.get("cells"); assertEquals("label1", cells.get(0).get("address").get("x").asText()); assertEquals("label2", cells.get(0).get("address").get("y").asText()); assertEquals("label4", cells.get(0).get("address").get("z").asText()); assertEquals("label3", cells.get(1).get("address").get("x").asText()); assertEquals(2.0, cells.get(0).get("value").asDouble(), 1e-6); assertEquals(3.0, cells.get(1).get("value").asDouble(), 1e-6); }
@Test public void requireThatUDFSupportsSimpleObjectFields() throws IOException { Schema objectSchema = new Schema(); Tuple objectTuple = TupleFactory.getInstance().newTuple(); addToTuple("id", DataType.LONG, 123456789L, objectSchema, objectTuple); addToTuple("url", DataType.CHARARRAY, "example.com", objectSchema, objectTuple); addToTuple("value", DataType.INTEGER, 123, objectSchema, objectTuple); Schema schema = new Schema(); Tuple tuple = TupleFactory.getInstance().newTuple(); addToTuple("object", DataType.TUPLE, objectTuple, objectSchema, schema, tuple); VespaDocumentOperation docOp = new VespaDocumentOperation("docid=empty", "simple-object-fields=object"); docOp.setInputSchema(schema); String json = docOp.exec(tuple); ObjectMapper m = new ObjectMapper(); JsonNode root = m.readTree(json); JsonNode fields = root.get("fields"); JsonNode objectNode = fields.get("object"); assertEquals(123456789L, objectNode.get("id").asLong()); assertEquals("example.com", objectNode.get("url").asText()); assertEquals(123, objectNode.get("value").asInt()); }
@Test public void requireThatUDFSupportsBagAsMapFields() throws IOException { DataBag bag = BagFactory.getInstance().newDefaultBag(); Schema objectSchema = new Schema(); Tuple objectTuple = TupleFactory.getInstance().newTuple(); addToTuple("key", DataType.CHARARRAY, "123456", objectSchema, objectTuple); addToTuple("value", DataType.INTEGER, 123456, objectSchema, objectTuple); bag.add(objectTuple); objectSchema = new Schema(); objectTuple = TupleFactory.getInstance().newTuple(); addToTuple("key", DataType.CHARARRAY, "234567", objectSchema, objectTuple); addToTuple("value", DataType.INTEGER, 234567, objectSchema, objectTuple); bag.add(objectTuple); Schema schema = new Schema(); Tuple tuple = TupleFactory.getInstance().newTuple(); addToTuple("bag", DataType.BAG, bag, objectSchema, schema, tuple); VespaDocumentOperation docOp = new VespaDocumentOperation("docid=empty", "bag-as-map-fields=bag"); docOp.setInputSchema(schema); String json = docOp.exec(tuple); ObjectMapper m = new ObjectMapper(); JsonNode root = m.readTree(json); JsonNode fields = root.get("fields"); JsonNode bagNode = fields.get("bag"); assertEquals(123456, bagNode.get("123456").asInt()); assertEquals(234567, bagNode.get("234567").asInt()); }
@Test public void requireThatUDFPrintIdWhenVerbose() throws IOException { DataBag bag = BagFactory.getInstance().newDefaultBag(); Schema objectSchema = new Schema(); Tuple objectTuple = TupleFactory.getInstance().newTuple(); addToTuple("key", DataType.CHARARRAY, "123456", objectSchema, objectTuple); addToTuple("value", DataType.INTEGER, 123456, objectSchema, objectTuple); bag.add(objectTuple); objectSchema = new Schema(); objectTuple = TupleFactory.getInstance().newTuple(); addToTuple("key", DataType.CHARARRAY, "234567", objectSchema, objectTuple); addToTuple("value", DataType.INTEGER, 234567, objectSchema, objectTuple); bag.add(objectTuple); Schema schema = new Schema(); Tuple tuple = TupleFactory.getInstance().newTuple(); addToTuple("bag", DataType.BAG, bag, objectSchema, schema, tuple); VespaDocumentOperation docOp = new VespaDocumentOperation("docid=7654321", "bag-as-map-fields=bag","verbose=true"); docOp.setInputSchema(schema); String json = docOp.exec(tuple); assertTrue(outContent.toString().contains("Processing docId: 7654321")); }
@Test public void requireThatUDFVerboseSetToFalseByDefault() throws IOException { DataBag bag = BagFactory.getInstance().newDefaultBag(); Schema objectSchema = new Schema(); Tuple objectTuple = TupleFactory.getInstance().newTuple(); addToTuple("key", DataType.CHARARRAY, "123456", objectSchema, objectTuple); addToTuple("value", DataType.INTEGER, 123456, objectSchema, objectTuple); bag.add(objectTuple); objectSchema = new Schema(); objectTuple = TupleFactory.getInstance().newTuple(); addToTuple("key", DataType.CHARARRAY, "234567", objectSchema, objectTuple); addToTuple("value", DataType.INTEGER, 234567, objectSchema, objectTuple); bag.add(objectTuple); Schema schema = new Schema(); Tuple tuple = TupleFactory.getInstance().newTuple(); addToTuple("bag", DataType.BAG, bag, objectSchema, schema, tuple); VespaDocumentOperation docOp = new VespaDocumentOperation("docid=7654321", "bag-as-map-fields=bag"); docOp.setInputSchema(schema); String json = docOp.exec(tuple); assertEquals("", outContent.toString()); }