Examples with DataSourceV2Relation org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation used on opensource projects

Example 1 with DataSourceV2Relation

use of org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation in project OpenLineage by OpenLineage.

the class AppendDataDatasetBuilderTest method testApplyWithDataSourceV2Relation.

@Test
public void testApplyWithDataSourceV2Relation() {
    AppendData appendData = mock(AppendData.class);
    DataSourceV2Relation relation = mock(DataSourceV2Relation.class);
    OpenLineage.OutputDataset dataset = mock(OpenLineage.OutputDataset.class);
    when(appendData.table()).thenReturn(relation);
    try (MockedStatic mockedPlanUtils3 = mockStatic(PlanUtils3.class)) {
        try (MockedStatic mockedFacetUtils = mockStatic(DatasetVersionDatasetFacetUtils.class)) {
            when(DatasetVersionDatasetFacetUtils.extractVersionFromDataSourceV2Relation(relation)).thenReturn(Optional.of("v2"));
            when(PlanUtils3.fromDataSourceV2Relation(eq(factory), eq(context), eq(relation), any())).thenReturn(Collections.singletonList(dataset));
            List<OpenLineage.OutputDataset> datasets = builder.apply(appendData);
            assertEquals(1, datasets.size());
            assertEquals(dataset, datasets.get(0));
        }
    }
}

Example 2 with DataSourceV2Relation

use of org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation in project OpenLineage by OpenLineage.

the class DataSourceV2ScanRelationInputDatasetBuilderTest method testApply.

@Test
public void testApply() {
    OpenLineage.DatasetFacetsBuilder datasetFacetsBuilder = mock(OpenLineage.DatasetFacetsBuilder.class);
    List<OpenLineage.InputDataset> datasets = mock(List.class);
    DataSourceV2ScanRelation scanRelation = mock(DataSourceV2ScanRelation.class);
    DataSourceV2Relation relation = mock(DataSourceV2Relation.class);
    when(openLineage.newDatasetFacetsBuilder()).thenReturn(datasetFacetsBuilder);
    when(context.getOpenLineage()).thenReturn(openLineage);
    when(scanRelation.relation()).thenReturn(relation);
    try (MockedStatic planUtils3MockedStatic = mockStatic(PlanUtils3.class)) {
        try (MockedStatic facetUtilsMockedStatic = mockStatic(DatasetVersionDatasetFacetUtils.class)) {
            when(PlanUtils3.fromDataSourceV2Relation(factory, context, relation, datasetFacetsBuilder)).thenReturn(datasets);
            assertEquals(datasets, builder.apply(scanRelation));
            facetUtilsMockedStatic.verify(() -> DatasetVersionDatasetFacetUtils.includeDatasetVersion(context, datasetFacetsBuilder, relation), times(1));
        }
    }
}

Example 3 with DataSourceV2Relation

use of org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation in project OpenLineage by OpenLineage.

the class DatasetVersionDatasetFacetUtilsTest method testIncludeDatasetVersion.

@Test
public void testIncludeDatasetVersion() {
    OpenLineage.DatasetFacetsBuilder datasetFacetsBuilder = new OpenLineage.DatasetFacetsBuilder();
    DataSourceV2Relation relation = mock(DataSourceV2Relation.class);
    OpenLineage.DatasetVersionDatasetFacet datasetVersionDatasetFacet = mock(OpenLineage.DatasetVersionDatasetFacet.class);
    when(datasetVersionDatasetFacet.getDatasetVersion()).thenReturn("v2");
    when(openLineage.newDatasetVersionDatasetFacet("v2")).thenReturn(datasetVersionDatasetFacet);
    when(relation.identifier()).thenReturn(Option.apply(identifier));
    when(relation.catalog()).thenReturn(Option.apply(tableCatalog));
    when(relation.table()).thenReturn(table);
    when(table.properties()).thenReturn(tableProperties);
    try (MockedStatic mocked = mockStatic(CatalogUtils3.class)) {
        when(CatalogUtils3.getDatasetVersion(tableCatalog, identifier, tableProperties)).thenReturn(Optional.of("v2"));
        DatasetVersionDatasetFacetUtils.includeDatasetVersion(openLineageContext, datasetFacetsBuilder, relation);
        assertEquals("v2", datasetFacetsBuilder.build().getVersion().getDatasetVersion());
    }
}

Example 4 with DataSourceV2Relation

use of org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation in project OpenLineage by OpenLineage.

the class DatasetVersionDatasetFacetUtilsTest method testIncludeDatasetVersionWhenNoDatasetVersion.

@Test
public void testIncludeDatasetVersionWhenNoDatasetVersion() {
    OpenLineage.DatasetFacetsBuilder datasetFacetsBuilder = new OpenLineage.DatasetFacetsBuilder();
    DataSourceV2Relation relation = mock(DataSourceV2Relation.class);
    when(relation.identifier()).thenReturn(Option.apply(identifier));
    when(relation.catalog()).thenReturn(Option.apply(tableCatalog));
    when(relation.table()).thenReturn(table);
    when(table.properties()).thenReturn(tableProperties);
    try (MockedStatic mocked = mockStatic(CatalogUtils3.class)) {
        when(CatalogUtils3.getDatasetVersion(tableCatalog, identifier, tableProperties)).thenReturn(Optional.empty());
        DatasetVersionDatasetFacetUtils.includeDatasetVersion(openLineageContext, datasetFacetsBuilder, relation);
        assertEquals(null, datasetFacetsBuilder.build().getVersion());
    }
}

Example 5 with DataSourceV2Relation

use of org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation in project iceberg by apache.

the class BaseProcedure method refreshSparkCache.

protected void refreshSparkCache(Identifier ident, Table table) {
    CacheManager cacheManager = spark.sharedState().cacheManager();
    DataSourceV2Relation relation = DataSourceV2Relation.create(table, Option.apply(tableCatalog), Option.apply(ident));
    cacheManager.recacheByPlan(spark, relation);
}