Search in sources :

Example 21 with SQLContext

use of org.apache.spark.sql.SQLContext in project Gaffer by gchq.

the class GetDataFrameOfElementsHandlerTest method checkGetCorrectElementsInDataFrame.

@Test
public void checkGetCorrectElementsInDataFrame() throws OperationException {
    final Graph graph = getGraph("/schema-DataFrame/dataSchema.json", getElements());
    final SQLContext sqlContext = getSqlContext("checkGetCorrectElementsInDataFrame");
    // Edges group - check get correct edges
    GetDataFrameOfElements dfOperation = new GetDataFrameOfElements.Builder().sqlContext(sqlContext).view(new View.Builder().edge(EDGE_GROUP).build()).build();
    Dataset<Row> dataFrame = graph.execute(dfOperation, new User());
    Set<Row> results = new HashSet<>(dataFrame.collectAsList());
    final Set<Row> expectedRows = new HashSet<>();
    for (int i = 0; i < NUM_ELEMENTS; i++) {
        final MutableList<Object> fields1 = new MutableList<>();
        fields1.appendElem(EDGE_GROUP);
        fields1.appendElem("" + i);
        fields1.appendElem("B");
        fields1.appendElem(1);
        fields1.appendElem(2);
        fields1.appendElem(3.0F);
        fields1.appendElem(4.0D);
        fields1.appendElem(5L);
        fields1.appendElem(100L);
        expectedRows.add(Row$.MODULE$.fromSeq(fields1));
        final MutableList<Object> fields2 = new MutableList<>();
        fields2.appendElem(EDGE_GROUP);
        fields2.appendElem("" + i);
        fields2.appendElem("C");
        fields2.appendElem(6);
        fields2.appendElem(7);
        fields2.appendElem(8.0F);
        fields2.appendElem(9.0D);
        fields2.appendElem(10L);
        fields2.appendElem(i * 200L);
        expectedRows.add(Row$.MODULE$.fromSeq(fields2));
    }
    assertEquals(expectedRows, results);
    // Entities group - check get correct entities
    dfOperation = new GetDataFrameOfElements.Builder().sqlContext(sqlContext).view(new View.Builder().entity(ENTITY_GROUP).build()).build();
    dataFrame = graph.execute(dfOperation, new User());
    results.clear();
    results.addAll(dataFrame.collectAsList());
    expectedRows.clear();
    for (int i = 0; i < NUM_ELEMENTS; i++) {
        final MutableList<Object> fields1 = new MutableList<>();
        fields1.clear();
        fields1.appendElem(ENTITY_GROUP);
        fields1.appendElem("" + i);
        fields1.appendElem(1);
        fields1.appendElem(i);
        fields1.appendElem(3.0F);
        fields1.appendElem(4.0D);
        fields1.appendElem(5L);
        fields1.appendElem(6);
        expectedRows.add(Row$.MODULE$.fromSeq(fields1));
    }
    assertEquals(expectedRows, results);
    sqlContext.sparkContext().stop();
}
Also used : GetDataFrameOfElements(uk.gov.gchq.gaffer.spark.operation.dataframe.GetDataFrameOfElements) User(uk.gov.gchq.gaffer.user.User) View(uk.gov.gchq.gaffer.data.elementdefinition.view.View) Graph(uk.gov.gchq.gaffer.graph.Graph) MutableList(scala.collection.mutable.MutableList) Row(org.apache.spark.sql.Row) SQLContext(org.apache.spark.sql.SQLContext) HashSet(java.util.HashSet) Test(org.junit.Test)

Example 22 with SQLContext

use of org.apache.spark.sql.SQLContext in project Gaffer by gchq.

the class GetDataFrameOfElementsHandlerTest method checkViewIsRespected.

