From b3bf0b60a91309addc4eeb0f0df4526ad7ae7dd1 Mon Sep 17 00:00:00 2001
From: Steeve Morin <steeve.morin@gmail.com>
Date: Fri, 30 Sep 2011 05:12:41 +0800
Subject: [PATCH] get correct key type in schema

---
 .../apache/cassandra/pig/CassandraStorage.java    | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/src/java/src/org/apache/cassandra/pig/CassandraStorage.java b/src/java/src/org/apache/cassandra/pig/CassandraStorage.java
index 9613de7..b07031e 100644
--- a/src/java/src/org/apache/cassandra/pig/CassandraStorage.java
+++ b/src/java/src/org/apache/cassandra/pig/CassandraStorage.java
@@ -187,10 +187,12 @@ private List<AbstractType> getDefaultMarshallers(CfDef cfDef) throws IOException
         ArrayList<AbstractType> marshallers = new ArrayList<AbstractType>();
         AbstractType comparator = null;
         AbstractType default_validator = null;
+        AbstractType key_validator = null;
         try
         {
             comparator = TypeParser.parse(cfDef.getComparator_type());
             default_validator = TypeParser.parse(cfDef.getDefault_validation_class());
+            key_validator = TypeParser.parse(cfDef.getKey_validation_class());
         }
         catch (ConfigurationException e)
         {
@@ -199,13 +201,14 @@ private List<AbstractType> getDefaultMarshallers(CfDef cfDef) throws IOException
 
         marshallers.add(comparator);
         marshallers.add(default_validator);
+        marshallers.add(key_validator);
         return marshallers;
     }
 
-    private Map<ByteBuffer, AbstractType> getValidatorMap(CfDef cfDef) throws  IOException
+    private Map<ByteBuffer, AbstractType> getValidatorMap(CfDef cfDef) throws IOException
     {
         Map<ByteBuffer, AbstractType> validators = new HashMap<ByteBuffer, AbstractType>();
-        for (ColumnDef cd : cfDef.column_metadata)
+        for (ColumnDef cd : cfDef.getColumn_metadata())
         {
             if (cd.getValidation_class() != null && !cd.getValidation_class().isEmpty())
             {
@@ -321,10 +324,14 @@ public ResourceSchema getSchema(String location, Job job) throws IOException
         // top-level schema, no type
         ResourceSchema schema = new ResourceSchema();
 
+        // get default marshallers and validators
+        List<AbstractType> marshallers = getDefaultMarshallers(cfDef);
+        Map<ByteBuffer,AbstractType> validators = getValidatorMap(cfDef);
+
         // add key
         ResourceFieldSchema keyFieldSchema = new ResourceFieldSchema();
         keyFieldSchema.setName("key");
-        keyFieldSchema.setType(DataType.CHARARRAY); //TODO: get key type
+        keyFieldSchema.setType(getPigType(marshallers.get(2)));
 
         // will become the bag of tuples
         ResourceFieldSchema bagFieldSchema = new ResourceFieldSchema();
@@ -332,8 +339,6 @@ public ResourceSchema getSchema(String location, Job job) throws IOException
         bagFieldSchema.setType(DataType.BAG);
         ResourceSchema bagSchema = new ResourceSchema();
 
-        List<AbstractType> marshallers = getDefaultMarshallers(cfDef);
-        Map<ByteBuffer,AbstractType> validators = getValidatorMap(cfDef);
         List<ResourceFieldSchema> tupleFields = new ArrayList<ResourceFieldSchema>();
 
         // default comparator/validator