@Override public int getInt(int columnIndex) throws SQLException { wasNull = getCurrentRow().isNullAt(columnIndex-1); return wasNull ? 0 : getCurrentRow().getInt(columnIndex-1); }
@SuppressWarnings("unchecked") @Test public void udf2Test() { spark.udf().register("stringLengthTest", (String str1, String str2) -> str1.length() + str2.length(), DataTypes.IntegerType); Row result = spark.sql("SELECT stringLengthTest('test', 'test2')").head(); Assert.assertEquals(9, result.getInt(0)); }
@SuppressWarnings("unchecked") @Test public void udf2Test() { spark.udf().register("stringLengthTest", (String str1, String str2) -> str1.length() + str2.length(), DataTypes.IntegerType); Row result = spark.sql("SELECT stringLengthTest('test', 'test2')").head(); Assert.assertEquals(9, result.getInt(0)); }
@SuppressWarnings("unchecked") @Test public void udf2Test() { spark.udf().register("stringLengthTest", (String str1, String str2) -> str1.length() + str2.length(), DataTypes.IntegerType); Row result = spark.sql("SELECT stringLengthTest('test', 'test2')").head(); Assert.assertEquals(9, result.getInt(0)); }
@SuppressWarnings("unchecked") @Test public void udf3Test() { spark.udf().registerJava("stringLengthTest", StringLengthTest.class.getName(), DataTypes.IntegerType); Row result = spark.sql("SELECT stringLengthTest('test', 'test2')").head(); Assert.assertEquals(9, result.getInt(0)); // returnType is not provided spark.udf().registerJava("stringLengthTest2", StringLengthTest.class.getName(), null); result = spark.sql("SELECT stringLengthTest('test', 'test2')").head(); Assert.assertEquals(9, result.getInt(0)); }
@SuppressWarnings("unchecked") @Test public void udf3Test() { spark.udf().registerJava("stringLengthTest", StringLengthTest.class.getName(), DataTypes.IntegerType); Row result = spark.sql("SELECT stringLengthTest('test', 'test2')").head(); Assert.assertEquals(9, result.getInt(0)); // returnType is not provided spark.udf().registerJava("stringLengthTest2", StringLengthTest.class.getName(), null); result = spark.sql("SELECT stringLengthTest('test', 'test2')").head(); Assert.assertEquals(9, result.getInt(0)); }
@SuppressWarnings("unchecked") @Test public void udf3Test() { spark.udf().registerJava("stringLengthTest", StringLengthTest.class.getName(), DataTypes.IntegerType); Row result = spark.sql("SELECT stringLengthTest('test', 'test2')").head(); Assert.assertEquals(9, result.getInt(0)); // returnType is not provided spark.udf().registerJava("stringLengthTest2", StringLengthTest.class.getName(), null); result = spark.sql("SELECT stringLengthTest('test', 'test2')").head(); Assert.assertEquals(9, result.getInt(0)); }
@Test public void pivot() { Dataset<Row> df = spark.table("courseSales"); List<Row> actual = df.groupBy("year") .pivot("course", Arrays.asList("dotNET", "Java")) .agg(sum("earnings")).orderBy("year").collectAsList(); Assert.assertEquals(2012, actual.get(0).getInt(0)); Assert.assertEquals(15000.0, actual.get(0).getDouble(1), 0.01); Assert.assertEquals(20000.0, actual.get(0).getDouble(2), 0.01); Assert.assertEquals(2013, actual.get(1).getInt(0)); Assert.assertEquals(48000.0, actual.get(1).getDouble(1), 0.01); Assert.assertEquals(30000.0, actual.get(1).getDouble(2), 0.01); }
@Test public void pivot() { Dataset<Row> df = spark.table("courseSales"); List<Row> actual = df.groupBy("year") .pivot("course", Arrays.asList("dotNET", "Java")) .agg(sum("earnings")).orderBy("year").collectAsList(); Assert.assertEquals(2012, actual.get(0).getInt(0)); Assert.assertEquals(15000.0, actual.get(0).getDouble(1), 0.01); Assert.assertEquals(20000.0, actual.get(0).getDouble(2), 0.01); Assert.assertEquals(2013, actual.get(1).getInt(0)); Assert.assertEquals(48000.0, actual.get(1).getDouble(1), 0.01); Assert.assertEquals(30000.0, actual.get(1).getDouble(2), 0.01); }
@Test public void pivot() { Dataset<Row> df = spark.table("courseSales"); List<Row> actual = df.groupBy("year") .pivot("course", Arrays.asList("dotNET", "Java")) .agg(sum("earnings")).orderBy("year").collectAsList(); Assert.assertEquals(2012, actual.get(0).getInt(0)); Assert.assertEquals(15000.0, actual.get(0).getDouble(1), 0.01); Assert.assertEquals(20000.0, actual.get(0).getDouble(2), 0.01); Assert.assertEquals(2013, actual.get(1).getInt(0)); Assert.assertEquals(48000.0, actual.get(1).getDouble(1), 0.01); Assert.assertEquals(30000.0, actual.get(1).getDouble(2), 0.01); }
@SuppressWarnings("unchecked") @Test public void udf1Test() { spark.udf().register("stringLengthTest", (String str) -> str.length(), DataTypes.IntegerType); Row result = spark.sql("SELECT stringLengthTest('test')").head(); Assert.assertEquals(4, result.getInt(0)); }
@SuppressWarnings("unchecked") @Test public void udf1Test() { spark.udf().register("stringLengthTest", (String str) -> str.length(), DataTypes.IntegerType); Row result = spark.sql("SELECT stringLengthTest('test')").head(); Assert.assertEquals(4, result.getInt(0)); }
@SuppressWarnings("unchecked") @Test public void udf1Test() { spark.udf().register("stringLengthTest", (String str) -> str.length(), DataTypes.IntegerType); Row result = spark.sql("SELECT stringLengthTest('test')").head(); Assert.assertEquals(4, result.getInt(0)); }
@Test public void isInCollectionWorksCorrectlyOnJava() { List<Row> rows = Arrays.asList( RowFactory.create(1, "x"), RowFactory.create(2, "y"), RowFactory.create(3, "z")); StructType schema = createStructType(Arrays.asList( createStructField("a", IntegerType, false), createStructField("b", StringType, false))); Dataset<Row> df = spark.createDataFrame(rows, schema); // Test with different types of collections Assert.assertTrue(Arrays.equals( (Row[]) df.filter(df.col("a").isInCollection(Arrays.asList(1, 2))).collect(), (Row[]) df.filter((FilterFunction<Row>) r -> r.getInt(0) == 1 || r.getInt(0) == 2).collect() )); Assert.assertTrue(Arrays.equals( (Row[]) df.filter(df.col("a").isInCollection(new HashSet<>(Arrays.asList(1, 2)))).collect(), (Row[]) df.filter((FilterFunction<Row>) r -> r.getInt(0) == 1 || r.getInt(0) == 2).collect() )); Assert.assertTrue(Arrays.equals( (Row[]) df.filter(df.col("a").isInCollection(new ArrayList<>(Arrays.asList(3, 1)))).collect(), (Row[]) df.filter((FilterFunction<Row>) r -> r.getInt(0) == 3 || r.getInt(0) == 1).collect() )); }
@Test public void isInCollectionWorksCorrectlyOnJava() { List<Row> rows = Arrays.asList( RowFactory.create(1, "x"), RowFactory.create(2, "y"), RowFactory.create(3, "z")); StructType schema = createStructType(Arrays.asList( createStructField("a", IntegerType, false), createStructField("b", StringType, false))); Dataset<Row> df = spark.createDataFrame(rows, schema); // Test with different types of collections Assert.assertTrue(Arrays.equals( (Row[]) df.filter(df.col("a").isInCollection(Arrays.asList(1, 2))).collect(), (Row[]) df.filter((FilterFunction<Row>) r -> r.getInt(0) == 1 || r.getInt(0) == 2).collect() )); Assert.assertTrue(Arrays.equals( (Row[]) df.filter(df.col("a").isInCollection(new HashSet<>(Arrays.asList(1, 2)))).collect(), (Row[]) df.filter((FilterFunction<Row>) r -> r.getInt(0) == 1 || r.getInt(0) == 2).collect() )); Assert.assertTrue(Arrays.equals( (Row[]) df.filter(df.col("a").isInCollection(new ArrayList<>(Arrays.asList(3, 1)))).collect(), (Row[]) df.filter((FilterFunction<Row>) r -> r.getInt(0) == 3 || r.getInt(0) == 1).collect() )); }
@SuppressWarnings("unchecked") @Test public void udf6Test() { spark.udf().register("returnOne", () -> 1, DataTypes.IntegerType); Row result = spark.sql("SELECT returnOne()").head(); Assert.assertEquals(1, result.getInt(0)); } }
@SuppressWarnings("unchecked") @Test public void udf6Test() { spark.udf().register("returnOne", () -> 1, DataTypes.IntegerType); Row result = spark.sql("SELECT returnOne()").head(); Assert.assertEquals(1, result.getInt(0)); } }
Assert.assertEquals(shortValue, simpleRow.getShort(3)); Assert.assertEquals(shortValue, simpleRow.get(3)); Assert.assertEquals(intValue, simpleRow.getInt(4)); Assert.assertEquals(intValue, simpleRow.get(4)); Assert.assertEquals(intValue, simpleRow.getInt(5)); Assert.assertEquals(intValue, simpleRow.get(5)); Assert.assertEquals(longValue, simpleRow.getLong(6));
Assert.assertEquals(shortValue, simpleRow.getShort(3)); Assert.assertEquals(shortValue, simpleRow.get(3)); Assert.assertEquals(intValue, simpleRow.getInt(4)); Assert.assertEquals(intValue, simpleRow.get(4)); Assert.assertEquals(intValue, simpleRow.getInt(5)); Assert.assertEquals(intValue, simpleRow.get(5)); Assert.assertEquals(longValue, simpleRow.getLong(6));