@Test
public void checkViewIsRespected() throws OperationException {
    final Graph graph = getGraph("/schema-DataFrame/dataSchema.json", getElements());
    final SQLContext sqlContext = getSqlContext("checkViewIsRespected");
    // Edges group - check get correct edges
    GetDataFrameOfElements dfOperation = new GetDataFrameOfElements.Builder().sqlContext(sqlContext).view(new View.Builder().edge(EDGE_GROUP, new ViewElementDefinition.Builder().preAggregationFilter(new ElementFilter.Builder().select("count").execute(new IsMoreThan(800L)).build()).build()).build()).build();
    Dataset<Row> dataFrame = graph.execute(dfOperation, new User());
    Set<Row> results = new HashSet<>(dataFrame.collectAsList());
    final Set<Row> expectedRows = new HashSet<>();
    for (int i = 0; i < NUM_ELEMENTS; i++) {
        if (i * 200L > 800L) {
            final MutableList<Object> fields2 = new MutableList<>();
            fields2.appendElem(EDGE_GROUP);
            fields2.appendElem("" + i);
            fields2.appendElem("C");
            fields2.appendElem(6);
            fields2.appendElem(7);
            fields2.appendElem(8.0F);
            fields2.appendElem(9.0D);
            fields2.appendElem(10L);
            fields2.appendElem(i * 200L);
            expectedRows.add(Row$.MODULE$.fromSeq(fields2));
        }
    }
    assertEquals(expectedRows, results);
    // Entities group - check get correct entities
    dfOperation = new GetDataFrameOfElements.Builder().sqlContext(sqlContext).view(new View.Builder().entity(ENTITY_GROUP, new ViewElementDefinition.Builder().postAggregationFilter(new ElementFilter.Builder().select("property1").execute(new IsMoreThan(1)).build()).build()).build()).build();
    dataFrame = graph.execute(dfOperation, new User());
    results.clear();
    results.addAll(dataFrame.collectAsList());
    expectedRows.clear();
    for (int i = 2; i < NUM_ELEMENTS; i++) {
        final MutableList<Object> fields1 = new MutableList<>();
        fields1.clear();
        fields1.appendElem(ENTITY_GROUP);
        fields1.appendElem("" + i);
        fields1.appendElem(1);
        fields1.appendElem(i);
        fields1.appendElem(3.0F);
        fields1.appendElem(4.0D);
        fields1.appendElem(5L);
        fields1.appendElem(6);
        expectedRows.add(Row$.MODULE$.fromSeq(fields1));
    }
    assertEquals(expectedRows, results);
    sqlContext.sparkContext().stop();
}
Also used : GetDataFrameOfElements(uk.gov.gchq.gaffer.spark.operation.dataframe.GetDataFrameOfElements) User(uk.gov.gchq.gaffer.user.User) ViewElementDefinition(uk.gov.gchq.gaffer.data.elementdefinition.view.ViewElementDefinition) View(uk.gov.gchq.gaffer.data.elementdefinition.view.View) Graph(uk.gov.gchq.gaffer.graph.Graph) MutableList(scala.collection.mutable.MutableList) ElementFilter(uk.gov.gchq.gaffer.data.element.function.ElementFilter) Row(org.apache.spark.sql.Row) IsMoreThan(uk.gov.gchq.gaffer.function.filter.IsMoreThan) SQLContext(org.apache.spark.sql.SQLContext) HashSet(java.util.HashSet) Test(org.junit.Test)

Example 23 with SQLContext

use of org.apache.spark.sql.SQLContext in project Gaffer by gchq.

the class GetDataFrameOfElementsHandlerTest method checkGetCorrectElementsInDataFrameWithProjection.

@Test
public void checkGetCorrectElementsInDataFrameWithProjection() throws OperationException {
    final Graph graph = getGraph("/schema-DataFrame/dataSchema.json", getElements());
    final SQLContext sqlContext = getSqlContext("checkGetCorrectElementsInDataFrameWithProjection");
    // Get all edges
    final GetDataFrameOfElements dfOperation = new GetDataFrameOfElements.Builder().sqlContext(sqlContext).view(new View.Builder().edge(EDGE_GROUP).build()).build();
    final Dataset<Row> dataFrame = graph.execute(dfOperation, new User());
    // Check get correct rows when ask for src, dst and property2 columns
    Set<Row> results = new HashSet<>(dataFrame.select("src", "dst", "property2").collectAsList());
    final Set<Row> expectedRows = new HashSet<>();
    for (int i = 0; i < NUM_ELEMENTS; i++) {
        MutableList<Object> fields1 = new MutableList<>();
        fields1.appendElem("" + i);
        fields1.appendElem("B");
        fields1.appendElem(3.0F);
        expectedRows.add(Row$.MODULE$.fromSeq(fields1));
        MutableList<Object> fields2 = new MutableList<>();
        fields2.appendElem("" + i);
        fields2.appendElem("C");
        fields2.appendElem(8.0F);
        expectedRows.add(Row$.MODULE$.fromSeq(fields2));
    }
    assertEquals(expectedRows, results);
    // Check get correct rows when ask for just property2 column
    results = new HashSet<>(dataFrame.select("property2").collectAsList());
    expectedRows.clear();
    for (int i = 0; i < NUM_ELEMENTS; i++) {
        MutableList<Object> fields1 = new MutableList<>();
        fields1.appendElem(3.0F);
        expectedRows.add(Row$.MODULE$.fromSeq(fields1));
        MutableList<Object> fields2 = new MutableList<>();
        fields2.appendElem(8.0F);
        expectedRows.add(Row$.MODULE$.fromSeq(fields2));
    }
    assertEquals(expectedRows, results);
    sqlContext.sparkContext().stop();
}
Also used : GetDataFrameOfElements(uk.gov.gchq.gaffer.spark.operation.dataframe.GetDataFrameOfElements) User(uk.gov.gchq.gaffer.user.User) Graph(uk.gov.gchq.gaffer.graph.Graph) MutableList(scala.collection.mutable.MutableList) Row(org.apache.spark.sql.Row) SQLContext(org.apache.spark.sql.SQLContext) HashSet(java.util.HashSet) Test(org.junit.Test)

