[GH-3493] Document the raster, type, map and GeoArrow Python APIs

python/sedona/spark/geoarrow/geoarrow.py

@@ -51,7 +51,7 @@ def dataframe_to_arrow(df, crs=None):
       `pyproj.CRS`). If provided, this will override any CRS present in the output
       geometries. If omitted, the CRS will be inferred from the values present in
       the output if exactly one CRS is present in the output.
-    :return:
+    :return: A ``pyarrow.Table``.
     """
     import pyarrow as pa
 
@@ -311,6 +311,19 @@ def infer_schema(gdf: "gpd.GeoDataFrame") -> StructType:
 # Modified backport from Spark 4.0
 # https://github.com/apache/spark/blob/3515b207c41d78194d11933cd04bddc21f8418dd/python/pyspark/sql/pandas/conversion.py#L632
 def create_spatial_dataframe(spark: SparkSession, gdf: "gpd.GeoDataFrame") -> DataFrame:
+    """
+    Create a Spark DataFrame from a GeoPandas GeoDataFrame, transferring the data with Arrow.
+
+    This is faster than ``spark.createDataFrame(gdf)`` for large GeoDataFrames.
+    Geometry columns become Sedona geometry columns, and the other columns get the
+    Spark types of their Arrow types.
+
+    :param spark: The Spark session.
+    :param gdf: The GeoPandas GeoDataFrame to convert. It must have at least one
+      geometry column.
+    :return: A Spark DataFrame with the columns of ``gdf``.
+    :raises ValueError: If ``gdf`` has no geometry column.
+    """
     from pyspark.sql.pandas.types import (
         to_arrow_type,
     )

python/sedona/spark/maps/SedonaKepler.py

@@ -20,6 +20,10 @@
 
 
 class SedonaKepler:
+    """Plot Sedona DataFrames on Kepler.gl maps in a Jupyter notebook.
+
+    Requires the ``keplergl`` package, installed with ``apache-sedona[kepler-map]``.
+    """
 
     @classmethod
     def create_map(
@@ -33,14 +37,15 @@ def create_map(
     ):
         """
         Creates a map visualization using kepler, optionally taking a sedona dataFrame as data input
+
         :param df: [Optional] SedonaDataFrame to plot on the map
         :param name: [Optional] Name to be associated with the given
-        dataframe, if a df is passed with no name, a default name of 'unnamed' is set for it.
-        param config: [Optional] A map config to be applied to the rendered map
-        param height: [Optional] Height of the map in pixels, default is 600
-        param use_arrow: [Optional] Whether to use arrow for data transfer, default is False
-        param show_docs: [Optional] Whether to show the documentation, default is False
-        :return: A map object
+            dataframe, if a df is passed with no name, a default name of 'unnamed' is set for it.
+        :param config: [Optional] A Kepler.gl map config to be applied to the rendered map
+        :param height: [Optional] Height of the map in pixels, default is 600
+        :param use_arrow: [Optional] Whether to use Arrow to transfer the data to the map, default is False
+        :param show_docs: [Optional] Whether to show the Kepler.gl documentation, default is False
+        :return: A ``keplergl.KeplerGl`` map object
         """
 
         try:
@@ -61,7 +66,8 @@ def create_map(
     @classmethod
     def add_df(cls, kepler_map, df, name="unnamed"):
         """
-        Adds a SedonaDataFrame to a given map object.
+        Adds a SedonaDataFrame to a given map object. Raster columns are left out.
+
         :param kepler_map: Map object to add SedonaDataFrame to
         :param df: SedonaDataFrame to add
         :param name: [Optional] Name to assign to the dataframe, default name assigned is 'unnamed'

python/sedona/spark/maps/SedonaPyDeck.py

@@ -31,6 +31,10 @@
 
 
 class SedonaPyDeck:
+    """Plot Sedona DataFrames on pydeck maps in a Jupyter notebook.
+
+    Requires the ``pydeck`` package, installed with ``apache-sedona[pydeck-map]``.
+    """
 
     # User Facing APIs
     @classmethod
@@ -127,14 +131,16 @@ def create_geometry_map(
     ):
         """
         Create a pydeck map with a GeoJsonLayer added for plotting given geometries.
-        :param line_color:
+
         :param df: SedonaDataFrame with polygons
         :param fill_color: Optional color for the plotted polygons
+        :param line_color: Optional color for the outlines of the plotted geometries
         :param elevation_col: Optional column/numeric value to determine elevation for plotted polygons
         :param initial_view_state: optional initial view state of the pydeck map
         :param map_style: optional map_style of the pydeck map
         :param map_provider: optional map_provider of the pydeck map
         :param api_keys: Optional dictionary of API keys for Map providers
+        :param stroked: Whether to draw the outlines of the polygons
         :return: A pydeck map with a GeoJsonLayer map added
         """
         pdk = _try_import_pydeck()
