sedona/python/tests/spatial_rdd/test_spatial_rdd_writer.py at 8440048b06643f6fdfeaeadac7e3f86acf9a217b · apache/sedona · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
#  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.

import os
import shutil

import pytest

from sedona.core.SpatialRDD import PointRDD
from sedona.core.enums import FileDataSplitter
from sedona.core.geom.envelope import Envelope
from tests.test_base import TestBase
from tests.tools import tests_resource

wkb_folder = "wkb"
wkt_folder = "wkt"

test_save_as_wkb_with_data = os.path.join(
    tests_resource, wkb_folder, "testSaveAsWKBWithData"
)
test_save_as_wkb = os.path.join(tests_resource, wkb_folder, "testSaveAsWKB")
test_save_as_empty_wkb = os.path.join(tests_resource, wkb_folder, "testSaveAsEmptyWKB")
test_save_as_wkt = os.path.join(tests_resource, wkt_folder, "testSaveAsWKT")
test_save_as_wkt_with_data = os.path.join(
    tests_resource, wkt_folder, "testSaveAsWKTWithData"
)

inputLocation = os.path.join(tests_resource, "arealm-small.csv")
queryWindowSet = os.path.join(tests_resource, "zcta510-small.csv")
offset = 1
splitter = FileDataSplitter.CSV
gridType = "rtree"
indexType = "rtree"
numPartitions = 11
distance = 0.01
queryPolygonSet = "primaryroads-polygon.csv"
inputCount = 3000
inputBoundary = Envelope(
    minx=-173.120769, maxx=-84.965961, miny=30.244859, maxy=71.355134
)
rectangleMatchCount = 103
rectangleMatchWithOriginalDuplicatesCount = 103
polygonMatchCount = 472
polygonMatchWithOriginalDuplicatesCount = 562


## todo add missing tests
def remove_directory(path: str) -> bool:
    try:
        shutil.rmtree(path)
    except Exception as e:
        return False
    return True


@pytest.fixture
def remove_wkb_directory():
    remove_directory(test_save_as_wkb_with_data)


class TestSpatialRDDWriter(TestBase):

    def test_save_as_geo_json_with_data(self):
        spatial_rdd = PointRDD(
            sparkContext=self.sc,
            InputLocation=inputLocation,
            Offset=offset,
            splitter=splitter,
            carryInputData=True,
            partitions=numPartitions,
        )

        spatial_rdd.saveAsGeoJSON(test_save_as_wkb_with_data)

        result_wkb = PointRDD(
            sparkContext=self.sc,
            InputLocation=test_save_as_wkb_with_data,
            splitter=FileDataSplitter.GEOJSON,
            carryInputData=True,
            partitions=numPartitions,
        )

        assert result_wkb.rawSpatialRDD.count() == spatial_rdd.rawSpatialRDD.count()