Hive - Zero to One [3.1.3]

package com.artists.hive.udf;

import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;

/**
 * @Author: ArtistS
 * @Site: https://github.com/ArtistSu
 * @Create: 2023/5/21 17:45
 * @Description: Custom UDF to calculate the length of a string
 **/
public class MyLength extends GenericUDF {
    
    @Override
    public ObjectInspector initialize(ObjectInspector[] objectInspectors) throws UDFArgumentException {
        // Return int type inspector because our function calculates the length of a field
        if (objectInspectors.length != 1) {
            throw new UDFArgumentException("Only accept one argument");
        }

        ObjectInspector argument = objectInspectors[0];
        if (ObjectInspector.Category.PRIMITIVE != argument.getCategory()) {
            throw new UDFArgumentException("Only accept primitive argument");
        }

        PrimitiveObjectInspector primitiveObjectInspector = (PrimitiveObjectInspector) argument;
        if (primitiveObjectInspector.getPrimitiveCategory() != PrimitiveObjectInspector.PrimitiveCategory.STRING) {
            throw new UDFArgumentException("Only accept string argument");
        }

        return PrimitiveObjectInspectorFactory.javaIntObjectInspector;
    }

    @Override
    public Object evaluate(DeferredObject[] deferredObjects) throws HiveException {
        // evaluate() is called once per row of data
        DeferredObject deferredObject = deferredObjects[0];
        Object obj = deferredObject.get();

        if (obj == null) {
            return 0;
        }

        return obj.toString().length();
    }

    @Override
    public String getDisplayString(String[] strings) {
        // This method is called for display in execution plans (e.g., when using explain)
        return "my_length(" + (strings.length > 0 ? strings[0] : "") + ")";
    }
}