@@ -187,12 +193,14 @@ def create_scatterplot_map(
     ):
         """
         Create a pydeck map with a scatterplot layer
-        :param radius_scale:
-        :param radius_max_pixels:
-        :param radius_min_pixels:
-        :param radius_col:
+
         :param df: SedonaDataFrame to plot
         :param fill_color: color of the points
+        :param radius_col: Column holding the radius of each point in meters, or a
+            number to use as the radius of every point
+        :param radius_min_pixels: Minimum radius of a point on screen, in pixels
+        :param radius_max_pixels: Maximum radius of a point on screen, in pixels
+        :param radius_scale: Multiplier applied to every radius
         :param initial_view_state: optional initial view state of a pydeck map
         :param map_style: optional map_style to be added to the pydeck map
         :param map_provider: optional map_provider to be added to the pydeck map
@@ -243,6 +251,7 @@ def create_heatmap(
     ):
         """
         Create a pydeck map with a heatmap layer added
+
         :param df: SedonaDataFrame to be used to plot the heatmap
         :param color_range: Optional color_range for the heatmap
         :param weight: Optional column to determine weight of each point while plotting the heatmap

python/sedona/spark/raster/meta.py

@@ -32,6 +32,32 @@ class PixelAnchor(Enum):
 
 
 class AffineTransform:
+    """Affine transform from pixel coordinates to world coordinates.
+
+    The pixel at column ``col`` and row ``row`` maps to
+    ``x = ip_x + col * scale_x + row * skew_x`` and
+    ``y = ip_y + col * skew_y + row * scale_y``. ``repr()`` shows the transform
+    as a 3x3 matrix.
+
+    Parameters
+    ----------
+    scale_x : float
+        Pixel width in world units.
+    skew_y : float
+        Change in y per column; 0 for a raster without rotation.
+    skew_x : float
+        Change in x per row; 0 for a raster without rotation.
+    scale_y : float
+        Pixel height in world units, negative for a north-up raster.
+    ip_x : float
+        X coordinate of the upper-left pixel.
+    ip_y : float
+        Y coordinate of the upper-left pixel.
+    pixel_anchor : PixelAnchor
+        Whether ``ip_x`` and ``ip_y`` refer to the center or the upper-left
+        corner of that pixel.
+    """
+
     scale_x: float
     skew_y: float
     skew_x: float
