Source code for langchain.document_loaders.pyspark_dataframe
import itertools
import logging
import sys
from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple
from langchain.docstore.document import Document
from langchain.document_loaders.base import BaseLoader
logger = logging.getLogger(__file__)
if TYPE_CHECKING:
    from pyspark.sql import SparkSession
[docs]class PySparkDataFrameLoader(BaseLoader):
    """Load `PySpark` DataFrames."""
[docs]    def __init__(
        self,
        spark_session: Optional["SparkSession"] = None,
        df: Optional[Any] = None,
        page_content_column: str = "text",
        fraction_of_memory: float = 0.1,
    ):
        """Initialize with a Spark DataFrame object.
        Args:
            spark_session: The SparkSession object.
            df: The Spark DataFrame object.
            page_content_column: The name of the column containing the page content.
             Defaults to "text".
            fraction_of_memory: The fraction of memory to use. Defaults to 0.1.
        """
        try:
            from pyspark.sql import DataFrame, SparkSession
        except ImportError:
            raise ImportError(
                "pyspark is not installed. "
                "Please install it with `pip install pyspark`"
            )
        self.spark = (
            spark_session if spark_session else SparkSession.builder.getOrCreate()
        )
        if not isinstance(df, DataFrame):
            raise ValueError(
                f"Expected data_frame to be a PySpark DataFrame, got {type(df)}"
            )
        self.df = df
        self.page_content_column = page_content_column
        self.fraction_of_memory = fraction_of_memory