Java Code Examples for org.apache.pig.impl.util.UDFContext#getUDFContext()

@Override
public void prepareToWrite(RecordWriter writer) throws IOException {
    // Store writer to use in putNext()
    this.writer = writer;

    // Get the schema string from the UDFContext object.
    UDFContext udfc = UDFContext.getUDFContext();
    Properties p = udfc.getUDFProperties(this.getClass(), new String[]{ udfContextSignature });
    String strSchema = p.getProperty(SCHEMA_SIGNATURE);
    if (strSchema == null) {
        throw new IOException("Could not find schema in UDF context");
    }

    schema = new ResourceSchema(Utils.getSchemaFromString(strSchema));
    fields = schema.getFields();
}
 

@Override
public void setLocation(String location, Job job) throws IOException {
    UDFContext udfContext = UDFContext.getUDFContext();
    Properties properties = udfContext.getUDFProperties(RegisteredJarVisibilityLoader.class);

    if (!properties.containsKey(REGISTERED_JAR_VISIBILITY_SCHEMA)) {
        LOG.info("Storing " + RegisteredJarVisibilitySchema.class.getName() + " in UDFContext.");
        properties.put(REGISTERED_JAR_VISIBILITY_SCHEMA, new RegisteredJarVisibilitySchema());
        LOG.info("Stored " + RegisteredJarVisibilitySchema.class.getName() + " in UDFContext.");
    } else {
        LOG.info("Retrieving " + REGISTERED_JAR_VISIBILITY_SCHEMA + " from UDFContext.");
        RegisteredJarVisibilitySchema registeredJarVisibilitySchema =
                (RegisteredJarVisibilitySchema) properties.get(REGISTERED_JAR_VISIBILITY_SCHEMA);
        LOG.info("Retrieved " + REGISTERED_JAR_VISIBILITY_SCHEMA + " from UDFContext.");
    }

    super.setLocation(location, job);
}
 

public ResourceSchema getSchema(String location, Job job)
        throws IOException {

    if (schema != null) {
        // Send schema to backend
        // Schema should have been passed as an argument (-> constructor)
        // or provided in the default constructor

        UDFContext udfc = UDFContext.getUDFContext();
        Properties p = udfc.getUDFProperties(this.getClass(), new String[]{ udfContextSignature });
        p.setProperty(SCHEMA_SIGNATURE, schema.toString());

        return schema;
    } else {
        // Should never get here
        throw new IllegalArgumentException(
            "No schema found: default schema was never created and no user-specified schema was found."
        );
    }
}
 

public ResourceSchema getSchema(String location, Job job)
throws IOException {

    ResourceSchema s;
    if (schema!=null) {
        s = schema;
    } else {
        // Parse the schema
        s = (new JsonMetadata()).getSchema(location, job, true);

        if (s == null) {
            throw new IOException("Unable to parse schema found in file in " + location);
        }
    }

    // Now that we have determined the schema, store it in our
    // UDFContext properties object so we have it when we need it on the
    // backend
    UDFContext udfc = UDFContext.getUDFContext();
    Properties p =
        udfc.getUDFProperties(this.getClass(), new String[]{udfcSignature});
    p.setProperty(SCHEMA_SIGNATURE, s.toString());

    return s;
}
 

@Override
public void prepareToWrite(RecordWriter writer) throws IOException {
    // Store the record writer reference so we can use it when it's time
    // to write tuples
    this.writer = writer;

    // Get the schema string from the UDFContext object.
    UDFContext udfc = UDFContext.getUDFContext();
    Properties p =
        udfc.getUDFProperties(this.getClass(), new String[]{udfcSignature});
    String strSchema = p.getProperty(SCHEMA_SIGNATURE);
    if (strSchema == null) {
        throw new IOException("Could not find schema in UDF context");
    }

    // Parse the schema from the string stored in the properties object.
    schema = new ResourceSchema(Utils.getSchemaFromString(strSchema));

    // Build a Json factory
    jsonFactory = new JsonFactory();
}
 

@Override
public void prepareToRead(RecordReader reader, PigSplit split) throws IOException {
    // Save reader to use in getNext()
    this.reader = reader;

    splitIndex = split.getSplitIndex();

    // Get schema from front-end
    UDFContext udfc = UDFContext.getUDFContext();
    Properties p = udfc.getUDFProperties(this.getClass(), new String[] { udfContextSignature });

    String strSchema = p.getProperty(SCHEMA_SIGNATURE);
    if (strSchema == null) {
        throw new IOException("Could not find schema in UDF context");
    }
    schema = new ResourceSchema(Utils.getSchemaFromString(strSchema));

    requiredFields = (boolean[]) ObjectSerializer.deserialize(p.getProperty(REQUIRED_FIELDS_SIGNATURE));
    if (requiredFields != null) {
        numRequiredFields = 0;
        for (int i = 0; i < requiredFields.length; i++) {
            if (requiredFields[i])
                numRequiredFields++;
        }
    }
}
 

@Override
public RequiredFieldResponse pushProjection(RequiredFieldList requiredFieldList) throws FrontendException {
    if (requiredFieldList == null)
        return null;

    if (fields != null && requiredFieldList.getFields() != null)
    {
        requiredFields = new boolean[fields.length];

        for (RequiredField f : requiredFieldList.getFields()) {
            requiredFields[f.getIndex()] = true;
        }

        UDFContext udfc = UDFContext.getUDFContext();
        Properties p = udfc.getUDFProperties(this.getClass(), new String[]{ udfContextSignature });
        try {
            p.setProperty(REQUIRED_FIELDS_SIGNATURE, ObjectSerializer.serialize(requiredFields));
        } catch (Exception e) {
            throw new RuntimeException("Cannot serialize requiredFields for pushProjection");
        }
    }

    return new RequiredFieldResponse(true);
}
 

