/
githubmirror
/
spark
Обзор
Документация
Войти
/
githubmirror
/
spark
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
python/pyspark/sql/connect/readwriter.py
1 102 строки
38 KB
Lavan Vivekanandasarma
[SPARK-56854][PYTHON] Filter None values in DataFrame[Stream]Reader/Writer .option(s)
15 май 2026, 00:50
15 май 2026, 00:50
f1cdb6d
Код
Авторство
О чём код?
# # 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. # from typing import Dict from typing import Optional, Union, List, overload, Tuple, cast, Callable from typing import TYPE_CHECKING from pyspark.sql.connect.plan import ( Read, RelationChanges, DataSource, LogicalPlan, WriteOperation, WriteOperationV2, Parse, ) import pyspark.sql.connect.proto as proto from pyspark.sql.types import StructType from pyspark.sql.utils import to_str from pyspark.sql.readwriter import ( DataFrameWriter as PySparkDataFrameWriter, DataFrameReader as PySparkDataFrameReader, DataFrameWriterV2 as PySparkDataFrameWriterV2, ) from pyspark.errors import ( AnalysisException, PySparkAttributeError, PySparkTypeError, PySparkValueError, ) from pyspark.sql.connect.functions import builtin as F if TYPE_CHECKING: from pyspark.sql.connect.dataframe import DataFrame from pyspark.sql.connect._typing import ColumnOrName, OptionalPrimitiveType from pyspark.sql.connect.session import SparkSession from pyspark.sql.metrics import ExecutionInfo __all__ = ["DataFrameReader", "DataFrameWriter"] PathOrPaths = Union[str, List[str]] TupleOrListOfString = Union[List[str], Tuple[str, ...]] class OptionUtils: def _set_opts( self, schema: Optional[Union[StructType, str]] = None, **options: "OptionalPrimitiveType", ) -> None: """ Set named options (filter out those the value is None) """ if schema is not None: self.schema(schema) # type: ignore[attr-defined] for k, v in options.items(): if v is not None: self.option(k, v) # type: ignore[attr-defined] class DataFrameReader(OptionUtils): # TODO(SPARK-40539) Achieve parity with PySpark. def __init__(self, client: "SparkSession"): self._client = client self._format: Optional[str] = None self._schema = "" self._options: Dict[str, str] = {} def format(self, source: str) -> "DataFrameReader": self._format = source return self format.__doc__ = PySparkDataFrameReader.format.__doc__ def schema(self, schema: Union[StructType, str]) -> "DataFrameReader": if isinstance(schema, StructType): self._schema = schema.json() elif isinstance(schema, str): self._schema = schema else: raise PySparkTypeError( errorClass="NOT_EXPECTED_TYPE", messageParameters={ "arg_name": "schema", "expected_type": "str or struct type", "arg_type": type(schema).__name__, }, ) return self schema.__doc__ = PySparkDataFrameReader.schema.__doc__ def option(self, key: str, value: "OptionalPrimitiveType") -> "DataFrameReader": self._options[key] = cast(str, to_str(value)) return self option.__doc__ = PySparkDataFrameReader.option.__doc__ def options(self, **options: "OptionalPrimitiveType") -> "DataFrameReader": for k in options: self.option(k, to_str(options[k])) return self options.__doc__ = PySparkDataFrameReader.options.__doc__ def load( self, path: Optional[PathOrPaths] = None, format: Optional[str] = None, schema: Optional[Union[StructType, str]] = None, **options: "OptionalPrimitiveType", ) -> "DataFrame": if format is not None: self.format(format) if schema is not None: self.schema(schema) self.options(**options) paths = path if isinstance(path, str): paths = [path] plan = DataSource( format=self._format, schema=self._schema, options=self._options, paths=paths, # type: ignore[arg-type] ) return self._df(plan) load.__doc__ = PySparkDataFrameReader.load.__doc__ def _df(self, plan: LogicalPlan) -> "DataFrame": from pyspark.sql.connect.dataframe import DataFrame return DataFrame(plan, self._client) def table(self, tableName: str) -> "DataFrame": return self._df(Read(tableName, self._options)) table.__doc__ = PySparkDataFrameReader.table.__doc__ def changes(self, tableName: str) -> "DataFrame": if self._schema: raise AnalysisException( "User specified schema not supported with `changes`.", errorClass="_LEGACY_ERROR_TEMP_1189", messageParameters={"operation": "changes"}, ) return self._df(RelationChanges(tableName, self._options)) changes.__doc__ = PySparkDataFrameReader.changes.__doc__ def json( self, path: Union[PathOrPaths, "DataFrame"], schema: Optional[Union[StructType, str]] = None, primitivesAsString: Optional[Union[bool, str]] = None, prefersDecimal: Optional[Union[bool, str]] = None, allowComments: Optional[Union[bool, str]] = None, allowUnquotedFieldNames: Optional[Union[bool, str]] = None, allowSingleQuotes: Optional[Union[bool, str]] = None, allowNumericLeadingZero: Optional[Union[bool, str]] = None, allowBackslashEscapingAnyCharacter: Optional[Union[bool, str]] = None, mode: Optional[str] = None, columnNameOfCorruptRecord: Optional[str] = None, dateFormat: Optional[str] = None, timestampFormat: Optional[str] = None, multiLine: Optional[Union[bool, str]] = None, allowUnquotedControlChars: Optional[Union[bool, str]] = None, lineSep: Optional[str] = None, samplingRatio: Optional[Union[float, str]] = None, dropFieldIfAllNull: Optional[Union[bool, str]] = None, encoding: Optional[str] = None, locale: Optional[str] = None, pathGlobFilter: Optional[Union[bool, str]] = None, recursiveFileLookup: Optional[Union[bool, str]] = None, modifiedBefore: Optional[Union[bool, str]] = None, modifiedAfter: Optional[Union[bool, str]] = None, allowNonNumericNumbers: Optional[Union[bool, str]] = None, useUnsafeRow: Optional[Union[bool, str]] = None, ) -> "DataFrame": self._set_opts( primitivesAsString=primitivesAsString, prefersDecimal=prefersDecimal, allowComments=allowComments, allowUnquotedFieldNames=allowUnquotedFieldNames, allowSingleQuotes=allowSingleQuotes, allowNumericLeadingZero=allowNumericLeadingZero, allowBackslashEscapingAnyCharacter=allowBackslashEscapingAnyCharacter, mode=mode, columnNameOfCorruptRecord=columnNameOfCorruptRecord, dateFormat=dateFormat, timestampFormat=timestampFormat, multiLine=multiLine, allowUnquotedControlChars=allowUnquotedControlChars, lineSep=lineSep, samplingRatio=samplingRatio, dropFieldIfAllNull=dropFieldIfAllNull, encoding=encoding, locale=locale, pathGlobFilter=pathGlobFilter, recursiveFileLookup=recursiveFileLookup, modifiedBefore=modifiedBefore, modifiedAfter=modifiedAfter, allowNonNumericNumbers=allowNonNumericNumbers, useUnsafeRow=useUnsafeRow, ) if isinstance(path, str): path = [path] if isinstance(path, list): return self.load(path=path, format="json", schema=schema) from pyspark.sql.connect.dataframe import DataFrame if isinstance(path, DataFrame): # Schema must be set explicitly here because the DataFrame path # bypasses load(), which normally calls self.schema(schema). if schema is not None: self.schema(schema) return self._df( Parse( child=path._plan, format=proto.Parse.ParseFormat.PARSE_FORMAT_JSON, schema=self._schema, options=self._options, ) ) raise PySparkTypeError( errorClass="NOT_EXPECTED_TYPE", messageParameters={ "arg_name": "path", "expected_type": "str, list, or DataFrame", "arg_type": type(path).__name__, }, ) json.__doc__ = PySparkDataFrameReader.json.__doc__ def parquet(self, *paths: str, **options: "OptionalPrimitiveType") -> "DataFrame": mergeSchema = options.get("mergeSchema", None) pathGlobFilter = options.get("pathGlobFilter", None) modifiedBefore = options.get("modifiedBefore", None) modifiedAfter = options.get("modifiedAfter", None) recursiveFileLookup = options.get("recursiveFileLookup", None) datetimeRebaseMode = options.get("datetimeRebaseMode", None) int96RebaseMode = options.get("int96RebaseMode", None) self._set_opts( mergeSchema=mergeSchema, pathGlobFilter=pathGlobFilter, recursiveFileLookup=recursiveFileLookup, modifiedBefore=modifiedBefore, modifiedAfter=modifiedAfter, datetimeRebaseMode=datetimeRebaseMode, int96RebaseMode=int96RebaseMode, ) return self.load(path=list(paths), format="parquet") parquet.__doc__ = PySparkDataFrameReader.parquet.__doc__ def text( self, paths: PathOrPaths, wholetext: bool = False, lineSep: Optional[str] = None, pathGlobFilter: Optional[Union[bool, str]] = None, recursiveFileLookup: Optional[Union[bool, str]] = None, modifiedBefore: Optional[Union[bool, str]] = None, modifiedAfter: Optional[Union[bool, str]] = None, ) -> "DataFrame": self._set_opts( wholetext=wholetext, lineSep=lineSep, pathGlobFilter=pathGlobFilter, recursiveFileLookup=recursiveFileLookup, modifiedBefore=modifiedBefore, modifiedAfter=modifiedAfter, ) if isinstance(paths, str): paths = [paths] return self.load(path=paths, format="text") text.__doc__ = PySparkDataFrameReader.text.__doc__ def csv( self, path: Union[PathOrPaths, "DataFrame"], schema: Optional[Union[StructType, str]] = None, sep: Optional[str] = None, encoding: Optional[str] = None, quote: Optional[str] = None, escape: Optional[str] = None, comment: Optional[str] = None, header: Optional[Union[bool, str]] = None, inferSchema: Optional[Union[bool, str]] = None, ignoreLeadingWhiteSpace: Optional[Union[bool, str]] = None, ignoreTrailingWhiteSpace: Optional[Union[bool, str]] = None, nullValue: Optional[str] = None, nanValue: Optional[str] = None, positiveInf: Optional[str] = None, negativeInf: Optional[str] = None, dateFormat: Optional[str] = None, timestampFormat: Optional[str] = None, maxColumns: Optional[Union[int, str]] = None, maxCharsPerColumn: Optional[Union[int, str]] = None, maxMalformedLogPerPartition: Optional[Union[int, str]] = None, mode: Optional[str] = None, columnNameOfCorruptRecord: Optional[str] = None, multiLine: Optional[Union[bool, str]] = None, charToEscapeQuoteEscaping: Optional[str] = None, samplingRatio: Optional[Union[float, str]] = None, enforceSchema: Optional[Union[bool, str]] = None, emptyValue: Optional[str] = None, locale: Optional[str] = None, lineSep: Optional[str] = None, pathGlobFilter: Optional[Union[bool, str]] = None, recursiveFileLookup: Optional[Union[bool, str]] = None, modifiedBefore: Optional[Union[bool, str]] = None, modifiedAfter: Optional[Union[bool, str]] = None, unescapedQuoteHandling: Optional[str] = None, ) -> "DataFrame": self._set_opts( sep=sep, encoding=encoding, quote=quote, escape=escape, comment=comment, header=header, inferSchema=inferSchema, ignoreLeadingWhiteSpace=ignoreLeadingWhiteSpace, ignoreTrailingWhiteSpace=ignoreTrailingWhiteSpace, nullValue=nullValue, nanValue=nanValue, positiveInf=positiveInf, negativeInf=negativeInf, dateFormat=dateFormat, timestampFormat=timestampFormat, maxColumns=maxColumns, maxCharsPerColumn=maxCharsPerColumn, maxMalformedLogPerPartition=maxMalformedLogPerPartition, mode=mode, columnNameOfCorruptRecord=columnNameOfCorruptRecord, multiLine=multiLine, charToEscapeQuoteEscaping=charToEscapeQuoteEscaping, samplingRatio=samplingRatio, enforceSchema=enforceSchema, emptyValue=emptyValue, locale=locale, lineSep=lineSep, pathGlobFilter=pathGlobFilter, recursiveFileLookup=recursiveFileLookup, modifiedBefore=modifiedBefore, modifiedAfter=modifiedAfter, unescapedQuoteHandling=unescapedQuoteHandling, ) if isinstance(path, str): path = [path] from pyspark.sql.connect.dataframe import DataFrame if isinstance(path, DataFrame): # Schema must be set explicitly here because the DataFrame path # bypasses load(), which normally calls self.schema(schema). if schema is not None: self.schema(schema) return self._df( Parse( child=path._plan, format=proto.Parse.ParseFormat.PARSE_FORMAT_CSV, schema=self._schema, options=self._options, ) ) return self.load(path=path, format="csv", schema=schema) csv.__doc__ = PySparkDataFrameReader.csv.__doc__ def xml( self, path: Union[PathOrPaths, "DataFrame"], rowTag: Optional[str] = None, schema: Optional[Union[StructType, str]] = None, excludeAttribute: Optional[Union[bool, str]] = None, attributePrefix: Optional[str] = None, valueTag: Optional[str] = None, ignoreSurroundingSpaces: Optional[Union[bool, str]] = None, rowValidationXSDPath: Optional[str] = None, ignoreNamespace: Optional[Union[bool, str]] = None, wildcardColName: Optional[str] = None, encoding: Optional[str] = None, inferSchema: Optional[Union[bool, str]] = None, nullValue: Optional[str] = None, dateFormat: Optional[str] = None, timestampFormat: Optional[str] = None, mode: Optional[str] = None, columnNameOfCorruptRecord: Optional[str] = None, multiLine: Optional[Union[bool, str]] = None, samplingRatio: Optional[Union[float, str]] = None, locale: Optional[str] = None, ) -> "DataFrame": self._set_opts( rowTag=rowTag, schema=schema, excludeAttribute=excludeAttribute, attributePrefix=attributePrefix, valueTag=valueTag, ignoreSurroundingSpaces=ignoreSurroundingSpaces, rowValidationXSDPath=rowValidationXSDPath, ignoreNamespace=ignoreNamespace, wildcardColName=wildcardColName, encoding=encoding, inferSchema=inferSchema, nullValue=nullValue, dateFormat=dateFormat, timestampFormat=timestampFormat, mode=mode, columnNameOfCorruptRecord=columnNameOfCorruptRecord, multiLine=multiLine, samplingRatio=samplingRatio, locale=locale, ) if isinstance(path, str): path = [path] if isinstance(path, list): return self.load(path=path, format="xml", schema=schema) from pyspark.sql.connect.dataframe import DataFrame if isinstance(path, DataFrame): # Schema must be set explicitly here because the DataFrame path # bypasses load(), which normally calls self.schema(schema). if schema is not None: self.schema(schema) return self._df( Parse( child=path._plan, format=proto.Parse.ParseFormat.PARSE_FORMAT_XML, schema=self._schema, options=self._options, ) ) raise PySparkTypeError( errorClass="NOT_EXPECTED_TYPE", messageParameters={ "arg_name": "path", "expected_type": "str, list, or DataFrame", "arg_type": type(path).__name__, }, ) xml.__doc__ = PySparkDataFrameReader.xml.__doc__ def orc( self, path: PathOrPaths, mergeSchema: Optional[bool] = None, pathGlobFilter: Optional[Union[bool, str]] = None, recursiveFileLookup: Optional[Union[bool, str]] = None, modifiedBefore: Optional[Union[bool, str]] = None, modifiedAfter: Optional[Union[bool, str]] = None, ) -> "DataFrame": self._set_opts( mergeSchema=mergeSchema, pathGlobFilter=pathGlobFilter, modifiedBefore=modifiedBefore, modifiedAfter=modifiedAfter, recursiveFileLookup=recursiveFileLookup, ) if isinstance(path, str): path = [path] return self.load(path=path, format="orc") orc.__doc__ = PySparkDataFrameReader.orc.__doc__ @overload def jdbc( self, url: str, table: str, *, properties: Optional[Dict[str, str]] = None ) -> "DataFrame": ... @overload def jdbc( self, url: str, table: str, column: str, lowerBound: Union[int, str], upperBound: Union[int, str], numPartitions: int, *, properties: Optional[Dict[str, str]] = None, ) -> "DataFrame": ... @overload def jdbc( self, url: str, table: str, *, predicates: List[str], properties: Optional[Dict[str, str]] = None, ) -> "DataFrame": ... def jdbc( self, url: str, table: str, column: Optional[str] = None, lowerBound: Optional[Union[int, str]] = None, upperBound: Optional[Union[int, str]] = None, numPartitions: Optional[int] = None, predicates: Optional[List[str]] = None, properties: Optional[Dict[str, str]] = None, ) -> "DataFrame": if properties is None: properties = dict() self.format("jdbc") if column is not None: assert lowerBound is not None, "lowerBound can not be None when ``column`` is specified" assert upperBound is not None, "upperBound can not be None when ``column`` is specified" assert numPartitions is not None, ( "numPartitions can not be None when ``column`` is specified" ) self.options( partitionColumn=column, lowerBound=lowerBound, upperBound=upperBound, numPartitions=numPartitions, ) self.options(**properties) self.options(url=url, dbtable=table) return self.load() else: self.options(**properties) self.options(url=url, dbtable=table) if predicates is not None: plan = DataSource( format=self._format, schema=self._schema, options=self._options, predicates=predicates, ) return self._df(plan) else: return self.load() jdbc.__doc__ = PySparkDataFrameReader.jdbc.__doc__ @property def _jreader(self) -> None: raise PySparkAttributeError( errorClass="JVM_ATTRIBUTE_NOT_SUPPORTED", messageParameters={"attr_name": "_jreader"} ) DataFrameReader.__doc__ = PySparkDataFrameReader.__doc__ class DataFrameWriter(OptionUtils): def __init__( self, plan: "LogicalPlan", session: "SparkSession", callback: Optional[Callable[["ExecutionInfo"], None]] = None, ): self._df: "LogicalPlan" = plan self._spark: "SparkSession" = session self._write: "WriteOperation" = WriteOperation(self._df) self._callback = callback if callback is not None else lambda _: None def mode(self, saveMode: Optional[str]) -> "DataFrameWriter": # At the JVM side, the default value of mode is already set to "error". # So, if the given saveMode is None, we will not call JVM-side's mode method. if saveMode is not None: self._write.mode = saveMode return self mode.__doc__ = PySparkDataFrameWriter.mode.__doc__ def format(self, source: str) -> "DataFrameWriter": self._write.source = source return self format.__doc__ = PySparkDataFrameWriter.format.__doc__ def option(self, key: str, value: "OptionalPrimitiveType") -> "DataFrameWriter": if value is None: return self self._write.options[key] = to_str(value) return self option.__doc__ = PySparkDataFrameWriter.option.__doc__ def options(self, **options: "OptionalPrimitiveType") -> "DataFrameWriter": for k in options: self.option(k, options[k]) return self options.__doc__ = PySparkDataFrameWriter.options.__doc__ @overload def partitionBy(self, *cols: str) -> "DataFrameWriter": ... @overload def partitionBy(self, *cols: List[str]) -> "DataFrameWriter": ... def partitionBy(self, *cols: Union[str, List[str]]) -> "DataFrameWriter": if len(cols) == 1 and isinstance(cols[0], (list, tuple)): cols = cols[0] # type: ignore[assignment] self._write.partitioning_cols = cast(List[str], cols) return self partitionBy.__doc__ = PySparkDataFrameWriter.partitionBy.__doc__ @overload def bucketBy(self, numBuckets: int, col: str, *cols: str) -> "DataFrameWriter": ... @overload def bucketBy(self, numBuckets: int, col: TupleOrListOfString) -> "DataFrameWriter": ... def bucketBy( self, numBuckets: int, col: Union[str, TupleOrListOfString], *cols: Optional[str] ) -> "DataFrameWriter": if not isinstance(numBuckets, int): raise PySparkTypeError( errorClass="NOT_EXPECTED_TYPE", messageParameters={ "expected_type": "int", "arg_name": "numBuckets", "arg_type": type(numBuckets).__name__, }, ) if isinstance(col, (list, tuple)): if cols: raise PySparkValueError( errorClass="CANNOT_SET_TOGETHER", messageParameters={ "arg_list": f"`col` of type {type(col).__name__} and `cols`", }, ) col, cols = col[0], col[1:] # type: ignore[assignment] for c in cols: if not isinstance(c, str): raise PySparkTypeError( errorClass="NOT_EXPECTED_TYPE", messageParameters={ "arg_name": "cols", "expected_type": "list[str]", "arg_type": type(c).__name__, }, ) if not isinstance(col, str): raise PySparkTypeError( errorClass="NOT_EXPECTED_TYPE", messageParameters={ "arg_name": "col", "expected_type": "list[str]", "arg_type": type(col).__name__, }, ) self._write.num_buckets = numBuckets self._write.bucket_cols = cast(List[str], [col, *cols]) return self bucketBy.__doc__ = PySparkDataFrameWriter.bucketBy.__doc__ @overload def sortBy(self, col: str, *cols: str) -> "DataFrameWriter": ... @overload def sortBy(self, col: TupleOrListOfString) -> "DataFrameWriter": ... def sortBy( self, col: Union[str, TupleOrListOfString], *cols: Optional[str] ) -> "DataFrameWriter": if isinstance(col, (list, tuple)): if cols: raise PySparkValueError( errorClass="CANNOT_SET_TOGETHER", messageParameters={ "arg_list": f"`col` of type {type(col).__name__} and `cols`", }, ) col, cols = col[0], col[1:] # type: ignore[assignment] for c in cols: if not isinstance(c, str): raise PySparkTypeError( errorClass="NOT_EXPECTED_TYPE", messageParameters={ "arg_name": "cols", "expected_type": "list[str]", "arg_type": type(c).__name__, }, ) if not isinstance(col, str): raise PySparkTypeError( errorClass="NOT_EXPECTED_TYPE", messageParameters={ "arg_name": "col", "expected_type": "list[str]", "arg_type": type(col).__name__, }, ) self._write.sort_cols = cast(List[str], [col, *cols]) return self sortBy.__doc__ = PySparkDataFrameWriter.sortBy.__doc__ @overload def clusterBy(self, *cols: str) -> "DataFrameWriter": ... @overload def clusterBy(self, *cols: List[str]) -> "DataFrameWriter": ... def clusterBy(self, *cols: Union[str, List[str]]) -> "DataFrameWriter": if len(cols) == 1 and isinstance(cols[0], (list, tuple)): cols = cols[0] # type: ignore[assignment] assert len(cols) > 0, "clusterBy needs one or more clustering columns." self._write.clustering_cols = cast(List[str], cols) return self clusterBy.__doc__ = PySparkDataFrameWriter.clusterBy.__doc__ def save( self, path: Optional[str] = None, format: Optional[str] = None, mode: Optional[str] = None, partitionBy: Optional[Union[str, List[str]]] = None, **options: "OptionalPrimitiveType", ) -> None: self.mode(mode).options(**options) if partitionBy is not None: self.partitionBy(partitionBy) if format is not None: self.format(format) self._write.path = path _, _, ei = self._spark.client.execute_command( self._write.command(self._spark.client), self._write.observations ) self._callback(ei) save.__doc__ = PySparkDataFrameWriter.save.__doc__ def insertInto(self, tableName: str, overwrite: Optional[bool] = None) -> None: if overwrite is not None: self.mode("overwrite" if overwrite else "append") self._write.table_name = tableName self._write.table_save_method = "insert_into" _, _, ei = self._spark.client.execute_command( self._write.command(self._spark.client), self._write.observations ) self._callback(ei) insertInto.__doc__ = PySparkDataFrameWriter.insertInto.__doc__ def saveAsTable( self, name: str, format: Optional[str] = None, mode: Optional[str] = None, partitionBy: Optional[Union[str, List[str]]] = None, **options: "OptionalPrimitiveType", ) -> None: self.mode(mode).options(**options) if partitionBy is not None: self.partitionBy(partitionBy) if format is not None: self.format(format) self._write.table_name = name self._write.table_save_method = "save_as_table" _, _, ei = self._spark.client.execute_command( self._write.command(self._spark.client), self._write.observations ) self._callback(ei) saveAsTable.__doc__ = PySparkDataFrameWriter.saveAsTable.__doc__ def json( self, path: str, mode: Optional[str] = None, compression: Optional[str] = None, dateFormat: Optional[str] = None, timestampFormat: Optional[str] = None, lineSep: Optional[str] = None, encoding: Optional[str] = None, ignoreNullFields: Optional[Union[bool, str]] = None, ) -> None: self.mode(mode) self._set_opts( compression=compression, dateFormat=dateFormat, timestampFormat=timestampFormat, lineSep=lineSep, encoding=encoding, ignoreNullFields=ignoreNullFields, ) self.format("json").save(path) json.__doc__ = PySparkDataFrameWriter.json.__doc__ def parquet( self, path: str, mode: Optional[str] = None, partitionBy: Optional[Union[str, List[str]]] = None, compression: Optional[str] = None, ) -> None: self.mode(mode) if partitionBy is not None: self.partitionBy(partitionBy) self._set_opts(compression=compression) self.format("parquet").save(path) parquet.__doc__ = PySparkDataFrameWriter.parquet.__doc__ def text( self, path: str, compression: Optional[str] = None, lineSep: Optional[str] = None ) -> None: self._set_opts(compression=compression, lineSep=lineSep) self.format("text").save(path) text.__doc__ = PySparkDataFrameWriter.text.__doc__ def csv( self, path: str, mode: Optional[str] = None, compression: Optional[str] = None, sep: Optional[str] = None, quote: Optional[str] = None, escape: Optional[str] = None, header: Optional[Union[bool, str]] = None, nullValue: Optional[str] = None, escapeQuotes: Optional[Union[bool, str]] = None, quoteAll: Optional[Union[bool, str]] = None, dateFormat: Optional[str] = None, timestampFormat: Optional[str] = None, ignoreLeadingWhiteSpace: Optional[Union[bool, str]] = None, ignoreTrailingWhiteSpace: Optional[Union[bool, str]] = None, charToEscapeQuoteEscaping: Optional[str] = None, encoding: Optional[str] = None, emptyValue: Optional[str] = None, lineSep: Optional[str] = None, ) -> None: self.mode(mode) self._set_opts( compression=compression, sep=sep, quote=quote, escape=escape, header=header, nullValue=nullValue, escapeQuotes=escapeQuotes, quoteAll=quoteAll, dateFormat=dateFormat, timestampFormat=timestampFormat, ignoreLeadingWhiteSpace=ignoreLeadingWhiteSpace, ignoreTrailingWhiteSpace=ignoreTrailingWhiteSpace, charToEscapeQuoteEscaping=charToEscapeQuoteEscaping, encoding=encoding, emptyValue=emptyValue, lineSep=lineSep, ) self.format("csv").save(path) csv.__doc__ = PySparkDataFrameWriter.csv.__doc__ def xml( self, path: str, rowTag: Optional[str] = None, mode: Optional[str] = None, attributePrefix: Optional[str] = None, valueTag: Optional[str] = None, rootTag: Optional[str] = None, declaration: Optional[bool] = None, arrayElementName: Optional[str] = None, nullValue: Optional[str] = None, dateFormat: Optional[str] = None, timestampFormat: Optional[str] = None, compression: Optional[str] = None, encoding: Optional[str] = None, validateName: Optional[bool] = None, ) -> None: self.mode(mode) self._set_opts( rowTag=rowTag, attributePrefix=attributePrefix, valueTag=valueTag, rootTag=rootTag, declaration=declaration, arrayElementName=arrayElementName, nullValue=nullValue, dateFormat=dateFormat, timestampFormat=timestampFormat, compression=compression, encoding=encoding, validateName=validateName, ) self.format("xml").save(path) xml.__doc__ = PySparkDataFrameWriter.xml.__doc__ def orc( self, path: str, mode: Optional[str] = None, partitionBy: Optional[Union[str, List[str]]] = None, compression: Optional[str] = None, ) -> None: self.mode(mode) if partitionBy is not None: self.partitionBy(partitionBy) self._set_opts(compression=compression) self.format("orc").save(path) orc.__doc__ = PySparkDataFrameWriter.orc.__doc__ def jdbc( self, url: str, table: str, mode: Optional[str] = None, properties: Optional[Dict[str, str]] = None, ) -> None: if properties is None: properties = dict() self.format("jdbc").mode(mode).options(**properties).options(url=url, dbtable=table).save() jdbc.__doc__ = PySparkDataFrameWriter.jdbc.__doc__ class DataFrameWriterV2(OptionUtils): def __init__( self, plan: "LogicalPlan", session: "SparkSession", table: str, callback: Optional[Callable[["ExecutionInfo"], None]] = None, ): self._df: "LogicalPlan" = plan self._spark: "SparkSession" = session self._table_name: str = table self._write: "WriteOperationV2" = WriteOperationV2(self._df, self._table_name) self._callback = callback if callback is not None else lambda _: None def using(self, provider: str) -> "DataFrameWriterV2": self._write.provider = provider return self using.__doc__ = PySparkDataFrameWriterV2.using.__doc__ def option(self, key: str, value: "OptionalPrimitiveType") -> "DataFrameWriterV2": if value is None: return self self._write.options[key] = to_str(value) return self option.__doc__ = PySparkDataFrameWriterV2.option.__doc__ def options(self, **options: "OptionalPrimitiveType") -> "DataFrameWriterV2": for k in options: self.option(k, options[k]) return self options.__doc__ = PySparkDataFrameWriterV2.options.__doc__ def tableProperty(self, property: str, value: str) -> "DataFrameWriterV2": self._write.table_properties[property] = value return self tableProperty.__doc__ = PySparkDataFrameWriterV2.tableProperty.__doc__ def partitionedBy(self, col: "ColumnOrName", *cols: "ColumnOrName") -> "DataFrameWriterV2": self._write.partitioning_columns = [F._to_col(c) for c in [col] + list(cols)] return self partitionedBy.__doc__ = PySparkDataFrameWriterV2.partitionedBy.__doc__ def clusterBy(self, col: str, *cols: str) -> "DataFrameWriterV2": self._write.clustering_columns = [col] self._write.clustering_columns.extend(cols) return self clusterBy.__doc__ = PySparkDataFrameWriterV2.clusterBy.__doc__ def create(self) -> None: self._write.mode = "create" _, _, ei = self._spark.client.execute_command( self._write.command(self._spark.client), self._write.observations ) self._callback(ei) create.__doc__ = PySparkDataFrameWriterV2.create.__doc__ def replace(self) -> None: self._write.mode = "replace" _, _, ei = self._spark.client.execute_command( self._write.command(self._spark.client), self._write.observations ) self._callback(ei) replace.__doc__ = PySparkDataFrameWriterV2.replace.__doc__ def createOrReplace(self) -> None: self._write.mode = "create_or_replace" _, _, ei = self._spark.client.execute_command( self._write.command(self._spark.client), self._write.observations ) self._callback(ei) createOrReplace.__doc__ = PySparkDataFrameWriterV2.createOrReplace.__doc__ def append(self) -> None: self._write.mode = "append" _, _, ei = self._spark.client.execute_command( self._write.command(self._spark.client), self._write.observations ) self._callback(ei) append.__doc__ = PySparkDataFrameWriterV2.append.__doc__ def overwrite(self, condition: "ColumnOrName") -> None: self._write.mode = "overwrite" self._write.overwrite_condition = F._to_col(condition) _, _, ei = self._spark.client.execute_command( self._write.command(self._spark.client), self._write.observations ) self._callback(ei) overwrite.__doc__ = PySparkDataFrameWriterV2.overwrite.__doc__ def overwritePartitions(self) -> None: self._write.mode = "overwrite_partitions" _, _, ei = self._spark.client.execute_command( self._write.command(self._spark.client), self._write.observations ) self._callback(ei) overwritePartitions.__doc__ = PySparkDataFrameWriterV2.overwritePartitions.__doc__ def _test() -> None: import sys import os import doctest from pyspark.sql import SparkSession as PySparkSession import pyspark.sql.connect.readwriter globs = pyspark.sql.connect.readwriter.__dict__.copy() globs["spark"] = ( PySparkSession.builder.appName("sql.connect.readwriter tests") .remote(os.environ.get("SPARK_CONNECT_TESTING_REMOTE", "local[4]")) .getOrCreate() ) failure_count, test_count = doctest.testmod( pyspark.sql.connect.readwriter, globs=globs, optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.IGNORE_EXCEPTION_DETAIL, ) globs["spark"].stop() if failure_count: sys.exit(-1) if __name__ == "__main__": _test()