from dask.dataframe.utils import (
meta_nonempty,
_nonempty_index,
meta_nonempty_dataframe,
)
from dask.dataframe.extensions import make_array_nonempty, make_scalar
from dask.base import normalize_token
import shapely.geometry
from shapely.geometry.base import BaseGeometry
import geopandas
from geopandas.array import GeometryArray, GeometryDtype, from_shapely
from .core import GeoSeries, GeoDataFrame
get_parallel_type.register(geopandas.GeoDataFrame, lambda _: GeoDataFrame)
get_parallel_type.register(geopandas.GeoSeries, lambda _: GeoSeries)
@make_meta.register(BaseGeometry)
def make_meta_shapely_geometry(x, index=None):
return x
@make_array_nonempty.register(GeometryDtype)
def _(dtype):
return from_shapely(
[shapely.geometry.LineString([(i, i), (i, i + 1)]) for i in range(2)])
@make_scalar.register(GeometryDtype.type)
tolist_dispatch,
)
from dask.dataframe.utils import (
UNKNOWN_CATEGORIES,
_nonempty_scalar,
_scalar_from_dtype,
is_arraylike,
is_scalar,
)
import cudf
from cudf.utils.dtypes import is_string_dtype
from .core import DataFrame, Index, Series
get_parallel_type.register(cudf.DataFrame, lambda _: DataFrame)
get_parallel_type.register(cudf.Series, lambda _: Series)
get_parallel_type.register(cudf.Index, lambda _: Index)
@meta_nonempty.register(cudf.Index)
def _nonempty_index(idx):
if isinstance(idx, cudf.core.index.RangeIndex):
return cudf.core.index.RangeIndex(2, name=idx.name)
elif isinstance(idx, cudf.core.index.DatetimeIndex):
start = "1970-01-01"
data = np.array([start, "1970-01-02"], dtype=idx.dtype)
values = cudf.core.column.as_column(data)
return cudf.core.index.DatetimeIndex(values, name=idx.name)
elif isinstance(idx, cudf.core.index.StringIndex):
return cudf.core.index.StringIndex(["cat", "dog"], name=idx.name)