/** get a list of index expression */
private List<IndexExpression> getIndexExpressions() throws IOException
{
    UDFContext context = UDFContext.getUDFContext();
    Properties property = context.getUDFProperties(AbstractCassandraStorage.class);
    if (property.getProperty(PARTITION_FILTER_SIGNATURE) != null)
        return indexExpressionsFromString(property.getProperty(PARTITION_FILTER_SIGNATURE));
    else
        return null;
}
 

/**
 * Initialise the database connection and prepared statement here.
 */
@SuppressWarnings("unchecked")
@Override
public void prepareToWrite(RecordWriter writer)
    throws IOException {
  ps = null;
  con = null;
  if (insertQuery == null) {
    throw new IOException("SQL Insert command not specified");
  }
  try {
    if (user == null || pass == null) {
      con = DriverManager.getConnection(jdbcURL);
    } else {
      con = DriverManager.getConnection(jdbcURL, user, pass);
    }
    con.setAutoCommit(false);
    ps = con.prepareStatement(insertQuery);
  } catch (SQLException e) {
    log.error("Unable to connect to JDBC @" + jdbcURL);
    throw new IOException("JDBC Error", e);
  }
  count = 0;

  // Try to get the schema from the UDFContext object.
  UDFContext udfc = UDFContext.getUDFContext();
  Properties p =
      udfc.getUDFProperties(this.getClass(), new String[]{udfcSignature});
  String strSchema = p.getProperty(SCHEMA_SIGNATURE);
  if (strSchema != null) {
      // Parse the schema from the string stored in the properties object.
      schema = new ResourceSchema(Utils.getSchemaFromString(strSchema));
  }
}
 

/** get the columnfamily definition for the signature */
protected CfInfo getCfInfo(String signature) throws IOException
{
    UDFContext context = UDFContext.getUDFContext();
    Properties property = context.getUDFProperties(AbstractCassandraStorage.class);
    String prop = property.getProperty(signature);
    CfInfo cfInfo = new CfInfo();
    cfInfo.cfDef = cfdefFromString(prop.substring(2));
    cfInfo.compactCqlTable = prop.charAt(0) == '1' ? true : false;
    cfInfo.cql3Table = prop.charAt(1) == '1' ? true : false;
    return cfInfo;
}
 

protected long getSeed()
{
  if(seed == null)
  {
    UDFContext context = UDFContext.getUDFContext();
    return Long.parseLong(context.getUDFProperties(this.getClass()).getProperty("seed"));
  }
  else
  {
    return seed;
  }
}
 

/** retrieve where clause for partition filter */
private String getWhereClauseForPartitionFilter()
{
    UDFContext context = UDFContext.getUDFContext();
    Properties property = context.getUDFProperties(AbstractCassandraStorage.class);
    return property.getProperty(PARTITION_FILTER_SIGNATURE);
}
 

/**
* The schema returned here is a bag containing pairs.  
* The pairs are the lsh_id (or the ID of the hash) and the hash value.
*/
public Schema outputSchema(Schema input) {
      try {
        validateInputSchema(input);
        long randomSeed = new Random().nextLong();
        UDFContext context = UDFContext.getUDFContext();
        context.getUDFProperties(this.getClass()).setProperty("seed", "" + randomSeed);
        Schema bagSchema = new Schema();
        bagSchema.add(new Schema.FieldSchema("lsh_id", DataType.INTEGER));
        bagSchema.add(new Schema.FieldSchema("hash", DataType.LONG));
        return new Schema(new Schema.FieldSchema("lsh", bagSchema, DataType.BAG));
      }catch (Exception e){
         throw new RuntimeException("Unable to create output schema", e);
      }
}
 

@Override
public void checkSchema(ResourceSchema s) throws IOException {
    // Not actually checking schema
    // Just storing it to use in the backend
    
    UDFContext udfc = UDFContext.getUDFContext();
    Properties p =
        udfc.getUDFProperties(this.getClass(), new String[]{ udfContextSignature });
    p.setProperty(SCHEMA_SIGNATURE, s.toString());
}
 

protected String getPropertyFromUDFContext(String key) {
  UDFContext udfContext = UDFContext.getUDFContext();
  return udfContext.getUDFProperties(this.getClass(), new String[]{signature}).getProperty(key);
}
 

protected Object getFromUDFContext(String key) {
  UDFContext udfContext = UDFContext.getUDFContext();
  return udfContext.getUDFProperties(this.getClass(), new String[]{signature}).get(key);
}
 

private void storeInUDFContext(final String signature,final String key,final String value) {
    final UDFContext udfContext = UDFContext.getUDFContext();
    final Properties props = udfContext.getUDFProperties(this.getClass(), new String[]{signature});
    props.put(key, value);
}
 

private void storeInUDFContext(String key, Object value) {
    UDFContext udfContext = UDFContext.getUDFContext();
    Properties props = udfContext.getUDFProperties(
            this.getClass(), new String[]{theUDFContextSignature});
    props.put(key, value);
}
 

public void setPartitionFilter(Expression partitionFilter) throws IOException
{
    UDFContext context = UDFContext.getUDFContext();
    Properties property = context.getUDFProperties(AbstractCassandraStorage.class);
    property.setProperty(PARTITION_FILTER_SIGNATURE, partitionFilterToWhereClauseString(partitionFilter));
}
 

/**
 * Helper method to return the context properties for this class
 * 
 * @return context properties
 */
protected Properties getContextProperties() {
  UDFContext context = UDFContext.getUDFContext();
  Properties properties = context.getUDFProperties(this.getClass());
  return properties;
}