Example 24 with SQLContext

use of org.apache.spark.sql.SQLContext in project Gaffer by gchq.

the class GetDataFrameOfElementsHandler method doOperation.

public Dataset<Row> doOperation(final GetDataFrameOfElements operation, final Context context, final AccumuloStore store) throws OperationException {
    final SQLContext sqlContext = operation.getSqlContext();
    final AccumuloStoreRelation relation = new AccumuloStoreRelation(sqlContext, operation.getConverters(), operation.getView(), store, context.getUser());
    return sqlContext.baseRelationToDataFrame(relation);
}
Also used : SQLContext(org.apache.spark.sql.SQLContext)

Example 25 with SQLContext

use of org.apache.spark.sql.SQLContext in project Gaffer by gchq.

the class FilterToOperationConverterTest method testIncompatibleGroups.

@Test
public void testIncompatibleGroups() throws OperationException {
    final Schema schema = getSchema();
    final SQLContext sqlContext = getSqlContext("testIncompatibleGroups");
    final Filter[] filters = new Filter[2];
    filters[0] = new EqualTo(SchemaToStructTypeConverter.GROUP, "A");
    filters[1] = new EqualTo(SchemaToStructTypeConverter.GROUP, "B");
    final FiltersToOperationConverter converter = new FiltersToOperationConverter(sqlContext, getViewFromSchema(schema), schema, filters);
    final AbstractGetRDD<?> operation = converter.getOperation();
    assertNull(operation);
    sqlContext.sparkContext().stop();
}
Also used : Filter(org.apache.spark.sql.sources.Filter) Schema(uk.gov.gchq.gaffer.store.schema.Schema) SQLContext(org.apache.spark.sql.SQLContext) EqualTo(org.apache.spark.sql.sources.EqualTo) Test(org.junit.Test)

Aggregations

SQLContext (org.apache.spark.sql.SQLContext)34 Test (org.junit.Test)20 HashSet (java.util.HashSet)15 Schema (uk.gov.gchq.gaffer.store.schema.Schema)14 Row (org.apache.spark.sql.Row)12 Filter (org.apache.spark.sql.sources.Filter)11 User (uk.gov.gchq.gaffer.user.User)11 View (uk.gov.gchq.gaffer.data.elementdefinition.view.View)10 EqualTo (org.apache.spark.sql.sources.EqualTo)9 Graph (uk.gov.gchq.gaffer.graph.Graph)9 GetDataFrameOfElements (uk.gov.gchq.gaffer.spark.operation.dataframe.GetDataFrameOfElements)8 JavaSparkContext (org.apache.spark.api.java.JavaSparkContext)7 MutableList (scala.collection.mutable.MutableList)7 ArrayList (java.util.ArrayList)6 DataFrame (org.apache.spark.sql.DataFrame)5 IsMoreThan (uk.gov.gchq.gaffer.function.filter.IsMoreThan)5 EntitySeed (uk.gov.gchq.gaffer.operation.data.EntitySeed)5 GetRDDOfElements (uk.gov.gchq.gaffer.spark.operation.scalardd.GetRDDOfElements)5 SparkConf (org.apache.spark.SparkConf)4 SparkContext (org.apache.spark.SparkContext)4