Java Code Examples for org.jruby.runtime.ThreadContext#getRuntime()

/**
 * This is a convenience method which will run the given block on the first element
 * of each tuple contained.
 *
 * @param context the context the method is being executed in
 * @param block   a block to call on the elements of the bag
 * @return        enumerator object if null block given, nil otherwise
 */
@JRubyMethod(name = {"flat_each", "flatten"})
public IRubyObject flatten(ThreadContext context, Block block) throws ExecException {
    Ruby runtime = context.getRuntime();

    if (!block.isGiven())
        return PigJrubyLibrary.enumeratorize(runtime, this, "flatten");
    /*  In a future release of JRuby when enumeratorize is made public (which is planned), should replace the above with the below
    if (!block.isGiven())
        return RubyEnumerator.enumeratorize(context.getRuntime(), this, "flatten");
    */

    for (Tuple t : this)
        block.yield(context, PigJrubyLibrary.pigToRuby(runtime, t.get(0)));

    return context.nil;
}
 

/**
 * This is an implementation of the each method which opens up the Enumerable interface,
 * and makes it very convenient to iterate over the elements of a DataBag. Note that currently,
 * due to a deficiency in JRuby, it is not possible to call each without a block given.
 *
 * @param context the context the method is being executed in
 * @param block   a block to call on the elements of the bag
 * @return        enumerator object if null block given, nil otherwise
 */
@JRubyMethod
public IRubyObject each(ThreadContext context, Block block) throws ExecException{
    Ruby runtime = context.getRuntime();

    if (!block.isGiven())
        return PigJrubyLibrary.enumeratorize(runtime, this, "each");
    /*  In a future release of JRuby when enumeratorize is made public (which is planned), should replace the above with the below
    if (!block.isGiven())
        return RubyEnumerator.enumeratorize(context.getRuntime(), this, "each");
    */

    for (Tuple t : this)
        block.yield(context, PigJrubyLibrary.pigToRuby(runtime, t));

    return context.nil;
}
 

/**
 * Given a field name, this will return the index of it in the schema.
 *
 * @param context the context the method is being executed in
 * @param arg     a field name to look for
 * @return        the index for that field name
 */
@JRubyMethod
public RubyFixnum index(ThreadContext context, IRubyObject arg) {
    if (arg instanceof RubyString) {
        try {
            return new RubyFixnum(context.getRuntime(), internalSchema.getPosition(arg.toString()));
        } catch (FrontendException e) {
            throw new RuntimeException("Unable to find position for argument: " + arg);
        }
    } else {
        throw new RuntimeException("Invalid arguement passed to index: " + arg);
    }
}
 

/**
 * This method returns a string representation of the RubyDataBag. If given an optional
 * argument, then if that argument is true, the contents of the bag will also be printed.
 *
 * @param context the context the method is being executed in
 * @param args    optional true/false argument passed to inspect
 * @return        string representation of the RubyDataBag
 */
@JRubyMethod(name = {"inspect", "to_s", "to_string"}, optional = 1)
public RubyString inspect(ThreadContext context, IRubyObject[] args) {
    Ruby runtime = context.getRuntime();
    StringBuilder sb = new StringBuilder();
    sb.append("[DataBag: size: ").append(internalDB.size());
    if (args.length > 0 && args[0].isTrue())
        sb.append(" = ").append(internalDB.toString());
    sb.append("]");
    return RubyString.newString(runtime, sb);
}
 

/**
 * This method returns a copy of the encapsulated DataBag.
 *
 * @param context the context the method is being executed in
 * @return        the copied RubyDataBag
 */
//TODO see if a deepcopy is necessary as well (and consider adding to DataBag and Tuple)
@JRubyMethod
public RubyDataBag clone(ThreadContext context) {
    DataBag b = mBagFactory.newDefaultBag();
    for (Tuple t : this)
        b.add(t);
    Ruby runtime = context.getRuntime();
    return new RubyDataBag(runtime, runtime.getClass("DataBag"), b);
}
 

/**
 * Overrides equality leveraging DataByteArray's equality.
 *
 * @param context the context the method is being executed in
 * @param arg     a RubyDataByteArray against which to test equality
 * @return        true if they are equal, false otherwise
 */
@JRubyMethod(name = {"eql?", "=="})
public RubyBoolean equals(ThreadContext context, IRubyObject arg) {
    Ruby runtime = context.getRuntime();
    if (arg instanceof RubyDataByteArray) {
        return RubyBoolean.newBoolean(runtime, internalDBA.equals(((RubyDataByteArray)arg).getDBA()));
    } else {
        return runtime.getFalse();
    }
}
 

/**
 * This method allows access into the Schema nested in the encapsulated Schema. For example,
 * if the encapsulated Schema is a bag Schema, this allows the user to access the schema of
 * the interior Tuple.
 *
 * @param context the context the method is being executed in
 * @return        a RubySchema encapsulating the nested Schema
 */