@@ -52,11 +78,39 @@ def __init__(
         self.pixel_anchor = pixel_anchor
 
     def with_anchor(self, pixel_anchor: PixelAnchor):
+        """Return the same transform expressed for another pixel anchor.
+
+        Parameters
+        ----------
+        pixel_anchor : PixelAnchor
+            The pixel anchor of the returned transform.
+
+        Returns
+        -------
+        AffineTransform
+            This transform if it already uses ``pixel_anchor``; otherwise a new
+            transform whose origin is shifted by half a pixel.
+        """
         if pixel_anchor == self.pixel_anchor:
             return self
         return self._do_change_pixel_anchor(self.pixel_anchor, pixel_anchor)
 
     def translate(self, offset_x: float, offset_y: float):
+        """Return the transform of a raster whose origin is offset by whole pixels.
+
+        Parameters
+        ----------
+        offset_x : float
+            Offset in pixel columns.
+        offset_y : float
+            Offset in pixel rows.
+
+        Returns
+        -------
+        AffineTransform
+            A new transform whose origin is pixel (``offset_x``, ``offset_y``)
+            of this one, for example the transform of a tile cut from the raster.
+        """
         new_ipx = self.ip_x + offset_x * self.scale_x + offset_y * self.skew_x
         new_ipy = self.ip_y + offset_x * self.skew_y + offset_y * self.scale_y
         return AffineTransform(
@@ -111,7 +165,19 @@ def __repr__(self):
 
 
 class SampleDimension:
-    """Raster band metadata."""
+    """Raster band metadata.
+
+    Parameters
+    ----------
+    description : str
+        Description of the band.
+    offset : float
+        Offset of the band: a stored pixel value means ``value * scale + offset``.
+    scale : float
+        Scale of the band.
+    nodata : float
+        NODATA value of the band, or NaN if the band has none.
+    """
 
     description: str
     offset: float

python/sedona/spark/raster/sedona_raster.py

@@ -206,6 +206,31 @@ def _generate_vrt_xml(
 
 
 class SedonaRaster(ABC):
+    """A raster value from a Sedona raster column.
+
+    ``DataFrame.collect()`` and Python UDFs receive raster column values as
+    ``SedonaRaster`` objects. Read the pixels with :meth:`as_numpy` or
+    :meth:`as_rasterio`, and the georeferencing with :attr:`affine_trans` and
+    :attr:`crs_wkt`. Call :meth:`close`, or use the raster as a context manager,
+    to release the resources it holds.
+
+    Sedona creates these objects when it reads a raster value; you do not
+    construct them yourself.
+
+    Parameters
+    ----------
+    width : int
+        Width of the raster in pixels.
+    height : int
+        Height of the raster in pixels.
+    bands_meta : list of SampleDimension
+        Metadata of each band.
+    affine_trans : AffineTransform
+        Transform from pixel coordinates to world coordinates.
+    crs_wkt : str
+        CRS of the raster as a WKT string.
+    """
+
     _width: int
     _height: int
     _bands_meta: List[SampleDimension]
@@ -297,7 +322,24 @@ def with_bands(
     ) -> "SedonaRaster":
         """Replace pixel data, preserving spatial metadata.
 
-        Only supported on InDbSedonaRaster.
+        Only supported on in-db rasters (``InDbSedonaRaster``); other rasters
+        raise ``TypeError``.
+
+        Parameters
+        ----------
+        new_data : np.ndarray
+            New pixel data, of shape (height, width) for a single band or
+            (bands, height, width). Height and width must match this raster;
+            the band count and dtype may differ.
+        nodata : float or sequence of float, optional
+            NODATA value of the output bands: a scalar for every band, or one
+            value per band. Use ``float('nan')`` for a band without NODATA. When
+            omitted, the output bands inherit NODATA from the source bands.
+
+        Returns
+        -------
+        SedonaRaster
+            A new raster with the given pixel data and the same spatial metadata.
         """
         raise TypeError(f"with_bands() is not supported on {type(self).__name__}.")
 

python/sedona/spark/raster_utils/SedonaUtils.py

@@ -19,15 +19,24 @@
 
 
 class SedonaUtils:
+    """Notebook display helpers for Sedona DataFrames."""
+
     @classmethod
     def display_image(cls, df):
         """Display raster images in a Jupyter notebook.
 
         Accepts DataFrames with either:
+
         - A raster column (GridCoverage2D) — auto-applies RS_AsImage
         - An HTML image column from RS_AsImage() — renders directly
 
         Falls back to the SedonaMapUtils HTML table path for other DataFrames.
+
+        Parameters
+        ----------
+        df : pyspark.sql.DataFrame
+            The DataFrame to display. All of its rows are collected to the
+            driver, so limit large DataFrames first.
         """
         from IPython.display import HTML, display
 

python/sedona/spark/sql/types.py

@@ -45,6 +45,21 @@
 
 
 class GeometryType(UserDefinedType):
+    """Spark SQL type of Sedona geometry columns.
+
+    Geometry values are Shapely geometries in Python: ``DataFrame.collect()``
+    returns Shapely geometry objects, and Shapely geometries passed to
+    ``createDataFrame`` are stored as geometries. Use ``GeometryType()`` in a
+    schema to declare a geometry column.
+
+    Examples
+    --------
+    >>> from pyspark.sql.types import StructField, StructType
+    >>> from shapely.geometry import Point
+    >>> from sedona.spark.sql.types import GeometryType
+    >>> schema = StructType([StructField("geom", GeometryType())])
+    >>> df = sedona.createDataFrame([(Point(1, 2),)], schema)  # doctest: +SKIP
+    """
 
     @classmethod
     def sqlType(cls):
@@ -70,6 +85,13 @@ def scalaUDT(cls):
 
 
 class GeographyType(UserDefinedType):
+    """Spark SQL type of Sedona geography columns.
+
+    Geography values are :class:`sedona.spark.core.geom.geography.Geography`
+    objects in Python, which hold a Shapely geometry in their ``geometry``
+    attribute. ``DataFrame.collect()`` returns them, and ``Geography`` objects
+    passed to ``createDataFrame`` are stored as geographies.
+    """
 
     @classmethod
     def sqlType(cls):
@@ -95,6 +117,12 @@ def scalaUDT(cls):
 
 
 class Box2DType(UserDefinedType):
+    """Spark SQL type of Sedona ``box2d`` columns, such as the result of ``ST_Box2D``.
+
+    Values are :class:`sedona.spark.core.geom.box2d.Box2D` objects in Python, with
+    ``xmin``, ``ymin``, ``xmax`` and ``ymax`` attributes. A column of this type is
+    stored as a struct of those four doubles.
+    """
 
     @classmethod
     def sqlType(cls):
@@ -126,6 +154,12 @@ def scalaUDT(cls):
 
 
 class Box3DType(UserDefinedType):
+    """Spark SQL type of Sedona ``box3d`` columns, such as the result of ``ST_Box3D``.
+
+    Values are :class:`sedona.spark.core.geom.box3d.Box3D` objects in Python, with
+    ``xmin``, ``ymin``, ``zmin``, ``xmax``, ``ymax`` and ``zmax`` attributes. A
+    column of this type is stored as a struct of those six doubles.
+    """
 
     @classmethod
     def sqlType(cls):
@@ -159,6 +193,14 @@ def scalaUDT(cls):
 
 
 class RasterType(UserDefinedType):
+    """Spark SQL type of Sedona raster columns.
+
+    Raster values are :class:`sedona.spark.raster.SedonaRaster` objects in
+    Python, which ``DataFrame.collect()`` and Python UDFs receive. Converting
+    raster values requires ``rasterio``; without it, they raise
+    ``NotImplementedError``. Only in-db rasters can be converted back to a
+    raster column, for example when a UDF returns one.
+    """
 
     @classmethod
     def sqlType(cls):