src/test/org/apache/sqoop/TestParquetImport.java - sqoop - Git at Google

 /**
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */

 package org.apache.sqoop;

 import org.apache.sqoop.testutil.CommonArgs;
 import org.apache.sqoop.testutil.HsqldbTestServer;
 import org.apache.sqoop.testutil.ImportJobTestCase;

 import org.apache.avro.Schema;
 import org.apache.avro.Schema.Field;
 import org.apache.avro.Schema.Type;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.junit.Test;
 import org.kitesdk.data.CompressionType;
 import org.kitesdk.data.Dataset;
 import org.kitesdk.data.DatasetReader;
 import org.kitesdk.data.Datasets;

 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;

 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;

 /**
  * Tests --as-parquetfile.
  */
 public class TestParquetImport extends ImportJobTestCase {

   public static final Log LOG = LogFactory
       .getLog(TestParquetImport.class.getName());

   /**
    * Create the argv to pass to Sqoop.
    *
    * @return the argv as an array of strings.
    */
   protected String[] getOutputArgv(boolean includeHadoopFlags,
           String[] extraArgs) {
     ArrayList<String> args = new ArrayList<String>();

     if (includeHadoopFlags) {
       CommonArgs.addHadoopFlags(args);
     }

     args.add("--table");
     args.add(getTableName());
     args.add("--connect");
     args.add(HsqldbTestServer.getUrl());
     args.add("--warehouse-dir");
     args.add(getWarehouseDir());
     args.add("--m");
     args.add("1");
     args.add("--split-by");
     args.add("INTFIELD1");
     args.add("--as-parquetfile");
     if (extraArgs != null) {
       args.addAll(Arrays.asList(extraArgs));
     }

     return args.toArray(new String[args.size()]);
   }

   protected String[] getOutputQueryArgv(boolean includeHadoopFlags, String[] extraArgs) {
     ArrayList<String> args = new ArrayList<String>();

     if (includeHadoopFlags) {
       CommonArgs.addHadoopFlags(args);
     }

     args.add("--query");
     args.add("SELECT * FROM " + getTableName() + " WHERE $CONDITIONS");
     args.add("--connect");
     args.add(HsqldbTestServer.getUrl());
     args.add("--target-dir");
     args.add(getWarehouseDir() + "/" + getTableName());
     args.add("--m");
     args.add("1");
     args.add("--split-by");
     args.add("INTFIELD1");
     args.add("--as-parquetfile");
     if (extraArgs != null) {
       args.addAll(Arrays.asList(extraArgs));
     }

     return args.toArray(new String[args.size()]);
   }

   @Test
   public void testSnappyCompression() throws IOException {
     runParquetImportTest("snappy");
   }

   @Test
   public void testDeflateCompression() throws IOException {
     runParquetImportTest("deflate");
   }

   private void runParquetImportTest(String codec) throws IOException {
     String[] types = {"BIT", "INTEGER", "BIGINT", "REAL", "DOUBLE", "VARCHAR(6)",
         "VARBINARY(2)",};
     String[] vals = {"true", "100", "200", "1.0", "2.0", "'s'", "'0102'", };
     createTableWithColTypes(types, vals);

     String [] extraArgs = { "--compression-codec", codec};
     runImport(getOutputArgv(true, extraArgs));

     assertEquals(CompressionType.forName(codec), getCompressionType());

     Schema schema = getSchema();
     assertEquals(Type.RECORD, schema.getType());
     List<Field> fields = schema.getFields();
     assertEquals(types.length, fields.size());
     checkField(fields.get(0), "DATA_COL0", Type.BOOLEAN);
     checkField(fields.get(1), "DATA_COL1", Type.INT);
     checkField(fields.get(2), "DATA_COL2", Type.LONG);
     checkField(fields.get(3), "DATA_COL3", Type.FLOAT);
     checkField(fields.get(4), "DATA_COL4", Type.DOUBLE);
     checkField(fields.get(5), "DATA_COL5", Type.STRING);
     checkField(fields.get(6), "DATA_COL6", Type.BYTES);

     DatasetReader<GenericRecord> reader = getReader();
     try {
       GenericRecord record1 = reader.next();
       assertNotNull(record1);
       assertEquals("DATA_COL0", true, record1.get("DATA_COL0"));
       assertEquals("DATA_COL1", 100, record1.get("DATA_COL1"));
       assertEquals("DATA_COL2", 200L, record1.get("DATA_COL2"));
       assertEquals("DATA_COL3", 1.0f, record1.get("DATA_COL3"));
       assertEquals("DATA_COL4", 2.0, record1.get("DATA_COL4"));
       assertEquals("DATA_COL5", "s", record1.get("DATA_COL5"));
       Object object = record1.get("DATA_COL6");
       assertTrue(object instanceof ByteBuffer);
       ByteBuffer b = ((ByteBuffer) object);
       assertEquals((byte) 1, b.get(0));
       assertEquals((byte) 2, b.get(1));
       assertFalse(reader.hasNext());
     } finally {
       reader.close();
     }
   }

   @Test
   public void testOverrideTypeMapping() throws IOException {
     String [] types = { "INT" };
     String [] vals = { "10" };
     createTableWithColTypes(types, vals);

     String [] extraArgs = { "--map-column-java", "DATA_COL0=String"};
     runImport(getOutputArgv(true, extraArgs));

     Schema schema = getSchema();
     assertEquals(Type.RECORD, schema.getType());
     List<Field> fields = schema.getFields();
     assertEquals(types.length, fields.size());
     checkField(fields.get(0), "DATA_COL0", Type.STRING);

     DatasetReader<GenericRecord> reader = getReader();
     try {
       assertTrue(reader.hasNext());
       GenericRecord record1 = reader.next();
       assertEquals("DATA_COL0", "10", record1.get("DATA_COL0"));
       assertFalse(reader.hasNext());
     } finally {
       reader.close();
     }
   }

   @Test
   public void testFirstUnderscoreInColumnName() throws IOException {
     String [] names = { "_NAME" };
     String [] types = { "INT" };
     String [] vals = { "1987" };
     createTableWithColTypesAndNames(names, types, vals);

     runImport(getOutputArgv(true, null));

     Schema schema = getSchema();
     assertEquals(Type.RECORD, schema.getType());
     List<Field> fields = schema.getFields();
     assertEquals(types.length, fields.size());
     checkField(fields.get(0), "__NAME", Type.INT);

     DatasetReader<GenericRecord> reader = getReader();
     try {
       assertTrue(reader.hasNext());
       GenericRecord record1 = reader.next();
       assertEquals("__NAME", 1987, record1.get("__NAME"));
       assertFalse(reader.hasNext());
     } finally {
       reader.close();
     }
   }

   @Test
   public void testNonIdentCharactersInColumnName() throws IOException {
     String [] names = { "test_p-a+r/quet" };
     String [] types = { "INT" };
     String [] vals = { "2015" };
     createTableWithColTypesAndNames(names, types, vals);

     runImport(getOutputArgv(true, null));

     Schema schema = getSchema();
     assertEquals(Type.RECORD, schema.getType());
     List<Field> fields = schema.getFields();
     assertEquals(types.length, fields.size());
     checkField(fields.get(0), "TEST_P_A_R_QUET", Type.INT);

     DatasetReader<GenericRecord> reader = getReader();
     try {
       assertTrue(reader.hasNext());
       GenericRecord record1 = reader.next();
       assertEquals("TEST_P_A_R_QUET", 2015, record1.get("TEST_P_A_R_QUET"));
       assertFalse(reader.hasNext());
     } finally {
       reader.close();
     }
   }

   @Test
   public void testNullableParquetImport() throws IOException, SQLException {
     String [] types = { "INT" };
     String [] vals = { null };
     createTableWithColTypes(types, vals);

     runImport(getOutputArgv(true, null));

     DatasetReader<GenericRecord> reader = getReader();
     try {
       assertTrue(reader.hasNext());
       GenericRecord record1 = reader.next();
       assertNull(record1.get("DATA_COL0"));
       assertFalse(reader.hasNext());
     } finally {
       reader.close();
     }
   }

   @Test
   public void testQueryImport() throws IOException, SQLException {
     String [] types = { "INT" };
     String [] vals = { "1" };
     createTableWithColTypes(types, vals);

     runImport(getOutputQueryArgv(true, null));

     DatasetReader<GenericRecord> reader = getReader();
     try {
       assertTrue(reader.hasNext());
       GenericRecord record1 = reader.next();
       assertEquals(1, record1.get("DATA_COL0"));
       assertFalse(reader.hasNext());
     } finally {
       reader.close();
     }
   }

   @Test
   public void testIncrementalParquetImport() throws IOException, SQLException {
     String [] types = { "INT" };
     String [] vals = { "1" };
     createTableWithColTypes(types, vals);

     runImport(getOutputArgv(true, null));
     runImport(getOutputArgv(true, new String[]{"--append"}));

     DatasetReader<GenericRecord> reader = getReader();
     try {
       assertTrue(reader.hasNext());
       GenericRecord record1 = reader.next();
       assertEquals(1, record1.get("DATA_COL0"));
       record1 = reader.next();
       assertEquals(1, record1.get("DATA_COL0"));
       assertFalse(reader.hasNext());
     } finally {
       reader.close();
     }
   }

   @Test
   public void testOverwriteParquetDatasetFail() throws IOException, SQLException {
     String [] types = { "INT" };
     String [] vals = {};
     createTableWithColTypes(types, vals);

     runImport(getOutputArgv(true, null));
     try {
       runImport(getOutputArgv(true, null));
       fail("");
     } catch (IOException ex) {
       // ok
     }
   }

   private CompressionType getCompressionType() {
     return getDataset().getDescriptor().getCompressionType();
   }

   private Schema getSchema() {
     return getDataset().getDescriptor().getSchema();
   }

   private DatasetReader<GenericRecord> getReader() {
     return getDataset().newReader();
   }

   private Dataset<GenericRecord> getDataset() {
     String uri = "dataset:file:" + getTablePath();
     return Datasets.load(uri, GenericRecord.class);
   }

   @Override
   public void tearDown() {
     super.tearDown();
     String uri = "dataset:file:" + getTablePath();
     if (Datasets.exists(uri)) {
       Datasets.delete(uri);
     }
   }

   private void checkField(Field field, String name, Type type) {
     assertEquals(name, field.name());
     assertEquals(Type.UNION, field.schema().getType());
     assertEquals(Type.NULL, field.schema().getTypes().get(0).getType());
     assertEquals(type, field.schema().getTypes().get(1).getType());
   }

 }
	/**
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/

	package org.apache.sqoop;

	import org.apache.sqoop.testutil.CommonArgs;
	import org.apache.sqoop.testutil.HsqldbTestServer;
	import org.apache.sqoop.testutil.ImportJobTestCase;

	import org.apache.avro.Schema;
	import org.apache.avro.Schema.Field;
	import org.apache.avro.Schema.Type;
	import org.apache.avro.generic.GenericRecord;
	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.junit.Test;
	import org.kitesdk.data.CompressionType;
	import org.kitesdk.data.Dataset;
	import org.kitesdk.data.DatasetReader;
	import org.kitesdk.data.Datasets;

	import java.io.IOException;
	import java.nio.ByteBuffer;
	import java.sql.SQLException;
	import java.util.ArrayList;
	import java.util.Arrays;
	import java.util.List;

	import static org.junit.Assert.assertEquals;
	import static org.junit.Assert.assertFalse;
	import static org.junit.Assert.assertNotNull;
	import static org.junit.Assert.assertNull;
	import static org.junit.Assert.assertTrue;
	import static org.junit.Assert.fail;

	/**
	* Tests --as-parquetfile.
	*/
	public class TestParquetImport extends ImportJobTestCase {

	public static final Log LOG = LogFactory
	.getLog(TestParquetImport.class.getName());

	/**
	* Create the argv to pass to Sqoop.
	*
	* @return the argv as an array of strings.
	*/
	protected String[] getOutputArgv(boolean includeHadoopFlags,
	String[] extraArgs) {
	ArrayList<String> args = new ArrayList<String>();

	if (includeHadoopFlags) {
	CommonArgs.addHadoopFlags(args);
	}

	args.add("--table");
	args.add(getTableName());
	args.add("--connect");
	args.add(HsqldbTestServer.getUrl());
	args.add("--warehouse-dir");
	args.add(getWarehouseDir());
	args.add("--m");
	args.add("1");
	args.add("--split-by");
	args.add("INTFIELD1");
	args.add("--as-parquetfile");
	if (extraArgs != null) {
	args.addAll(Arrays.asList(extraArgs));
	}

	return args.toArray(new String[args.size()]);
	}

	protected String[] getOutputQueryArgv(boolean includeHadoopFlags, String[] extraArgs) {
	ArrayList<String> args = new ArrayList<String>();

	if (includeHadoopFlags) {
	CommonArgs.addHadoopFlags(args);
	}

	args.add("--query");
	args.add("SELECT * FROM " + getTableName() + " WHERE $CONDITIONS");
	args.add("--connect");
	args.add(HsqldbTestServer.getUrl());
	args.add("--target-dir");
	args.add(getWarehouseDir() + "/" + getTableName());
	args.add("--m");
	args.add("1");
	args.add("--split-by");
	args.add("INTFIELD1");
	args.add("--as-parquetfile");
	if (extraArgs != null) {
	args.addAll(Arrays.asList(extraArgs));
	}

	return args.toArray(new String[args.size()]);
	}

	@Test
	public void testSnappyCompression() throws IOException {
	runParquetImportTest("snappy");
	}

	@Test
	public void testDeflateCompression() throws IOException {
	runParquetImportTest("deflate");
	}

	private void runParquetImportTest(String codec) throws IOException {
	String[] types = {"BIT", "INTEGER", "BIGINT", "REAL", "DOUBLE", "VARCHAR(6)",
	"VARBINARY(2)",};
	String[] vals = {"true", "100", "200", "1.0", "2.0", "'s'", "'0102'", };
	createTableWithColTypes(types, vals);

	String [] extraArgs = { "--compression-codec", codec};
	runImport(getOutputArgv(true, extraArgs));

	assertEquals(CompressionType.forName(codec), getCompressionType());

	Schema schema = getSchema();
	assertEquals(Type.RECORD, schema.getType());
	List<Field> fields = schema.getFields();
	assertEquals(types.length, fields.size());
	checkField(fields.get(0), "DATA_COL0", Type.BOOLEAN);
	checkField(fields.get(1), "DATA_COL1", Type.INT);
	checkField(fields.get(2), "DATA_COL2", Type.LONG);
	checkField(fields.get(3), "DATA_COL3", Type.FLOAT);
	checkField(fields.get(4), "DATA_COL4", Type.DOUBLE);
	checkField(fields.get(5), "DATA_COL5", Type.STRING);
	checkField(fields.get(6), "DATA_COL6", Type.BYTES);

	DatasetReader<GenericRecord> reader = getReader();
	try {
	GenericRecord record1 = reader.next();
	assertNotNull(record1);
	assertEquals("DATA_COL0", true, record1.get("DATA_COL0"));
	assertEquals("DATA_COL1", 100, record1.get("DATA_COL1"));
	assertEquals("DATA_COL2", 200L, record1.get("DATA_COL2"));
	assertEquals("DATA_COL3", 1.0f, record1.get("DATA_COL3"));
	assertEquals("DATA_COL4", 2.0, record1.get("DATA_COL4"));
	assertEquals("DATA_COL5", "s", record1.get("DATA_COL5"));
	Object object = record1.get("DATA_COL6");
	assertTrue(object instanceof ByteBuffer);
	ByteBuffer b = ((ByteBuffer) object);
	assertEquals((byte) 1, b.get(0));
	assertEquals((byte) 2, b.get(1));
	assertFalse(reader.hasNext());
	} finally {
	reader.close();
	}
	}

	@Test
	public void testOverrideTypeMapping() throws IOException {
	String [] types = { "INT" };
	String [] vals = { "10" };
	createTableWithColTypes(types, vals);

	String [] extraArgs = { "--map-column-java", "DATA_COL0=String"};
	runImport(getOutputArgv(true, extraArgs));

	Schema schema = getSchema();
	assertEquals(Type.RECORD, schema.getType());
	List<Field> fields = schema.getFields();
	assertEquals(types.length, fields.size());
	checkField(fields.get(0), "DATA_COL0", Type.STRING);

	DatasetReader<GenericRecord> reader = getReader();
	try {
	assertTrue(reader.hasNext());
	GenericRecord record1 = reader.next();
	assertEquals("DATA_COL0", "10", record1.get("DATA_COL0"));
	assertFalse(reader.hasNext());
	} finally {
	reader.close();
	}
	}

	@Test
	public void testFirstUnderscoreInColumnName() throws IOException {
	String [] names = { "_NAME" };
	String [] types = { "INT" };
	String [] vals = { "1987" };
	createTableWithColTypesAndNames(names, types, vals);

	runImport(getOutputArgv(true, null));

	Schema schema = getSchema();
	assertEquals(Type.RECORD, schema.getType());
	List<Field> fields = schema.getFields();
	assertEquals(types.length, fields.size());
	checkField(fields.get(0), "__NAME", Type.INT);

	DatasetReader<GenericRecord> reader = getReader();
	try {
	assertTrue(reader.hasNext());
	GenericRecord record1 = reader.next();
	assertEquals("__NAME", 1987, record1.get("__NAME"));
	assertFalse(reader.hasNext());
	} finally {
	reader.close();
	}
	}

	@Test
	public void testNonIdentCharactersInColumnName() throws IOException {
	String [] names = { "test_p-a+r/quet" };
	String [] types = { "INT" };
	String [] vals = { "2015" };
	createTableWithColTypesAndNames(names, types, vals);

	runImport(getOutputArgv(true, null));

	Schema schema = getSchema();
	assertEquals(Type.RECORD, schema.getType());
	List<Field> fields = schema.getFields();
	assertEquals(types.length, fields.size());
	checkField(fields.get(0), "TEST_P_A_R_QUET", Type.INT);

	DatasetReader<GenericRecord> reader = getReader();
	try {
	assertTrue(reader.hasNext());
	GenericRecord record1 = reader.next();
	assertEquals("TEST_P_A_R_QUET", 2015, record1.get("TEST_P_A_R_QUET"));
	assertFalse(reader.hasNext());
	} finally {
	reader.close();
	}
	}

	@Test
	public void testNullableParquetImport() throws IOException, SQLException {
	String [] types = { "INT" };
	String [] vals = { null };
	createTableWithColTypes(types, vals);

	runImport(getOutputArgv(true, null));

	DatasetReader<GenericRecord> reader = getReader();
	try {
	assertTrue(reader.hasNext());
	GenericRecord record1 = reader.next();
	assertNull(record1.get("DATA_COL0"));
	assertFalse(reader.hasNext());
	} finally {
	reader.close();
	}
	}

	@Test
	public void testQueryImport() throws IOException, SQLException {
	String [] types = { "INT" };
	String [] vals = { "1" };
	createTableWithColTypes(types, vals);

	runImport(getOutputQueryArgv(true, null));

	DatasetReader<GenericRecord> reader = getReader();
	try {
	assertTrue(reader.hasNext());
	GenericRecord record1 = reader.next();
	assertEquals(1, record1.get("DATA_COL0"));
	assertFalse(reader.hasNext());
	} finally {
	reader.close();
	}
	}

	@Test
	public void testIncrementalParquetImport() throws IOException, SQLException {
	String [] types = { "INT" };
	String [] vals = { "1" };
	createTableWithColTypes(types, vals);

	runImport(getOutputArgv(true, null));
	runImport(getOutputArgv(true, new String[]{"--append"}));

	DatasetReader<GenericRecord> reader = getReader();
	try {
	assertTrue(reader.hasNext());
	GenericRecord record1 = reader.next();
	assertEquals(1, record1.get("DATA_COL0"));
	record1 = reader.next();
	assertEquals(1, record1.get("DATA_COL0"));
	assertFalse(reader.hasNext());
	} finally {
	reader.close();
	}
	}

	@Test
	public void testOverwriteParquetDatasetFail() throws IOException, SQLException {
	String [] types = { "INT" };
	String [] vals = {};
	createTableWithColTypes(types, vals);

	runImport(getOutputArgv(true, null));
	try {
	runImport(getOutputArgv(true, null));
	fail("");
	} catch (IOException ex) {
	// ok
	}
	}

	private CompressionType getCompressionType() {
	return getDataset().getDescriptor().getCompressionType();
	}

	private Schema getSchema() {
	return getDataset().getDescriptor().getSchema();
	}

	private DatasetReader<GenericRecord> getReader() {
	return getDataset().newReader();
	}

	private Dataset<GenericRecord> getDataset() {
	String uri = "dataset:file:" + getTablePath();
	return Datasets.load(uri, GenericRecord.class);
	}

	@Override
	public void tearDown() {
	super.tearDown();
	String uri = "dataset:file:" + getTablePath();
	if (Datasets.exists(uri)) {
	Datasets.delete(uri);
	}
	}

	private void checkField(Field field, String name, Type type) {
	assertEquals(name, field.name());
	assertEquals(Type.UNION, field.schema().getType());
	assertEquals(Type.NULL, field.schema().getTypes().get(0).getType());
	assertEquals(type, field.schema().getTypes().get(1).getType());
	}

	}