@JRubyMethod(name = {"get", "inner", "in"})
public RubySchema get(ThreadContext context) {
    if (internalSchema.size() != 1)
        throw new RuntimeException("Can only return nested schema if there is one schema to get");
    Ruby runtime = context.getRuntime();
    try {
        return new RubySchema(runtime, runtime.getClass("Schema"), internalSchema.getField(0).schema, false);
    } catch (FrontendException e) {
        throw new RuntimeException("Schema does not have a nested FieldScema", e);
    }
}
 

private ConcurrentHashMap<IRubyObject, IRubyObject> toCHM(ThreadContext context, IRubyObject options) {
    Ruby runtime = context.getRuntime();
    if (!options.isNil() && options.respondsTo("[]")) {
        IRubyObject rInitialCapacity = options.callMethod(context, "[]", runtime.newSymbol("initial_capacity"));
        IRubyObject rLoadFactor      = options.callMethod(context, "[]", runtime.newSymbol("load_factor"));
        int initialCapacity = !rInitialCapacity.isNil() ? RubyNumeric.num2int(rInitialCapacity.convertToInteger()) : DEFAULT_INITIAL_CAPACITY;
        float loadFactor    = !rLoadFactor.isNil() ?      (float)RubyNumeric.num2dbl(rLoadFactor.convertToFloat()) : DEFAULT_LOAD_FACTOR;
        return newCHM(initialCapacity, loadFactor);
    } else {
        return newCHM();
    }
}
 

/**
 * Given a field name this string will search the RubySchema for a FieldSchema
 * with that name and return it encapsulated in a Schema.
 *
 * @param context the context the method is being executed in
 * @param arg     a RubyString serving as an alias to look
 *                for in the Schema
 * @return        the found RubySchema
 */
@JRubyMethod
public RubySchema find(ThreadContext context, IRubyObject arg) {
    if (arg instanceof RubyString) {
        Ruby runtime = context.getRuntime();
        return new RubySchema(runtime, runtime.getClass("Schema"), RubySchema.find(internalSchema, arg.toString()), false);
    } else {
        throw new RuntimeException("Invalid arguement passed to find: " + arg);
    }
}
 

/**
 * This method provides addition semantics, modifying the original Schema in place.
 * This method can be given any number of arguments, much as with the constructor.
 *
 * @param context the context the method is being executed in
 * @param args    a varargs which can be any valid set of arguments that
 *                can initialize a RubySchema
 */
@JRubyMethod(name = "add!", rest = true)
public void addInPlace(ThreadContext context, IRubyObject[] args) {
    Ruby runtime = context.getRuntime();
    List<Schema.FieldSchema> lfs = internalSchema.getFields();
    RubySchema rs = new RubySchema(runtime, runtime.getClass("Schema")).initialize(args);
    for (Schema.FieldSchema fs : rs.getInternalSchema().getFields())
        lfs.add(fs);
    RubySchema.fixSchemaNames(internalSchema);
}
 

/**
 * This is a version of [] which allows the range to be specified as such: [1,2].
 *
 * @param context the context the method is being executed in
 * @param arg1    a Fixnum start index
 * @param arg2    a Fixnum end index
 * @return        the RubySchema object encapsulated the found Schema
 */
@JRubyMethod(name = {"[]", "slice"})
public RubySchema get(ThreadContext context, IRubyObject arg1, IRubyObject arg2) {
    if (arg1 instanceof RubyFixnum && arg2 instanceof RubyFixnum) {
        Ruby runtime = context.getRuntime();
        int min = (int)((RubyFixnum)arg1).getLongValue();
        int max = (int)((RubyFixnum)arg2).getLongValue() - 1;
        return new RubySchema(runtime, runtime.getClass("Schema"), new Schema(internalSchema.getFields().subList(min, max + 1)), false);
    } else {
        throw new RuntimeException("Bad arguments given to get function: ( " + arg1.toString() + " , " + arg2.toString()+ " )");
    }
}
 

/**
 * This is a ruby method which takes an array of arguments and constructs a Bag schema from them. The name
 * will be set automatically.
 *
 * @param context the context the method is being executed in
 * @param self    the RubyClass for the Class object this was invoked on
 * @param arg     a list of arguments to instantiate the new RubySchema
 * @return        the new RubySchema
 */
@JRubyMethod(meta = true, name = {"b", "bag"})
public static RubySchema bag(ThreadContext context, IRubyObject self, IRubyObject arg) {
    Schema s = tuple(context, self, arg).getInternalSchema();
    Ruby runtime = context.getRuntime();
    try {
        return new RubySchema(runtime, runtime.getClass("Schema"), new Schema(new Schema.FieldSchema("bag_0", s, DataType.BAG)));
    } catch (FrontendException e) {
        throw new RuntimeException("Error making map", e);
    }
}
 

/**
 * This is a ruby method which takes an array of arguments and constructs a Map schema from them. The name
 * will be set automatically.
 *
 * @param context the context the method is being executed in
 * @param self    the RubyClass for the Class object this was invoked on
 * @param arg     a list of arguments to instantiate the new RubySchema
 * @return        the new RubySchema
 */
@JRubyMethod(meta = true, name = {"m", "map"})
public static RubySchema map(ThreadContext context, IRubyObject self, IRubyObject arg) {
    Schema s = tuple(context, self, arg).getInternalSchema();
    Ruby runtime = context.getRuntime();
    try {
        return new RubySchema(runtime, runtime.getClass("Schema"), new Schema(new Schema.FieldSchema("map_0", s.getField(0).schema, DataType.MAP)));
    } catch (FrontendException e) {
        throw new RuntimeException("Error making map", e);
    }
}
 

/**
 * This is a ruby method which takes an array of arguments and constructs a Tuple schema from them. The name
 * will be set automatically.
 *
 * @param context the context the method is being executed in
 * @param self    the RubyClass for the Class object this was invoked on
 * @param arg     a list of arguments to instantiate the new RubySchema
 * @return        the new RubySchema
 */
@JRubyMethod(meta = true, name = {"t", "tuple"})
public static RubySchema tuple(ThreadContext context, IRubyObject self, IRubyObject arg) {
    if (arg instanceof RubyArray) {
        Schema s = rubyArgToSchema(arg);
        Ruby runtime = context.getRuntime();
        return new RubySchema(runtime, runtime.getClass("Schema"), s);
    } else {
        throw new RuntimeException("Bad argument given to Schema.tuple");
    }
}
 

/**
 * @param context the context the method is being executed in
 * @return        a RubySchema copy of the Schema
 */
@JRubyMethod
public RubySchema clone(ThreadContext context) {
    Ruby runtime = context.getRuntime();
    return new RubySchema(runtime, runtime.getClass("Schema"), internalSchema);
}
 

/**
 * @param context the context the method is being executed in
 * @return        the size of the encapsulated Schema
 */
@JRubyMethod(name = {"size", "length"})
public RubyFixnum size(ThreadContext context) {
    return new RubyFixnum(context.getRuntime(), internalSchema.size());
}
 

@JRubyMethod(name = "docinfo?", required = 1)
public IRubyObject hasDocInfo(ThreadContext context, IRubyObject location) {
  Ruby runtime = context.getRuntime();
  LocationType locationType = "footer".equals(location.asJavaString()) ? LocationType.FOOTER : LocationType.HEADER;
  return runtime.newBoolean(delegate.hasDocInfo(locationType));
}
 

/**
 * This calls to the append function of the underlying DataByteArray. This accepts
 * the same types that set accepts.
 *
 * @param context the context the method is being executed in
 * @param arg a value to append to the encpasulated DataByteArray. In the case of a
 *            RubyDataByteArray, the bytes will be copied and appended; in the case
 *            of a String, the bits will be added; in the case of a byte array,
 *            the bytes will be appended directly.
 */
@JRubyMethod(name = {"add!", "append"})
public void append(ThreadContext context, IRubyObject arg) {
    Ruby runtime = context.getRuntime();
    RubyDataByteArray toAppend = new RubyDataByteArray(runtime, runtime.getClass("DataByteArray")).initialize(arg);
    internalDBA.append(toAppend.getDBA());
}
 

/**
 * This method accepts as an argument anything that could be passed to set (ie a
 * RubyDataByteArray, RubyString, or byte array). It then adds those values together.
 *
 * @param context the context the method is being executed in
 * @param arg     any argument that can validly initialize a RubyDataByteArray
 * @return        a <u>new</u> RubyDataByteArray that is the concatentation of
 *                the encapsulated DataByteArray and arg
 */
@JRubyMethod(name = "+")
public IRubyObject plus(ThreadContext context, IRubyObject arg) {
    Ruby runtime = context.getRuntime();
    RubyDataByteArray toAdd = new RubyDataByteArray(runtime, runtime.getClass("DataByteArray")).initialize(arg);
    return new RubyDataByteArray(runtime, runtime.getClass("DataByteArray"), internalDBA, toAdd.getDBA());
}
 

/**
 * This is a helper method to generate a RubySchema of the given type without an alias.
 *
 * @param context the context the method is being executed in
 * @param type    the DataType.PIGTYPE value to make the Schema from
 * @return        a RubySchema object encapsulated a Schema of the specified type
 */
private static RubySchema makeNullAliasRubySchema(ThreadContext context, byte type) {
   Ruby runtime = context.getRuntime();
   return new RubySchema(runtime, runtime.getClass("Schema"), new Schema(new Schema.FieldSchema(null, type)));
}