diff --git a/.github/workflows/build_python_connect.yml b/.github/workflows/build_python_connect.yml index 993eccf3..6dfdafb5 100644 --- a/.github/workflows/build_python_connect.yml +++ b/.github/workflows/build_python_connect.yml @@ -17,7 +17,7 @@ # under the License. # -name: Python Connect Parity (Spark 4.2.0) +name: Python Connect Parity (Spark 4.3.0-rc3) # Cancel an in-progress run when a newer commit supersedes it (per workflow + ref). concurrency: @@ -25,8 +25,8 @@ concurrency: cancel-in-progress: true # Mirrors Apache Spark's build_python_connect.yml approach, adapted for the -# Rust-backed client: download the Spark 4.2.0 Connect server, start it, pull the -# official v4.2.0 connect test files, and run them through our client - the gate is +# Rust-backed client: download the Spark 4.3.0-rc3 Connect server, start it, pull the +# official v4.3.0-rc3 connect test files, and run them through our client - the gate is # "every test passes except the checked-in known-environmental-failures manifest". # A separate job enforces Rust test coverage across all crates. @@ -38,7 +38,7 @@ on: workflow_dispatch: env: - SPARK_VERSION: "4.2.0" + SPARK_VERSION: "4.3.0" RUST_PARITY_LOG_DIR: /tmp/parity-test-logs SPARK_CONNECT_TESTING_REMOTE: sc://localhost:15002 SPARK_TESTING: "1" @@ -47,7 +47,7 @@ env: jobs: # --------------------------------------------------------------------------- # Parity gate: the official connect suite through our client, against a real - # Spark 4.2.0 Connect server (see scripts/run_official_tests.py). + # Spark 4.3.0-rc3 Connect server (see scripts/run_official_tests.py). # --------------------------------------------------------------------------- connect-parity: name: Connect parity (official suite via our client, Python ${{ matrix.python-version }}) @@ -118,12 +118,12 @@ jobs: path: ~/spark-dist key: spark-dist-${{ env.SPARK_VERSION }}-hadoop3 - - name: Download Apache Spark ${{ env.SPARK_VERSION }} (Connect server) + - name: Download Apache Spark ${{ env.SPARK_VERSION }} (Connect server, RC) id: spark-download run: | - PKG="spark-${SPARK_VERSION}-bin-hadoop3" - if [ ! -d "${HOME}/spark-dist/${PKG}" ]; then - URL="https://archive.apache.org/dist/spark/spark-${SPARK_VERSION}/${PKG}.tgz" + PKG="spark-${SPARK_VERSION}-bin-hadoop3-connect" + if [ ! -d "${HOME}/spark-dist/spark-4.3.0-bin-hadoop3"* ]; then + URL="https://dist.apache.org/repos/dist/dev/spark/v4.3.0-rc3-bin/${PKG}.tgz" mkdir -p "${HOME}/spark-dist" echo "Downloading ${URL}" wget -q -O "${HOME}/spark-dist/${PKG}.tgz" "${URL}" @@ -132,7 +132,8 @@ jobs: else echo "Using cached Spark distribution." fi - echo "SPARK_HOME=${HOME}/spark-dist/${PKG}" >> "$GITHUB_OUTPUT" + SPARK_HOME=$(find "${HOME}/spark-dist" -maxdepth 1 -type d -name 'spark-4.3.0-bin-hadoop3*' | head -1) + echo "SPARK_HOME=${SPARK_HOME}" >> "$GITHUB_OUTPUT" - name: Cache Spark ${{ env.SPARK_VERSION }} source (official test files) uses: actions/cache@v4 @@ -140,11 +141,11 @@ jobs: path: ~/spark-source key: spark-source-${{ env.SPARK_VERSION }} - - name: Clone Apache Spark v${{ env.SPARK_VERSION }} (official test files + reference client) + - name: Clone Apache Spark v4.3.0-rc3 (official test files + reference client) id: spark-source run: | if [ ! -d "${HOME}/spark-source/python" ]; then - git clone --depth 1 --branch "v${SPARK_VERSION}" \ + git clone --depth 1 --branch "v4.3.0-rc3" \ https://github.com/apache/spark.git "${HOME}/spark-source" else echo "Using cached Spark source." @@ -164,7 +165,6 @@ jobs: # this one server, and a 1g default degrades late in the run (GC pressure), # which is what made the streaming listener/observation files flake. "${SPARK_HOME}/sbin/start-connect-server.sh" \ - --packages "org.apache.spark:spark-connect_2.13:${SPARK_VERSION}" \ --driver-memory 3g \ --conf spark.sql.session.localRelationSizeLimit=1073741824 \ > /tmp/connect-server.log 2>&1 & @@ -232,11 +232,9 @@ jobs: sleep 1 done echo "Starting Connect server in local-cluster mode..." - # Keep --packages: local-cluster launches real executor JVMs that each need the - # spark-connect jar on their classpath, so it is not redundant here. + # Connect jar is bundled in RC distribution, no --packages needed. "${SPARK_HOME}/sbin/start-connect-server.sh" \ --master "local-cluster[2, 4, 1024]" \ - --packages "org.apache.spark:spark-connect_2.13:${SPARK_VERSION}" \ --driver-memory 3g \ --conf spark.sql.session.localRelationSizeLimit=1073741824 \ > /tmp/connect-server-localcluster.log 2>&1 & @@ -282,7 +280,7 @@ jobs: if-no-files-found: warn # --------------------------------------------------------------------------- - # Behavioral e2e tests (Rust drop-in against a live server), in their OWN job + # Behavioral e2e tests (Rust drop-in against a live 4.3.0-rc3 server), in their OWN job # with a separate Connect server so they never share server state with the # parity gate. The parity gate then starts against a server nothing else has # touched, which avoids the SPARK-59032 arrow-UDF hang (a UDF that crashes a @@ -338,12 +336,12 @@ jobs: path: ~/spark-dist key: spark-dist-${{ env.SPARK_VERSION }}-hadoop3 - - name: Download Apache Spark ${{ env.SPARK_VERSION }} (Connect server) + - name: Download Apache Spark ${{ env.SPARK_VERSION }} (Connect server, RC) id: spark-download run: | - PKG="spark-${SPARK_VERSION}-bin-hadoop3" - if [ ! -d "${HOME}/spark-dist/${PKG}" ]; then - URL="https://archive.apache.org/dist/spark/spark-${SPARK_VERSION}/${PKG}.tgz" + PKG="spark-${SPARK_VERSION}-bin-hadoop3-connect" + if [ ! -d "${HOME}/spark-dist/spark-4.3.0-bin-hadoop3"* ]; then + URL="https://dist.apache.org/repos/dist/dev/spark/v4.3.0-rc3-bin/${PKG}.tgz" mkdir -p "${HOME}/spark-dist" echo "Downloading ${URL}" wget -q -O "${HOME}/spark-dist/${PKG}.tgz" "${URL}" @@ -352,7 +350,8 @@ jobs: else echo "Using cached Spark distribution." fi - echo "SPARK_HOME=${HOME}/spark-dist/${PKG}" >> "$GITHUB_OUTPUT" + SPARK_HOME=$(find "${HOME}/spark-dist" -maxdepth 1 -type d -name 'spark-4.3.0-bin-hadoop3*' | head -1) + echo "SPARK_HOME=${SPARK_HOME}" >> "$GITHUB_OUTPUT" - name: Start Spark Connect server env: @@ -361,7 +360,6 @@ jobs: run: | SPARK_HOME="${{ steps.spark-download.outputs.SPARK_HOME }}" "${SPARK_HOME}/sbin/start-connect-server.sh" \ - --packages "org.apache.spark:spark-connect_2.13:${SPARK_VERSION}" \ --driver-memory 3g \ --conf spark.sql.session.localRelationSizeLimit=1073741824 \ > /tmp/connect-server.log 2>&1 & @@ -383,7 +381,7 @@ jobs: -- --test-threads=1 # --------------------------------------------------------------------------- - # E2E performance benchmark: our client vs reference, same 4.2.0 server. + # E2E performance benchmark: our client vs reference, same 4.3.0-rc3 server. # --------------------------------------------------------------------------- connect-benchmark: name: E2E benchmark vs reference @@ -415,7 +413,7 @@ jobs: run: | sudo apt-get update && sudo apt-get install -y protobuf-compiler python -m pip install --upgrade pip - python -m pip install 'pandas>=2.2,<3.0' pyarrow numpy grpcio grpcio-status googleapis-common-protos cloudpickle zstandard py4j "pyspark==${SPARK_VERSION}" + python -m pip install 'pandas>=2.2,<3.0' pyarrow numpy grpcio grpcio-status googleapis-common-protos cloudpickle zstandard py4j "https://dist.apache.org/repos/dist/dev/spark/v4.3.0-rc3-bin/pyspark-4.3.0.tar.gz" - name: Build the Rust extension into the skin # Build the extension into ./python/pyspark so our client is usable via # PYTHONPATH. We do NOT install a wheel: our package and the reference @@ -433,24 +431,24 @@ jobs: with: path: ~/spark-source key: spark-source-${{ env.SPARK_VERSION }} - - name: Download & start Spark ${{ env.SPARK_VERSION }} server + - name: Download & start Spark ${{ env.SPARK_VERSION }} server (RC) run: | - PKG="spark-${SPARK_VERSION}-bin-hadoop3" - if [ ! -d "${HOME}/spark-dist/${PKG}" ]; then + PKG="spark-${SPARK_VERSION}-bin-hadoop3-connect" + if [ ! -d "${HOME}/spark-dist/spark-4.3.0-bin-hadoop3"* ]; then mkdir -p "${HOME}/spark-dist" - wget -q -O "/tmp/${PKG}.tgz" "https://archive.apache.org/dist/spark/spark-${SPARK_VERSION}/${PKG}.tgz" + wget -q -O "/tmp/${PKG}.tgz" "https://dist.apache.org/repos/dist/dev/spark/v4.3.0-rc3-bin/${PKG}.tgz" tar -xzf "/tmp/${PKG}.tgz" -C "${HOME}/spark-dist" rm -f "/tmp/${PKG}.tgz" fi if [ ! -d "${HOME}/spark-source/python" ]; then - git clone --depth 1 --branch "v${SPARK_VERSION}" https://github.com/apache/spark.git "${HOME}/spark-source" + git clone --depth 1 --branch "v4.3.0-rc3" https://github.com/apache/spark.git "${HOME}/spark-source" fi - "${HOME}/spark-dist/${PKG}/sbin/start-connect-server.sh" \ - --packages "org.apache.spark:spark-connect_2.13:${SPARK_VERSION}" > /tmp/server.log 2>&1 & + SPARK_HOME=$(find "${HOME}/spark-dist" -maxdepth 1 -type d -name 'spark-4.3.0-bin-hadoop3*' | head -1) + "${SPARK_HOME}/sbin/start-connect-server.sh" > /tmp/server.log 2>&1 & for i in $(seq 1 90); do (echo > /dev/tcp/localhost/15002) 2>/dev/null && break; sleep 1; done - name: Benchmark reference client run: | - # Reference uses the pip-installed pyspark 4.2.0 (no ./python on the path). + # Reference uses the pip-installed pyspark 4.3.0-rc3 (no ./python on the path). python scripts/benchmark_e2e.py --isolate --label reference --out /tmp/bench_reference.json - name: Benchmark our client run: | diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 7f017ad7..bda75b06 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -32,7 +32,7 @@ on: env: CARGO_TERM_COLOR: always - SPARK_VERSION: "4.2.0" + SPARK_VERSION: "4.3.0" jobs: coverage: @@ -98,12 +98,12 @@ jobs: path: ~/spark-dist key: spark-dist-${{ env.SPARK_VERSION }}-hadoop3 - - name: Download Apache Spark ${{ env.SPARK_VERSION }} (Connect server) + - name: Download Apache Spark ${{ env.SPARK_VERSION }} (Connect server, RC) id: spark-download run: | - PKG="spark-${SPARK_VERSION}-bin-hadoop3" - if [ ! -d "${HOME}/spark-dist/${PKG}" ]; then - URL="https://archive.apache.org/dist/spark/spark-${SPARK_VERSION}/${PKG}.tgz" + PKG="spark-${SPARK_VERSION}-bin-hadoop3-connect" + if [ ! -d "${HOME}/spark-dist/spark-4.3.0-bin-hadoop3"* ]; then + URL="https://dist.apache.org/repos/dist/dev/spark/v4.3.0-rc3-bin/${PKG}.tgz" mkdir -p "${HOME}/spark-dist" echo "Downloading ${URL}" wget -q -O "${HOME}/spark-dist/${PKG}.tgz" "${URL}" @@ -112,7 +112,8 @@ jobs: else echo "Using cached Spark distribution." fi - echo "SPARK_HOME=${HOME}/spark-dist/${PKG}" >> "$GITHUB_OUTPUT" + SPARK_HOME=$(find "${HOME}/spark-dist" -maxdepth 1 -type d -name 'spark-4.3.0-bin-hadoop3*' | head -1) + echo "SPARK_HOME=${SPARK_HOME}" >> "$GITHUB_OUTPUT" - name: Start Spark Connect server env: @@ -124,7 +125,6 @@ jobs: SPARK_HOME="${{ steps.spark-download.outputs.SPARK_HOME }}" echo "Starting Connect server on sc://localhost:15002" "${SPARK_HOME}/sbin/start-connect-server.sh" \ - --packages "org.apache.spark:spark-connect_2.13:${SPARK_VERSION}" \ --driver-memory 3g \ --conf spark.sql.session.localRelationSizeLimit=1073741824 \ > /tmp/connect-server.log 2>&1 & diff --git a/.github/workflows/pandas-parity.yml b/.github/workflows/pandas-parity.yml index bdf10f6a..ccfd8733 100644 --- a/.github/workflows/pandas-parity.yml +++ b/.github/workflows/pandas-parity.yml @@ -23,7 +23,7 @@ # full visibility into pass/fail status. If this becomes a critical blocker, add # continue-on-error: true back in; currently it provides honest signal on pandas parity. -name: Pandas Connect Parity (Spark 4.2.0) +name: Pandas Connect Parity (Spark 4.3.0-rc3) # Cancel an in-progress run when a newer commit supersedes it (per workflow + ref). concurrency: @@ -38,7 +38,7 @@ on: workflow_dispatch: env: - SPARK_VERSION: "4.2.0" + SPARK_VERSION: "4.3.0" SPARK_CONNECT_TESTING_REMOTE: sc://localhost:15002 SPARK_TESTING: "1" SPARK_SKIP_CONNECT_COMPAT_TESTS: "1" # force remote-only (pure connect, no local JVM) @@ -46,7 +46,7 @@ env: jobs: # --------------------------------------------------------------------------- # Pandas-on-Spark Connect parity: run the official pandas tests through our - # Rust client against a real Spark 4.2.0 Connect server. This is a separate, + # Rust client against a real Spark 4.3.0-rc3 Connect server. This is a separate, # long-running job (30-60 min) to avoid blocking the main parity gate. # --------------------------------------------------------------------------- pandas-parity: @@ -122,12 +122,12 @@ jobs: path: ~/spark-dist key: spark-dist-${{ env.SPARK_VERSION }}-hadoop3 - - name: Download Apache Spark ${{ env.SPARK_VERSION }} (Connect server) + - name: Download Apache Spark ${{ env.SPARK_VERSION }} (Connect server, RC) id: spark-download run: | - PKG="spark-${SPARK_VERSION}-bin-hadoop3" - if [ ! -d "${HOME}/spark-dist/${PKG}" ]; then - URL="https://archive.apache.org/dist/spark/spark-${SPARK_VERSION}/${PKG}.tgz" + PKG="spark-${SPARK_VERSION}-bin-hadoop3-connect" + if [ ! -d "${HOME}/spark-dist/spark-4.3.0-bin-hadoop3"* ]; then + URL="https://dist.apache.org/repos/dist/dev/spark/v4.3.0-rc3-bin/${PKG}.tgz" mkdir -p "${HOME}/spark-dist" echo "Downloading ${URL}" wget -q -O "${HOME}/spark-dist/${PKG}.tgz" "${URL}" @@ -136,7 +136,8 @@ jobs: else echo "Using cached Spark distribution." fi - echo "SPARK_HOME=${HOME}/spark-dist/${PKG}" >> "$GITHUB_OUTPUT" + SPARK_HOME=$(find "${HOME}/spark-dist" -maxdepth 1 -type d -name 'spark-4.3.0-bin-hadoop3*' | head -1) + echo "SPARK_HOME=${SPARK_HOME}" >> "$GITHUB_OUTPUT" - name: Cache Spark ${{ env.SPARK_VERSION }} source (official test files) uses: actions/cache@v4 @@ -144,11 +145,11 @@ jobs: path: ~/spark-source key: spark-source-${{ env.SPARK_VERSION }} - - name: Clone Apache Spark v${{ env.SPARK_VERSION }} (official test files + reference client) + - name: Clone Apache Spark v4.3.0-rc3 (official test files + reference client) id: spark-source run: | if [ ! -d "${HOME}/spark-source/python" ]; then - git clone --depth 1 --branch "v${SPARK_VERSION}" \ + git clone --depth 1 --branch "v4.3.0-rc3" \ https://github.com/apache/spark.git "${HOME}/spark-source" else echo "Using cached Spark source." @@ -177,7 +178,6 @@ jobs: # how Apache's PYTHON connect tests run (the runner's UTC); the pom.xml America/Los_Angeles # is for the JVM/Scala maven-surefire tests, not the separately-launched python connect server. "${SPARK_HOME}/sbin/start-connect-server.sh" \ - --packages "org.apache.spark:spark-connect_2.13:${SPARK_VERSION}" \ --driver-memory 3g \ --conf spark.sql.session.localRelationSizeLimit=1073741824 \ --conf spark.sql.session.timeZone=UTC \ diff --git a/.github/workflows/pipelines-parity.yml b/.github/workflows/pipelines-parity.yml index 3280c264..ca819fe9 100644 --- a/.github/workflows/pipelines-parity.yml +++ b/.github/workflows/pipelines-parity.yml @@ -24,7 +24,7 @@ # 2. The full official upstream pipeline connect suite (incl. server-backed tests like # test_spark_connect.py) run through OUR Rust client via the transport plugin. -name: Pipelines Connect Parity (Spark 4.2.0) +name: Pipelines Connect Parity (Spark 4.3.0-rc3) # Cancel an in-progress run when a newer commit supersedes it (per workflow + ref). concurrency: @@ -39,7 +39,7 @@ on: workflow_dispatch: env: - SPARK_VERSION: "4.2.0" + SPARK_VERSION: "4.3.0" SPARK_CONNECT_TESTING_REMOTE: sc://localhost:15002 SPARK_TESTING: "1" SPARK_SKIP_CONNECT_COMPAT_TESTS: "1" # force remote-only (pure connect, no local JVM) @@ -47,7 +47,7 @@ env: jobs: # --------------------------------------------------------------------------- # Pipelines Connect parity: run the official pipeline connect tests through our - # Rust client against a real Spark 4.2.0 Connect server. Two blocking groups: pure-Python + # Rust client against a real Spark 4.3.0-rc3 Connect server. Two blocking groups: pure-Python # pipeline tests and the full server-backed official suite. # --------------------------------------------------------------------------- pipelines-parity: @@ -117,12 +117,12 @@ jobs: path: ~/spark-dist key: spark-dist-${{ env.SPARK_VERSION }}-hadoop3 - - name: Download Apache Spark ${{ env.SPARK_VERSION }} (Connect server) + - name: Download Apache Spark ${{ env.SPARK_VERSION }} (Connect server, RC) id: spark-download run: | - PKG="spark-${SPARK_VERSION}-bin-hadoop3" - if [ ! -d "${HOME}/spark-dist/${PKG}" ]; then - URL="https://archive.apache.org/dist/spark/spark-${SPARK_VERSION}/${PKG}.tgz" + PKG="spark-${SPARK_VERSION}-bin-hadoop3-connect" + if [ ! -d "${HOME}/spark-dist/spark-4.3.0-bin-hadoop3"* ]; then + URL="https://dist.apache.org/repos/dist/dev/spark/v4.3.0-rc3-bin/${PKG}.tgz" mkdir -p "${HOME}/spark-dist" echo "Downloading ${URL}" wget -q -O "${HOME}/spark-dist/${PKG}.tgz" "${URL}" @@ -131,7 +131,8 @@ jobs: else echo "Using cached Spark distribution." fi - echo "SPARK_HOME=${HOME}/spark-dist/${PKG}" >> "$GITHUB_OUTPUT" + SPARK_HOME=$(find "${HOME}/spark-dist" -maxdepth 1 -type d -name 'spark-4.3.0-bin-hadoop3*' | head -1) + echo "SPARK_HOME=${SPARK_HOME}" >> "$GITHUB_OUTPUT" - name: Cache Spark ${{ env.SPARK_VERSION }} source (official test files) uses: actions/cache@v4 @@ -139,11 +140,11 @@ jobs: path: ~/spark-source key: spark-source-${{ env.SPARK_VERSION }} - - name: Clone Apache Spark v${{ env.SPARK_VERSION }} (official test files + reference client) + - name: Clone Apache Spark v4.3.0-rc3 (official test files + reference client) id: spark-source run: | if [ ! -d "${HOME}/spark-source/python" ]; then - git clone --depth 1 --branch "v${SPARK_VERSION}" \ + git clone --depth 1 --branch "v4.3.0-rc3" \ https://github.com/apache/spark.git "${HOME}/spark-source" else echo "Using cached Spark source." @@ -178,7 +179,6 @@ jobs: echo "Starting Connect server on ${SPARK_CONNECT_TESTING_REMOTE}" # Give the driver extra heap for the pipeline tests which may be I/O heavy. "${SPARK_HOME}/sbin/start-connect-server.sh" \ - --packages "org.apache.spark:spark-connect_2.13:${SPARK_VERSION}" \ --driver-memory 3g \ --conf spark.sql.session.localRelationSizeLimit=1073741824 \ > /tmp/connect-server.log 2>&1 & diff --git a/.github/workflows/vendor-drift.yml b/.github/workflows/vendor-drift.yml index ad4a63be..10103f91 100644 --- a/.github/workflows/vendor-drift.yml +++ b/.github/workflows/vendor-drift.yml @@ -18,7 +18,7 @@ name: Vendored upstream drift # Guards that directories vendored verbatim from Apache Spark (e.g. cloudpickle, -# brought in via git subtree) stay byte-identical to the pinned v4.2.0 tag. Catches +# brought in via git subtree) stay byte-identical to the pinned v4.3.0-rc3 tag. Catches # accidental local edits to upstream-tracked code. concurrency: group: ${{ github.workflow }}-${{ github.ref }} @@ -32,7 +32,7 @@ on: jobs: vendored-drift: - name: Vendored files match Apache Spark v4.2.0 + name: Vendored files match Apache Spark v4.3.0-rc3 runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 diff --git a/README.md b/README.md index 7f64da75..6b961c52 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ # Spark Connect Rust Client A fast, native **Rust** client for **Apache Spark Connect** - and a drop-in -`pyspark` replacement with **100% public-API parity with PySpark 4.2.0**. It +`pyspark` replacement with **100% public-API parity with PySpark 4.3.0**. It builds `spark.connect` protobuf plans, manages the gRPC channel, and decodes Arrow results in Rust, speaking the same protocol and returning the same results as the reference client. @@ -28,7 +28,7 @@ as the reference client. [![PyPI](https://img.shields.io/pypi/v/pyspark-client-rust?color=c2410c&label=pyspark-client-rust)](https://pypi.org/project/pyspark-client-rust/) [![Crates.io](https://img.shields.io/crates/v/apache-spark-connect?color=c2410c&label=apache-spark-connect)](https://crates.io/crates/apache-spark-connect) [![docs.rs](https://img.shields.io/docsrs/apache-spark-connect?color=c2410c&label=docs.rs)](https://docs.rs/apache-spark-connect) -![Spark](https://img.shields.io/badge/Apache%20Spark-4.2.0%2B-c2410c) +![Spark](https://img.shields.io/badge/Apache%20Spark-4.3.0%2B-c2410c) ![License](https://img.shields.io/badge/license-Apache--2.0-blue) [![Rust coverage](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/apache/spark-connect-rust/badges/coverage-rust.json)](https://github.com/apache/spark-connect-rust/actions/workflows/coverage.yml) @@ -63,7 +63,7 @@ Your Spark Connect code then runs unchanged; use it exactly like ```toml [dependencies] -apache-spark-connect = "4.2" +apache-spark-connect = "4.3" ``` ## Quickstart (Rust) diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index 7101aadf..a5edcdb7 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -414,6 +414,14 @@ impl PyDataFrame { PyDataFrame::new(self.dataframe.zip_with_index(indexColName)) } + /// Zip this DataFrame with another column-wise (side-by-side). Mirrors + /// `DataFrame.zip(other)`. + fn zip(&self, other: &PyDataFrame) -> PyResult { + Ok(PyDataFrame::new( + self.dataframe.zip(&other.dataframe).to_pyerr()?, + )) + } + /// Union with another DataFrame. fn union(&self, other: &PyDataFrame) -> PyDataFrame { PyDataFrame::new(self.dataframe.union(&other.dataframe)) @@ -1680,6 +1688,25 @@ impl PyMergeIntoWriter { #[pymethods] impl PyMergeIntoWriter { + /// Nested builder-class references, mirroring pyspark's `MergeIntoWriter.WhenMatched` + /// etc. (the objects returned by `whenMatched()` / `whenNotMatched()` / + /// `whenNotMatchedBySource()`). + #[classattr] + #[allow(non_snake_case)] + fn WhenMatched(py: Python<'_>) -> Py { + py.get_type::().unbind() + } + #[classattr] + #[allow(non_snake_case)] + fn WhenNotMatched(py: Python<'_>) -> Py { + py.get_type::().unbind() + } + #[classattr] + #[allow(non_snake_case)] + fn WhenNotMatchedBySource(py: Python<'_>) -> Py { + py.get_type::().unbind() + } + #[pyo3(name = "whenMatched", signature = (condition=None))] fn when_matched(&mut self, condition: Option<&PyColumn>) -> PyResult { let cond = condition.map(|c| c.column.clone()); diff --git a/crates/pyspark-rs/src/dispatch_generated.rs b/crates/pyspark-rs/src/dispatch_generated.rs index 6523d54d..2d4352fb 100644 --- a/crates/pyspark-rs/src/dispatch_generated.rs +++ b/crates/pyspark-rs/src/dispatch_generated.rs @@ -141,6 +141,14 @@ pub fn call_builtin(name: &str, args: Vec) -> PyR if args.is_empty() { return Err(PyErr::new::("Missing required column argument")); } Ok(spark_funcs::base64(args[0].clone())) }, + "to_base32" => { + if args.is_empty() { return Err(PyErr::new::("Missing required column argument")); } + Ok(spark_funcs::to_base32(args[0].clone())) + }, + "from_base32" => { + if args.is_empty() { return Err(PyErr::new::("Missing required column argument")); } + Ok(spark_funcs::from_base32(args[0].clone())) + }, "bin" => { if args.is_empty() { return Err(PyErr::new::("Missing required column argument")); } Ok(spark_funcs::bin(args[0].clone())) @@ -407,6 +415,14 @@ pub fn call_builtin(name: &str, args: Vec) -> PyR if args.is_empty() { return Err(PyErr::new::("Missing required column argument")); } Ok(spark_funcs::is_variant_null(args[0].clone())) }, + "collect_union" => { + if args.is_empty() { return Err(PyErr::new::("Missing required column argument")); } + Ok(spark_funcs::collect_union(args[0].clone())) + }, + "variant_delete" => { + if args.is_empty() { return Err(PyErr::new::("Missing required column argument")); } + Ok(spark_funcs::variant_delete(args[0].clone(), args[1..].to_vec())) + }, "isnan" => { if args.is_empty() { return Err(PyErr::new::("Missing required column argument")); } Ok(spark_funcs::isnan(args[0].clone())) @@ -1363,7 +1379,8 @@ pub fn call_builtin(name: &str, args: Vec) -> PyR }, "instr" => { if args.len() < 2 { return Err(PyErr::new::(format!("Missing required arguments for instr: expected at least 2, got {}", args.len()))); } - Ok(spark_funcs::instr(args[0].clone(), args[1].clone())) + // Forward the optional 4.3.0 start/occurrence args too (not just the first 2). + Ok(spark_funcs::call_function("instr", args.clone())) }, "json_tuple" => Ok(spark_funcs::json_tuple(args.clone())), "least" => { @@ -1712,7 +1729,8 @@ pub fn call_builtin(name: &str, args: Vec) -> PyR }, "regexp_replace" => { if args.len() < 3 { return Err(PyErr::new::(format!("Missing required arguments for regexp_replace: expected at least 3, got {}", args.len()))); } - Ok(spark_funcs::regexp_replace(args[0].clone(), args[1].clone(), args[2].clone())) + // Forward the optional 4.3.0 position arg too (not just the first 3). + Ok(spark_funcs::call_function("regexp_replace", args.clone())) }, "rpad" => { if args.len() < 3 { return Err(PyErr::new::(format!("Missing required arguments for rpad: expected at least 3, got {}", args.len()))); } diff --git a/crates/pyspark-rs/src/functions.rs b/crates/pyspark-rs/src/functions.rs index 2e9e9cc0..3b628ad0 100644 --- a/crates/pyspark-rs/src/functions.rs +++ b/crates/pyspark-rs/src/functions.rs @@ -120,6 +120,7 @@ pub fn register_functions(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult< m.add_function(wrap_pyfunction!(pyfunc_call_named_function, m)?)?; m.add_function(wrap_pyfunction!(pyfunc_invoke_function, m)?)?; m.add_function(wrap_pyfunction!(pyfunc_sha2, m)?)?; + m.add_function(wrap_pyfunction!(pyfunc_variant_strip_nulls, m)?)?; m.add_function(wrap_pyfunction!(pyfunc_window, m)?)?; m.add_function(wrap_pyfunction!(pyfunc_window_with_slide_and_start, m)?)?; m.add_function(wrap_pyfunction!(pyfunc_from_avro, m)?)?; @@ -402,6 +403,15 @@ fn pyfunc_sha2(col: &PyColumn, numbits: i32) -> PyColumn { PyColumn::new(spark_funcs::sha2(col.column.clone(), numbits)) } +/// Wrapper for variant_strip_nulls(v, include_arrays) - bool second arg. +#[pyfunction] +fn pyfunc_variant_strip_nulls(v: &PyColumn, include_arrays: bool) -> PyColumn { + PyColumn::new(spark_funcs::variant_strip_nulls( + v.column.clone(), + include_arrays, + )) +} + /// Wrapper for window(col, window_duration). #[pyfunction] fn pyfunc_window(col: &PyColumn, window_duration: &str) -> PyColumn { diff --git a/crates/pyspark-rs/src/lib.rs b/crates/pyspark-rs/src/lib.rs index 43a850f4..33d23209 100644 --- a/crates/pyspark-rs/src/lib.rs +++ b/crates/pyspark-rs/src/lib.rs @@ -71,7 +71,7 @@ use resource::{ PyExecutorResourceRequests, PyResourceProfile, PyResourceProfileBuilder, PyTaskResourceRequests, }; use row::PyRow; -use session::{PySparkSession, PySparkSessionBuilder}; +use session::{PyHook, PySparkSession, PySparkSessionBuilder}; use stat::PyStatFunctions; use streaming::{ PyDataStreamReader, PyDataStreamWriter, PyListenerEventStream, PyStreamingQuery, @@ -91,6 +91,7 @@ use window::{PyFrameBound, PyWindow, PyWindowSpec}; #[pymodule(gil_used = false)] fn _pyspark(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 5650e606..4fa2ad22 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -50,9 +50,11 @@ impl PySparkSessionBuilder { } /// Set the remote Spark Connect server URL. Returns the builder (chainable). - fn remote(&self, url: &str) -> PySparkSessionBuilder { + /// Mirrors `SparkSession.Builder.remote(location="sc://localhost")`. + #[pyo3(signature = (location="sc://localhost"))] + fn remote(&self, location: &str) -> PySparkSessionBuilder { let mut b = self.clone(); - b.remote_url = Some(url.to_string()); + b.remote_url = Some(location.to_string()); b } @@ -99,13 +101,25 @@ impl PySparkSessionBuilder { /// already-running remote server, so this is a no-op (chainable), matching /// pyspark's Connect behavior. #[pyo3(name = "appName")] - fn app_name(&self, _name: &str) -> PySparkSessionBuilder { + fn app_name(&self, name: &str) -> PySparkSessionBuilder { + let _ = name; self.clone() } /// `master` - not applicable to a remote Connect session (no local cluster to /// point at); accepted and ignored for API parity, like pyspark Connect. - fn master(&self, _url: &str) -> PySparkSessionBuilder { + fn master(&self, master: &str) -> PySparkSessionBuilder { + let _ = master; + self.clone() + } + + /// `_registerHook` - register an execution-hook factory (mirrors pyspark's Builder). + /// Accepted for API parity; this native-transport client sends ExecutePlanRequests + /// from Rust, so registered `Hook.on_execute_plan` callbacks are not invoked (see + /// `PyHook`). Chainable, matching the builder pattern. + #[pyo3(name = "_registerHook")] + fn register_hook(&self, hook_factory: &Bound<'_, PyAny>) -> PySparkSessionBuilder { + let _ = hook_factory; self.clone() } @@ -232,6 +246,28 @@ fn log_client_banner(py: Python<'_>) { }); } +/// `SparkSession.Hook` - injects behavior around plan execution (mirrors pyspark's +/// `SparkSession.Hook`). In pyspark, `on_execute_plan` rewrites the Python +/// `ExecutePlanRequest` before the gRPC client sends it. This client uses a native Rust +/// transport that builds and sends the request in Rust (there is no Python request to +/// rewrite), so the hook is a passthrough - the same architectural reason +/// `SparkSession.client` is absent. Exposed for API parity and subclassing. +#[pyclass(name = "Hook", module = "pyspark.sql.connect.session", subclass)] +pub struct PyHook; + +#[pymethods] +impl PyHook { + #[new] + fn new() -> Self { + PyHook + } + + /// Called before sending an ExecutePlanRequest; returns the (unchanged) request. + fn on_execute_plan<'py>(&self, request: Bound<'py, PyAny>) -> Bound<'py, PyAny> { + request + } +} + /// Python wrapper for a Spark session. #[pyclass(name = "SparkSession", module = "pyspark.sql.session")] pub struct PySparkSession { @@ -652,6 +688,13 @@ impl PySparkSession { py.get_type::().unbind() } + /// The execution hook base class, reached as `SparkSession.Hook` (mirrors pyspark). + #[classattr] + #[allow(non_snake_case)] + fn Hook(py: Python<'_>) -> Py { + py.get_type::().unbind() + } + /// UDTF registration accessor, mirroring `SparkSession.udtf.register(name, cls)`. /// Returns a registration object whose `register` cloudpickles the Python UDTF /// class and returns a bound, name-registered table function. diff --git a/crates/pyspark-rs/src/types.rs b/crates/pyspark-rs/src/types.rs index b58db7cb..4aab99a1 100644 --- a/crates/pyspark-rs/src/types.rs +++ b/crates/pyspark-rs/src/types.rs @@ -932,6 +932,9 @@ pub struct PyDateType; #[pymethods] impl PyDateType { + /// Number of days from 0001-01-01 to 1970-01-01 (mirrors pyspark's DateType.EPOCH_ORDINAL). + #[classattr] + const EPOCH_ORDINAL: i32 = 719163; // --- DataType object-model methods (v4.2.0 parity) --- #[pyo3(name = "json")] fn __obj_json(slf: &Bound<'_, Self>) -> PyResult { @@ -2143,6 +2146,10 @@ pub struct PyYearMonthIntervalType { #[pymethods] impl PyYearMonthIntervalType { + #[classattr] + const YEAR: i32 = 0; + #[classattr] + const MONTH: i32 = 1; // --- DataType object-model methods (v4.2.0 parity) --- #[pyo3(name = "json")] fn __obj_json(slf: &Bound<'_, Self>) -> PyResult { @@ -2222,6 +2229,14 @@ pub struct PyDayTimeIntervalType { #[pymethods] impl PyDayTimeIntervalType { + #[classattr] + const DAY: i32 = 0; + #[classattr] + const HOUR: i32 = 1; + #[classattr] + const MINUTE: i32 = 2; + #[classattr] + const SECOND: i32 = 3; // --- DataType object-model methods (v4.2.0 parity) --- #[pyo3(name = "json")] fn __obj_json(slf: &Bound<'_, Self>) -> PyResult { @@ -2350,6 +2365,16 @@ impl PyVariantType { } } +// Class constants shared by all spatial types (SpatialType is macro-generated with no +// #[pymethods], so this companion block adds the reference SpatialType.MIXED_* constants). +#[pymethods] +impl PySpatialType { + #[classattr] + const MIXED_CRS: &'static str = "SRID:ANY"; + #[classattr] + const MIXED_SRID: i32 = -1; +} + #[pyclass(name = "GeometryType", module = "pyspark.sql.types", extends = PySpatialType)] pub struct PyGeometryType { pub srid: i32, @@ -2357,6 +2382,14 @@ pub struct PyGeometryType { #[pymethods] impl PyGeometryType { + #[classattr] + const DEFAULT_CRS: &'static str = "OGC:CRS84"; + #[classattr] + const DEFAULT_SRID: i32 = 4326; + #[classattr] + const MIXED_CRS: &'static str = "SRID:ANY"; + #[classattr] + const MIXED_SRID: i32 = -1; #[new] #[pyo3(signature = (srid=0))] fn new(srid: i32) -> pyo3::PyClassInitializer { @@ -2391,6 +2424,16 @@ pub struct PyGeographyType { #[pymethods] impl PyGeographyType { + #[classattr] + const DEFAULT_ALG: &'static str = "SPHERICAL"; + #[classattr] + const DEFAULT_CRS: &'static str = "OGC:CRS84"; + #[classattr] + const DEFAULT_SRID: i32 = 4326; + #[classattr] + const MIXED_CRS: &'static str = "SRID:ANY"; + #[classattr] + const MIXED_SRID: i32 = -1; #[new] #[pyo3(signature = (srid=0))] fn new(srid: i32) -> pyo3::PyClassInitializer { diff --git a/crates/spark-connect-core/src/channel.rs b/crates/spark-connect-core/src/channel.rs index 79a3063a..43a9e28b 100644 --- a/crates/spark-connect-core/src/channel.rs +++ b/crates/spark-connect-core/src/channel.rs @@ -12,7 +12,7 @@ use std::collections::BTreeMap; use crate::error::{Result, SparkError}; /// Client version reported in the user-agent (mirrors `pyspark.__version__`). -pub const SPARK_VERSION: &str = "4.2.0"; +pub const SPARK_VERSION: &str = "4.3.0"; /// Default Spark Connect server port. pub const DEFAULT_PORT: u16 = 15002; diff --git a/crates/spark-connect-proto/proto/PROTO_VERSION.txt b/crates/spark-connect-proto/proto/PROTO_VERSION.txt index 51be872e..9ca23982 100644 --- a/crates/spark-connect-proto/proto/PROTO_VERSION.txt +++ b/crates/spark-connect-proto/proto/PROTO_VERSION.txt @@ -1 +1 @@ -v4.2.0 +v4.3.0 diff --git a/crates/spark-connect/src/functions.rs b/crates/spark-connect/src/functions.rs index 3926fbb1..bdd049c9 100644 --- a/crates/spark-connect/src/functions.rs +++ b/crates/spark-connect/src/functions.rs @@ -628,6 +628,14 @@ pub fn base64(col1: Column) -> Column { func("base64", vec![col1.expression().clone()]) } +pub fn to_base32(col1: Column) -> Column { + func("to_base32", vec![col1.expression().clone()]) +} + +pub fn from_base32(col1: Column) -> Column { + func("from_base32", vec![col1.expression().clone()]) +} + /// Mirrors `pyspark.sql.functions.bin`. pub fn bin(col1: Column) -> Column { func("bin", vec![col1.expression().clone()]) @@ -1012,6 +1020,17 @@ pub fn is_variant_null(col1: Column) -> Column { func("is_variant_null", vec![col1.expression().clone()]) } +pub fn collect_union(col1: Column) -> Column { + func("collect_union", vec![col1.expression().clone()]) +} + +pub fn variant_strip_nulls(v: Column, include_arrays: bool) -> Column { + func( + "variant_strip_nulls", + vec![v.expression().clone(), lit_bool(include_arrays)], + ) +} + /// Mirrors `pyspark.sql.functions.isnan`. pub fn isnan(col1: Column) -> Column { func("isnan", vec![col1.expression().clone()]) diff --git a/crates/spark-connect/src/lib.rs b/crates/spark-connect/src/lib.rs index 5f1a8faf..d6a6e3fd 100644 --- a/crates/spark-connect/src/lib.rs +++ b/crates/spark-connect/src/lib.rs @@ -63,5 +63,8 @@ pub use storage::StorageLevelExt; pub use streaming::{ DataStreamReader, DataStreamWriter, StreamingQuery, StreamingQueryManager, Trigger, }; -pub use types::{DataType, StructField}; +pub use types::{ + DataType, DateType, DayTimeIntervalType, GeographyType, GeometryType, SpatialType, StructField, + YearMonthIntervalType, +}; pub use window::{FrameBound, FrameType, Window, WindowSpec}; diff --git a/crates/spark-connect/src/types.rs b/crates/spark-connect/src/types.rs index 1f06c73f..fbc55a56 100644 --- a/crates/spark-connect/src/types.rs +++ b/crates/spark-connect/src/types.rs @@ -88,6 +88,58 @@ pub enum DataType { Unparsed { data_type_string: String }, } +/// Field-index constants for [`DataType::YearMonthInterval`], mirroring +/// `pyspark.sql.types.YearMonthIntervalType`. +pub struct YearMonthIntervalType; +impl YearMonthIntervalType { + pub const YEAR: i32 = 0; + pub const MONTH: i32 = 1; +} + +/// Field-index constants for [`DataType::DayTimeInterval`], mirroring +/// `pyspark.sql.types.DayTimeIntervalType`. +pub struct DayTimeIntervalType; +impl DayTimeIntervalType { + pub const DAY: i32 = 0; + pub const HOUR: i32 = 1; + pub const MINUTE: i32 = 2; + pub const SECOND: i32 = 3; +} + +/// Constants for [`DataType::Date`], mirroring `pyspark.sql.types.DateType`. +pub struct DateType; +impl DateType { + /// Number of days from 0001-01-01 to 1970-01-01. + pub const EPOCH_ORDINAL: i32 = 719163; +} + +/// CRS/SRID constants shared by the spatial types, mirroring +/// `pyspark.sql.types.SpatialType`. +pub struct SpatialType; +impl SpatialType { + pub const MIXED_CRS: &'static str = "SRID:ANY"; + pub const MIXED_SRID: i32 = -1; +} + +/// Constants for [`DataType::Geometry`], mirroring `pyspark.sql.types.GeometryType`. +pub struct GeometryType; +impl GeometryType { + pub const DEFAULT_CRS: &'static str = "OGC:CRS84"; + pub const DEFAULT_SRID: i32 = 4326; + pub const MIXED_CRS: &'static str = "SRID:ANY"; + pub const MIXED_SRID: i32 = -1; +} + +/// Constants for [`DataType::Geography`], mirroring `pyspark.sql.types.GeographyType`. +pub struct GeographyType; +impl GeographyType { + pub const DEFAULT_ALG: &'static str = "SPHERICAL"; + pub const DEFAULT_CRS: &'static str = "OGC:CRS84"; + pub const DEFAULT_SRID: i32 = 4326; + pub const MIXED_CRS: &'static str = "SRID:ANY"; + pub const MIXED_SRID: i32 = -1; +} + /// A field in a StructType, mirroring `pyspark.sql.types.StructField`. // Not `Eq`: metadata values are arbitrary JSON (`serde_json::Value`, which is only `PartialEq` // because of floats), matching pyspark's `Dict[str, Any]` field metadata. diff --git a/dev/parity/4.3.0-introspection-diff.md b/dev/parity/4.3.0-introspection-diff.md new file mode 100644 index 00000000..ca46a652 --- /dev/null +++ b/dev/parity/4.3.0-introspection-diff.md @@ -0,0 +1,74 @@ +# PySpark 4.3.0-rc3 introspection diff (connect classes) + +Runtime `dir()` diff of our built `pyspark._pyspark` classes vs the reference +`pyspark-client==4.3.0` (installed from the `v4.3.0-rc3` staging sdist). Method = +public callables + properties; signatures compared on parameter name + default +(annotations and the `self, /` marker normalized away). + +Re-run against RC3: the results are the same as against RC1. Between `v4.3.0-rc1` and +`v4.3.0-rc3`, upstream changed no connect-client class or `__all__` surface. + +## Method-name parity: 100% (0 gaps) + +| Class | ref | ours | missing | +|---|---|---|---| +| DataFrame | 118 | 119 | 0 | +| Column | 37 | 36 | 0 (justified: `to_plan`) | +| SparkSession | 40 | 42 | 0 | +| SparkSession.Builder | 8 | 8 | 0 | +| DataFrameReader | 14 | 14 | 0 | +| DataFrameWriter | 18 | 18 | 0 | +| DataFrameWriterV2 | 12 | 12 | 0 | +| GroupedData | 15 | 15 | 0 | +| Catalog | 37 | 37 | 0 | +| Window | 7 | 7 | 0 | +| WindowSpec | 4 | 4 | 0 | +| Observation | 1 | 2 | 0 | + +Justified non-gaps (native Rust transport, not the Python gRPC client): `Column.to_plan`, +`SparkSession.client`. + +## `SparkSession.Hook` (new in 4.3.0) — implemented + +Nested `SparkSession.Hook` base class (`on_execute_plan`) + the `_registerHook` factory, +letting user code inject behavior around plan execution. In pyspark, `on_execute_plan` +rewrites the Python `ExecutePlanRequest` before the gRPC client sends it. This client uses +a native Rust transport that builds/sends the request in Rust — there is no Python request +to rewrite — so `Hook.on_execute_plan` is a passthrough and `_registerHook` is an accepted +chainable no-op, the same architectural rationale as the justified-absent +`SparkSession.client` / `Column.to_plan`. Both are present for API parity + subclassing. + +## Signature notes (benign — ours accepts a superset of valid pyspark calls) + +Fixed here: `SparkSession.Builder.appName/master/remote` param names now match pyspark +(`name`/`master`/`location`), and `remote` defaults to `sc://localhost`. + +Remaining, intentionally benign: `Builder.config` accepts an extra optional `conf`; +`emptyDataFrame`, `DataFrameReader.text`, `freqItems`, `persist`, `replace`, `melt`, +`unpivot`, `nearestByJoin` use `None`-sentinel defaults (resolved to the documented value +internally) or default args where the reference marks them required — every valid +reference call is accepted. + +## Full-surface sweep (beyond the 12 core classes) + +A second pass diffed the authoritative public API (`__all__`) of `pyspark.sql.functions`, +`pyspark.sql`, `pyspark.sql.types`, `pyspark.sql.streaming`, plus datatype class constants +and `MergeIntoWriter`'s nested classes. It found — and this PR closed — real gaps the +class-only diff missed: + +- `functions.__all__`: `variant_delete` (wrapper + dispatch arm), `PandasUDFType`, + `ArrowUDFType`, `UserDefinedFunction`, `UserDefinedTableFunction`, `AnalyzeArgument`, + `AnalyzeResult`, `SkipRestOfInputTableException`. +- `types.__all__`: `Geometry`, `GeometryType`, `Geography`, `GeographyType`, `VariantVal`. +- Datatype constants: `DayTimeIntervalType.{DAY,HOUR,MINUTE,SECOND}`, + `YearMonthIntervalType.{YEAR,MONTH}`, `DateType.EPOCH_ORDINAL`, and the + `Geometry/Geography/SpatialType` CRS/SRID/ALG constants (added as PyO3 `#[classattr]`s + AND as native `apache-spark-connect` constants). +- `MergeIntoWriter.{WhenMatched,WhenNotMatched,WhenNotMatchedBySource}` nested classes. + +After the fixes the `__all__` + constants sweep is fully clean (0 missing). The new +functions and constants exist on the native Rust side too (`spark_connect::functions::*`, +`spark_connect::{DayTimeIntervalType, GeometryType, ...}`). + +Note: this is a NAME + SIGNATURE proof, not a behavior proof. Behavioral parity (the +official 4.3.0 connect test suite through our transport) runs on CI against the RC. diff --git a/docs/architecture.md b/docs/architecture.md index f7fdb920..571f1672 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -36,7 +36,7 @@ flowchart TD ARROW["Arrow result decode"] end - SERVER["Spark Connect server (Spark 4.2+, :15002)"] + SERVER["Spark Connect server (Spark 4.3+, :15002)"] OUT["Rows / Arrow - optionally DataFusion or Polars (Rust); pandas (Python)"] PY --> SKIN --> EXT --> API diff --git a/docs/compatibility.md b/docs/compatibility.md index cc5fb6a0..9f7d684f 100644 --- a/docs/compatibility.md +++ b/docs/compatibility.md @@ -9,8 +9,8 @@ on PyPI - and how to tell which one you have installed - see ## Supported Spark -Apache Spark **4.2.0 and later**. The crate and wheel version tracks the Spark -release it targets (starting at `4.2.0`), so the version number tells you which +Apache Spark **4.3.0 and later**. The crate and wheel version tracks the Spark +release it targets (starting at `4.3.0`), so the version number tells you which Spark it speaks. ## Rust migration for the 4.3 line diff --git a/docs/configuration.md b/docs/configuration.md index 053059d4..58e25c32 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -51,13 +51,13 @@ export SPARK_HOME=/path/to/spark # Start the server on sc://localhost:15002 $SPARK_HOME/sbin/start-connect-server.sh \ - --packages "org.apache.spark:spark-connect_2.13:4.2.0" + --packages "org.apache.spark:spark-connect_2.13:4.3.0" ``` The server listens on port `15002` by default. To use a different port, add `--conf spark.connect.grpc.binding.port=`. !!! tip - The Spark Connect server requires a JVM and Apache Spark 4.2.0 or later. Stop it with `$SPARK_HOME/sbin/stop-connect-server.sh`. + The Spark Connect server requires a JVM and Apache Spark 4.3.0 or later. Stop it with `$SPARK_HOME/sbin/stop-connect-server.sh`. ## Remote Connections diff --git a/docs/index.md b/docs/index.md index 6722fa55..ae32b273 100644 --- a/docs/index.md +++ b/docs/index.md @@ -13,7 +13,7 @@ results as the reference client.
![PyPI](https://img.shields.io/pypi/v/pyspark-client-rust?color=c2410c&label=pyspark-client-rust) -![Spark](https://img.shields.io/badge/Apache%20Spark-4.2.0%2B-c2410c) +![Spark](https://img.shields.io/badge/Apache%20Spark-4.3.0%2B-c2410c) ![License](https://img.shields.io/badge/license-Apache--2.0-blue)
@@ -74,7 +74,7 @@ runs unchanged.
-### Spark 4.2.0+ +### Spark 4.3.0+ The crate and wheel version tracks the Spark release it targets, so the version number tells you which Spark it speaks.
diff --git a/docs/installation.md b/docs/installation.md index d6f9fb01..311ed307 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -8,7 +8,7 @@ Add the crate to your `Cargo.toml`: ```toml [dependencies] -apache-spark-connect = "4.2" +apache-spark-connect = "4.3" ``` ### Optional Features @@ -17,7 +17,7 @@ Enable optional features for additional functionality: ```toml [dependencies] -apache-spark-connect = { version = "4.2", features = ["datafusion", "polars"] } +apache-spark-connect = { version = "4.3", features = ["datafusion", "polars"] } ``` **Requirements:** @@ -91,7 +91,7 @@ pip install dist/pyspark_client_rust-*.whl ## Version and Compatibility -The crate and package versions track Apache Spark: version `4.2.x` supports Spark 4.2.0 and later. The wire protocol is identical to the reference client, so existing Spark Connect code works unchanged. +The crate and package versions track Apache Spark: version `4.3.x` supports Spark 4.3.0 and later. The wire protocol is identical to the reference client, so existing Spark Connect code works unchanged. ## Next Steps diff --git a/docs/quickstart.md b/docs/quickstart.md index 8f37b405..31060b3f 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -7,7 +7,7 @@ Get up and running with Spark Connect in minutes. This guide walks you through y Before you start, ensure you have a running Spark Connect server. See [Configuration and Connection](configuration.md) for how to start one locally: ```bash -$SPARK_HOME/sbin/start-connect-server.sh --packages "org.apache.spark:spark-connect_2.13:4.2.0" +$SPARK_HOME/sbin/start-connect-server.sh --packages "org.apache.spark:spark-connect_2.13:4.3.0" ``` The server listens on `sc://localhost:15002` by default. diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index f15c70fc..94b268c5 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -7,13 +7,13 @@ Common issues and solutions when using Spark Connect Rust. **Symptom:** `Error: Failed to connect to sc://localhost:15002` or connection timeout. **Solution:** -1. Verify the server is running: `$SPARK_HOME/sbin/start-connect-server.sh --packages "org.apache.spark:spark-connect_2.13:4.2.0"` +1. Verify the server is running: `$SPARK_HOME/sbin/start-connect-server.sh --packages "org.apache.spark:spark-connect_2.13:4.3.0"` 2. Check the port is correct. By default it's `15002`. To verify, look for `spark.connect.grpc.binding.port` in the server logs. 3. If connecting to a remote server, verify network connectivity: `ping spark.example.com` and check firewalls. -4. Ensure Apache Spark 4.2.0 or later is installed: `$SPARK_HOME/bin/spark-shell --version` +4. Ensure Apache Spark 4.3.0 or later is installed: `$SPARK_HOME/bin/spark-shell --version` !!! tip - Run the server in the foreground to see startup messages: `$SPARK_HOME/sbin/start-connect-server.sh --packages "org.apache.spark:spark-connect_2.13:4.2.0" 2>&1` + Run the server in the foreground to see startup messages: `$SPARK_HOME/sbin/start-connect-server.sh --packages "org.apache.spark:spark-connect_2.13:4.3.0" 2>&1` ### TLS / Authentication Error @@ -61,6 +61,6 @@ This is normal. The JVM and Spark runtime warm up on first use. Subsequent queri **Symptom:** Errors about incompatible schema or unexpected null values. **Solution:** -1. Verify the Spark Connect server and client are compatible versions: both should be Spark 4.2.0 or later. +1. Verify the Spark Connect server and client are compatible versions: both should be Spark 4.3.0 or later. 2. If you recently upgraded Apache Spark, restart the Spark Connect server to apply the new version. 3. Check that data on disk matches the expected schema (e.g., after schema changes, re-write the data). diff --git a/docs/udfs.md b/docs/udfs.md index 85aebe97..580cd2fa 100644 --- a/docs/udfs.md +++ b/docs/udfs.md @@ -62,15 +62,15 @@ fn main() { ```toml # Cargo.toml [dependencies] -spark-connect-macros = { package = "apache-spark-connect-macros", version = "4.2" } +spark-connect-macros = { package = "apache-spark-connect-macros", version = "4.3" } # host-only: the client is not pulled into the wasm build [target.'cfg(not(target_arch = "wasm32"))'.dependencies] -spark-connect = { package = "apache-spark-connect", version = "4.2", features = ["wasm-udf"] } +spark-connect = { package = "apache-spark-connect", version = "4.3", features = ["wasm-udf"] } [build-dependencies] -spark-connect-build = { package = "apache-spark-connect-build", version = "4.2" } -spark-connect-macros = { package = "apache-spark-connect-macros", version = "4.2" } +spark-connect-build = { package = "apache-spark-connect-build", version = "4.3" } +spark-connect-macros = { package = "apache-spark-connect-macros", version = "4.3" } ``` Run it: @@ -134,7 +134,7 @@ Nothing here is needed unless the `wasm-udf` feature is enabled: there is nothing extra to install or put on a path; point `SPARK_CONNECT_PYTHON` at the interpreter if it isn't your default `python3`. - **Executors** - the `wasmtime` Python package. -- **Spark** - 4.2.0+. +- **Spark** - 4.3.0+. ## Advanced diff --git a/docs/which-client.md b/docs/which-client.md index 36745137..977dd856 100644 --- a/docs/which-client.md +++ b/docs/which-client.md @@ -40,7 +40,7 @@ also why the two are separate implementations that can have separate bugs. import pyspark pyspark.__rust_client__ # True only in pyspark-client-rust pyspark.__engine__ # "rust" - pyspark.__version__ # the Spark version this client targets, e.g. "4.2.0" + pyspark.__version__ # the Spark version this client targets, e.g. "4.3.0" ``` === "Log on connect" @@ -49,7 +49,7 @@ also why the two are separate implementations that can have separate bugs. `pyspark` logger: ``` - pyspark-client-rust 4.2.0 active: Spark Connect client backed by the native Rust + pyspark-client-rust 4.3.0 active: Spark Connect client backed by the native Rust engine (tonic), a drop-in for pyspark-client -- not the reference Python client. ``` @@ -79,6 +79,6 @@ bug we want to hear about, since the goal is byte-for-byte parity. ## Server compatibility -This client speaks to Apache Spark Connect servers **4.2.0 and later** (see -[Compatibility](compatibility.md)). To talk to servers older than 4.2, keep using the +This client speaks to Apache Spark Connect servers **4.3.0 and later** (see +[Compatibility](compatibility.md)). To talk to servers older than 4.3, keep using the reference client; the two can coexist in separate environments. diff --git a/examples/README.md b/examples/README.md index e84da409..8c3ace82 100644 --- a/examples/README.md +++ b/examples/README.md @@ -69,7 +69,7 @@ Read a file into a dataframe, save the result as a deltalake table, and append a The Delta Lake jars need to be added onto the server before it starts: ```bash -$ $SPARK_HOME/sbin/start-connect-server.sh --packages "org.apache.spark:spark-connect_2.13:4.2.0,io.delta:delta-spark_2.13:4.2.0" \ +$ $SPARK_HOME/sbin/start-connect-server.sh --packages "org.apache.spark:spark-connect_2.13:4.3.0,io.delta:delta-spark_2.13:4.3.0" \ --conf "spark.driver.extraJavaOptions=-Divy.cache.dir=/tmp -Divy.home=/tmp" \ --conf "spark.sql.extensions=io.delta.sql.DeltaSparkSessionExtension" \ --conf "spark.sql.catalog.spark_catalog=org.apache.spark.sql.delta.catalog.DeltaCatalog" diff --git a/examples/wasm-udf-inline/Cargo.toml b/examples/wasm-udf-inline/Cargo.toml index 68f5dcd7..addeb88f 100644 --- a/examples/wasm-udf-inline/Cargo.toml +++ b/examples/wasm-udf-inline/Cargo.toml @@ -31,16 +31,16 @@ edition = "2021" publish = false [dependencies] -spark-connect-macros = { package = "apache-spark-connect-macros", path = "../../crates/spark-connect-macros", version = "4.2.0" } +spark-connect-macros = { package = "apache-spark-connect-macros", path = "../../crates/spark-connect-macros", version = "4.3.0" } # Host-only: the client + the macro-generated constructors (which need the # `wasm-udf` feature). Never compiled for wasm32. [target.'cfg(not(target_arch = "wasm32"))'.dependencies] -spark-connect = { package = "apache-spark-connect", path = "../../crates/spark-connect", version = "4.2.0", features = ["wasm-udf"] } +spark-connect = { package = "apache-spark-connect", path = "../../crates/spark-connect", version = "4.3.0", features = ["wasm-udf"] } [build-dependencies] # The build helper that compiles + embeds the wasm module. -spark-connect-build = { package = "apache-spark-connect-build", path = "../../crates/spark-connect-build", version = "4.2.0" } +spark-connect-build = { package = "apache-spark-connect-build", path = "../../crates/spark-connect-build", version = "4.3.0" } # Forces the proc-macro to be built before build.rs runs, so the helper can find # its dylib to feed the wasm `rustc` invocation. -spark-connect-macros = { package = "apache-spark-connect-macros", path = "../../crates/spark-connect-macros", version = "4.2.0" } +spark-connect-macros = { package = "apache-spark-connect-macros", path = "../../crates/spark-connect-macros", version = "4.3.0" } diff --git a/examples/wasm-udfs/Cargo.toml b/examples/wasm-udfs/Cargo.toml index 73e1e3f4..93f2949c 100644 --- a/examples/wasm-udfs/Cargo.toml +++ b/examples/wasm-udfs/Cargo.toml @@ -29,13 +29,13 @@ edition = "2021" publish = false [dependencies] -spark-connect-macros = { package = "apache-spark-connect-macros", path = "../../crates/spark-connect-macros", version = "4.2.0" } +spark-connect-macros = { package = "apache-spark-connect-macros", path = "../../crates/spark-connect-macros", version = "4.3.0" } # Host-only: the macro-generated constructors (which need the `wasm-udf` # feature). Never compiled for wasm32. [target.'cfg(not(target_arch = "wasm32"))'.dependencies] -spark-connect = { package = "apache-spark-connect", path = "../../crates/spark-connect", version = "4.2.0", features = ["wasm-udf"] } +spark-connect = { package = "apache-spark-connect", path = "../../crates/spark-connect", version = "4.3.0", features = ["wasm-udf"] } [build-dependencies] -spark-connect-build = { package = "apache-spark-connect-build", path = "../../crates/spark-connect-build", version = "4.2.0" } -spark-connect-macros = { package = "apache-spark-connect-macros", path = "../../crates/spark-connect-macros", version = "4.2.0" } +spark-connect-build = { package = "apache-spark-connect-build", path = "../../crates/spark-connect-build", version = "4.3.0" } +spark-connect-macros = { package = "apache-spark-connect-macros", path = "../../crates/spark-connect-macros", version = "4.3.0" } diff --git a/python/.coveragerc b/python/.coveragerc index bbb7556c..4fe8f831 100644 --- a/python/.coveragerc +++ b/python/.coveragerc @@ -28,6 +28,7 @@ omit = */pyspark/pipelines/* */pyspark/logger/* */pyspark/errors/* + */pyspark/messages/* */pyspark/sql/streaming/* */pyspark/sql/worker/* */pyspark/sql/plot/* diff --git a/python/pyspark/__init__.py b/python/pyspark/__init__.py index 5a10e957..20f56a3c 100644 --- a/python/pyspark/__init__.py +++ b/python/pyspark/__init__.py @@ -83,4 +83,5 @@ def getAll(self): "Row", "SparkConf", "StorageLevel", + "__version__", ] diff --git a/python/pyspark/accumulators.py b/python/pyspark/accumulators.py index fcfa3470..0e68db70 100644 --- a/python/pyspark/accumulators.py +++ b/python/pyspark/accumulators.py @@ -24,7 +24,7 @@ import threading from typing import Callable, Dict, Generic, Tuple, Type, TYPE_CHECKING, TypeVar, Union, Optional -from pyspark.serializers import read_int, CPickleSerializer +from pyspark.serializers import read_int, CPickleSerializer, CloudPickleSerializer from pyspark.errors import PySparkRuntimeError if TYPE_CHECKING: @@ -39,6 +39,10 @@ pickleSer = CPickleSerializer() +# Safe serializer for special accumulators. The client will unpickle the metrics from the server +# unconditionally, so we need to restrict the allowed classes. +specialAccumulatorSer = CloudPickleSerializer(allowed_names=[("pstats", "Stats")]) + # Holds accumulators registered on the current machine, keyed by ID. This is then used to send # the local accumulator updates back to the driver program at the end of a task. _accumulatorRegistry: Dict[int, "Accumulator"] = {} diff --git a/python/pyspark/errors/error-conditions.json b/python/pyspark/errors/error-conditions.json index 38417cbf..756d55f1 100644 --- a/python/pyspark/errors/error-conditions.json +++ b/python/pyspark/errors/error-conditions.json @@ -72,7 +72,7 @@ }, "CANNOT_CONVERT_COLUMN_INTO_BOOL": { "message": [ - "Cannot convert column into bool: please use '&' for 'and', '|' for 'or', '~' for 'not' when building DataFrame boolean expressions." + "Cannot convert column into bool (offending column: ): please use '&' for 'and', '|' for 'or', '~' for 'not' when building DataFrame boolean expressions." ] }, "CANNOT_CONVERT_TYPE": { @@ -194,7 +194,7 @@ }, "DATA_SOURCE_EXTRANEOUS_FILTERS": { "message": [ - ".pushFilters() returned filters that are not part of the input. Make sure that each returned filter is one of the input filters by reference." + ".pushFilters() returned filters that are not part of the input. Make sure that each returned filter is one of the input filters by reference. Input filters: . Extraneous filters: ." ] }, "DATA_SOURCE_INVALID_RETURN_TYPE": { @@ -430,6 +430,11 @@ "Invalid number of dataframes in group ." ] }, + "INVALID_OPERATION_UUID_ID": { + "message": [ + "Parameter value must be a valid UUID format: " + ] + }, "INVALID_PANDAS_UDF": { "message": [ "Invalid function: " @@ -546,6 +551,16 @@ " and should be of the same length, got and ." ] }, + "LOCAL_CONNECT_RUNTIME_DIR_UNAVAILABLE": { + "message": [ + "Cannot claim the per-user runtime directory (was it created by another user?); remove it or point SPARK_LOCAL_CONNECT_DISCOVERY at a path you own." + ] + }, + "LOCAL_CONNECT_SERVER_START_FAILED": { + "message": [ + "Failed to start a persistent local Spark Connect server: ." + ] + }, "LOCAL_RELATION_SIZE_LIMIT_EXCEEDED": { "message": [ "Local relation size ( bytes) exceeds the limit ( bytes)." @@ -851,11 +866,7 @@ "An Observation can be used with a DataFrame only once." ] }, - "SCHEMA_MISMATCH_FOR_PANDAS_UDF": { - "message": [ - "Result vector from was not the required length: expected , got ." - ] - }, + "SESSION_ALREADY_EXIST": { "message": [ "Cannot start a remote Spark session because there is a regular Spark session already running." @@ -1083,11 +1094,6 @@ "Unexpected tuple with StructType." ] }, - "UNKNOWN_EXPLAIN_MODE": { - "message": [ - "Unknown explain mode: ''. Accepted explain modes are 'simple', 'extended', 'codegen', 'cost', 'formatted'." - ] - }, "UNKNOWN_INTERRUPT_TYPE": { "message": [ "Unknown interrupt type: ''. Accepted interrupt types are 'all'." @@ -1118,6 +1124,12 @@ " is not supported in conversion to Arrow." ] }, + "UNSUPPORTED_FILTER": { + "message": [ + "Unsupported filter ." + ], + "sqlState": "0A000" + }, "UNSUPPORTED_JOIN_TYPE": { "message": [ "Unsupported join type: ''. Supported join types include: ." @@ -1198,11 +1210,6 @@ "Value for `` has to be amongst the following values: ." ] }, - "VALUE_NOT_ANY_OR_ALL": { - "message": [ - "Value for `` must be 'any' or 'all', got ''." - ] - }, "VALUE_NOT_BETWEEN": { "message": [ "Value for `` must be between and ." @@ -1213,11 +1220,6 @@ "Value for `` must be a non-empty string, got ''." ] }, - "VALUE_NOT_PEARSON": { - "message": [ - "Value for `` only supports 'pearson', got ''." - ] - }, "VALUE_NOT_PLAIN_COLUMN_REFERENCE": { "message": [ "Value `` in `` should be a plain column reference such as `df.col` or `col('column')`." @@ -1228,16 +1230,17 @@ "Value for `` must be positive, got ''." ] }, - "VALUE_NOT_TRUE": { - "message": [ - "Value for `` must be True, got ''." - ] - }, "VALUE_OUT_OF_BOUNDS": { "message": [ "Value for `` must be between and (inclusive), got " ] }, + "VARIANT_SIZE_LIMIT": { + "message": [ + "Cannot build variant bigger than bytes." + ], + "sqlState": "22023" + }, "WKB_PARSE_ERROR" : { "message" : [ "Error parsing WKB: at position " diff --git a/python/pyspark/errors/error_classes.py b/python/pyspark/errors/error_classes.py index 12094d61..4a840baf 100644 --- a/python/pyspark/errors/error_classes.py +++ b/python/pyspark/errors/error_classes.py @@ -23,6 +23,8 @@ # For more information, please see: https://issues.apache.org/jira/browse/SPARK-46810 # This discrepancy will be resolved as part of: https://issues.apache.org/jira/browse/SPARK-47429 ERROR_CLASSES_JSON = ( - importlib.resources.files("pyspark.errors").joinpath("error-conditions.json").read_text() + importlib.resources.files("pyspark.errors") + .joinpath("error-conditions.json") + .read_text(encoding="utf-8") ) ERROR_CLASSES_MAP = json.loads(ERROR_CLASSES_JSON) diff --git a/python/pyspark/errors/utils.py b/python/pyspark/errors/utils.py index c796e37a..53a5e189 100644 --- a/python/pyspark/errors/utils.py +++ b/python/pyspark/errors/utils.py @@ -98,20 +98,16 @@ def __init__(self) -> None: def get_sqlstate(self, errorClass: Optional[str]) -> Optional[str]: """ Returns the SQL state for the given error class. + + The SQL state is declared on the main class, so a sub-class inherits it. """ if errorClass is None: return None - error_classes = errorClass.split(".") - try: - if len(error_classes) == 1: - return self.error_info_map[errorClass]["sqlState"] - else: - return self.error_info_map[error_classes[0]]["sub_class"][error_classes[1]][ - "sqlState" - ] - except KeyError: + main_class_info = self.error_info_map.get(errorClass.split(".")[0]) + if main_class_info is None: return None + return main_class_info.get("sqlState") def get_error_message(self, errorClass: str, messageParameters: Dict[str, str]) -> str: """ diff --git a/python/pyspark/instrumentation_utils.py b/python/pyspark/instrumentation_utils.py index 9ad76c42..1f822207 100644 --- a/python/pyspark/instrumentation_utils.py +++ b/python/pyspark/instrumentation_utils.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with diff --git a/python/pyspark/logger/logger.py b/python/pyspark/logger/logger.py index 8f5cfb86..8d08275f 100644 --- a/python/pyspark/logger/logger.py +++ b/python/pyspark/logger/logger.py @@ -1,4 +1,3 @@ -# -*- encoding: utf-8 -*- # # Licensed to the Apache Software Foundation (ASF) under one or more # contributor license agreements. See the NOTICE file distributed with diff --git a/python/pyspark/messages/README.md b/python/pyspark/messages/README.md new file mode 100644 index 00000000..95e5962e --- /dev/null +++ b/python/pyspark/messages/README.md @@ -0,0 +1,23 @@ +## PySpark <-> Spark message interface + +This module implements message-based communication between PySpark and the Spark. +It introduces abstraction layers, which handle the receiving and sending of messages. +Through these abstractions, the underlying data transport channel (Unix domain socket, gRPC, etc.) + can be *decoupled from the core PySpark logic*, which processes the data. + +Overall, introducing these abstractions allows the same PySpark code to work with +different underlying data transport channels transparently. + +This module defines the following message types: + +### Spark -> PySpark + +1. Initialization - UDF Payload, parameters, ... +2. Data - Data to invoke the UDF on +3. Finish - UDF has been invoked on all available data + +### PySpark -> Spark + +1. Response data +2. Exceptions +3. Finish - All processing is done diff --git a/python/pyspark/messages/__init__.py b/python/pyspark/messages/__init__.py new file mode 100644 index 00000000..69cfbf6b --- /dev/null +++ b/python/pyspark/messages/__init__.py @@ -0,0 +1,26 @@ +# +# 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 pyspark.messages.spark_message_receiver import SparkMessageReceiver +from pyspark.messages.zero_copy_byte_stream import ZeroCopyByteStream +from pyspark.messages.socket.spark_socket_message_receiver import SparkSocketMessageReceiver + +__all__ = [ + "SparkMessageReceiver", + "SparkSocketMessageReceiver", + "ZeroCopyByteStream", +] diff --git a/python/pyspark/messages/socket/__init__.py b/python/pyspark/messages/socket/__init__.py new file mode 100644 index 00000000..cce3acad --- /dev/null +++ b/python/pyspark/messages/socket/__init__.py @@ -0,0 +1,16 @@ +# +# 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. +# diff --git a/python/pyspark/messages/socket/spark_socket_message_receiver.py b/python/pyspark/messages/socket/spark_socket_message_receiver.py new file mode 100644 index 00000000..7099f020 --- /dev/null +++ b/python/pyspark/messages/socket/spark_socket_message_receiver.py @@ -0,0 +1,67 @@ +# +# 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 BinaryIO + +from pyspark.serializers import read_int, SpecialLengths +from pyspark.messages.zero_copy_byte_stream import ZeroCopyByteStream +from pyspark.messages.spark_message_receiver import ( + SparkMessageReceiver, +) + + +def _assert_message_id(message_id: int, expected: int) -> None: + assert message_id == expected, ( + f"Expected message with id {expected} but got message with id {message_id} instead." + ) + + +class SparkSocketMessageReceiver(SparkMessageReceiver): + def __init__(self, infile: BinaryIO): + super().__init__() + self._infile = infile + + def _do_get_init_message(self) -> ZeroCopyByteStream: + message_id = read_int(self._infile) + _assert_message_id(message_id, SpecialLengths.START_OF_INIT_MESSAGE) + + # Read the length and init content + message_length = read_int(self._infile) + message_content = self._infile.read(message_length) + + # The init message is fully materialized, so mark the stream finished: a parse that + # runs past the declared message length (e.g. a writer/reader protocol mismatch) must + # fail fast with EOFError instead of blocking forever on a chunk that will never come. + stream = ZeroCopyByteStream(memoryview(message_content)) + stream.finish() + return stream + + def _do_get_data_stream(self) -> BinaryIO: + # For socket communication, we just pass along the underlying socket + # for the data channel. We already stripped the initialization data + # at this state. Therefore, any bytes following this are data bytes. + # + # Note: We deliberately did not introduce a message header for + # data messages to reduce the overhead, especially for small + # batch sizes and real-time-mode (RTM). + return self._infile + + def _do_get_finish_signal_from_stream(self) -> None: + # If everything finished properly, we should read END_OF_STREAM. + # Anything else means something went wrong during processing. + message_id = read_int(self._infile) + _assert_message_id(message_id, SpecialLengths.END_OF_STREAM) diff --git a/python/pyspark/messages/spark_message_receiver.py b/python/pyspark/messages/spark_message_receiver.py new file mode 100644 index 00000000..ec6b6fc3 --- /dev/null +++ b/python/pyspark/messages/spark_message_receiver.py @@ -0,0 +1,122 @@ +# +# 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 enum import Enum +from functools import wraps +from typing import BinaryIO, Callable, TypeVar +from abc import ABC, abstractmethod + +from pyspark.messages.zero_copy_byte_stream import ZeroCopyByteStream + + +T = TypeVar("T", bound="SparkMessageReceiver") +R = TypeVar("R") + + +class MessageState(Enum): + WAITING_FOR_INIT = 1 + WAITING_FOR_DATA = 2 + WAITING_FOR_FINISH = 3 + DONE = 4 + + +class SparkMessageReceiver(ABC): + """ + Generic class that implements receiving messages from Spark. + Caution: This class is STATEFUL. It is expected, that the + methods of this class are called in the following order: + + 1. Init -> 2. Data stream -> 3. Finish + + This order is verified using assertions in the class. Each function + can be called EXACTLY ONCE in the specified order. + """ + + def __init__(self) -> None: + self._state = MessageState.WAITING_FOR_INIT + + @staticmethod + def _state_transition( + required_state: MessageState, next_state: MessageState + ) -> Callable[[Callable[[T], R]], Callable[[T], R]]: + """Decorator to enforce state transitions.""" + + def decorator(func: Callable[[T], R]) -> Callable[[T], R]: + @wraps(func) + def wrapper(self: T) -> R: + assert self._state == required_state + result = func(self) + self._state = next_state + return result + + return wrapper + + return decorator + + @_state_transition(MessageState.WAITING_FOR_INIT, MessageState.WAITING_FOR_DATA) + def get_init_message(self) -> ZeroCopyByteStream: + """ + Returns: + the binary contents of the initial message as a ZeroCopyByteStream. + """ + return self._do_get_init_message() + + @_state_transition(MessageState.WAITING_FOR_DATA, MessageState.WAITING_FOR_FINISH) + def get_data_stream(self) -> BinaryIO: + """ + Returns: + A binary stream containing the data to invoke the UDF on. + """ + return self._do_get_data_stream() + + @_state_transition(MessageState.WAITING_FOR_FINISH, MessageState.DONE) + def get_finish_signal_from_stream(self) -> None: + """ + Consumes the finish message from the JVM and transitions to the DONE state. + The finish message marks the end of the stream. Raises an `AssertionError` if + the finish signal was not received correctly. + """ + self._do_get_finish_signal_from_stream() + + @abstractmethod + def _do_get_init_message(self) -> ZeroCopyByteStream: + """ + Returns the contents of the init message + as a 'ZeroCopyByteStream'. + + To be implemented by child classes. + """ + ... + + @abstractmethod + def _do_get_data_stream(self) -> BinaryIO: + """ + Returns the Spark data stream. + + To be implemented by child classes. + """ + ... + + @abstractmethod + def _do_get_finish_signal_from_stream(self) -> None: + """ + Consumes the finish signal from the stream. + Raises an `AssertionError` if the signal is not received correctly. + + To be implemented by child classes. + """ + ... diff --git a/python/pyspark/messages/zero_copy_byte_stream.py b/python/pyspark/messages/zero_copy_byte_stream.py new file mode 100644 index 00000000..611b4f92 --- /dev/null +++ b/python/pyspark/messages/zero_copy_byte_stream.py @@ -0,0 +1,176 @@ +# +# 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 threading +from typing import Optional +from collections import deque + + +class ZeroCopyByteStream: + """ + Accepts chunks of bytes as zero-copy memory views. Implements + a file-like interface on top of the received chunks. + + read() calls that access bytes from a single chunk are served + as zero-copy reads. If a read() call crosses chunk boundaries, + memory copies are required. The later case is unexpected and only + implemented for correctness. + + This implementation is thread-safe. + """ + + def __init__(self, initial_view: Optional[memoryview] = None): + if not isinstance(initial_view, memoryview) and initial_view is not None: + raise TypeError( + "Only memoryview and None are allowed as the initial " + + f"ZeroCopyByteStream view. Recveied type {type(initial_view)} instead." + ) + + self._chunks = deque[memoryview]() + self._current_chunk = initial_view + self._current_position = 0 + self._eof = False + self._condition = threading.Condition() + + def add_next_chunk(self, chunk: memoryview) -> None: + """ + Adds the next chunk as a read source. + + Chunks can only be added if the stream has not + been finished before. + + The chunk to be added cannot be None. + """ + if not isinstance(chunk, memoryview): + raise TypeError( + "Only memoryviews can be added to the ZeroCopyByteStreams. " + + f"Received {type(chunk)} instead." + ) + with self._condition: + if self._eof: + raise ValueError("Cannot add chunk after ZeroCopyByteStream has been finished") + self._chunks.append(chunk) + self._condition.notify() + + def finish(self) -> None: + """ + Marks the stream as ended. + + Idempotent: can be called multiple times. + """ + with self._condition: + self._eof = True + self._condition.notify() + + @property + def finished(self) -> bool: + """ + Returns whether the stream has been marked as finished and was fully + consumed. If finished() == True, any attempts to read the stream + will raise an `EOFError`. + """ + with self._condition: + # It is finished if: we read all content of the current chunk, + # there are no remaining chunks, and the input has been marked as done + return self._current_chunk is None and len(self._chunks) == 0 and self._eof + + def _try_read_bytes(self, size: int) -> Optional[memoryview]: + """ + Reads up to ``size`` bytes from the current or next available chunk. + Returns a zero-copy memoryview slice (may be shorter than ``size`` + if the current chunk doesn't have enough data), or None on EOF. + + Blocks until at least some data is available or EOF is reached. + + Internal, assumes to be run inside locked self._condition! + """ + # Ensure we have a current chunk + while self._current_chunk is None: + try: + self._current_chunk = self._chunks.popleft() + self._current_position = 0 + except IndexError: + # No chunks available - check for EOF + if self._eof: + return None + # Block until data arrives or EOF is signaled + self._condition.wait() + + remaining = len(self._current_chunk) - self._current_position + to_read = min(remaining, size) + + # Read slice from current chunk (zero-copy) + result = self._current_chunk[self._current_position : self._current_position + to_read] + self._current_position += to_read + + # If entire chunk consumed, clear it for next chunk + assert self._current_position <= len(self._current_chunk), ( + f"Current position {self._current_position} was unexpectedly " + + f"larger than max position {len(self._current_chunk)}" + ) + if self._current_position == len(self._current_chunk): + self._current_chunk = None + self._current_position = 0 + + return result + + def read(self, size: int) -> memoryview: + """ + Reads size bytes. If the read fails because the + stream was exhausted before size bytes could be read + an `EOFError` will be raised. + + It is required that size >= 0. + """ + if size < 0: + raise ValueError( + "ZeroCopyByteStream.read() cannot be called" + + f" with negative size. Received {size}" + ) + + if size == 0: + return memoryview(b"") + + with self._condition: + first = self._try_read_bytes(size) + + # Zero-copy: either EOF or the read fits in a single chunk + if first is None: + raise EOFError( + f"ZeroCopyByteStream.read() tried to read {size} byte(s), " + + "however, the stream was exhausted after reading 0 byte." + ) + elif len(first) == size: + return first + + # Slow path: read crosses chunk boundaries (requires copy) + buf = bytearray(size) + buf[0 : len(first)] = first + bytes_copied = len(first) + + while bytes_copied < size: + chunk = self._try_read_bytes(size - bytes_copied) + if chunk is None: + raise EOFError( + f"ZeroCopyByteStream.read() tried to read {size} byte(s), " + + "however, the stream was exhausted after reading " + + f"{bytes_copied} byte." + ) + buf[bytes_copied : bytes_copied + len(chunk)] = chunk + bytes_copied += len(chunk) + + return memoryview(buf) diff --git a/python/pyspark/ml/connect/__init__.py b/python/pyspark/ml/connect/__init__.py index 98f809c7..651fba4d 100644 --- a/python/pyspark/ml/connect/__init__.py +++ b/python/pyspark/ml/connect/__init__.py @@ -17,11 +17,27 @@ BinaryClassificationEvaluator, MulticlassClassificationEvaluator, ) -from pyspark.ml.connect.pipeline import Pipeline +from pyspark.ml.connect.pipeline import Pipeline, PipelineModel from pyspark.ml.connect.tuning import CrossValidator, CrossValidatorModel from pyspark._pyspark import MLModel +# Base ABCs (vendored from upstream base.py) and the submodules, matching the reference +# `pyspark.ml.connect.__all__`. +from pyspark.ml.connect.base import Estimator, Transformer, Evaluator, Model +from pyspark.ml.connect import feature, evaluation, tuning + __all__ = [ + # Reference `pyspark.ml.connect` public surface. + "Estimator", + "Transformer", + "Evaluator", + "Model", + "feature", + "evaluation", + "Pipeline", + "PipelineModel", + "tuning", + # Fork-additional concrete estimators/evaluators (Rust-backed). "StandardScaler", "MaxAbsScaler", "StringIndexer", @@ -29,7 +45,6 @@ "LogisticRegression", "RegressionEvaluator", "BinaryClassificationEvaluator", - "Pipeline", "MulticlassClassificationEvaluator", "CrossValidator", "CrossValidatorModel", diff --git a/python/pyspark/ml/connect/base.py b/python/pyspark/ml/connect/base.py new file mode 100644 index 00000000..92fe1ae4 --- /dev/null +++ b/python/pyspark/ml/connect/base.py @@ -0,0 +1,360 @@ +# +# 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 abc import ABCMeta, abstractmethod +from typing import ( + Any, + Generic, + List, + Optional, + TypeVar, + Union, + TYPE_CHECKING, + Tuple, + Callable, +) + +import pandas as pd + +from pyspark import since +from pyspark.ml.common import inherit_doc +from pyspark.sql.dataframe import DataFrame +from pyspark.ml.param import Params +from pyspark.ml.param.shared import ( + HasLabelCol, + HasFeaturesCol, + HasPredictionCol, +) + +if TYPE_CHECKING: + from pyspark.ml._typing import ParamMap + +M = TypeVar("M", bound="Transformer") + + +@inherit_doc +class Estimator(Params, Generic[M], metaclass=ABCMeta): + """ + Abstract class for estimators that fit models to data. + + .. versionadded:: 3.5.0 + + .. deprecated:: 4.0.0 + """ + + @abstractmethod + def _fit(self, dataset: Union[DataFrame, pd.DataFrame]) -> M: + """ + Fits a model to the input dataset. This is called by the default implementation of fit. + + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + input dataset + + Returns + ------- + :class:`Transformer` + fitted model + """ + raise NotImplementedError() + + def fit( + self, + dataset: Union[DataFrame, pd.DataFrame], + params: Optional["ParamMap"] = None, + ) -> Union[M, List[M]]: + """ + Fits a model to the input dataset with optional parameters. + + .. versionadded:: 3.5.0 + + .. deprecated:: 4.0.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` or py:class:`pandas.DataFrame` + input dataset, it can be either pandas dataframe or spark dataframe. + params : a dict of param values, optional + an optional param map that overrides embedded params. + + Returns + ------- + :py:class:`Transformer` + fitted model + """ + if params is None: + params = dict() + + if isinstance(params, dict): + if params: + return self.copy(params)._fit(dataset) + else: + return self._fit(dataset) + else: + raise TypeError( + "Params must be either a param map or a list/tuple of param maps, " + "but got %s." % type(params) + ) + + +_SPARKML_TRANSFORMER_TMP_OUTPUT_COLNAME = "_sparkML_transformer_tmp_output" + + +@inherit_doc +class Transformer(Params, metaclass=ABCMeta): + """ + Abstract class for transformers that transform one dataset into another. + + .. versionadded:: 3.5.0 + + .. deprecated:: 4.0.0 + """ + + def _input_columns(self) -> List[str]: + """ + Return a list of input column names which are used as inputs of transformation. + """ + raise NotImplementedError() + + def _output_columns(self) -> List[Tuple[str, str]]: + """ + Return a list of output transformed columns, each elements in the list + is a tuple of (column_name, column_spark_type) + """ + raise NotImplementedError() + + def _get_transform_fn(self) -> Callable[..., Any]: + """ + Return a transformation function that accepts one or more `pd.Series` instances as inputs + and returns transformed result as an instance of `pd.Series` or `pd.DataFrame`. + If there's only one output column, the transformed result must be an + instance of `pd.Series`, if there are multiple output columns, the transformed result + must be an instance of `pd.DataFrame` with column names matching output schema + returned by `_output_columns` interface. + """ + raise NotImplementedError() + + def transform( + self, dataset: Union[DataFrame, pd.DataFrame], params: Optional["ParamMap"] = None + ) -> Union[DataFrame, pd.DataFrame]: + """ + Transforms the input dataset. + The dataset can be either pandas dataframe or spark dataframe, + if it is a spark DataFrame, the result of transformation is a new spark DataFrame + that contains all existing columns and output columns with names, + If it is a pandas DataFrame, the result of transformation is a shallow copy + of the input pandas dataframe with output columns with names. + + Note: Transformers does not allow output column having the same name with + existing columns. + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` or py:class:`pandas.DataFrame` + input dataset. + + params : dict, optional + an optional param map that overrides embedded params. + + Returns + ------- + :py:class:`pyspark.sql.DataFrame` or py:class:`pandas.DataFrame` + transformed dataset, the type of output dataframe is consistent with + input dataframe. + """ + if params is None: + params = dict() + if isinstance(params, dict): + if params: + return self.copy(params)._transform(dataset) + else: + return self._transform(dataset) + + def _transform(self, dataset: Union[DataFrame, pd.DataFrame]) -> Union[DataFrame, pd.DataFrame]: + from pyspark.ml.connect.util import transform_dataframe_column + + input_cols = self._input_columns() + transform_fn = self._get_transform_fn() + output_cols = self._output_columns() + + existing_cols = list(dataset.columns) + for col_name, _ in output_cols: + if col_name in existing_cols: + raise ValueError( + "Transformers does not allow output column having the same name with " + "existing columns." + ) + + return transform_dataframe_column( + dataset, + input_cols=input_cols, + transform_fn=transform_fn, + output_cols=output_cols, + ) + + +@inherit_doc +class Evaluator(Params, metaclass=ABCMeta): + """ + Base class for evaluators that compute metrics from predictions. + + .. versionadded:: 3.5.0 + + .. deprecated:: 4.0.0 + """ + + @abstractmethod + def _evaluate(self, dataset: Union["DataFrame", "pd.DataFrame"]) -> float: + """ + Evaluates the output. + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + a dataset that contains labels/observations and predictions + + Returns + ------- + float + metric + """ + raise NotImplementedError() + + def evaluate(self, dataset: DataFrame, params: Optional["ParamMap"] = None) -> float: + """ + Evaluates the output with optional parameters. + + .. versionadded:: 3.5.0 + + .. deprecated:: 4.0.0 + + Parameters + ---------- + dataset : :py:class:`pyspark.sql.DataFrame` + a dataset that contains labels/observations and predictions + params : dict, optional + an optional param map that overrides embedded params + + Returns + ------- + float + metric + """ + if params is None: + params = dict() + if isinstance(params, dict): + if params: + return self.copy(params)._evaluate(dataset) + else: + return self._evaluate(dataset) + else: + raise TypeError("Params must be a param map but got %s." % type(params)) + + @since("1.5.0") + def isLargerBetter(self) -> bool: + """ + Indicates whether the metric returned by :py:meth:`evaluate` should be maximized + (True, default) or minimized (False). + A given evaluator may support multiple metrics which may be maximized or minimized. + """ + raise NotImplementedError() + + +@inherit_doc +class Model(Transformer, metaclass=ABCMeta): + """ + Abstract class for models that are fitted by estimators. + + .. versionadded:: 3.5.0 + + .. deprecated:: 4.0.0 + """ + + pass + + +@inherit_doc +class _PredictorParams(HasLabelCol, HasFeaturesCol, HasPredictionCol): + """ + Params for :py:class:`Predictor` and :py:class:`PredictorModel`. + + .. versionadded:: 3.5.0 + + .. deprecated:: 4.0.0 + """ + + pass + + +@inherit_doc +class Predictor(Estimator[M], _PredictorParams, metaclass=ABCMeta): + """ + Estimator for prediction tasks (regression and classification). + """ + + @since("3.5.0") + def setLabelCol(self, value: str) -> "Predictor": + """ + Sets the value of :py:attr:`labelCol`. + """ + return self._set(labelCol=value) + + @since("3.5.0") + def setFeaturesCol(self, value: str) -> "Predictor": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("3.5.0") + def setPredictionCol(self, value: str) -> "Predictor": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + +@inherit_doc +class PredictionModel(Model, _PredictorParams, metaclass=ABCMeta): + """ + Model for prediction tasks (regression and classification). + """ + + @since("3.5.0") + def setFeaturesCol(self, value: str) -> "PredictionModel": + """ + Sets the value of :py:attr:`featuresCol`. + """ + return self._set(featuresCol=value) + + @since("3.5.0") + def setPredictionCol(self, value: str) -> "PredictionModel": + """ + Sets the value of :py:attr:`predictionCol`. + """ + return self._set(predictionCol=value) + + @property + @abstractmethod + @since("3.5.0") + def numFeatures(self) -> int: + """ + Returns the number of features the model was trained on. If unknown, returns -1 + """ + raise NotImplementedError() diff --git a/python/pyspark/ml/deepspeed/deepspeed_distributor.py b/python/pyspark/ml/deepspeed/deepspeed_distributor.py index 3fd1d3bb..162cd2ad 100644 --- a/python/pyspark/ml/deepspeed/deepspeed_distributor.py +++ b/python/pyspark/ml/deepspeed/deepspeed_distributor.py @@ -100,7 +100,9 @@ def __init__( @staticmethod def _get_deepspeed_config_path(deepspeed_config: Union[str, Dict[str, Any]]) -> str: if isinstance(deepspeed_config, dict): - with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".json") as file: + with tempfile.NamedTemporaryFile( + mode="w+", delete=False, suffix=".json", encoding="utf-8" + ) as file: json.dump(deepspeed_config, file) return file.name deepspeed_config_path = deepspeed_config diff --git a/python/pyspark/ml/dl_util.py b/python/pyspark/ml/dl_util.py index 26538c32..89cf2ea8 100644 --- a/python/pyspark/ml/dl_util.py +++ b/python/pyspark/ml/dl_util.py @@ -119,7 +119,7 @@ def create_fn_run_script( with open("{fn_output_path}", "wb") as f: cloudpickle.dump(output, f) """) - with open(script_path, "w") as f: + with open(script_path, "w", encoding="utf-8") as f: if prefix_code != "": f.write(prefix_code) f.write(code_snippet) diff --git a/python/pyspark/ml/functions.py b/python/pyspark/ml/functions.py new file mode 100644 index 00000000..b3cd3978 --- /dev/null +++ b/python/pyspark/ml/functions.py @@ -0,0 +1,872 @@ +# +# 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 __future__ import annotations + +import inspect +import uuid +from typing import Any, Callable, Iterator, List, Mapping, TYPE_CHECKING, Tuple, Union, Optional + +import numpy as np + +try: + import pandas as pd +except ImportError: + pass # Let it throw a better error message later when the API is invoked. + +from pyspark.sql.functions import pandas_udf +from pyspark.sql.column import Column +from pyspark.sql.types import ( + ArrayType, + ByteType, + DataType, + DoubleType, + FloatType, + IntegerType, + LongType, + ShortType, + StringType, + StructType, +) +from pyspark.ml.util import try_remote_functions + +if TYPE_CHECKING: + from pyspark.sql._typing import UserDefinedFunctionLike + +supported_scalar_types = ( + ByteType, + ShortType, + IntegerType, + LongType, + FloatType, + DoubleType, + StringType, +) + +# Callable type for end user predict functions that take a variable number of ndarrays as +# input and returns one of the following as output: +# - single ndarray (single output) +# - dictionary of named ndarrays (multiple outputs represented in columnar form) +# - list of dictionaries of named ndarrays (multiple outputs represented in row form) +PredictBatchFunction = Callable[ + [np.ndarray], Union[np.ndarray, Mapping[str, np.ndarray], List[Mapping[str, np.dtype]]] +] + + +@try_remote_functions +def vector_to_array(col: Column, dtype: str = "float64") -> Column: + """ + Converts a column of MLlib sparse/dense vectors into a column of dense arrays. + + .. versionadded:: 3.0.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + col : :py:class:`pyspark.sql.Column` or str + Input column + dtype : str, optional + The data type of the output array. Valid values: "float64" or "float32". + + Returns + ------- + :py:class:`pyspark.sql.Column` + The converted column of dense arrays. + + Examples + -------- + >>> from pyspark.ml.linalg import Vectors + >>> from pyspark.ml.functions import vector_to_array + >>> from pyspark.mllib.linalg import Vectors as OldVectors + >>> df = spark.createDataFrame([ + ... (Vectors.dense(1.0, 2.0, 3.0), OldVectors.dense(10.0, 20.0, 30.0)), + ... (Vectors.sparse(3, [(0, 2.0), (2, 3.0)]), + ... OldVectors.sparse(3, [(0, 20.0), (2, 30.0)]))], + ... ["vec", "oldVec"]) + >>> df1 = df.select(vector_to_array("vec").alias("vec"), + ... vector_to_array("oldVec").alias("oldVec")) + >>> df1.collect() + [Row(vec=[1.0, 2.0, 3.0], oldVec=[10.0, 20.0, 30.0]), + Row(vec=[2.0, 0.0, 3.0], oldVec=[20.0, 0.0, 30.0])] + >>> df2 = df.select(vector_to_array("vec", "float32").alias("vec"), + ... vector_to_array("oldVec", "float32").alias("oldVec")) + >>> df2.collect() + [Row(vec=[1.0, 2.0, 3.0], oldVec=[10.0, 20.0, 30.0]), + Row(vec=[2.0, 0.0, 3.0], oldVec=[20.0, 0.0, 30.0])] + >>> df1.schema.fields + [StructField('vec', ArrayType(DoubleType(), False), False), + StructField('oldVec', ArrayType(DoubleType(), False), False)] + >>> df2.schema.fields + [StructField('vec', ArrayType(FloatType(), False), False), + StructField('oldVec', ArrayType(FloatType(), False), False)] + """ + from pyspark.core.context import SparkContext + from pyspark.sql.classic.column import Column, _to_java_column + + sc = SparkContext._active_spark_context + assert sc is not None and sc._jvm is not None + return Column( + getattr(sc._jvm, "org.apache.spark.ml.functions").vector_to_array( + _to_java_column(col), dtype + ) + ) + + +@try_remote_functions +def array_to_vector(col: Column) -> Column: + """ + Converts a column of array of numeric type into a column of pyspark.ml.linalg.DenseVector + instances + + .. versionadded:: 3.1.0 + + .. versionchanged:: 3.5.0 + Supports Spark Connect. + + Parameters + ---------- + col : :py:class:`pyspark.sql.Column` or str + Input column + + Returns + ------- + :py:class:`pyspark.sql.Column` + The converted column of dense vectors. + + Examples + -------- + >>> from pyspark.ml.functions import array_to_vector + >>> df1 = spark.createDataFrame([([1.5, 2.5],),], schema='v1 array') + >>> df1.select(array_to_vector('v1').alias('vec1')).collect() + [Row(vec1=DenseVector([1.5, 2.5]))] + >>> df2 = spark.createDataFrame([([1.5, 3.5],),], schema='v1 array') + >>> df2.select(array_to_vector('v1').alias('vec1')).collect() + [Row(vec1=DenseVector([1.5, 3.5]))] + >>> df3 = spark.createDataFrame([([1, 3],),], schema='v1 array') + >>> df3.select(array_to_vector('v1').alias('vec1')).collect() + [Row(vec1=DenseVector([1.0, 3.0]))] + """ + from pyspark.core.context import SparkContext + from pyspark.sql.classic.column import Column, _to_java_column + + sc = SparkContext._active_spark_context + assert sc is not None and sc._jvm is not None + return Column( + getattr(sc._jvm, "org.apache.spark.ml.functions").array_to_vector(_to_java_column(col)) + ) + + +def _batched( + data: Union[pd.Series, pd.DataFrame, Tuple[pd.Series]], batch_size: int +) -> Iterator[pd.DataFrame]: + """Generator that splits a pandas dataframe/series into batches.""" + if isinstance(data, pd.DataFrame): + df = data + elif isinstance(data, pd.Series): + df = pd.concat((data,), axis=1) + else: # isinstance(data, Tuple[pd.Series]): + df = pd.concat(data, axis=1) + + index = 0 + data_size = len(df) + while index < data_size: + yield df.iloc[index : index + batch_size] + index += batch_size + + +def _is_tensor_col(data: Union[pd.Series, pd.DataFrame]) -> bool: + if isinstance(data, pd.Series): + return data.dtype == np.object_ and isinstance(data.iloc[0], (np.ndarray, list)) + elif isinstance(data, pd.DataFrame): + return any(data.dtypes == np.object_) and any( + [isinstance(d, (np.ndarray, list)) for d in data.iloc[0]] + ) + else: + raise ValueError( + "Unexpected data type: {}, expected pd.Series or pd.DataFrame.".format(type(data)) + ) + + +def _has_tensor_cols(data: Union[pd.Series, pd.DataFrame, Tuple[pd.Series]]) -> bool: + """Check if input Series/DataFrame/Tuple contains any tensor-valued columns.""" + if isinstance(data, (pd.Series, pd.DataFrame)): + return _is_tensor_col(data) + else: # isinstance(data, Tuple): + return any(_is_tensor_col(elem) for elem in data) + + +def _validate_and_transform_multiple_inputs( + batch: pd.DataFrame, input_shapes: List[Optional[List[int]]], num_input_cols: int +) -> List[np.ndarray]: + multi_inputs = [batch[col].to_numpy() for col in batch.columns] + if input_shapes: + if len(input_shapes) == num_input_cols: + multi_inputs = [ + ( + np.vstack(v).reshape([-1] + input_shapes[i]) # type: ignore + if input_shapes[i] + else v + ) + for i, v in enumerate(multi_inputs) + ] + if not all([len(x) == len(batch) for x in multi_inputs]): + raise ValueError("Input data does not match expected shape.") + else: + raise ValueError("input_tensor_shapes must match columns") + + return multi_inputs + + +def _validate_and_transform_single_input( + batch: pd.DataFrame, + input_shapes: List[List[int] | None], + has_tensors: bool, + has_tuple: bool, +) -> np.ndarray: + # multiple input columns for single expected input + if has_tensors: + # tensor columns + if len(batch.columns) == 1: + # one tensor column and one expected input, vstack rows + single_input = np.vstack(batch.iloc[:, 0]) # type: ignore[call-overload] + else: + raise ValueError( + "Multiple input columns found, but model expected a single " + "input, use `array` to combine columns into tensors." + ) + else: + # scalar columns + if len(batch.columns) == 1: + # single scalar column, remove extra dim + np_batch = batch.to_numpy() + single_input = np.squeeze(np_batch, -1) if len(np_batch.shape) > 1 else np_batch + if input_shapes and input_shapes[0] not in [None, [], [1]]: + raise ValueError("Invalid input_tensor_shape for scalar column.") + elif not has_tuple: + # columns grouped via `array`, convert to single tensor + single_input = batch.to_numpy() + if input_shapes and input_shapes[0] != [len(batch.columns)]: + raise ValueError("Input data does not match expected shape.") + else: + raise ValueError( + "Multiple input columns found, but model expected a single " + "input, use `array` to combine columns into tensors." + ) + + # if input_tensor_shapes provided, try to reshape input + if input_shapes: + if len(input_shapes) == 1: + single_input = single_input.reshape([-1] + input_shapes[0]) # type: ignore + if len(single_input) != len(batch): + raise ValueError("Input data does not match expected shape.") + else: + raise ValueError("Multiple input_tensor_shapes found, but model expected one input") + + return single_input + + +def _validate_and_transform_prediction_result( + preds: np.ndarray | Mapping[str, np.ndarray] | List[Mapping[str, Any]], + num_input_rows: int, + return_type: DataType, +) -> pd.DataFrame | pd.Series: + """Validate numpy-based model predictions against the expected pandas_udf return_type and + transforms the predictions into an equivalent pandas DataFrame or Series.""" + if isinstance(return_type, StructType): + struct_rtype: StructType = return_type + fieldNames = struct_rtype.names + if isinstance(preds, dict): + # dictionary of columns + predNames = list(preds.keys()) + for field in struct_rtype.fields: + if isinstance(field.dataType, ArrayType): + if len(preds[field.name].shape) == 2: + preds[field.name] = list(preds[field.name]) + else: + raise ValueError( + "Prediction results for ArrayType must be two-dimensional." + ) + elif isinstance(field.dataType, supported_scalar_types): + if len(preds[field.name].shape) != 1: + raise ValueError( + "Prediction results for scalar types must be one-dimensional." + ) + else: + raise ValueError("Unsupported field type in return struct type.") + + if len(preds[field.name]) != num_input_rows: + raise ValueError("Prediction results must have same length as input data") + + elif isinstance(preds, list) and isinstance(preds[0], dict): + # rows of dictionaries + predNames = list(preds[0].keys()) + if len(preds) != num_input_rows: + raise ValueError("Prediction results must have same length as input data.") + for field in struct_rtype.fields: + if isinstance(field.dataType, ArrayType): + if len(preds[0][field.name].shape) != 1: + raise ValueError( + "Prediction results for ArrayType must be one-dimensional." + ) + elif isinstance(field.dataType, supported_scalar_types): + if not np.isscalar(preds[0][field.name]): + raise ValueError("Invalid scalar prediction result.") + else: + raise ValueError("Unsupported field type in return struct type.") + else: + raise ValueError( + "Prediction results for StructType must be a dictionary or " + "a list of dictionary, got: {}".format(type(preds)) + ) + + # check column names + if set(predNames) != set(fieldNames): + raise ValueError( + "Prediction result columns did not match expected return_type " + "columns: expected {}, got: {}".format(fieldNames, predNames) + ) + + return pd.DataFrame(preds) + elif isinstance(return_type, ArrayType): + if isinstance(preds, np.ndarray): + if len(preds) != num_input_rows: + raise ValueError("Prediction results must have same length as input data.") + if len(preds.shape) != 2: + raise ValueError("Prediction results for ArrayType must be two-dimensional.") + else: + raise ValueError("Prediction results for ArrayType must be an ndarray.") + + return pd.Series(list(preds)) + elif isinstance(return_type, supported_scalar_types): + preds_array: np.ndarray = preds # type: ignore + if len(preds_array) != num_input_rows: + raise ValueError("Prediction results must have same length as input data.") + if not ( + (len(preds_array.shape) == 2 and preds_array.shape[1] == 1) + or len(preds_array.shape) == 1 + ): + raise ValueError("Invalid shape for scalar prediction result.") + + output = np.squeeze(preds_array, -1) if len(preds_array.shape) > 1 else preds_array + return pd.Series(output).astype(output.dtype) + else: + raise ValueError("Unsupported return type") + + +def predict_batch_udf( + make_predict_fn: Callable[ + [], + PredictBatchFunction, + ], + *, + return_type: DataType, + batch_size: int, + input_tensor_shapes: Optional[Union[List[Optional[List[int]]], Mapping[int, List[int]]]] = None, +) -> UserDefinedFunctionLike: + """Given a function which loads a model and returns a `predict` function for inference over a + batch of numpy inputs, returns a Pandas UDF wrapper for inference over a Spark DataFrame. + + The returned Pandas UDF does the following on each DataFrame partition: + + * calls the `make_predict_fn` to load the model and cache its `predict` function. + * batches the input records as numpy arrays and invokes `predict` on each batch. + + Note: this assumes that the `make_predict_fn` encapsulates all of the necessary dependencies for + running the model, or the Spark executor environment already satisfies all runtime requirements. + + For the conversion of the Spark DataFrame to numpy arrays, there is a one-to-one mapping between + the input arguments of the `predict` function (returned by the `make_predict_fn`) and the input + columns sent to the Pandas UDF (returned by the `predict_batch_udf`) at runtime. Each input + column will be converted as follows: + + * scalar column -> 1-dim np.ndarray + * tensor column + tensor shape -> N-dim np.ndarray + + Note that any tensor columns in the Spark DataFrame must be represented as a flattened + one-dimensional array, and multiple scalar columns can be combined into a single tensor column + using the standard :py:func:`pyspark.sql.functions.array()` function. + + .. versionadded:: 3.4.0 + + Parameters + ---------- + make_predict_fn : callable + Function which is responsible for loading a model and returning a + :py:class:`PredictBatchFunction` which takes one or more numpy arrays as input and returns + one of the following: + + * a numpy array (for a single output) + * a dictionary of named numpy arrays (for multiple outputs) + * a row-oriented list of dictionaries (for multiple outputs). + + For a dictionary of named numpy arrays, the arrays can only be one or two dimensional, since + higher dimensional arrays are not supported. For a row-oriented list of dictionaries, each + element in the dictionary must be either a scalar or one-dimensional array. + return_type : :py:class:`pyspark.sql.types.DataType` or str. + Spark SQL datatype for the expected output: + + * Scalar (e.g. IntegerType, FloatType) --> 1-dim numpy array. + * ArrayType --> 2-dim numpy array. + * StructType --> dict with keys matching struct fields. + * StructType --> list of dict with keys matching struct fields, for models like the + `Huggingface pipeline for sentiment analysis + `_. + + batch_size : int + Batch size to use for inference. This is typically a limitation of the model + and/or available hardware resources and is usually smaller than the Spark partition size. + input_tensor_shapes : list, dict, optional. + A list of ints or a dictionary of ints (key) and list of ints (value). + Input tensor shapes for models with tensor inputs. This can be a list of shapes, + where each shape is a list of integers or None (for scalar inputs). Alternatively, this + can be represented by a "sparse" dictionary, where the keys are the integer indices of the + inputs, and the values are the shapes. Each tensor input value in the Spark DataFrame must + be represented as a single column containing a flattened 1-D array. The provided + `input_tensor_shapes` will be used to reshape the flattened array into the expected tensor + shape. For the list form, the order of the tensor shapes must match the order of the + selected DataFrame columns. The batch dimension (typically -1 or None in the first + dimension) should not be included, since it will be determined by the batch_size argument. + Tabular datasets with scalar-valued columns should not provide this argument. + + Returns + ------- + :py:class:`UserDefinedFunctionLike` + A Pandas UDF for model inference on a Spark DataFrame. + + Examples + -------- + For a pre-trained TensorFlow MNIST model with two-dimensional input images represented as a + flattened tensor value stored in a single Spark DataFrame column of type `array`. + + .. code-block:: python + + from pyspark.ml.functions import predict_batch_udf + + def make_mnist_fn(): + # load/init happens once per python worker + import tensorflow as tf + model = tf.keras.models.load_model('/path/to/mnist_model') + + # predict on batches of tasks/partitions, using cached model + def predict(inputs: np.ndarray) -> np.ndarray: + # inputs.shape = [batch_size, 784], see input_tensor_shapes + # outputs.shape = [batch_size, 10], see return_type + return model.predict(inputs) + + return predict + + mnist_udf = predict_batch_udf(make_mnist_fn, + return_type=ArrayType(FloatType()), + batch_size=100, + input_tensor_shapes=[[784]]) + + df = spark.read.parquet("/path/to/mnist_data") + df.show(5) + # +--------------------+ + # | data| + # +--------------------+ + # |[0.0, 0.0, 0.0, 0...| + # |[0.0, 0.0, 0.0, 0...| + # |[0.0, 0.0, 0.0, 0...| + # |[0.0, 0.0, 0.0, 0...| + # |[0.0, 0.0, 0.0, 0...| + # +--------------------+ + + df.withColumn("preds", mnist_udf("data")).show(5) + # +--------------------+--------------------+ + # | data| preds| + # +--------------------+--------------------+ + # |[0.0, 0.0, 0.0, 0...|[-13.511008, 8.84...| + # |[0.0, 0.0, 0.0, 0...|[-5.3957458, -2.2...| + # |[0.0, 0.0, 0.0, 0...|[-7.2014456, -8.8...| + # |[0.0, 0.0, 0.0, 0...|[-19.466187, -13....| + # |[0.0, 0.0, 0.0, 0...|[-5.7757926, -7.8...| + # +--------------------+--------------------+ + + To demonstrate usage with different combinations of input and output types, the following + examples just use simple mathematical transforms as the models. + + * Single scalar column + Input DataFrame has a single scalar column, which will be passed to the `predict` + function as a 1-D numpy array. + + >>> import numpy as np + >>> import pandas as pd + >>> from pyspark.ml.functions import predict_batch_udf + >>> from pyspark.sql.types import FloatType + >>> + >>> df = spark.createDataFrame(pd.DataFrame(np.arange(100))) + >>> df.show(5) + +---+ + | 0| + +---+ + | 0| + | 1| + | 2| + | 3| + | 4| + +---+ + only showing top 5 rows + + >>> def make_times_two_fn(): + ... def predict(inputs: np.ndarray) -> np.ndarray: + ... # inputs.shape = [batch_size] + ... # outputs.shape = [batch_size] + ... return inputs * 2 + ... return predict + ... + >>> times_two_udf = predict_batch_udf(make_times_two_fn, + ... return_type=FloatType(), + ... batch_size=10) + >>> df = spark.createDataFrame(pd.DataFrame(np.arange(100))) + >>> df.withColumn("x2", times_two_udf("0")).show(5) + +---+---+ + | 0| x2| + +---+---+ + | 0|0.0| + | 1|2.0| + | 2|4.0| + | 3|6.0| + | 4|8.0| + +---+---+ + only showing top 5 rows + + * Multiple scalar columns + Input DataFrame has multiple columns of scalar values. If the user-provided `predict` + function expects a single input, then the user must combine the multiple columns into a + single tensor using `pyspark.sql.functions.array`. + + >>> import numpy as np + >>> import pandas as pd + >>> from pyspark.ml.functions import predict_batch_udf + >>> from pyspark.sql.functions import array + >>> + >>> data = np.arange(0, 1000, dtype=np.float64).reshape(-1, 4) + >>> pdf = pd.DataFrame(data, columns=['a','b','c','d']) + >>> df = spark.createDataFrame(pdf) + >>> df.show(5) + +----+----+----+----+ + | a| b| c| d| + +----+----+----+----+ + | 0.0| 1.0| 2.0| 3.0| + | 4.0| 5.0| 6.0| 7.0| + | 8.0| 9.0|10.0|11.0| + |12.0|13.0|14.0|15.0| + |16.0|17.0|18.0|19.0| + +----+----+----+----+ + only showing top 5 rows + + >>> def make_sum_fn(): + ... def predict(inputs: np.ndarray) -> np.ndarray: + ... # inputs.shape = [batch_size, 4] + ... # outputs.shape = [batch_size] + ... return np.sum(inputs, axis=1) + ... return predict + ... + >>> sum_udf = predict_batch_udf(make_sum_fn, + ... return_type=FloatType(), + ... batch_size=10, + ... input_tensor_shapes=[[4]]) + >>> df.withColumn("sum", sum_udf(array("a", "b", "c", "d"))).show(5) + +----+----+----+----+----+ + | a| b| c| d| sum| + +----+----+----+----+----+ + | 0.0| 1.0| 2.0| 3.0| 6.0| + | 4.0| 5.0| 6.0| 7.0|22.0| + | 8.0| 9.0|10.0|11.0|38.0| + |12.0|13.0|14.0|15.0|54.0| + |16.0|17.0|18.0|19.0|70.0| + +----+----+----+----+----+ + only showing top 5 rows + + If the `predict` function expects multiple inputs, then the number of selected input columns + must match the number of expected inputs. + + >>> def make_sum_fn(): + ... def predict(x1: np.ndarray, + ... x2: np.ndarray, + ... x3: np.ndarray, + ... x4: np.ndarray) -> np.ndarray: + ... # xN.shape = [batch_size] + ... # outputs.shape = [batch_size] + ... return x1 + x2 + x3 + x4 + ... return predict + ... + >>> sum_udf = predict_batch_udf(make_sum_fn, + ... return_type=FloatType(), + ... batch_size=10) + >>> df.withColumn("sum", sum_udf("a", "b", "c", "d")).show(5) + +----+----+----+----+----+ + | a| b| c| d| sum| + +----+----+----+----+----+ + | 0.0| 1.0| 2.0| 3.0| 6.0| + | 4.0| 5.0| 6.0| 7.0|22.0| + | 8.0| 9.0|10.0|11.0|38.0| + |12.0|13.0|14.0|15.0|54.0| + |16.0|17.0|18.0|19.0|70.0| + +----+----+----+----+----+ + only showing top 5 rows + + * Multiple tensor columns + Input DataFrame has multiple columns, where each column is a tensor. The number of columns + should match the number of expected inputs for the user-provided `predict` function. + + >>> import numpy as np + >>> import pandas as pd + >>> from pyspark.ml.functions import predict_batch_udf + >>> from pyspark.sql.types import ArrayType, FloatType, StructType, StructField + >>> from typing import Mapping + >>> + >>> data = np.arange(0, 1000, dtype=np.float64).reshape(-1, 4) + >>> pdf = pd.DataFrame(data, columns=['a','b','c','d']) + >>> pdf_tensor = pd.DataFrame() + >>> pdf_tensor['t1'] = pdf.values.tolist() + >>> pdf_tensor['t2'] = pdf.drop(columns='d').values.tolist() + >>> df = spark.createDataFrame(pdf_tensor) + >>> df.show(5) + +--------------------+------------------+ + | t1| t2| + +--------------------+------------------+ + |[0.0, 1.0, 2.0, 3.0]| [0.0, 1.0, 2.0]| + |[4.0, 5.0, 6.0, 7.0]| [4.0, 5.0, 6.0]| + |[8.0, 9.0, 10.0, ...| [8.0, 9.0, 10.0]| + |[12.0, 13.0, 14.0...|[12.0, 13.0, 14.0]| + |[16.0, 17.0, 18.0...|[16.0, 17.0, 18.0]| + +--------------------+------------------+ + only showing top 5 rows + + >>> def make_multi_sum_fn(): + ... def predict(x1: np.ndarray, x2: np.ndarray) -> np.ndarray: + ... # x1.shape = [batch_size, 4] + ... # x2.shape = [batch_size, 3] + ... # outputs.shape = [batch_size] + ... return np.sum(x1, axis=1) + np.sum(x2, axis=1) + ... return predict + ... + >>> multi_sum_udf = predict_batch_udf( + ... make_multi_sum_fn, + ... return_type=FloatType(), + ... batch_size=5, + ... input_tensor_shapes=[[4], [3]], + ... ) + >>> df.withColumn("sum", multi_sum_udf("t1", "t2")).show(5) + +--------------------+------------------+-----+ + | t1| t2| sum| + +--------------------+------------------+-----+ + |[0.0, 1.0, 2.0, 3.0]| [0.0, 1.0, 2.0]| 9.0| + |[4.0, 5.0, 6.0, 7.0]| [4.0, 5.0, 6.0]| 37.0| + |[8.0, 9.0, 10.0, ...| [8.0, 9.0, 10.0]| 65.0| + |[12.0, 13.0, 14.0...|[12.0, 13.0, 14.0]| 93.0| + |[16.0, 17.0, 18.0...|[16.0, 17.0, 18.0]|121.0| + +--------------------+------------------+-----+ + only showing top 5 rows + + * Multiple outputs + Some models can provide multiple outputs. These can be returned as a dictionary of named + values, which can be represented in either columnar or row-based formats. + + >>> def make_multi_sum_fn(): + ... def predict_columnar(x1: np.ndarray, x2: np.ndarray) -> Mapping[str, np.ndarray]: + ... # x1.shape = [batch_size, 4] + ... # x2.shape = [batch_size, 3] + ... return { + ... "sum1": np.sum(x1, axis=1), + ... "sum2": np.sum(x2, axis=1) + ... } + ... return predict_columnar + ... + >>> multi_sum_udf = predict_batch_udf( + ... make_multi_sum_fn, + ... return_type=StructType([ + ... StructField("sum1", FloatType(), True), + ... StructField("sum2", FloatType(), True) + ... ]), + ... batch_size=5, + ... input_tensor_shapes=[[4], [3]], + ... ) + >>> df.withColumn("preds", multi_sum_udf("t1", "t2")).select("t1", "t2", "preds.*").show(5) + +--------------------+------------------+----+----+ + | t1| t2|sum1|sum2| + +--------------------+------------------+----+----+ + |[0.0, 1.0, 2.0, 3.0]| [0.0, 1.0, 2.0]| 6.0| 3.0| + |[4.0, 5.0, 6.0, 7.0]| [4.0, 5.0, 6.0]|22.0|15.0| + |[8.0, 9.0, 10.0, ...| [8.0, 9.0, 10.0]|38.0|27.0| + |[12.0, 13.0, 14.0...|[12.0, 13.0, 14.0]|54.0|39.0| + |[16.0, 17.0, 18.0...|[16.0, 17.0, 18.0]|70.0|51.0| + +--------------------+------------------+----+----+ + only showing top 5 rows + + >>> def make_multi_sum_fn(): + ... def predict_row(x1: np.ndarray, x2: np.ndarray) -> list[Mapping[str, float]]: + ... # x1.shape = [batch_size, 4] + ... # x2.shape = [batch_size, 3] + ... return [{'sum1': np.sum(x1[i]), 'sum2': np.sum(x2[i])} for i in range(len(x1))] + ... return predict_row + ... + >>> multi_sum_udf = predict_batch_udf( + ... make_multi_sum_fn, + ... return_type=StructType([ + ... StructField("sum1", FloatType(), True), + ... StructField("sum2", FloatType(), True) + ... ]), + ... batch_size=5, + ... input_tensor_shapes=[[4], [3]], + ... ) + >>> df.withColumn("sum", multi_sum_udf("t1", "t2")).select("t1", "t2", "sum.*").show(5) + +--------------------+------------------+----+----+ + | t1| t2|sum1|sum2| + +--------------------+------------------+----+----+ + |[0.0, 1.0, 2.0, 3.0]| [0.0, 1.0, 2.0]| 6.0| 3.0| + |[4.0, 5.0, 6.0, 7.0]| [4.0, 5.0, 6.0]|22.0|15.0| + |[8.0, 9.0, 10.0, ...| [8.0, 9.0, 10.0]|38.0|27.0| + |[12.0, 13.0, 14.0...|[12.0, 13.0, 14.0]|54.0|39.0| + |[16.0, 17.0, 18.0...|[16.0, 17.0, 18.0]|70.0|51.0| + +--------------------+------------------+----+----+ + only showing top 5 rows + + Note that the multiple outputs can be arrays as well. + + >>> def make_multi_times_two_fn(): + ... def predict(x1: np.ndarray, x2: np.ndarray) -> Mapping[str, np.ndarray]: + ... # x1.shape = [batch_size, 4] + ... # x2.shape = [batch_size, 3] + ... return {"t1x2": x1 * 2, "t2x2": x2 * 2} + ... return predict + ... + >>> multi_times_two_udf = predict_batch_udf( + ... make_multi_times_two_fn, + ... return_type=StructType([ + ... StructField("t1x2", ArrayType(FloatType()), True), + ... StructField("t2x2", ArrayType(FloatType()), True) + ... ]), + ... batch_size=5, + ... input_tensor_shapes=[[4], [3]], + ... ) + >>> df.withColumn("x2", multi_times_two_udf("t1", "t2")).select("t1", "t2", "x2.*").show(5) + +--------------------+------------------+--------------------+------------------+ + | t1| t2| t1x2| t2x2| + +--------------------+------------------+--------------------+------------------+ + |[0.0, 1.0, 2.0, 3.0]| [0.0, 1.0, 2.0]|[0.0, 2.0, 4.0, 6.0]| [0.0, 2.0, 4.0]| + |[4.0, 5.0, 6.0, 7.0]| [4.0, 5.0, 6.0]|[8.0, 10.0, 12.0,...| [8.0, 10.0, 12.0]| + |[8.0, 9.0, 10.0, ...| [8.0, 9.0, 10.0]|[16.0, 18.0, 20.0...|[16.0, 18.0, 20.0]| + |[12.0, 13.0, 14.0...|[12.0, 13.0, 14.0]|[24.0, 26.0, 28.0...|[24.0, 26.0, 28.0]| + |[16.0, 17.0, 18.0...|[16.0, 17.0, 18.0]|[32.0, 34.0, 36.0...|[32.0, 34.0, 36.0]| + +--------------------+------------------+--------------------+------------------+ + only showing top 5 rows + """ + # generate a new uuid each time this is invoked on the driver to invalidate executor-side cache. + model_uuid = uuid.uuid4() + + def predict(data: Iterator[Union[pd.Series, pd.DataFrame]]) -> Iterator[pd.DataFrame]: + # TODO: adjust return type hint when Iterator[Union[pd.Series, pd.DataFrame]] is supported + from pyspark.ml.model_cache import ModelCache + + # get predict function (from cache or from running user-provided make_predict_fn) + predict_fn = ModelCache.get(model_uuid) + if not predict_fn: + predict_fn = make_predict_fn() + ModelCache.add(model_uuid, predict_fn) + + # get number of expected parameters for predict function + signature = inspect.signature(predict_fn) + num_expected_cols = len(signature.parameters) + + # convert sparse input_tensor_shapes to dense if needed + input_shapes: List[List[int] | None] + if isinstance(input_tensor_shapes, Mapping): + input_shapes = [None] * num_expected_cols + for index, shape in input_tensor_shapes.items(): + input_shapes[index] = shape + else: + input_shapes = input_tensor_shapes # type: ignore + + # iterate over pandas batch, invoking predict_fn with ndarrays + for pandas_batch in data: + has_tuple = isinstance(pandas_batch, Tuple) # type: ignore + has_tensors = _has_tensor_cols(pandas_batch) + + # require input_tensor_shapes for any tensor columns + if has_tensors and not input_shapes: + raise ValueError("Tensor columns require input_tensor_shapes") + + for batch in _batched(pandas_batch, batch_size): + num_input_rows = len(batch) + num_input_cols = len(batch.columns) + if num_input_cols == num_expected_cols and num_expected_cols > 1: + # input column per expected input for multiple inputs + multi_inputs = _validate_and_transform_multiple_inputs( + batch, input_shapes, num_input_cols + ) + # run model prediction function on multiple (numpy) inputs + preds = predict_fn(*multi_inputs) + elif num_expected_cols == 1: + # one or more input columns for single expected input + single_input = _validate_and_transform_single_input( + batch, input_shapes, has_tensors, has_tuple + ) + # run model prediction function on single (numpy) inputs + preds = predict_fn(single_input) + else: + msg = "Model expected {} inputs, but received {} columns" + raise ValueError(msg.format(num_expected_cols, num_input_cols)) + + # return transformed predictions to Spark + yield _validate_and_transform_prediction_result(preds, num_input_rows, return_type) # type: ignore + + return pandas_udf(predict, return_type) # type: ignore[call-overload] + + +def _test() -> None: + import doctest + from pyspark.sql import SparkSession + import pyspark.ml.functions + import sys + + from pyspark.sql.pandas.utils import ( + require_minimum_pandas_version, + require_minimum_pyarrow_version, + ) + + try: + require_minimum_pandas_version() + require_minimum_pyarrow_version() + except Exception as e: + print( + f"Skipping pyspark.ml.functions doctests: {e}", + file=sys.stderr, + ) + sys.exit(0) + + globs = pyspark.ml.functions.__dict__.copy() + spark = SparkSession.builder.master("local[2]").appName("ml.functions tests").getOrCreate() + sc = spark.sparkContext + globs["sc"] = sc + globs["spark"] = spark + + failure_count, test_count = doctest.testmod( + pyspark.ml.functions, + globs=globs, + optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, + ) + spark.stop() + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/python/pyspark/ml/param/_shared_params_code_gen.py b/python/pyspark/ml/param/_shared_params_code_gen.py index bbcaa208..a885d6f1 100644 --- a/python/pyspark/ml/param/_shared_params_code_gen.py +++ b/python/pyspark/ml/param/_shared_params_code_gen.py @@ -333,6 +333,12 @@ def get{name[0].upper()}{name[1:]}(self) -> {paramType}: "0.0", "TypeConverters.toFloat", ), + ( + "intermediateStorageLevel", + "StorageLevel for intermediate datasets. Cannot be 'NONE'.", + '"MEMORY_AND_DISK"', + "TypeConverters.toString", + ), ( "numTrainWorkers", "number of training workers", diff --git a/python/pyspark/ml/param/shared.py b/python/pyspark/ml/param/shared.py index e60f2a74..ae8100b1 100644 --- a/python/pyspark/ml/param/shared.py +++ b/python/pyspark/ml/param/shared.py @@ -789,6 +789,29 @@ def getMaxBlockSizeInMB(self) -> float: return self.getOrDefault(self.maxBlockSizeInMB) +class HasIntermediateStorageLevel(Params): + """ + Mixin for param intermediateStorageLevel: StorageLevel for intermediate datasets. Cannot be 'NONE'. + """ + + intermediateStorageLevel: "Param[str]" = Param( + Params._dummy(), + "intermediateStorageLevel", + "StorageLevel for intermediate datasets. Cannot be 'NONE'.", + typeConverter=TypeConverters.toString, + ) + + def __init__(self) -> None: + super().__init__() + self._setDefault(intermediateStorageLevel="MEMORY_AND_DISK") + + def getIntermediateStorageLevel(self) -> str: + """ + Gets the value of intermediateStorageLevel or its default value. + """ + return self.getOrDefault(self.intermediateStorageLevel) + + class HasNumTrainWorkers(Params): """ Mixin for param numTrainWorkers: number of training workers diff --git a/python/pyspark/ml/torch/distributor.py b/python/pyspark/ml/torch/distributor.py index dc0eda96..f9e7f12a 100644 --- a/python/pyspark/ml/torch/distributor.py +++ b/python/pyspark/ml/torch/distributor.py @@ -431,12 +431,20 @@ def _get_torchrun_args(local_mode: bool, num_processes: int) -> Tuple[List[Any], else: processes_per_node = 1 node_rank = os.environ["RANK"] + # Set by set_torch_config; no constant fallback, since concurrent runs sharing a + # rendezvous endpoint would collide on a fixed id. + rdzv_id = os.environ.get("PYSPARK_TORCH_DISTRIBUTOR_RDZV_ID") + if not rdzv_id: + raise RuntimeError( + "Missing PYSPARK_TORCH_DISTRIBUTOR_RDZV_ID environment variable: a unique " + "per-run rendezvous id is required for distributed training." + ) torchrun_args = [ f"--nnodes={num_processes // processes_per_node}", f"--node_rank={node_rank}", f"--rdzv_endpoint={master_addr}:{master_port}", - "--rdzv_id=0", # TODO: setup random ID that is gleaned from env variables + f"--rdzv_id={rdzv_id}", ] return torchrun_args, processes_per_node @@ -482,7 +490,7 @@ def _execute_command( tail: collections.deque = collections.deque(maxlen=_TAIL_LINES_TO_KEEP) try: for line in task.stdout: # type: ignore - decoded = line.decode() + decoded = line.decode("utf-8") tail.append(decoded) if redirect_to_stdout: if ( @@ -660,6 +668,7 @@ def _get_spark_task_function( # Spark task program def wrapped_train_fn(iterator): # type: ignore[no-untyped-def] import os + import secrets import pandas as pd import pyarrow from pyspark import BarrierTaskContext @@ -687,6 +696,15 @@ def set_torch_config(context: "BarrierTaskContext") -> None: os.environ["MASTER_ADDR"] = str(addrs[0]) os.environ["MASTER_PORT"] = str(get_free_port(addrs[0], context)) + # Unique per run so that concurrent runs sharing a rendezvous endpoint + # do not collide. + rdzv_id = secrets.token_hex(16) if context.partitionId() == 0 else "" + rdzv_id = context.allGather(str(rdzv_id))[0] + if not rdzv_id: + raise RuntimeError( + "Failed to generate a shared rendezvous id for distributed training." + ) + os.environ["PYSPARK_TORCH_DISTRIBUTOR_RDZV_ID"] = rdzv_id os.environ["WORLD_SIZE"] = str(num_processes) os.environ["NODE_RANK"] = str(context.partitionId()) os.environ["RANK"] = str(context.partitionId()) @@ -883,7 +901,7 @@ def _setup_spark_partition_data( schema_file_path = os.path.join(save_dir, "schema.json") schema_json_string = json.dumps(input_schema_json) - with open(schema_file_path, "w") as f: + with open(schema_file_path, "w", encoding="utf-8") as f: f.write(schema_json_string) os.environ[SPARK_PARTITION_ARROW_DATA_FILE] = arrow_file_path @@ -1088,7 +1106,7 @@ def _get_spark_partition_data_loader( arrow_file = os.environ[SPARK_PARTITION_ARROW_DATA_FILE] schema_file = os.environ[SPARK_DATAFRAME_SCHEMA_FILE] - with open(schema_file, "r") as fp: + with open(schema_file, "r", encoding="utf-8") as fp: schema = StructType.fromJson(json.load(fp)) dataset = _SparkPartitionTorchDataset(arrow_file, schema, num_samples) diff --git a/python/pyspark/ml/torch/torch_run_process_wrapper.py b/python/pyspark/ml/torch/torch_run_process_wrapper.py index 67ec4923..88415a2c 100644 --- a/python/pyspark/ml/torch/torch_run_process_wrapper.py +++ b/python/pyspark/ml/torch/torch_run_process_wrapper.py @@ -69,7 +69,7 @@ def sigterm_handler(*args: Any) -> None: task.stdin.close() # type: ignore[union-attr] try: for line in task.stdout: # type: ignore[union-attr] - decoded = line.decode() + decoded = line.decode("utf-8") print(decoded.rstrip()) task.wait() finally: diff --git a/python/pyspark/mllib/clustering.py b/python/pyspark/mllib/clustering.py index 08d32da2..ca733599 100644 --- a/python/pyspark/mllib/clustering.py +++ b/python/pyspark/mllib/clustering.py @@ -135,7 +135,7 @@ def computeCost(self, x: Union["VectorLike", RDD["VectorLike"]]) -> float: Parameters ---------- - point : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` + x : :py:class:`pyspark.mllib.linalg.Vector` or :py:class:`pyspark.RDD` A data point (or RDD of points) to compute the cost(s). :py:class:`pyspark.mllib.linalg.Vector` can be replaced with equivalent objects (list, tuple, numpy.ndarray). diff --git a/python/pyspark/pandas/__init__.py b/python/pyspark/pandas/__init__.py index ac749c19..bb05ec3b 100644 --- a/python/pyspark/pandas/__init__.py +++ b/python/pyspark/pandas/__init__.py @@ -75,6 +75,7 @@ "options", "option_context", "NamedAgg", + "show_versions", ] diff --git a/python/pyspark/pandas/data_type_ops/timedelta_ops.py b/python/pyspark/pandas/data_type_ops/timedelta_ops.py index aa3e8f64..6f17474f 100644 --- a/python/pyspark/pandas/data_type_ops/timedelta_ops.py +++ b/python/pyspark/pandas/data_type_ops/timedelta_ops.py @@ -18,6 +18,7 @@ from datetime import timedelta from typing import Any, Union +import numpy as np import pandas as pd from pandas.api.types import CategoricalDtype @@ -72,6 +73,38 @@ def restore(self, col: pd.Series) -> pd.Series: else: return col.astype(self.dtype) + def _with_inferred_unit( + self, result: SeriesOrIndex, left: IndexOpsLike, right: Union[IndexOpsMixin, timedelta] + ) -> SeriesOrIndex: + # pandas 3.0.0+ promotes timedelta arithmetic to the finer resolution of the + # operands; before that, pandas-on-Spark represented timedelta as nanoseconds. + if LooseVersion(pd.__version__) < "3.0.0": + return result + + def unit_of(obj: Union[IndexOpsMixin, timedelta]) -> str: + if isinstance(obj, IndexOpsMixin): + dtype = obj.dtype + if isinstance(dtype, np.dtype) and np.issubdtype(dtype, np.timedelta64): + return np.datetime_data(dtype)[0] + elif isinstance(obj, pd.Timedelta): + return obj.unit + # datetime.timedelta scalars and object-backed interval columns map to microseconds. + return "us" + + promoted = np.promote_types( + np.dtype(f"timedelta64[{unit_of(left)}]"), + np.dtype(f"timedelta64[{unit_of(right)}]"), + ) + # DayTimeIntervalType stores microseconds and cannot represent finer resolutions. + if np.datetime_data(promoted)[0] == "ns": + promoted = np.dtype("timedelta64[us]") + + field = result._internal.data_fields[0] + if field.dtype == promoted: + # Already the right resolution (the common us case); avoid rebuilding the field. + return result + return result._with_new_scol(result.spark.column, field=field.copy(dtype=promoted)) + def sub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: _sanitize_list_like(right) @@ -80,7 +113,8 @@ def sub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: and isinstance(right.spark.data_type, DayTimeIntervalType) or isinstance(right, timedelta) ): - return pyspark_column_op("__sub__", left, right) + result = pyspark_column_op("__sub__", left, right) + return self._with_inferred_unit(result, left, right) else: raise TypeError("Timedelta subtraction can only be applied to timedelta series.") @@ -88,7 +122,8 @@ def rsub(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: _sanitize_list_like(right) if isinstance(right, timedelta): - return pyspark_column_op("__rsub__", left, right) + result = pyspark_column_op("__rsub__", left, right) + return self._with_inferred_unit(result, left, right) else: raise TypeError("Timedelta subtraction can only be applied to timedelta series.") diff --git a/python/pyspark/pandas/datetimes.py b/python/pyspark/pandas/datetimes.py index aee37e11..e0562749 100644 --- a/python/pyspark/pandas/datetimes.py +++ b/python/pyspark/pandas/datetimes.py @@ -688,8 +688,6 @@ def round(self, freq: Union[str, DateOffset], *args: Any, **kwargs: Any) -> "ps. - 'raise' will raise an NonExistentTimeError if there are nonexistent times - .. note:: this option only works with pandas 0.24.0+ - Returns ------- Series @@ -748,8 +746,6 @@ def floor(self, freq: Union[str, DateOffset], *args: Any, **kwargs: Any) -> "ps. - 'raise' will raise an NonExistentTimeError if there are nonexistent times - .. note:: this option only works with pandas 0.24.0+ - Returns ------- Series @@ -808,8 +804,6 @@ def ceil(self, freq: Union[str, DateOffset], *args: Any, **kwargs: Any) -> "ps.S - 'raise' will raise an NonExistentTimeError if there are nonexistent times - .. note:: this option only works with pandas 0.24.0+ - Returns ------- Series diff --git a/python/pyspark/pandas/frame.py b/python/pyspark/pandas/frame.py index 812f66b8..d196a850 100644 --- a/python/pyspark/pandas/frame.py +++ b/python/pyspark/pandas/frame.py @@ -2633,8 +2633,7 @@ def to_latex( By default the value will be read from the pandas config module. When set to False prevents from escaping latex special characters in column names. encoding : str, optional - A string representing the encoding to use in the output file, defaults to 'ascii' on - Python 2 and 'utf-8' on Python 3. + A string representing the encoding to use in the output file, defaults to 'utf-8'. decimal : str, default '.' Character recognized as decimal separator, e.g. ',' in Europe. multicolumn : bool, default True @@ -3584,8 +3583,9 @@ def pop(self, item: Name) -> "DataFrame": self._update_internal_frame(self.drop(columns=item)._internal) return result - # TODO(SPARK-46158): add axis parameter can work when '1' or 'columns' - def xs(self, key: Name, axis: Axis = 0, level: Optional[int] = None) -> DataFrameOrSeries: + def xs( + self, key: Name, axis: Axis = 0, level: Optional[Union[int, Name]] = None + ) -> DataFrameOrSeries: """ Return cross-section from the DataFrame. @@ -3596,9 +3596,8 @@ def xs(self, key: Name, axis: Axis = 0, level: Optional[int] = None) -> DataFram ---------- key : label or tuple of label Label contained in the index, or partially in a MultiIndex. - axis : 0 or 'index', default 0 + axis : {0 or 'index', 1 or 'columns'}, default 0 Axis to retrieve cross-section on. - currently only support 0 or 'index' level : object, defaults to first n levels (n=1 or len(key)) In case of a key partially contained in a MultiIndex, indicate which levels are used. Levels can be referred by label or position. @@ -3660,6 +3659,16 @@ class animal locomotion num_legs num_wings class locomotion mammal walks 4 0 + + Get values at specified column + + >>> df.xs('num_legs', axis=1) # doctest: +NORMALIZE_WHITESPACE + class animal locomotion + mammal cat walks 4 + dog walks 4 + bat flies 2 + bird penguin walks 2 + Name: num_legs, dtype: int64 """ from pyspark.pandas.series import first_series @@ -3670,8 +3679,85 @@ class locomotion raise KeyError(key) axis = validate_axis(axis) - if axis != 0: - raise NotImplementedError('axis should be either 0 or "index" currently.') + if axis == 1: + column_level = self._internal.column_labels_level + if column_level == 1: + if level is not None: + raise TypeError("Index must be a MultiIndex") + if is_name_like_tuple(key): + raise KeyError(key) + + column_key = cast(Label, key if is_name_like_tuple(key) else (key,)) + if len(column_key) > column_level: + raise KeyError( + "Key length ({}) exceeds index depth ({})".format(len(column_key), column_level) + ) + + if level is None: + level_idx = 0 + elif isinstance(level, int): + level_idx = level + else: + level_name = cast(Label, level if is_name_like_tuple(level) else (level,)) + try: + level_idx = self._internal.column_label_names.index(level_name) + except ValueError: + raise KeyError("Level {} not found".format(name_like_string(level_name))) + + original_level_idx = level_idx + if level_idx < 0: + level_idx += column_level + if level_idx < 0: + raise IndexError( + "Too many levels: Index has only {} levels, " + "{} is not a valid level number".format(column_level, original_level_idx) + ) + if level_idx >= column_level: + raise IndexError( + "Too many levels: Index has only {} levels, not {}".format( + column_level, level_idx + 1 + ) + ) + + drop_levels = range(level_idx, level_idx + len(column_key)) + selected = [ + (label, scol, field) + for label, scol, field in zip( + self._internal.column_labels, + self._internal.data_spark_columns, + self._internal.data_fields, + ) + if tuple(label[i] for i in drop_levels) == column_key + ] + if not selected: + raise KeyError(key) + + selected_column_labels, data_spark_columns, data_fields = zip(*selected) + # Match pandas by dropping the selected column levels from the result. + new_column_labels = [ + tuple(label[i] for i in range(column_level) if i not in drop_levels) + for label in selected_column_labels + ] + if len(new_column_labels[0]) == 0: + if len(selected_column_labels) == 1: + return self._psser_for(selected_column_labels[0]) + else: + new_column_labels = list(selected_column_labels) + column_label_names = self._internal.column_label_names + else: + column_label_names = [ + name + for i, name in enumerate(self._internal.column_label_names) + if i not in drop_levels + ] + + internal = self._internal.copy( + column_labels=list(new_column_labels), + data_spark_columns=list(data_spark_columns), + data_fields=list(data_fields), + column_label_names=list(column_label_names), + ) + return DataFrame(internal) if not is_name_like_tuple(key): key = (key,) @@ -3682,10 +3768,25 @@ class locomotion ) ) if level is None: - level = 0 + index_level = 0 + elif isinstance(level, int): + index_level = level + else: + level_name = cast(Label, level if is_name_like_tuple(level) else (level,)) + if self._internal.index_names.count(level_name) > 1: + raise ValueError( + "The name {} occurs multiple times, use a level number".format( + name_like_string(level_name) + ) + ) + try: + index_level = self._internal.index_names.index(level_name) + except ValueError: + raise KeyError("Level {} not found".format(name_like_string(level_name))) rows = [ - self._internal.index_spark_columns[lvl] == index for lvl, index in enumerate(key, level) + self._internal.index_spark_columns[lvl] == index + for lvl, index in enumerate(key, index_level) ] internal = self._internal.with_filter(reduce(lambda x, y: x & y, rows)) @@ -3700,11 +3801,16 @@ class locomotion return first_series(DataFrame(pdf.transpose())) else: index_spark_columns = ( - internal.index_spark_columns[:level] - + internal.index_spark_columns[level + len(key) :] + internal.index_spark_columns[:index_level] + + internal.index_spark_columns[index_level + len(key) :] + ) + index_names = ( + internal.index_names[:index_level] + internal.index_names[index_level + len(key) :] + ) + index_fields = ( + internal.index_fields[:index_level] + + internal.index_fields[index_level + len(key) :] ) - index_names = internal.index_names[:level] + internal.index_names[level + len(key) :] - index_fields = internal.index_fields[:level] + internal.index_fields[level + len(key) :] internal = internal.copy( index_spark_columns=index_spark_columns, @@ -14207,6 +14313,10 @@ def _set_axis_fallback(self, *args: Any, **kwargs: Any) -> "DataFrame": _f = self._build_fallback_method("set_axis") return _f(*args, **kwargs) + def _combine_fallback(self, *args: Any, **kwargs: Any) -> "DataFrame": + _f = self._build_fallback_method("combine") + return _f(*args, **kwargs) + def __getattr__(self, key: str) -> Any: if key.startswith("__"): raise AttributeError(key) diff --git a/python/pyspark/pandas/groupby.py b/python/pyspark/pandas/groupby.py index b0879c94..9907d955 100644 --- a/python/pyspark/pandas/groupby.py +++ b/python/pyspark/pandas/groupby.py @@ -1136,7 +1136,7 @@ def nth(self, n: int) -> FrameLike: There is a behavior difference between pandas-on-Spark and pandas: * when there is no aggregation column, and `n` not equal to 0 or -1, - the returned empty dataframe may have an index with different lenght `__len__`. + the returned empty dataframe may have an index with different length `__len__`. Examples -------- @@ -2459,8 +2459,10 @@ def idxmax(self, skipna: bool = True) -> FrameLike: groupkey_names = ["__groupkey_{}__".format(i) for i in range(len(self._groupkeys))] sdf = self._psdf._internal.spark_frame - for s, name in zip(self._groupkeys, groupkey_names): - sdf = sdf.withColumn(name, s.spark.column) + if groupkey_names: + sdf = sdf.withColumns( + {name: s.spark.column for s, name in zip(self._groupkeys, groupkey_names)} + ) index = self._psdf._internal.index_spark_column_names[0] index_spark_type = self._psdf._internal.index_fields[0].spark_type @@ -2554,8 +2556,10 @@ def idxmin(self, skipna: bool = True) -> FrameLike: groupkey_names = ["__groupkey_{}__".format(i) for i in range(len(self._groupkeys))] sdf = self._psdf._internal.spark_frame - for s, name in zip(self._groupkeys, groupkey_names): - sdf = sdf.withColumn(name, s.spark.column) + if groupkey_names: + sdf = sdf.withColumns( + {name: s.spark.column for s, name in zip(self._groupkeys, groupkey_names)} + ) index = self._psdf._internal.index_spark_column_names[0] index_spark_type = self._psdf._internal.index_fields[0].spark_type @@ -4067,12 +4071,14 @@ def describe(self) -> DataFrame: # Split "quartiles" columns into first, second, and third quartiles. for label in agg_column_labels: quartiles_col = name_like_string(tuple(list(label) + ["quartiles"])) - for i, percentile in enumerate(formatted_percentiles): - sdf = sdf.withColumn( - name_like_string(tuple(list(label) + [percentile])), - scol_for(sdf, quartiles_col)[i], - ) - sdf = sdf.drop(quartiles_col) + sdf = sdf.withColumns( + { + name_like_string(tuple(list(label) + [percentile])): scol_for( + sdf, quartiles_col + )[i] + for i, percentile in enumerate(formatted_percentiles) + } + ).drop(quartiles_col) # Reorder columns lexicographically by agg column followed by stats. stats = ["count", "mean", "std", "min"] + formatted_percentiles + ["max"] @@ -4303,6 +4309,216 @@ def corr( ) ) + @with_ansi_mode_context + def cov( + self, + min_periods: Optional[int] = None, + ddof: int = 1, + numeric_only: bool = False, + ) -> "DataFrame": + """ + Compute pairwise covariance of columns, excluding NA/null values. + + The returned DataFrame is the covariance matrix of the columns + of the DataFrame within each group. + + Both NA and null values are automatically excluded from the + calculation. A threshold can be set for the minimum number of + observations for each value created. Comparisons with observations + below this threshold will be returned as ``NaN``. + + .. versionadded:: 4.3.0 + + Parameters + ---------- + min_periods : int, optional + Minimum number of observations required per pair of columns + to have a valid result. + + ddof : int, default 1 + Delta degrees of freedom. The divisor used in calculations + is ``N - ddof``, where ``N`` represents the number of elements. + + numeric_only : bool, default False + Include only `float`, `int` or `boolean` data. + + Returns + ------- + DataFrame + The covariance matrix of the series of the DataFrame within each group. + + See Also + -------- + DataFrame.cov + Series.cov + + Examples + -------- + >>> df = ps.DataFrame( + ... {"A": [1, 1, 2, 2, 2], "B": [1, 2, 3, 4, 5], "C": [4, 6, 7, 9, 11]}, + ... columns=["A", "B", "C"]) + >>> df.groupby("A").cov().sort_index() + B C + A + 1 B 0.5 1.0 + C 1.0 2.0 + 2 B 1.0 2.0 + C 2.0 4.0 + + >>> df.groupby("A").cov(ddof=0).sort_index() + B C + A + 1 B 0.250000 0.500000 + C 0.500000 1.000000 + 2 B 0.666667 1.333333 + C 1.333333 2.666667 + + >>> df.groupby("A").cov(min_periods=3).sort_index() + B C + A + 1 B NaN NaN + C NaN NaN + 2 B 1.0 2.0 + C 2.0 4.0 + """ + if not isinstance(ddof, int): + raise TypeError("ddof must be integer") + if LooseVersion(pd.__version__) >= "3.0.0": + if not isinstance(numeric_only, bool): + raise ValueError("numeric_only accepts only Boolean values") + min_periods = 1 if min_periods is None else min_periods + + groupkey_names: List[str] = [str(key.name) for key in self._groupkeys] + internal, _, _ = self._prepare_reduce( + groupkey_names=groupkey_names, + accepted_spark_types=(NumericType, BooleanType) if numeric_only else None, + bool_to_numeric=False, + ) + + numeric_labels = [ + label + for label in internal.column_labels + if isinstance(internal.spark_type_for(label), (NumericType, BooleanType)) + ] + numeric_scols: List[Column] = [ + internal.spark_column_for(label).cast("double") for label in numeric_labels + ] + numeric_col_names: List[str] = [name_like_string(label) for label in numeric_labels] + num_scols = len(numeric_scols) + + sdf = internal.spark_frame + index_1_col_name = verify_temp_column_name(sdf, "__groupby_cov_index_1_temp_column__") + index_2_col_name = verify_temp_column_name(sdf, "__groupby_cov_index_2_temp_column__") + value_1_col_name = verify_temp_column_name(sdf, "__groupby_cov_value_1_temp_column__") + value_2_col_name = verify_temp_column_name(sdf, "__groupby_cov_value_2_temp_column__") + cov_output_col_name = verify_temp_column_name(sdf, "__groupby_cov_output_temp_column__") + count_output_col_name = verify_temp_column_name(sdf, "__groupby_cov_count_temp_column__") + + pair_scols: List[Column] = [] + for i in range(0, num_scols): + for j in range(i, num_scols): + pair_scols.append( + F.struct( + F.lit(i).alias(index_1_col_name), + F.lit(j).alias(index_2_col_name), + numeric_scols[i].alias(value_1_col_name), + numeric_scols[j].alias(value_2_col_name), + ) + ) + + sdf = sdf.select(*[F.col(key) for key in groupkey_names], *[F.inline(F.array(*pair_scols))]) + + # Null both values when either is null so pairwise non-null counts are accurate. + sdf = sdf.select( + *[F.col(key) for key in groupkey_names + [index_1_col_name, index_2_col_name]], + F.when(F.isnull(value_1_col_name) | F.isnull(value_2_col_name), F.lit(None)) + .otherwise(F.col(value_1_col_name)) + .alias(value_1_col_name), + F.when(F.isnull(value_1_col_name) | F.isnull(value_2_col_name), F.lit(None)) + .otherwise(F.col(value_2_col_name)) + .alias(value_2_col_name), + ) + + sdf = sdf.groupby(groupkey_names + [index_1_col_name, index_2_col_name]).agg( + SF.covar(F.col(value_1_col_name), F.col(value_2_col_name), ddof).alias( + cov_output_col_name + ), + F.count(F.when(~F.isnull(value_1_col_name), 1)).alias(count_output_col_name), + ) + + sdf = sdf.withColumn( + cov_output_col_name, + F.when(F.col(count_output_col_name) < min_periods, F.lit(None)).otherwise( + F.col(cov_output_col_name) + ), + ) + + # Mirror the (i, j) pair to (j, i) to fill in the lower triangle of the matrix. + auxiliary_col_name = verify_temp_column_name(sdf, "__groupby_cov_auxiliary_temp_column__") + sdf = sdf.withColumn( + auxiliary_col_name, + F.explode( + F.when( + F.col(index_1_col_name) == F.col(index_2_col_name), + F.lit([0]), + ).otherwise(F.lit([0, 1])) + ), + ).select( + *[F.col(key) for key in groupkey_names], + *[ + F.when(F.col(auxiliary_col_name) == 0, F.col(index_1_col_name)) + .otherwise(F.col(index_2_col_name)) + .alias(index_1_col_name), + F.when(F.col(auxiliary_col_name) == 0, F.col(index_2_col_name)) + .otherwise(F.col(index_1_col_name)) + .alias(index_2_col_name), + F.col(cov_output_col_name), + ], + ) + + array_col_name = verify_temp_column_name(sdf, "__groupby_cov_array_temp_column__") + sdf = sdf.groupby(groupkey_names + [index_1_col_name]).agg( + F.array_sort( + F.collect_list( + F.struct( + F.col(index_2_col_name), + F.col(cov_output_col_name), + ) + ) + ).alias(array_col_name) + ) + + for i in range(0, num_scols): + sdf = sdf.withColumn(auxiliary_col_name, F.get(F.col(array_col_name), i)).withColumn( + numeric_col_names[i], + F.col(f"{auxiliary_col_name}.{cov_output_col_name}"), + ) + + sdf = sdf.orderBy(groupkey_names + [index_1_col_name]) + + sdf = sdf.select( + *[F.col(col) for col in groupkey_names + numeric_col_names], + *[ + F.get(F.lit(numeric_col_names), F.col(index_1_col_name)).alias(auxiliary_col_name), + F.monotonically_increasing_id().alias(NATURAL_ORDER_COLUMN_NAME), + ], + ) + + return DataFrame( + InternalFrame( + spark_frame=sdf, + index_spark_columns=[ + scol_for(sdf, key) for key in groupkey_names + [auxiliary_col_name] + ], + index_names=( + [psser._column_label for psser in self._groupkeys] + + self._psdf._internal.index_names + ), + column_labels=numeric_labels, + column_label_names=internal.column_label_names, + ) + ) + class SeriesGroupBy(GroupBy[Series]): @staticmethod @@ -4371,6 +4587,76 @@ def _handle_output( else: return first_series(psdf).rename(self._psser.name) + def describe(self) -> DataFrame: + """ + Generate descriptive statistics that summarize the central tendency, + dispersion and shape of a dataset's distribution, excluding + ``NaN`` values. + + .. versionadded:: 4.3.0 + + .. note:: Unlike pandas, the percentiles in pandas-on-Spark are based upon + approximate percentile computation because computing percentiles + across a large dataset is extremely expensive. + + Returns + ------- + DataFrame + Summary statistics for each group and the series values. + + See Also + -------- + DataFrame.describe + + Examples + -------- + >>> df = ps.DataFrame({'a': [1, 1, 3], 'b': [4, 5, 6]}) + >>> df.groupby('a')['b'].describe().sort_index() # doctest: +NORMALIZE_WHITESPACE + count mean std min 25% 50% 75% max + a + 1 2.0 4.5 0.707107 4.0 4.0 4.0 5.0 5.0 + 3 1.0 6.0 NaN 6.0 6.0 6.0 6.0 6.0 + """ + # NOTE: keep in sync with GroupBy.describe + if isinstance(self._agg_columns[0].spark.data_type, StringType): + raise NotImplementedError( + "SeriesGroupBy.describe() doesn't support for string type for now" + ) + + # Use the base GroupBy.aggregate to bypass SeriesGroupBy's unsupported override. + psdf = GroupBy.aggregate(self, ["count", "mean", "std", "min", "quartiles", "max"]) + sdf = psdf._internal.spark_frame + agg_column_labels = [col._column_label for col in self._agg_columns] + formatted_percentiles = ["25%", "50%", "75%"] + + # Split "quartiles" columns into first, second, and third quartiles. + for label in agg_column_labels: + quartiles_col = name_like_string(tuple(list(label) + ["quartiles"])) + for i, percentile in enumerate(formatted_percentiles): + sdf = sdf.withColumn( + name_like_string(tuple(list(label) + [percentile])), + scol_for(sdf, quartiles_col)[i], + ) + sdf = sdf.drop(quartiles_col) + + # Build single-level column labels: ["count", "mean", "std", "min", "25%", ...] + stats = ["count", "mean", "std", "min"] + formatted_percentiles + ["max"] + # For SeriesGroupBy there is only one agg column label. + label = agg_column_labels[0] + multi_column_labels = [tuple(list(label) + [s]) for s in stats] + data_columns = [name_like_string(cl) for cl in multi_column_labels] + + internal = psdf._internal.copy( + spark_frame=sdf, + column_labels=[(s,) for s in stats], + data_spark_columns=[scol_for(sdf, col) for col in data_columns], + data_fields=None, + column_label_names=None, + ) + + # Cast columns to "float64" to match pandas.SeriesGroupBy.describe(). + return DataFrame(internal).astype("float64") + def agg(self, *args: Any, **kwargs: Any) -> None: return MissingPandasLikeSeriesGroupBy.agg(self, *args, **kwargs) diff --git a/python/pyspark/pandas/missing/groupby.py b/python/pyspark/pandas/missing/groupby.py index a6b672df..9e655384 100644 --- a/python/pyspark/pandas/missing/groupby.py +++ b/python/pyspark/pandas/missing/groupby.py @@ -42,7 +42,6 @@ class MissingPandasLikeDataFrameGroupBy: # Properties corrwith = _unsupported_property("corrwith") - cov = _unsupported_property("cov") dtypes = _unsupported_property("dtypes") groups = _unsupported_property("groups") hist = _unsupported_property("hist") @@ -84,7 +83,6 @@ class MissingPandasLikeSeriesGroupBy: # Functions agg = _unsupported_function("agg") aggregate = _unsupported_function("aggregate") - describe = _unsupported_function("describe") ngroup = _unsupported_function("ngroup") ohlc = _unsupported_function("ohlc") pct_change = _unsupported_function("pct_change") diff --git a/python/pyspark/pandas/missing/window.py b/python/pyspark/pandas/missing/window.py index a6d423d0..5342ed51 100644 --- a/python/pyspark/pandas/missing/window.py +++ b/python/pyspark/pandas/missing/window.py @@ -81,7 +81,6 @@ class MissingPandasLikeExpanding: apply = _unsupported_function_expanding("apply") corr = _unsupported_function_expanding("corr") cov = _unsupported_function_expanding("cov") - median = _unsupported_function_expanding("median") validate = _unsupported_function_expanding("validate") exclusions = _unsupported_property_expanding("exclusions") @@ -99,7 +98,6 @@ class MissingPandasLikeRolling: apply = _unsupported_function_rolling("apply") corr = _unsupported_function_rolling("corr") cov = _unsupported_function_rolling("cov") - median = _unsupported_function_rolling("median") validate = _unsupported_function_rolling("validate") exclusions = _unsupported_property_rolling("exclusions") @@ -117,7 +115,6 @@ class MissingPandasLikeExpandingGroupby: apply = _unsupported_function_expanding("apply") corr = _unsupported_function_expanding("corr") cov = _unsupported_function_expanding("cov") - median = _unsupported_function_expanding("median") validate = _unsupported_function_expanding("validate") exclusions = _unsupported_property_expanding("exclusions") @@ -135,7 +132,6 @@ class MissingPandasLikeRollingGroupby: apply = _unsupported_function_rolling("apply") corr = _unsupported_function_rolling("corr") cov = _unsupported_function_rolling("cov") - median = _unsupported_function_rolling("median") validate = _unsupported_function_rolling("validate") exclusions = _unsupported_property_rolling("exclusions") diff --git a/python/pyspark/pandas/namespace.py b/python/pyspark/pandas/namespace.py index 7be107ec..05875347 100644 --- a/python/pyspark/pandas/namespace.py +++ b/python/pyspark/pandas/namespace.py @@ -141,6 +141,7 @@ "broadcast", "read_orc", "json_normalize", + "show_versions", ] @@ -2731,8 +2732,10 @@ def resolve_func(psdf, this_column_labels, that_column_labels): # TODO: NaN and None difference for missing values. pandas seems to be filling NaN. sdf = psdf._internal.resolved_copy.spark_frame - for label in columns_to_add: - sdf = sdf.withColumn(name_like_string(label), F.lit(None)) + if columns_to_add: + sdf = sdf.withColumns( + {name_like_string(label): F.lit(None) for label in columns_to_add} + ) data_columns = psdf._internal.data_spark_column_names + [ name_like_string(label) for label in columns_to_add @@ -3879,6 +3882,134 @@ def flatten_schema(schema: StructType, prefix: str = "") -> Tuple[List[str], Lis return ps.DataFrame(internal) +def _get_sys_info() -> Dict[str, Any]: + """Returns system information as a dictionary.""" + import locale + import os + import platform + import struct + import sys as _sys + + uname_result = platform.uname() + try: + language_code, encoding = locale.getlocale() + except (TypeError, ValueError): + language_code, encoding = (None, None) + return { + "python": platform.python_version(), + "python-bits": struct.calcsize("P") * 8, + "OS": uname_result.system, + "OS-release": uname_result.release, + "Version": uname_result.version, + "machine": uname_result.machine, + "processor": uname_result.processor, + "byteorder": _sys.byteorder, + "LC_ALL": os.environ.get("LC_ALL"), + "LANG": os.environ.get("LANG"), + "LOCALE": {"language-code": language_code, "encoding": encoding}, + } + + +def _get_dependency_info() -> Dict[str, Optional[str]]: + """Returns dependency information as a dictionary.""" + import importlib + + import pyspark + + deps = [ + "pyspark", + "pandas", + "numpy", + "pyarrow", + "grpc", + "google.protobuf", + "matplotlib", + "IPython", + "sphinx", + "plotly", + "tabulate", + "scipy", + "mlflow", + ] + result: Dict[str, Optional[str]] = {} + for modname in deps: + if modname == "pyspark": + result[modname] = pyspark.__version__ + continue + try: + mod = importlib.import_module(modname) + except ImportError: + result[modname] = None + except Exception: + # Dependency conflicts may cause non-ImportError failures. + result[modname] = "N/A" + else: + result[modname] = getattr(mod, "__version__", None) + return result + + +def show_versions(as_json: Union[str, bool] = False) -> None: + """ + Provide useful information, important for bug reports. + + It comprises info about hosting operation system, pyspark.pandas version, + and versions of other installed relative packages. + + .. versionadded:: 4.3.0 + + Parameters + ---------- + as_json : str or bool, default False + * If False, outputs info in a human readable form to the console. + * If str, it will be considered as a path to a file. + Info will be written to that file in JSON format. + * If True, outputs info in JSON format to the console. + + Examples + -------- + >>> ps.show_versions() # doctest: +SKIP + INSTALLED VERSIONS + ------------------ + python : 3.10.6.final.0 + python-bits : 64 + ... + pyspark : 4.3.0.dev0 + pandas : 2.2.0 + numpy : 1.24.3 + pyarrow : 15.0.0 + ... + """ + sys_info = _get_sys_info() + deps = _get_dependency_info() + + if as_json: + import sys as _sys + + j = {"system": sys_info, "dependencies": deps} + + if as_json is True: + _sys.stdout.writelines(json.dumps(j, indent=2)) + else: + assert isinstance(as_json, str) + with open(as_json, "w", encoding="utf-8") as f: + json.dump(j, f, indent=2) + return + + locale_info = sys_info["LOCALE"] + sys_info["LOCALE"] = "{language_code}.{encoding}".format( + language_code=locale_info["language-code"], encoding=locale_info["encoding"] + ) + + maxlen = max(len(x) for x in deps) + print("\nINSTALLED VERSIONS") + print("------------------") + for k, v in sys_info.items(): + print(f"{k:<{maxlen}}: {v}") + print("") + for k, v in deps.items(): + print(f"{k:<{maxlen}}: {v}") + + def _get_index_map( sdf: PySparkDataFrame, index_col: Optional[Union[str, List[str]]] = None ) -> Tuple[Optional[List[PySparkColumn]], Optional[List[Label]]]: diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index 913089ca..1653121d 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -20,18 +20,19 @@ from pyspark.sql import functions as F from pyspark.sql.pandas.functions import pandas_udf -from pyspark.sql.types import DoubleType, LongType, BooleanType +from pyspark.sql.types import DoubleType, BooleanType from pyspark.pandas.base import IndexOpsMixin + unary_np_spark_mappings = { "abs": F.abs, "absolute": F.abs, "arccos": F.acos, - "arccosh": pandas_udf(lambda s: np.arccosh(s), DoubleType()), # type: ignore[call-overload] + "arccosh": F.acosh, "arcsin": F.asin, - "arcsinh": pandas_udf(lambda s: np.arcsinh(s), DoubleType()), # type: ignore[call-overload] + "arcsinh": F.asinh, "arctan": F.atan, - "arctanh": pandas_udf(lambda s: np.arctanh(s), DoubleType()), # type: ignore[call-overload] + "arctanh": F.atanh, "bitwise_not": F.bitwiseNOT, "cbrt": F.cbrt, "ceil": F.ceil, @@ -39,45 +40,54 @@ "conj": lambda _: NotImplemented, "conjugate": lambda _: NotImplemented, # It requires complex type "cos": F.cos, - "cosh": pandas_udf(lambda s: np.cosh(s), DoubleType()), # type: ignore[call-overload] - "deg2rad": pandas_udf(lambda s: np.deg2rad(s), DoubleType()), # type: ignore[call-overload] + "cosh": F.cosh, + "deg2rad": F.radians, "degrees": F.degrees, "exp": F.exp, - "exp2": pandas_udf(lambda s: np.exp2(s), DoubleType()), # type: ignore[call-overload] + "exp2": lambda c: F.pow(F.lit(2.0), c), "expm1": F.expm1, - "fabs": pandas_udf(lambda s: np.fabs(s), DoubleType()), # type: ignore[call-overload] + "fabs": lambda c: F.abs(c.cast("double")), "floor": F.floor, "frexp": lambda _: NotImplemented, # 'frexp' output lengths become different # and it cannot be supported via pandas UDF. - "invert": pandas_udf(lambda s: np.invert(s), DoubleType()), # type: ignore[call-overload] - "isfinite": lambda c: c != float("inf"), - "isinf": lambda c: c == float("inf"), + "invert": F.bitwise_not, + "isfinite": lambda c: F.coalesce( + ~(F.isnan(c) | (c == float("inf")) | (c == float("-inf"))), F.lit(False) + ), + "isinf": lambda c: F.coalesce((c == float("inf")) | (c == float("-inf")), F.lit(False)), "isnan": F.isnan, "isnat": lambda c: NotImplemented, # pandas-on-Spark and PySpark does not have Nat concept. "log": F.log, "log10": F.log10, "log1p": F.log1p, - "log2": pandas_udf(lambda s: np.log2(s), DoubleType()), # type: ignore[call-overload] "logical_not": lambda c: ~(c.cast(BooleanType())), "matmul": lambda _: NotImplemented, # Can return a NumPy array in pandas. - "negative": lambda c: c * -1, - "positive": lambda c: c, - "rad2deg": pandas_udf(lambda s: np.rad2deg(s), DoubleType()), # type: ignore[call-overload] + "negative": F.negative, + "positive": F.positive, + "rad2deg": F.degrees, "radians": F.radians, "reciprocal": pandas_udf( # type: ignore[call-overload] lambda s: np.reciprocal(s), DoubleType() ), - "rint": pandas_udf(lambda s: np.rint(s), DoubleType()), # type: ignore[call-overload] - "sign": lambda c: F.when(c == 0, 0).when(c < 0, -1).otherwise(1), + "rint": lambda c: F.rint(c.cast("double")), + "sign": F.signum, "signbit": lambda c: F.when(c < 0, True).otherwise(False), "sin": F.sin, - "sinh": pandas_udf(lambda s: np.sinh(s), DoubleType()), # type: ignore[call-overload] + "sinh": F.sinh, "spacing": pandas_udf(lambda s: np.spacing(s), DoubleType()), # type: ignore[call-overload] "sqrt": F.sqrt, - "square": pandas_udf(lambda s: np.square(s), DoubleType()), # type: ignore[call-overload] + "square": lambda c: c.cast("double") * c, "tan": F.tan, - "tanh": pandas_udf(lambda s: np.tanh(s), DoubleType()), # type: ignore[call-overload] - "trunc": pandas_udf(lambda s: np.trunc(s), DoubleType()), # type: ignore[call-overload] + "tanh": F.tanh, + "trunc": lambda c: F.when( + c.cast("double").isNull() + | F.isnan(c.cast("double")) + | c.cast("double").isin(float("-inf"), float("inf")), + c.cast("double"), + ).otherwise( + F.signum(c.cast("double")) + * (F.abs(c.cast("double")) - (F.abs(c.cast("double")) % F.lit(1.0))) + ), } binary_np_spark_mappings = { @@ -106,8 +116,10 @@ "ldexp": pandas_udf( # type: ignore[call-overload] lambda s1, s2: np.ldexp(s1, s2), DoubleType() ), - "left_shift": pandas_udf( # type: ignore[call-overload] - lambda s1, s2: np.left_shift(s1, s2), LongType() + # F.shiftleft accepts literal counts only; call_function also accepts a column. + # NumPy returns zero for counts outside an int64's bit width, unlike JVM shifts. + "left_shift": lambda c1, c2: F.when((c2 < 0) | (c2 >= 64), F.lit(0)).otherwise( + F.call_function("shiftleft", c1, c2) ), "logaddexp": pandas_udf( # type: ignore[call-overload] lambda s1, s2: np.logaddexp(s1, s2), DoubleType() @@ -128,9 +140,11 @@ "nextafter": pandas_udf( # type: ignore[call-overload] lambda s1, s2: np.nextafter(s1, s2), DoubleType() ), - "right_shift": pandas_udf( # type: ignore[call-overload] - lambda s1, s2: np.right_shift(s1, s2), LongType() - ), + # F.shiftright accepts literal counts only; call_function also accepts a column. + # NumPy sign-extends counts outside an int64's bit width, unlike JVM shifts. + "right_shift": lambda c1, c2: F.when( + (c2 < 0) | (c2 >= 64), F.call_function("shiftright", c1, F.lit(63)) + ).otherwise(F.call_function("shiftright", c1, c2)), } diff --git a/python/pyspark/pandas/series.py b/python/pyspark/pandas/series.py index 64961378..9f953679 100644 --- a/python/pyspark/pandas/series.py +++ b/python/pyspark/pandas/series.py @@ -4316,7 +4316,7 @@ def _rank( window = ( Window.orderBy( asc_func(self.spark.column), - asc_func(F.col(NATURAL_ORDER_COLUMN_NAME)), + F.col(NATURAL_ORDER_COLUMN_NAME).asc(), ) .partitionBy(*part_cols) .rowsBetween(Window.unboundedPreceding, Window.currentRow) @@ -6000,7 +6000,7 @@ def asof(self, where: Union[Any, List]) -> Union[Scalar, "Series"]: ----- Indices are assumed to be sorted. Raises if this is not the case and config 'compute.eager_check' is True. If 'compute.eager_check' is False pandas-on-Spark just - proceeds and performs by ignoring the indeces's order + proceeds and performs by ignoring the indices' order Examples -------- diff --git a/python/pyspark/pandas/supported_api_gen.py b/python/pyspark/pandas/supported_api_gen.py index c39b8880..3a9c4f1d 100644 --- a/python/pyspark/pandas/supported_api_gen.py +++ b/python/pyspark/pandas/supported_api_gen.py @@ -419,7 +419,7 @@ def _write_rst( all_supported_status : Dict Collected support status data. """ - with open(output_rst_file_path, "w") as w_fd: + with open(output_rst_file_path, "w", encoding="utf-8") as w_fd: w_fd.write(RST_HEADER) for module_info, supported_status in all_supported_status.items(): module, module_path = module_info diff --git a/python/pyspark/pandas/typedef/typehints.py b/python/pyspark/pandas/typedef/typehints.py index 1c3d8dd2..8d99584f 100644 --- a/python/pyspark/pandas/typedef/typehints.py +++ b/python/pyspark/pandas/typedef/typehints.py @@ -21,11 +21,10 @@ import datetime import decimal -import sys import typing from collections.abc import Iterable from inspect import isclass -from typing import Any, Callable, Generic, List, Tuple, Union, Type, get_type_hints +from typing import Any, Callable, Generic, List, Optional, Tuple, Union, Type, get_type_hints import numpy as np import pandas as pd @@ -127,15 +126,15 @@ def __repr__(self) -> str: class IndexNameTypeHolder: - name = None - tpe = None - short_name = "IndexNameType" + name: Optional[str] = None + tpe: Optional[Union[type, Dtype]] = None + short_name: str = "IndexNameType" class NameTypeHolder: - name = None - tpe = None - short_name = "NameType" + name: Optional[str] = None + tpe: Optional[Union[type, Dtype]] = None + short_name: str = "NameType" def as_spark_type( @@ -150,28 +149,24 @@ def as_spark_type( - dictionaries of field_name -> type - Python3's typing system """ - # For NumPy typing, NumPy version should be 1.21+ - if LooseVersion(np.__version__) >= LooseVersion("1.21"): - if ( - hasattr(tpe, "__origin__") - and tpe.__origin__ is np.ndarray - and hasattr(tpe, "__args__") - and len(tpe.__args__) > 1 - ): - # numpy.typing.NDArray for numpy < 2.5 - return types.ArrayType( - as_spark_type(tpe.__args__[1].__args__[0], raise_error=raise_error) - ) - elif ( - hasattr(tpe, "__origin__") - and hasattr(tpe.__origin__, "__value__") - and getattr(tpe.__origin__.__value__, "__origin__", None) is np.ndarray - and hasattr(tpe, "__args__") - and len(tpe.__args__) > 0 - ): - # numpy.typing.NDArray for numpy >= 2.5: a PEP 695 type alias whose __value__ - # resolves to np.ndarray[shape, dtype[scalar]], with the scalar at __args__[0] - return types.ArrayType(as_spark_type(tpe.__args__[0], raise_error=raise_error)) + if ( + hasattr(tpe, "__origin__") + and tpe.__origin__ is np.ndarray + and hasattr(tpe, "__args__") + and len(tpe.__args__) > 1 + ): + # numpy.typing.NDArray for numpy < 2.5 + return types.ArrayType(as_spark_type(tpe.__args__[1].__args__[0], raise_error=raise_error)) + elif ( + hasattr(tpe, "__origin__") + and hasattr(tpe.__origin__, "__value__") + and getattr(tpe.__origin__.__value__, "__origin__", None) is np.ndarray + and hasattr(tpe, "__args__") + and len(tpe.__args__) > 0 + ): + # numpy.typing.NDArray for numpy >= 2.5: a PEP 695 type alias whose __value__ + # resolves to np.ndarray[shape, dtype[scalar]], with the scalar at __args__[0] + return types.ArrayType(as_spark_type(tpe.__args__[0], raise_error=raise_error)) if isinstance(tpe, np.dtype) and tpe == np.dtype("object"): pass @@ -736,7 +731,7 @@ def create_type_for_series_type(param: Any) -> Type[SeriesType]: new_class: Type[NameTypeHolder] if isinstance(param, ExtensionDtype): new_class = type(NameTypeHolder.short_name, (NameTypeHolder,), {}) - new_class.tpe = param # type: ignore[assignment] + new_class.tpe = param else: if LooseVersion(pd.__version__) < "3.0.0": new_class = param.type if isinstance(param, np.dtype) else param @@ -877,21 +872,16 @@ def _new_type_holders( isinstance(param, slice) and param.step is None and param.stop is not None for param in params ) - if sys.version_info < (3, 11): - is_unnamed_params = all( - not isinstance(param, slice) and not isinstance(param, Iterable) for param in params - ) - else: - # PEP 646 changes `GenericAlias` instances into iterable ones at Python 3.11 - is_unnamed_params = all( - not isinstance(param, slice) - and ( - not isinstance(param, Iterable) - or isinstance(param, typing.GenericAlias) # type: ignore[attr-defined] - or isinstance(param, typing._GenericAlias) # type: ignore[attr-defined] - ) - for param in params + # PEP 646 changes `GenericAlias` instances into iterable ones at Python 3.11+ + is_unnamed_params = all( + not isinstance(param, slice) + and ( + not isinstance(param, Iterable) + or isinstance(param, typing.GenericAlias) # type: ignore[attr-defined] + or isinstance(param, typing._GenericAlias) # type: ignore[attr-defined] ) + for param in params + ) if is_named_params: # DataFrame["id": int, "A": int] @@ -903,13 +893,11 @@ def _new_type_holders( new_param.name = param.start if LooseVersion(pd.__version__) < "3.0.0": if isinstance(param.stop, ExtensionDtype): - new_param.tpe = param.stop # type: ignore[assignment] + new_param.tpe = param.stop else: # When the given argument is a numpy's dtype instance. new_param.tpe = ( - param.stop.type # type: ignore[assignment] - if isinstance(param.stop, np.dtype) - else param.stop + param.stop.type if isinstance(param.stop, np.dtype) else param.stop ) else: new_param.tpe = param.stop @@ -924,13 +912,9 @@ def _new_type_holders( ) if LooseVersion(pd.__version__) < "3.0.0": if isinstance(param, ExtensionDtype): - new_type.tpe = param # type: ignore[assignment] + new_type.tpe = param else: - new_type.tpe = ( - param.type # type: ignore[assignment] - if isinstance(param, np.dtype) - else param - ) + new_type.tpe = param.type if isinstance(param, np.dtype) else param else: new_type.tpe = param new_types.append(new_type) diff --git a/python/pyspark/pandas/window.py b/python/pyspark/pandas/window.py index 294923d5..20ebbddd 100644 --- a/python/pyspark/pandas/window.py +++ b/python/pyspark/pandas/window.py @@ -19,7 +19,9 @@ from typing import Any, Callable, Generic, List, Optional import numpy as np +import pandas as pd +from pyspark.loose_version import LooseVersion from pyspark.sql import Window from pyspark.sql import functions as F from pyspark.sql.internal import InternalFunction as SF @@ -110,6 +112,9 @@ def quantile(scol: Column) -> Column: return self._apply_as_series_or_frame(quantile) + def median(self) -> FrameLike: + return self.quantile(0.5) + def std(self) -> FrameLike: def std(scol: Column) -> Column: return F.when( @@ -128,6 +133,25 @@ def var(scol: Column) -> Column: return self._apply_as_series_or_frame(var) + def sem(self, ddof: int = 1) -> FrameLike: + if not isinstance(ddof, int): + raise TypeError("ddof must be integer") + + def sem(scol: Column) -> Column: + count = F.count(scol).over(self._window) + if LooseVersion(pd.__version__) >= "3.0.0": + # pandas 3.0 fixed sem to std(ddof) / sqrt(count) (pandas-dev/pandas#63181); + # earlier pandas divided by sqrt(count - ddof) instead. + sem_scol = SF.stddev(scol, ddof).over(self._window) / F.sqrt(count) + else: + sem_scol = F.stddev(scol).over(self._window) / F.sqrt(count - ddof) + return F.when( + (F.row_number().over(self._unbounded_window) >= self._min_periods) & (count > ddof), + sem_scol, + ).otherwise(F.lit(None)) + + return self._apply_as_series_or_frame(sem) + def skew(self) -> FrameLike: def skew(scol: Column) -> Column: return F.when( @@ -674,6 +698,64 @@ def quantile(self, quantile: float, accuracy: int = 10000) -> FrameLike: """ return super().quantile(quantile, accuracy) + def median(self) -> FrameLike: + """ + Calculate the rolling median. + + .. versionadded:: 4.3.0 + + .. note:: the current implementation of this API uses Spark's Window without + specifying partition specification. This leads to move all data into + single partition in single machine and could cause serious + performance degradation. Avoid this method against very large dataset. + + Returns + ------- + Series or DataFrame + Returned object type is determined by the caller of the rolling + calculation. + + See Also + -------- + pyspark.pandas.Series.rolling : Calling object with Series data. + pyspark.pandas.DataFrame.rolling : Calling object with DataFrames. + pyspark.pandas.Series.median : Equivalent method for Series. + pyspark.pandas.DataFrame.median : Equivalent method for DataFrame. + + Notes + ----- + `median` in pandas-on-Spark uses distributed percentile approximation + algorithm unlike pandas, the result might be different with pandas. + + Examples + -------- + >>> s = ps.Series([4, 3, 5, 2, 6]) + >>> s + 0 4 + 1 3 + 2 5 + 3 2 + 4 6 + dtype: int64 + + >>> s.rolling(2).median() + 0 NaN + 1 3.0 + 2 3.0 + 3 2.0 + 4 2.0 + dtype: float64 + + >>> s.rolling(3).median() + 0 NaN + 1 NaN + 2 4.0 + 3 3.0 + 4 5.0 + dtype: float64 + """ + return super().median() + def std(self) -> FrameLike: """ Calculate rolling standard deviation. @@ -774,6 +856,63 @@ def var(self) -> FrameLike: """ return super().var() + def sem(self, ddof: int = 1) -> FrameLike: + """ + Calculate the rolling standard error of the mean. + + .. versionadded:: 4.3.0 + + .. note:: the current implementation of this API uses Spark's Window without + specifying partition specification. This leads to move all data into + single partition in single machine and could cause serious + performance degradation. Avoid this method against very large dataset. + + Parameters + ---------- + ddof : int, default 1 + Delta Degrees of Freedom. The divisor used in calculations is N - ddof, + where N represents the number of elements. + + Returns + ------- + Series or DataFrame + Returns the same object type as the caller of the rolling calculation. + + See Also + -------- + pyspark.pandas.Series.rolling : Calling object with Series data. + pyspark.pandas.DataFrame.rolling : Calling object with DataFrames. + pyspark.pandas.Series.sem : Equivalent method for Series. + pyspark.pandas.DataFrame.sem : Equivalent method for DataFrame. + + Examples + -------- + >>> s = ps.Series([5, 5, 6, 7, 5, 5, 5]) + >>> s.rolling(3).sem() + 0 NaN + 1 NaN + 2 0.408248 + 3 0.707107 + 4 0.707107 + 5 0.816497 + 6 0.000000 + dtype: float64 + + For DataFrame, each rolling standard error of the mean is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.rolling(2).sem() + A B + 0 NaN NaN + 1 0.000000 0.000000 + 2 0.707107 7.778175 + 3 0.707107 9.192388 + 4 1.414214 16.970563 + 5 0.000000 0.000000 + 6 0.000000 0.000000 + """ + return super().sem(ddof) + def skew(self) -> FrameLike: """ Calculate unbiased rolling skew. @@ -1320,6 +1459,67 @@ def quantile(self, quantile: float, accuracy: int = 10000) -> FrameLike: """ return super().quantile(quantile, accuracy) + def median(self) -> FrameLike: + """ + Calculate the rolling median. + + .. versionadded:: 4.3.0 + + Returns + ------- + Series or DataFrame + Returned object type is determined by the caller of the rolling + calculation. + + See Also + -------- + pyspark.pandas.Series.rolling : Calling object with Series data. + pyspark.pandas.DataFrame.rolling : Calling object with DataFrames. + pyspark.pandas.Series.median : Equivalent method for Series. + pyspark.pandas.DataFrame.median : Equivalent method for DataFrame. + + Notes + ----- + `median` in pandas-on-Spark uses distributed percentile approximation + algorithm unlike pandas, the result might be different with pandas. + + Examples + -------- + >>> s = ps.Series([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5]) + >>> s.groupby(s).rolling(3).median().sort_index() + 2 0 NaN + 1 NaN + 3 2 NaN + 3 NaN + 4 3.0 + 4 5 NaN + 6 NaN + 7 4.0 + 8 4.0 + 5 9 NaN + 10 NaN + dtype: float64 + + For DataFrame, each rolling median is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.groupby(df.A).rolling(2).median().sort_index() + B + A + 2 0 NaN + 1 4.0 + 3 2 NaN + 3 9.0 + 4 9.0 + 4 5 NaN + 6 16.0 + 7 16.0 + 8 16.0 + 5 9 NaN + 10 25.0 + """ + return super().median() + def std(self) -> FrameLike: """ Calculate rolling standard deviation. @@ -1727,6 +1927,64 @@ def quantile(self, quantile: float, accuracy: int = 10000) -> FrameLike: """ return super().quantile(quantile, accuracy) + def median(self) -> FrameLike: + """ + Calculate the expanding median. + + .. versionadded:: 4.3.0 + + .. note:: the current implementation of this API uses Spark's Window without + specifying partition specification. This leads to move all data into + single partition in single machine and could cause serious + performance degradation. Avoid this method against very large dataset. + + Returns + ------- + Series or DataFrame + Returned object type is determined by the caller of the expanding + calculation. + + See Also + -------- + pyspark.pandas.Series.expanding : Calling object with Series data. + pyspark.pandas.DataFrame.expanding : Calling object with DataFrames. + pyspark.pandas.Series.median : Equivalent method for Series. + pyspark.pandas.DataFrame.median : Equivalent method for DataFrame. + + Notes + ----- + `median` in pandas-on-Spark uses distributed percentile approximation + algorithm unlike pandas, the result might be different with pandas. + + Examples + -------- + >>> s = ps.Series([4, 3, 5, 2, 6]) + >>> s + 0 4 + 1 3 + 2 5 + 3 2 + 4 6 + dtype: int64 + + >>> s.expanding(2).median() + 0 NaN + 1 3.0 + 2 4.0 + 3 3.0 + 4 4.0 + dtype: float64 + + >>> s.expanding(3).median() + 0 NaN + 1 NaN + 2 4.0 + 3 3.0 + 4 4.0 + dtype: float64 + """ + return super().median() + def std(self) -> FrameLike: """ Calculate expanding standard deviation. @@ -1827,6 +2085,63 @@ def var(self) -> FrameLike: """ return super().var() + def sem(self, ddof: int = 1) -> FrameLike: + """ + Calculate the expanding standard error of the mean. + + .. versionadded:: 4.3.0 + + .. note:: the current implementation of this API uses Spark's Window without + specifying partition specification. This leads to move all data into + single partition in single machine and could cause serious + performance degradation. Avoid this method against very large dataset. + + Parameters + ---------- + ddof : int, default 1 + Delta Degrees of Freedom. The divisor used in calculations is N - ddof, + where N represents the number of elements. + + Returns + ------- + Series or DataFrame + Returns the same object type as the caller of the expanding calculation. + + See Also + -------- + pyspark.pandas.Series.expanding : Calling object with Series data. + pyspark.pandas.DataFrame.expanding : Calling object with DataFrames. + pyspark.pandas.Series.sem : Equivalent method for Series. + pyspark.pandas.DataFrame.sem : Equivalent method for DataFrame. + + Examples + -------- + >>> s = ps.Series([5, 5, 6, 7, 5, 5, 5]) + >>> s.expanding(3).sem() + 0 NaN + 1 NaN + 2 0.408248 + 3 0.552771 + 4 0.447214 + 5 0.374166 + 6 0.321208 + dtype: float64 + + For DataFrame, each expanding standard error of the mean is computed column-wise. + + >>> df = ps.DataFrame({"A": s.to_numpy(), "B": s.to_numpy() ** 2}) + >>> df.expanding(2).sem() + A B + 0 NaN NaN + 1 0.000000 0.000000 + 2 0.408248 4.490731 + 3 0.552771 6.589132 + 4 0.447214 5.315073 + 5 0.374166 4.439970 + 6 0.321208 3.807887 + """ + return super().sem(ddof) + def skew(self) -> FrameLike: """ Calculate unbiased expanding skew. @@ -2293,6 +2608,49 @@ def quantile(self, quantile: float, accuracy: int = 10000) -> FrameLike: """ return super().quantile(quantile, accuracy) + def median(self) -> FrameLike: + """ + Calculate the expanding median. + + .. versionadded:: 4.3.0 + + Returns + ------- + Series or DataFrame + Returned object type is determined by the caller of the expanding + calculation. + + See Also + -------- + pyspark.pandas.Series.expanding : Calling expanding with Series data. + pyspark.pandas.DataFrame.expanding : Calling expanding with DataFrames. + pyspark.pandas.Series.median : Equivalent method for Series. + pyspark.pandas.DataFrame.median : Equivalent method for DataFrame. + + Notes + ----- + `median` in pandas-on-Spark uses distributed percentile approximation + algorithm unlike pandas, the result might be different with pandas. + + Examples + -------- + >>> s = ps.Series([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5]) + >>> s.groupby(s).expanding(3).median().sort_index() + 2 0 NaN + 1 NaN + 3 2 NaN + 3 NaN + 4 3.0 + 4 5 NaN + 6 NaN + 7 4.0 + 8 4.0 + 5 9 NaN + 10 NaN + dtype: float64 + """ + return super().median() + def std(self) -> FrameLike: """ Calculate expanding standard deviation. diff --git a/python/pyspark/pipelines/flow.py b/python/pyspark/pipelines/flow.py index b1922454..237bcf14 100644 --- a/python/pyspark/pipelines/flow.py +++ b/python/pyspark/pipelines/flow.py @@ -59,8 +59,14 @@ class AutoCdcFlow: :param apply_as_deletes: Optional delete condition for the merge operation. :param column_list: Optional columns to include in the output table. :param except_column_list: Optional columns to exclude from the output table. - :param stored_as_scd_type: Optional SCD type for the target table. Only 1 (or "1") is \ - supported. + :param stored_as_scd_type: Optional SCD type for the target table. 1 (or "1") and 2 (or "2") \ + are supported. + :param track_history_column_list: Optional SCD2-only columns whose value change opens a new \ + history record. + :param track_history_except_column_list: Optional SCD2-only columns excluded from history \ + tracking. + :param spark_conf: A dict where the keys are the Spark configuration property names and the + values are the property values. These properties will be set on the flow. :param source_code_location: The location of the source code that created this flow. """ @@ -72,5 +78,8 @@ class AutoCdcFlow: apply_as_deletes: Optional[Column] column_list: Optional[List[Column]] except_column_list: Optional[List[Column]] - stored_as_scd_type: Optional[Literal[1, "1"]] + stored_as_scd_type: Optional[Literal[1, 2, "1", "2"]] + track_history_column_list: Optional[List[Column]] + track_history_except_column_list: Optional[List[Column]] + spark_conf: Dict[str, str] source_code_location: SourceCodeLocation diff --git a/python/pyspark/pipelines/tests/test_graph_element_registry.py b/python/pyspark/pipelines/tests/test_graph_element_registry.py index fd8ed439..1e6fcf22 100644 --- a/python/pyspark/pipelines/tests/test_graph_element_registry.py +++ b/python/pyspark/pipelines/tests/test_graph_element_registry.py @@ -17,14 +17,11 @@ import unittest -from pyspark.errors import PySparkException, PySparkTypeError +from pyspark.errors import PySparkException from pyspark.pipelines.graph_element_registry import graph_element_registration_context from pyspark import pipelines as dp -from pyspark.pipelines.flow import AutoCdcFlow from pyspark.pipelines.output import Sink from pyspark.pipelines.tests.local_graph_element_registry import LocalGraphElementRegistry -from pyspark.sql import Column -from pyspark.sql.connect.functions.builtin import col, expr from typing import cast @@ -100,147 +97,6 @@ def flow2(): self.assertEqual(sink_obj.options["key1"], "value1") assert sink_obj.source_code_location.filename.endswith("test_graph_element_registry.py") - def test_create_auto_cdc_flow(self): - registry = LocalGraphElementRegistry() - with graph_element_registration_context(registry): - dp.create_streaming_table("target") - dp.create_auto_cdc_flow( - target="target", - source="source", - keys=[col("key")], - sequence_by=expr("seq"), - ) - - self.assertEqual(len(registry.outputs), 1) - self.assertEqual(len(registry.auto_cdc_flows), 1) - - flow = cast(AutoCdcFlow, registry.auto_cdc_flows[0]) - self.assertEqual(flow.target, "target") - self.assertEqual(flow.source, "source") - - # When name is not specified, it inherits the target's name at construction time. - self.assertEqual(flow.name, "target") - self.assertIsNone(flow.stored_as_scd_type) - self.assertIsNone(flow.apply_as_deletes) - assert flow.source_code_location.filename.endswith("test_graph_element_registry.py") - - def test_create_auto_cdc_flow_with_all_args(self): - registry = LocalGraphElementRegistry() - with graph_element_registration_context(registry): - dp.create_streaming_table("tgt") - dp.create_auto_cdc_flow( - target="tgt", - source="src", - keys=[col("id")], - sequence_by=expr("ts"), - apply_as_deletes=expr("op = 'DELETE'"), - column_list=[col("id"), col("val")], - stored_as_scd_type=1, - name="my_flow", - ) - - flow = cast(AutoCdcFlow, registry.auto_cdc_flows[0]) - self.assertEqual(flow.name, "my_flow") - self.assertEqual(flow.stored_as_scd_type, 1) - - def test_create_auto_cdc_flow_with_string_args(self): - # Verify that string forms of column / expression arguments are normalized to - # PySpark Columns, equivalent to passing col(...) / expr(...) directly. - registry = LocalGraphElementRegistry() - with graph_element_registration_context(registry): - dp.create_streaming_table("tgt") - dp.create_auto_cdc_flow( - target="tgt", - source="src", - keys=["id"], - sequence_by="ts", - apply_as_deletes="op = 'DELETE'", - column_list=["id", "val"], - ) - - flow = cast(AutoCdcFlow, registry.auto_cdc_flows[0]) - for k in flow.keys: - self.assertIsInstance(k, Column) - self.assertIsInstance(flow.sequence_by, Column) - self.assertIsInstance(flow.apply_as_deletes, Column) - assert flow.column_list is not None - for c in flow.column_list: - self.assertIsInstance(c, Column) - - def test_create_auto_cdc_flow_stored_as_scd_type_string(self): - registry = LocalGraphElementRegistry() - with graph_element_registration_context(registry): - dp.create_streaming_table("t") - dp.create_auto_cdc_flow( - target="t", - source="s", - keys=[col("k")], - sequence_by=expr("seq"), - stored_as_scd_type="1", - ) - - flow = cast(AutoCdcFlow, registry.auto_cdc_flows[0]) - self.assertEqual(flow.stored_as_scd_type, "1") - - def test_create_auto_cdc_flow_invalid_scd_type(self): - registry = LocalGraphElementRegistry() - with graph_element_registration_context(registry): - dp.create_streaming_table("t") - with self.assertRaises(PySparkTypeError) as ctx: - dp.create_auto_cdc_flow( - target="t", - source="s", - keys=[col("k")], - sequence_by=expr("seq"), - stored_as_scd_type=2, # type: ignore[arg-type] - ) - self.assertEqual(ctx.exception.getCondition(), "NOT_EXPECTED_TYPE") - - def test_create_auto_cdc_flow_with_except_column_list(self): - registry = LocalGraphElementRegistry() - with graph_element_registration_context(registry): - dp.create_streaming_table("tgt") - dp.create_auto_cdc_flow( - target="tgt", - source="src", - keys=[col("id")], - sequence_by=expr("ts"), - except_column_list=["op", "ts"], - ) - - flow = cast(AutoCdcFlow, registry.auto_cdc_flows[0]) - self.assertIsNone(flow.column_list) - assert flow.except_column_list is not None - self.assertEqual(len(flow.except_column_list), 2) - for c in flow.except_column_list: - self.assertIsInstance(c, Column) - - def test_create_auto_cdc_flow_rejects_non_str_target(self): - registry = LocalGraphElementRegistry() - with graph_element_registration_context(registry): - dp.create_streaming_table("tgt") - with self.assertRaises(PySparkTypeError) as ctx: - dp.create_auto_cdc_flow( - target=123, # type: ignore[arg-type] - source="src", - keys=[col("id")], - sequence_by=expr("ts"), - ) - self.assertEqual(ctx.exception.getCondition(), "NOT_EXPECTED_TYPE") - - def test_create_auto_cdc_flow_rejects_invalid_key_element(self): - registry = LocalGraphElementRegistry() - with graph_element_registration_context(registry): - dp.create_streaming_table("tgt") - with self.assertRaises(PySparkTypeError) as ctx: - dp.create_auto_cdc_flow( - target="tgt", - source="src", - keys=[123], # type: ignore[list-item] - sequence_by=expr("ts"), - ) - self.assertEqual(ctx.exception.getCondition(), "NOT_EXPECTED_TYPE") - def test_definition_without_graph_element_registry(self): for decorator in [dp.table, dp.temporary_view, dp.materialized_view]: with self.assertRaises(PySparkException) as context: @@ -273,19 +129,6 @@ def b(): "GRAPH_ELEMENT_DEFINED_OUTSIDE_OF_DECLARATIVE_PIPELINE", ) - with self.assertRaises(PySparkException) as context: - dp.create_auto_cdc_flow( - target="t", - source="s", - keys=["k"], - sequence_by="seq", - ) - - self.assertEqual( - context.exception.getCondition(), - "GRAPH_ELEMENT_DEFINED_OUTSIDE_OF_DECLARATIVE_PIPELINE", - ) - if __name__ == "__main__": from pyspark.testing import main diff --git a/python/pyspark/profiler.py b/python/pyspark/profiler.py index 8f797b5f..45ab1e5a 100644 --- a/python/pyspark/profiler.py +++ b/python/pyspark/profiler.py @@ -410,7 +410,7 @@ def dump(self, id: int, path: str) -> None: stats = self.stats() # dict if stats: p = os.path.join(path, "udf_%d_memory.txt" % id) - with open(p, "w+") as f: + with open(p, "w+", encoding="utf-8") as f: self._show_results(stats, stream=f) diff --git a/python/pyspark/serializers.py b/python/pyspark/serializers.py index 6de64a10..1307b695 100644 --- a/python/pyspark/serializers.py +++ b/python/pyspark/serializers.py @@ -55,6 +55,7 @@ import sys import os +import io from itertools import chain, product import marshal import struct @@ -416,6 +417,17 @@ def namedtuple(*args, **kwargs): _hijack_namedtuple() +class _RestrictedUnpickler(pickle.Unpickler): + def __init__(self, file, allowed_names, *, encoding="bytes"): + super().__init__(file, encoding=encoding) + self.allowed_names = allowed_names + + def find_class(self, module, name): + if (module, name) not in self.allowed_names: + raise pickle.UnpicklingError(f"Unpickling {module}.{name} is not allowed") + return super().find_class(module, name) + + class PickleSerializer(FramedSerializer): """ Serializes objects using Python's pickle serializer: @@ -434,6 +446,18 @@ def loads(self, obj, encoding="bytes"): class CloudPickleSerializer(FramedSerializer): + # Class-level default for `allowed_names`. Instances of this serializer are pickled and + # shipped across the driver and workers, which may run different Spark versions. Pickle + # restores an instance via __new__ + __dict__ without calling __init__, so an instance + # pickled by a version that predates `allowed_names` would otherwise lack the attribute + # entirely. Declaring the default here keeps `self.allowed_names` resolvable in that + # cross-version case; only a version-matched, explicitly restricted instance sets it. + allowed_names = None + + def __init__(self, allowed_names=None): + super().__init__() + self.allowed_names = allowed_names + def dumps(self, obj): from pyspark.util import print_exec @@ -451,6 +475,9 @@ def dumps(self, obj): raise pickle.PicklingError(msg) def loads(self, obj, encoding="bytes"): + if self.allowed_names is not None: + unpickler = _RestrictedUnpickler(io.BytesIO(obj), self.allowed_names, encoding=encoding) + return unpickler.load() return cloudpickle.loads(obj, encoding=encoding) diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py index 0550ff4b..20d37def 100644 --- a/python/pyspark/sql/conversion.py +++ b/python/pyspark/sql/conversion.py @@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Any, Callable, List, Optional, Sequence, Union, overload import pyspark -from pyspark.errors import PySparkTypeError, PySparkValueError +from pyspark.errors import PySparkNotImplementedError, PySparkRuntimeError, PySparkValueError from pyspark.sql.pandas.types import ( _dedup_names, _deduplicate_field_names, @@ -63,6 +63,7 @@ VariantType, VariantVal, _create_row, + _has_type, ) if TYPE_CHECKING: @@ -82,7 +83,6 @@ def flatten_struct(batch: "pa.RecordBatch", column_index: int = 0) -> "pa.Record Flatten a struct column at given index into a RecordBatch. Used by: - - ArrowStreamUDFSerializer.load_stream - SQL_GROUPED_MAP_ARROW_UDF mapper - SQL_GROUPED_MAP_ARROW_ITER_UDF mapper """ @@ -91,12 +91,27 @@ def flatten_struct(batch: "pa.RecordBatch", column_index: int = 0) -> "pa.Record struct = batch.column(column_index) return pa.RecordBatch.from_arrays(struct.flatten(), schema=pa.schema(struct.type)) + @classmethod + def select_columns(cls, batch: "pa.RecordBatch", column_indices: list[int]) -> "pa.RecordBatch": + """ + Select a subset of columns from a RecordBatch by index. + + Used by: SQL_COGROUPED_MAP_ARROW_UDF handler in worker.py + """ + import pyarrow as pa + + return pa.RecordBatch.from_arrays( + [batch.columns[i] for i in column_indices], + [batch.schema.names[i] for i in column_indices], + ) + @staticmethod def wrap_struct(batch: "pa.RecordBatch") -> "pa.RecordBatch": """ Wrap a RecordBatch's columns into a single struct column. - Used by: ArrowStreamUDFSerializer.dump_stream + Used by: Arrow UDF mappers in worker.py to re-wrap flattened batches + before serialization. """ import pyarrow as pa @@ -111,19 +126,20 @@ def wrap_struct(batch: "pa.RecordBatch") -> "pa.RecordBatch": @classmethod def enforce_schema( cls, - batch: "pa.RecordBatch", + batch: Union["pa.RecordBatch", "pa.Table"], arrow_schema: "pa.Schema", *, arrow_cast: bool = True, safecheck: bool = True, - ) -> "pa.RecordBatch": + reorder_by_name: bool = True, + ) -> Union["pa.RecordBatch", "pa.Table"]: """ - Enforce target schema on a RecordBatch by reordering columns and coercing types. + Enforce a target schema on an Arrow RecordBatch or Table. Parameters ---------- - batch : pa.RecordBatch - Input RecordBatch to transform. + batch : pa.RecordBatch or pa.Table + Input to transform. Output is of the same container type. arrow_schema : pa.Schema Target Arrow schema. Callers should pre-compute this once via to_arrow_schema() to avoid repeated conversion. @@ -132,11 +148,28 @@ def enforce_schema( If False, raise an error on type mismatch instead of casting. safecheck : bool, default True If True, use safe casting (fails on overflow/truncation). + reorder_by_name : bool, default True + If True, match columns by name and reorder to the target order; any + missing or extra names raise ``RESULT_COLUMN_NAMES_MISMATCH``. Output + columns are renamed to target names. + If False, match columns by position (ignore names) and preserve the + original column names in the output. Returns ------- - pa.RecordBatch - RecordBatch with columns reordered and types coerced to match target schema. + pa.RecordBatch or pa.Table + Same container type as ``batch``, with columns matched (and possibly + reordered/cast) per the target schema. + + Raises + ------ + PySparkRuntimeError + ``RESULT_COLUMN_NAMES_MISMATCH`` when ``reorder_by_name=True`` and the + batch has missing or extra column names. + ``RESULT_COLUMN_TYPES_MISMATCH`` when any column's type does not match + the target (and either ``arrow_cast=False`` or the cast itself fails). + ``RESULT_COLUMN_SCHEMA_MISMATCH`` when ``reorder_by_name=False`` and the + batch has a different number of columns than the target schema. """ import pyarrow as pa @@ -147,37 +180,68 @@ def enforce_schema( if batch.schema.equals(arrow_schema, check_metadata=False): return batch - # Check if columns are in the same order (by name) as the target schema. - # If so, use index-based access (faster than name lookup). - batch_names = [batch.schema.field(i).name for i in range(batch.num_columns)] target_names = [field.name for field in arrow_schema] - use_index = batch_names == target_names - coerced_arrays = [] - for i, field in enumerate(arrow_schema): - try: - arr = batch.column(i) if use_index else batch.column(field.name) - except KeyError: - raise PySparkTypeError( - f"Result column '{field.name}' does not exist in the output. " - f"Expected schema: {arrow_schema}, got: {batch.schema}." + # Step 1: pick source columns from batch to align with target schema + if reorder_by_name: + batch_names = [batch.schema.field(i).name for i in range(batch.num_columns)] + missing = sorted(set(target_names) - set(batch_names)) + extra = sorted(set(batch_names) - set(target_names)) + if missing or extra: + raise PySparkRuntimeError( + errorClass="RESULT_COLUMN_NAMES_MISMATCH", + messageParameters={ + "missing": f" Missing: {', '.join(missing)}." if missing else "", + "extra": f" Unexpected: {', '.join(extra)}." if extra else "", + }, ) - if arr.type != field.type: - if not arrow_cast: - raise PySparkTypeError( - f"Result type of column '{field.name}' does not match " - f"the expected type. Expected: {field.type}, got: {arr.type}." - ) + source_columns = [batch.column(name) for name in target_names] + output_names = target_names + else: + # Positional: require exact column-count match, then take columns by + # index, preserving the batch's original column names. + if batch.num_columns != len(arrow_schema): + raise PySparkRuntimeError( + errorClass="RESULT_COLUMN_SCHEMA_MISMATCH", + messageParameters={ + "expected": str(len(arrow_schema)), + "actual": str(batch.num_columns), + }, + ) + source_columns = [batch.column(i) for i in range(len(arrow_schema))] + output_names = [batch.schema.field(i).name for i in range(len(arrow_schema))] + + # Step 2: check types / cast, collect all mismatches + type_mismatches = [] + coerced_arrays = [] + for field, arr in zip(arrow_schema, source_columns): + if arr.type == field.type: + coerced_arrays.append(arr) + elif not arrow_cast: + type_mismatches.append((field.name, field.type, arr.type)) + coerced_arrays.append(arr) + else: try: - arr = arr.cast(target_type=field.type, safe=safecheck) - except (pa.ArrowInvalid, pa.ArrowTypeError) as e: - raise PySparkTypeError( - f"Result type of column '{field.name}' does not match " - f"the expected type. Expected: {field.type}, got: {arr.type}." - ) from e - coerced_arrays.append(arr) + coerced_arrays.append(arr.cast(target_type=field.type, safe=safecheck)) + except (pa.ArrowInvalid, pa.ArrowTypeError): + type_mismatches.append((field.name, field.type, arr.type)) + coerced_arrays.append(arr) + + if type_mismatches: + raise PySparkRuntimeError( + errorClass="RESULT_COLUMN_TYPES_MISMATCH", + messageParameters={ + "mismatch": ", ".join( + f"column '{name}' (expected {expected}, actual {actual})" + for name, expected, actual in type_mismatches + ) + }, + ) - return pa.RecordBatch.from_arrays(coerced_arrays, names=target_names) + # Preserve input container type (Table vs RecordBatch) + if isinstance(batch, pa.Table): + return pa.Table.from_arrays(coerced_arrays, names=output_names) + return pa.RecordBatch.from_arrays(coerced_arrays, names=output_names) @classmethod def to_pandas( @@ -618,6 +682,27 @@ def convert_array(value: Any) -> Any: assert isinstance(value, (list, array.array)) return list(value) + elif isinstance(dataType.elementType, (StringType, BinaryType)): + # Inline the scalar identity fast path so elements that are + # already the target Python type skip the per-element converter + # call entirely: `convert_string`/`convert_binary` return such + # elements unchanged. `str` and immutable `bytes` are the two + # element types whose converter is a no-op on a matching value. + # Any other element -- including `None` (whose nullability is + # enforced by `element_conv`) and values that need coercion (e.g. + # a bool to string) -- falls back to `element_conv`, reused + # unchanged. + fast_type = str if isinstance(dataType.elementType, StringType) else bytes + + def convert_array(value: Any) -> Any: + if value is None: + if not nullable: + raise PySparkValueError(f"input for {dataType} must not be None") + return None + else: + assert isinstance(value, (list, array.array)) + return [v if type(v) is fast_type else element_conv(v) for v in value] + else: def convert_array(value: Any) -> Any: @@ -675,6 +760,12 @@ def convert_binary(value: Any) -> Any: if not nullable: raise PySparkValueError(f"input for {dataType} must not be None") return None + elif type(value) is bytes: + # Fast path: `bytes(value)` returns `value` itself for a `bytes` + # input (no copy, as `bytes` is immutable), but still pays the + # constructor dispatch per element. Returning it directly skips + # that. `bytearray` falls through and is copied into `bytes`. + return value else: assert isinstance(value, (bytes, bytearray)) return bytes(value) @@ -737,13 +828,19 @@ def convert_string(value: Any) -> Any: if not nullable: raise PySparkValueError(f"input for {dataType} must not be None") return None + elif type(value) is str: + # Fast path: `str(value)` returns `value` itself for a `str` + # input (no copy), but still pays the constructor dispatch per + # element. Returning it directly skips that and the bool check. + return value + elif value is True: + # To match the PySpark Classic which convert bool to string in + # the JVM side (python.EvaluatePython.makeFromJava) + return "true" + elif value is False: + return "false" else: - if isinstance(value, bool): - # To match the PySpark Classic which convert bool to string in - # the JVM side (python.EvaluatePython.makeFromJava) - return str(value).lower() - else: - return str(value) + return str(value) return convert_string @@ -1356,6 +1453,24 @@ def convert( assert schema is not None and isinstance(schema, StructType) + # YearMonthIntervalType is serialized by the JVM as an Arrow YEAR_MONTH interval, which + # PyArrow cannot materialize into Python values: `to_pylist()` raises an opaque + # `KeyError: ` from `get_array_class_from_type`. That lookup fails for an + # empty column too (it resolves the array class before reading any element), so the check + # below is intentionally unconditional in the row count -- it covers empty results as well, + # surfacing a clean NOT_IMPLEMENTED instead of the opaque KeyError. Collecting such a value + # is therefore not supported in the Spark Connect client; raise the same NOT_IMPLEMENTED + # error as the classic PySpark path (YearMonthIntervalType.fromInternal). Note that, unlike + # classic, PYSPARK_YM_INTERVAL_LEGACY (returning the integer months) cannot be honored here, + # and an empty result raises rather than returning [] as classic would. + if any(_has_type(f.dataType, YearMonthIntervalType) for f in schema.fields): + raise PySparkNotImplementedError( + errorClass="NOT_IMPLEMENTED", + messageParameters={ + "feature": "Collecting a year-month interval value in Spark Connect" + }, + ) + fields = schema.fieldNames() if len(fields) > 0: diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py index 47d56448..cf847aa2 100644 --- a/python/pyspark/sql/functions.py +++ b/python/pyspark/sql/functions.py @@ -22,6 +22,7 @@ _call_function = _functions.pyfunc_call_function # Mixed function bindings _pyfunc_sha2 = _functions.pyfunc_sha2 +_pyfunc_variant_strip_nulls = _functions.pyfunc_variant_strip_nulls _pyfunc_window = _functions.pyfunc_window _pyfunc_window_with_slide_and_start = _functions.pyfunc_window_with_slide_and_start _pyfunc_from_avro = _functions.pyfunc_from_avro @@ -202,7 +203,10 @@ def _wrap(obj): 'translate': ((0,), None), 'trunc': ((0,), None), 'try_to_date': ((0,), None), + 'try_variant_array_append': ((0, 2), None), 'try_variant_get': ((0,), None), + 'try_variant_insert': ((0, 2), None), + 'try_variant_set': ((0, 2), None), 'tuple_intersection_agg_double': ((0,), None), 'tuple_intersection_agg_integer': ((0,), None), 'tuple_intersection_double': ((0, 1), None), @@ -223,7 +227,11 @@ def _wrap(obj): 'unix_timestamp': ((0,), None), 'user': ((), None), 'uuid': ((), None), + 'variant_array_append': ((0, 2), None), + 'variant_delete': ((0,), None), 'variant_get': ((0,), None), + 'variant_insert': ((0, 2), None), + 'variant_set': ((0, 2), None), 'version': ((), None), 'when': ((), None), } @@ -347,6 +355,28 @@ def sha2(col, numBits): """ return _pyfunc_sha2(_to_col(col), numBits) + +def variant_strip_nulls(v, include_arrays=True): + """Recursively removes variant-null object fields and array elements. + + Args: + v: a variant column + include_arrays: also strip nulls inside arrays (default True) + """ + return _pyfunc_variant_strip_nulls(_to_col(v), include_arrays) + + +def variant_delete(v, *paths): + """Deletes the value(s) at the given path(s) from a variant column. + + Args: + v: a variant column + *paths: one or more path expressions (Column or str) + """ + _a = [v] + _a.extend(paths) + return _dispatch("variant_delete", _a) + def window(timeColumn, windowDuration, slideDuration=None, startTime=None): """Buckets rows into one or more time windows specified by the given parameters. @@ -666,6 +696,27 @@ def broadcast(df): ] _sys.modules["pyspark.sql.functions.partitioning"] = partitioning +# Public classes/enums that pyspark.sql.functions re-exports (so `from +# pyspark.sql.functions import PandasUDFType` etc. work, matching reference pyspark). +# Guarded so a partial import never breaks the functions module. +try: + from pyspark.sql.udf import UserDefinedFunction # noqa: F401 +except Exception: # pragma: no cover + pass +try: + from pyspark.sql.udtf import ( # noqa: F401 + UserDefinedTableFunction, + AnalyzeArgument, + AnalyzeResult, + SkipRestOfInputTableException, + ) +except Exception: # pragma: no cover + pass +try: + from pyspark.sql.pandas.functions import PandasUDFType, ArrowUDFType # noqa: F401 +except Exception: # pragma: no cover + pass + # Build __all__ with all function names __all__ = [ "col", diff --git a/python/pyspark/sql/functions_generated.py b/python/pyspark/sql/functions_generated.py index e983a57c..f96110ff 100644 --- a/python/pyspark/sql/functions_generated.py +++ b/python/pyspark/sql/functions_generated.py @@ -345,6 +345,10 @@ def collect_set(col): _a = [] _a.append(col) return _dispatch("collect_set", _a) +def collect_union(col): + _a = [] + _a.append(col) + return _dispatch("collect_union", _a) def concat(*cols): _a = [] _a.extend(cols) @@ -409,7 +413,11 @@ def count_min_sketch(col, eps, confidence, seed=_UNSET): _a.append(confidence) if seed is not _UNSET: _a.append(seed) return _dispatch("count_min_sketch", _a) -counter_diff = _create_wrapper("counter_diff") +def counter_diff(value, startTime=_UNSET): + _a = [] + _a.append(value) + if startTime is not _UNSET: _a.append(startTime) + return _dispatch("counter_diff", _a) def covar_pop(col1, col2): _a = [] _a.append(col1) @@ -650,6 +658,10 @@ def format_string(format, *cols): _a.append(format) _a.extend(cols) return _dispatch("format_string", _a) +def from_base32(col): + _a = [] + _a.append(col) + return _dispatch("from_base32", _a) def from_csv(col, schema, options=_UNSET): _a = [] _a.append(col) @@ -738,7 +750,12 @@ def hll_union_agg(col, allowDifferentLgConfigK=_UNSET): _a.append(col) if allowDifferentLgConfigK is not _UNSET: _a.append(allowDifferentLgConfigK) return _dispatch("hll_union_agg", _a) -hmac = _create_wrapper("hmac") +def hmac(key, message, algorithm=_UNSET): + _a = [] + _a.append(key) + _a.append(message) + if algorithm is not _UNSET: _a.append(algorithm) + return _dispatch("hmac", _a) def hour(col): _a = [] _a.append(col) @@ -784,10 +801,12 @@ def input_file_block_start(): def input_file_name(): _a = [] return _dispatch("input_file_name", _a) -def instr(str, substr): +def instr(str, substr, start=_UNSET, occurrence=_UNSET): _a = [] _a.append(str) _a.append(substr) + if start is not _UNSET: _a.append(start) + if occurrence is not _UNSET: _a.append(occurrence) return _dispatch("instr", _a) def is_valid_utf8(str): _a = [] @@ -813,7 +832,11 @@ def isnull(col): _a = [] _a.append(col) return _dispatch("isnull", _a) -jaro_winkler_similarity = _create_wrapper("jaro_winkler_similarity") +def jaro_winkler_similarity(left, right): + _a = [] + _a.append(left) + _a.append(right) + return _dispatch("jaro_winkler_similarity", _a) def java_method(*cols): _a = [] _a.extend(cols) @@ -1412,11 +1435,12 @@ def regexp_like(str, regexp): _a.append(str) _a.append(regexp) return _dispatch("regexp_like", _a) -def regexp_replace(string, pattern, replacement): +def regexp_replace(string, pattern, replacement, position=_UNSET): _a = [] _a.append(string) _a.append(pattern) _a.append(replacement) + if position is not _UNSET: _a.append(position) return _dispatch("regexp_replace", _a) def regexp_substr(str, regexp): _a = [] @@ -1869,7 +1893,10 @@ def timestamp_millis(col): _a = [] _a.append(col) return _dispatch("timestamp_millis", _a) -timestamp_nanos = _create_wrapper("timestamp_nanos") +def timestamp_nanos(col): + _a = [] + _a.append(col) + return _dispatch("timestamp_nanos", _a) def timestamp_seconds(col): _a = [] _a.append(col) @@ -1883,6 +1910,10 @@ def toRadians(col): _a.append(col) return _dispatch("toRadians", _a) to_avro = _create_wrapper("to_avro") +def to_base32(col): + _a = [] + _a.append(col) + return _dispatch("to_base32", _a) def to_binary(col, format=_UNSET): _a = [] _a.append(col) @@ -2109,15 +2140,31 @@ def try_validate_utf8(str): _a = [] _a.append(str) return _dispatch("try_validate_utf8", _a) -try_variant_array_append = _create_wrapper("try_variant_array_append") +def try_variant_array_append(v, path, value): + _a = [] + _a.append(v) + _a.append(path) + _a.append(value) + return _dispatch("try_variant_array_append", _a) def try_variant_get(v, path, targetType): _a = [] _a.append(v) _a.append(path) _a.append(targetType) return _dispatch("try_variant_get", _a) -try_variant_insert = _create_wrapper("try_variant_insert") -try_variant_set = _create_wrapper("try_variant_set") +def try_variant_insert(v, path, value): + _a = [] + _a.append(v) + _a.append(path) + _a.append(value) + return _dispatch("try_variant_insert", _a) +def try_variant_set(v, path, value, create_if_missing=_UNSET): + _a = [] + _a.append(v) + _a.append(path) + _a.append(value) + if create_if_missing is not _UNSET: _a.append(create_if_missing) + return _dispatch("try_variant_set", _a) def tuple_difference_double(col1, col2): _a = [] _a.append(col1) @@ -2286,7 +2333,10 @@ def unix_millis(col): _a = [] _a.append(col) return _dispatch("unix_millis", _a) -unix_nanos = _create_wrapper("unix_nanos") +def unix_nanos(col): + _a = [] + _a.append(col) + return _dispatch("unix_nanos", _a) def unix_seconds(col): _a = [] _a.append(col) @@ -2335,22 +2385,64 @@ def variance(col): _a = [] _a.append(col) return _dispatch("variance", _a) -variant_array_append = _create_wrapper("variant_array_append") +def variant_array_append(v, path, value): + _a = [] + _a.append(v) + _a.append(path) + _a.append(value) + return _dispatch("variant_array_append", _a) def variant_get(v, path, targetType): _a = [] _a.append(v) _a.append(path) _a.append(targetType) return _dispatch("variant_get", _a) -variant_insert = _create_wrapper("variant_insert") -variant_set = _create_wrapper("variant_set") -vector_avg = _create_wrapper("vector_avg") -vector_cosine_similarity = _create_wrapper("vector_cosine_similarity") -vector_inner_product = _create_wrapper("vector_inner_product") -vector_l2_distance = _create_wrapper("vector_l2_distance") -vector_norm = _create_wrapper("vector_norm") -vector_normalize = _create_wrapper("vector_normalize") -vector_sum = _create_wrapper("vector_sum") +def variant_insert(v, path, value): + _a = [] + _a.append(v) + _a.append(path) + _a.append(value) + return _dispatch("variant_insert", _a) +def variant_set(v, path, value, create_if_missing=_UNSET): + _a = [] + _a.append(v) + _a.append(path) + _a.append(value) + if create_if_missing is not _UNSET: _a.append(create_if_missing) + return _dispatch("variant_set", _a) +def vector_avg(col): + _a = [] + _a.append(col) + return _dispatch("vector_avg", _a) +def vector_cosine_similarity(left, right): + _a = [] + _a.append(left) + _a.append(right) + return _dispatch("vector_cosine_similarity", _a) +def vector_inner_product(left, right): + _a = [] + _a.append(left) + _a.append(right) + return _dispatch("vector_inner_product", _a) +def vector_l2_distance(left, right): + _a = [] + _a.append(left) + _a.append(right) + return _dispatch("vector_l2_distance", _a) +def vector_norm(vector, degree=_UNSET): + _a = [] + _a.append(vector) + if degree is not _UNSET: _a.append(degree) + return _dispatch("vector_norm", _a) +def vector_normalize(vector, degree=_UNSET): + _a = [] + _a.append(vector) + if degree is not _UNSET: _a.append(degree) + return _dispatch("vector_normalize", _a) +def vector_sum(col): + _a = [] + _a.append(col) + return _dispatch("vector_sum", _a) def version(): _a = [] return _dispatch("version", _a) diff --git a/python/pyspark/sql/pandas/conversion.py b/python/pyspark/sql/pandas/conversion.py index 3f5d68d1..48998a6a 100644 --- a/python/pyspark/sql/pandas/conversion.py +++ b/python/pyspark/sql/pandas/conversion.py @@ -88,6 +88,7 @@ def create_arrow_array_from_pandas( """ import pyarrow as pa import pandas as pd + from pyspark.loose_version import LooseVersion from pyspark.sql.pandas.types import to_arrow_type, _create_converter_from_pandas if isinstance(series.dtype, pd.CategoricalDtype): @@ -113,7 +114,23 @@ def create_arrow_array_from_pandas( else: mask = series.isnull() try: - return pa.Array.from_pandas(series, mask=mask, type=arrow_type, safe=safecheck) + result = pa.Array.from_pandas(series, mask=mask, type=arrow_type, safe=safecheck) + # SPARK-46776: pyarrow < 19.0.0 ignores the requested ``type`` in the + # ``__arrow_array__`` protocol used by pyarrow-backed extension dtypes, so a + # ``string[pyarrow]`` series (backed by ``large_string`` since pandas 2.2.0) can + # come back as ``large_string`` even when ``string`` was requested, silently + # corrupting data on the JVM side. Cast back only for this exact (large_)string / + # (large_)binary offset-width mismatch; pyarrow >= 19.0.0 already honors the type. + if ( + arrow_type is not None + and LooseVersion(pa.__version__) < LooseVersion("19.0.0") + and ( + (pa.types.is_large_string(result.type) and pa.types.is_string(arrow_type)) + or (pa.types.is_large_binary(result.type) and pa.types.is_binary(arrow_type)) + ) + ): + result = result.cast(arrow_type) + return result except TypeError as e: error_msg = ( "Exception thrown when converting pandas.Series (%s) " @@ -135,15 +152,25 @@ def create_arrow_array_from_pandas( raise PySparkValueError(error_msg % (series.dtype, series.name, arrow_type)) from e -def create_arrow_batch_from_pandas( +def create_arrow_table_from_pandas( series_with_types: Iterable[Tuple["pd.Series", Optional[DataType]]], *, timezone: Optional[str] = None, safecheck: bool = False, prefers_large_types: bool = False, -) -> "pa.RecordBatch": +) -> "pa.Table": """ - Create an Arrow record batch from the given iterable of (series, spark_type) tuples. + Create an Arrow ``Table`` from the given iterable of (series, spark_type) tuples. + + A ``pa.Table`` is used (rather than a single ``pa.RecordBatch``) because + ``pa.Array.from_pandas`` may return a ``pa.ChunkedArray`` when the input + pandas Series is backed by a chunked Arrow array (e.g. pyarrow-backed + extension dtypes such as ``string[pyarrow]``) or when the data exceeds + the maximum size of a single Arrow array (e.g. string data larger than + 2 GB). ``pa.RecordBatch.from_arrays`` does not accept ``ChunkedArray``, + but ``pa.Table.from_arrays`` does. Call ``.to_batches()`` on the result + to obtain a zero-copy list of ``pa.RecordBatch`` aligned on a common + chunk boundary. Parameters ---------- @@ -158,8 +185,7 @@ def create_arrow_batch_from_pandas( Returns ------- - pyarrow.RecordBatch - Arrow RecordBatch + pyarrow.Table """ import pyarrow as pa @@ -173,7 +199,7 @@ def create_arrow_batch_from_pandas( ) for s, spark_type in series_with_types ] - return pa.RecordBatch.from_arrays(arrs, ["_%d" % i for i in range(len(arrs))]) + return pa.Table.from_arrays(arrs, names=["_%d" % i for i in range(len(arrs))]) def _convert_arrow_table_to_pandas( @@ -507,7 +533,7 @@ def _collect_as_arrow( jsocket_auth_server, ) = self._jdf.collectAsArrowToPython() - # Collect list of un-ordered batches where last element is a list of correct order indices + # Collect the batches already reordered by ArrowCollectSerializer. try: with _load_from_socket((port, auth_secret), ArrowCollectSerializer()) as batch_stream: if split_batches: @@ -519,32 +545,23 @@ def _collect_as_arrow( # converted. import pyarrow as pa - results = [] - for batch_or_indices in batch_stream: - if isinstance(batch_or_indices, pa.RecordBatch): - batch_or_indices = pa.RecordBatch.from_arrays( - [ - # This call actually reallocates the array - pa.concat_arrays([array]) - for array in batch_or_indices - ], - schema=batch_or_indices.schema, - ) - results.append(batch_or_indices) + batches = [ + pa.RecordBatch.from_arrays( + # This call actually reallocates the array + [pa.concat_arrays([array]) for array in batch], + schema=batch.schema, + ) + for batch in batch_stream + ] else: - results = list(batch_stream) + batches = list(batch_stream) finally: with unwrap_spark_exception(): # Join serving thread and raise any exceptions from collectAsArrowToPython jsocket_auth_server.getResult() - # Separate RecordBatches from batch order indices in results - batches = results[:-1] - batch_order = results[-1] - if len(batches) or empty_list_if_zero_records: - # Re-order the batch list using the correct order - return [batches[i] for i in batch_order] + return batches else: from pyspark.sql.pandas.types import to_arrow_schema import pyarrow as pa @@ -999,15 +1016,17 @@ def _create_from_pandas_with_arrow( if len(pdf.columns) == 0: arrow_batches = [pa.RecordBatch.from_pandas(pdf_slice) for pdf_slice in pdf_slices] else: - # Create Arrow batches directly using the standalone function + # Each slice may produce more than one RecordBatch when a column is + # backed by a ChunkedArray, so flatten the per-slice tables. arrow_batches = [ - create_arrow_batch_from_pandas( + b + for pdf_slice in pdf_slices + for b in create_arrow_table_from_pandas( [(c, t) for (_, c), t in zip(pdf_slice.items(), spark_types)], timezone=timezone, safecheck=safecheck, prefers_large_types=prefers_large_var_types, - ) - for pdf_slice in pdf_slices + ).to_batches() ] jsparkSession = self._jsparkSession diff --git a/python/pyspark/sql/pandas/functions.py b/python/pyspark/sql/pandas/functions.py index 71599fcb..289f0b1d 100644 --- a/python/pyspark/sql/pandas/functions.py +++ b/python/pyspark/sql/pandas/functions.py @@ -14,16 +14,33 @@ # See the License for the specific language governing permissions and # limitations under the License. # -"""Rust-backed pandas UDF entry points under the official import path.""" +"""Rust-backed pandas UDF entry points under the official import path. + +Canonical home of the ``PandasUDFType`` / ``ArrowUDFType`` enums (mirroring reference +pyspark); ``pyspark.sql.functions`` re-exports them from here. +""" from pyspark.sql.functions import pandas_udf # noqa: F401 -try: - from pyspark.sql.functions import PandasUDFType # noqa: F401 -except ImportError: # pragma: no cover - class PandasUDFType: # minimal stand-in; enum values set below - SCALAR = 200 - GROUPED_MAP = 201 - GROUPED_AGG = 202 - SCALAR_ITER = 204 - MAP_ITER = 205 - COGROUPED_MAP = 206 -__all__ = ["pandas_udf", "PandasUDFType"] + + +class PandasUDFType: + """Pandas UDF Types. See :func:`pyspark.sql.functions.pandas_udf`.""" + + SCALAR = 200 + SCALAR_ITER = 204 + GROUPED_MAP = 201 + GROUPED_AGG = 202 + GROUPED_AGG_ITER = 217 + MAP_ITER = 205 + COGROUPED_MAP = 206 + + +class ArrowUDFType: + """Arrow UDF Types. See :func:`pyspark.sql.functions.pandas_udf`.""" + + SCALAR = 250 + SCALAR_ITER = 251 + GROUPED_AGG = 252 + GROUPED_AGG_ITER = 254 + + +__all__ = ["pandas_udf", "PandasUDFType", "ArrowUDFType"] diff --git a/python/pyspark/sql/pandas/serializers.py b/python/pyspark/sql/pandas/serializers.py index 96784053..2ee42a9d 100644 --- a/python/pyspark/sql/pandas/serializers.py +++ b/python/pyspark/sql/pandas/serializers.py @@ -19,8 +19,7 @@ Serializers for PyArrow and pandas conversions. See `pyspark.serializers` for more details. """ -from itertools import groupby -from typing import IO, TYPE_CHECKING, Any, Iterable, Iterator, List, Optional, Tuple +from typing import IO, TYPE_CHECKING, Iterable, Iterator, List, Tuple from pyspark.errors import PySparkRuntimeError, PySparkValueError from pyspark.serializers import ( @@ -28,45 +27,12 @@ read_int, write_int, UTF8Deserializer, - CPickleSerializer, -) -from pyspark.sql import Row -from pyspark.sql.conversion import ( - ArrowBatchTransformer, - PandasToArrowConversion, -) -from pyspark.sql.pandas.types import ( - from_arrow_schema, - to_arrow_type, -) -from pyspark.sql.types import ( - DataType, - StringType, - StructType, - BinaryType, - StructField, - LongType, - IntegerType, ) if TYPE_CHECKING: - import pandas as pd import pyarrow as pa -def _normalize_packed(packed): - """ - Normalize UDF output to a uniform tuple-of-tuples form. - - Iterator UDFs yield a single (series, spark_type) tuple directly, - while batched UDFs return a tuple of tuples ((s1, t1), (s2, t2), ...). - This function normalizes both forms to a tuple of tuples. - """ - if len(packed) == 2 and isinstance(packed[1], DataType): - return (packed,) - return tuple(packed) - - class SpecialLengths: END_OF_DATA_SECTION = -1 PYTHON_EXCEPTION_THROWN = -2 @@ -76,49 +42,6 @@ class SpecialLengths: START_ARROW_STREAM = -6 -class ArrowCollectSerializer(Serializer): - """ - Deserialize a stream of batches followed by batch order information. Used in - PandasConversionMixin._collect_as_arrow() after invoking Dataset.collectAsArrowToPython() - in the JVM. - """ - - def __init__(self): - self.serializer = ArrowStreamSerializer() - - def dump_stream(self, iterator, stream): - return self.serializer.dump_stream(iterator, stream) - - def load_stream(self, stream): - """ - Load a stream of un-ordered Arrow RecordBatches, where the last iteration yields - a list of indices that can be used to put the RecordBatches in the correct order. - """ - # load the batches - for batch in self.serializer.load_stream(stream): - yield batch - - # load the batch order indices or propagate any error that occurred in the JVM - num = read_int(stream) - if num == -1: - error_msg = UTF8Deserializer().loads(stream) - raise PySparkRuntimeError( - errorClass="ERROR_OCCURRED_WHILE_CALLING", - messageParameters={ - "func_name": "ArrowCollectSerializer.load_stream", - "error_msg": error_msg, - }, - ) - batch_order = [] - for i in range(num): - index = read_int(stream) - batch_order.append(index) - yield batch_order - - def __repr__(self): - return "ArrowCollectSerializer(%s)" % self.serializer - - class ArrowStreamSerializer(Serializer): """ Serializes Arrow record batches as a plain stream. @@ -130,9 +53,10 @@ class ArrowStreamSerializer(Serializer): output batch. Default False. """ - def __init__(self, write_start_stream: bool = False) -> None: + def __init__(self, write_start_stream: bool = False, flush_per_batch: bool = False) -> None: super().__init__() self._write_start_stream: bool = write_start_stream + self._flush_per_batch: bool = flush_per_batch def dump_stream(self, iterator: Iterable["pa.RecordBatch"], stream: IO[bytes]) -> None: """Optionally prepend START_ARROW_STREAM, then write batches.""" @@ -147,6 +71,10 @@ def dump_stream(self, iterator: Iterable["pa.RecordBatch"], stream: IO[bytes]) - if writer is None: writer = pa.RecordBatchStreamWriter(stream, batch.schema) writer.write_batch(batch) + # In pipelined mode, flush after each batch so the JVM can read output + # while still sending input, rather than buffering all output. + if self._flush_per_batch: + stream.flush() finally: if writer is not None: writer.close() @@ -178,6 +106,44 @@ def __repr__(self) -> str: return "ArrowStreamSerializer(write_start_stream=%s)" % self._write_start_stream +class ArrowCollectSerializer(ArrowStreamSerializer): + """ + Extends :class:`ArrowStreamSerializer` to load Arrow RecordBatches that the JVM + sends out of order, followed by the indices giving their correct order, and yields + the batches already reordered. Used in PandasConversionMixin._collect_as_arrow() + after invoking Dataset.collectAsArrowToPython() in the JVM. + """ + + def load_stream(self, stream: IO[bytes]) -> Iterator["pa.RecordBatch"]: + """Load the out-of-order batches, then yield them in the correct order.""" + batches = list(super().load_stream(stream)) + + # Load the batch order indices, or propagate any error that occurred in the JVM. + num = read_int(stream) + if num == -1: + error_msg = UTF8Deserializer().loads(stream) + raise PySparkRuntimeError( + errorClass="ERROR_OCCURRED_WHILE_CALLING", + messageParameters={ + "func_name": "ArrowCollectSerializer.load_stream", + "error_msg": error_msg, + }, + ) + # Yield the batches in order, dropping our reference to each as it goes so + # that, when selfDestruct is enabled, the caller's reallocated copy is the + # only remaining reference and the original batch can be freed immediately + # rather than staying pinned here until the stream is fully consumed. The + # indices are a permutation, so each batch is yielded exactly once. + for _ in range(num): + i = read_int(stream) + batch = batches[i] + batches[i] = None + yield batch + + def __repr__(self) -> str: + return "ArrowCollectSerializer()" + + class ArrowStreamGroupSerializer(ArrowStreamSerializer): """ Extends :class:`ArrowStreamSerializer` with group-count protocol for loading @@ -218,1457 +184,3 @@ def load_stream( errorClass="INVALID_NUMBER_OF_DATAFRAMES_IN_GROUP", messageParameters={"dataframes_in_group": str(dataframes_in_group)}, ) - - -class ArrowStreamUDFSerializer(ArrowStreamSerializer): - """ - Same as :class:`ArrowStreamSerializer` but it flattens the struct to Arrow record batch - for applying each function with the raw record arrow batch. See also `DataFrame.mapInArrow`. - """ - - def load_stream(self, stream): - """ - Flatten the struct into Arrow's record batches. - """ - batches = super().load_stream(stream) - flattened = map(ArrowBatchTransformer.flatten_struct, batches) - return map(lambda b: [b], flattened) - - def dump_stream(self, iterator, stream): - """ - Override because Pandas UDFs require a START_ARROW_STREAM before the Arrow stream is sent. - """ - batches = self._write_stream_start( - (ArrowBatchTransformer.wrap_struct(x[0]) for x in iterator), stream - ) - return super().dump_stream(batches, stream) - - -class ArrowStreamUDTFSerializer(ArrowStreamUDFSerializer): - """ - Same as :class:`ArrowStreamUDFSerializer` but it does not flatten when loading batches. - """ - - def load_stream(self, stream): - return ArrowStreamSerializer.load_stream(self, stream) - - -class ArrowStreamArrowUDTFSerializer(ArrowStreamUDTFSerializer): - """ - Serializer for PyArrow-native UDTFs that work directly with PyArrow RecordBatches and Arrays. - """ - - def __init__(self, *, table_arg_offsets=None): - super().__init__() - self.table_arg_offsets = table_arg_offsets if table_arg_offsets else [] - - def load_stream(self, stream): - """ - Flatten the struct into Arrow's record batches. - """ - for batch in super().load_stream(stream): - # For each column: flatten struct columns at table_arg_offsets into RecordBatch, - # keep other columns as Array - yield [ - ( - ArrowBatchTransformer.flatten_struct(batch, column_index=i) - if i in self.table_arg_offsets - else batch.column(i) - ) - for i in range(batch.num_columns) - ] - - def dump_stream(self, iterator, stream): - """ - Override to handle type coercion for ArrowUDTF outputs. - ArrowUDTF returns iterator of (pa.RecordBatch, arrow_return_type) tuples. - """ - import pyarrow as pa - - def apply_type_coercion(): - for batch, arrow_return_type in iterator: - assert isinstance(arrow_return_type, pa.StructType), ( - f"Expected pa.StructType, got {type(arrow_return_type)}" - ) - coerced_batch = ArrowBatchTransformer.enforce_schema( - batch, pa.schema(arrow_return_type), safecheck=True - ) - yield coerced_batch, arrow_return_type - - return super().dump_stream(apply_type_coercion(), stream) - - -class ArrowStreamGroupUDFSerializer(ArrowStreamUDFSerializer): - """ - Serializer for grouped Arrow UDFs. - - Deserializes: - ``Iterator[Iterator[pa.RecordBatch]]`` - one inner iterator per group. - Each batch contains a single struct column. - - Serializes: - ``Iterator[Tuple[Iterator[pa.RecordBatch], pa.DataType]]`` - Each tuple contains iterator of flattened batches and their Arrow type. - - Used by: - - SQL_GROUPED_MAP_ARROW_UDF - - SQL_GROUPED_MAP_ARROW_ITER_UDF - - Parameters - ---------- - assign_cols_by_name : bool - If True, reorder serialized columns by schema name. - """ - - def __init__(self, *, assign_cols_by_name): - super().__init__() - self._assign_cols_by_name = assign_cols_by_name - - def load_stream(self, stream): - """ - Load grouped Arrow record batches from stream. - """ - for batches in ArrowStreamGroupSerializer.load_stream(self, stream): - yield batches - # Make sure the batches are fully iterated before getting the next group - for _ in batches: - pass - - def dump_stream(self, iterator, stream): - import pyarrow as pa - - # flatten inner list [([pa.RecordBatch], arrow_type)] into [(pa.RecordBatch, arrow_type)] - # so strip off inner iterator induced by ArrowStreamUDFSerializer.load_stream - batch_iter = ( - (batch, arrow_type) - for batches, arrow_type in iterator # tuple constructed in wrap_grouped_map_arrow_udf - for batch in batches - ) - - if self._assign_cols_by_name: - batch_iter = ( - ( - pa.RecordBatch.from_arrays( - [batch.column(field.name) for field in arrow_type], - names=[field.name for field in arrow_type], - ), - arrow_type, - ) - for batch, arrow_type in batch_iter - ) - - super().dump_stream(batch_iter, stream) - - -class ArrowStreamPandasSerializer(ArrowStreamSerializer): - """ - Serializes pandas.Series as Arrow data with Arrow streaming format. - - Parameters - ---------- - timezone : str - A timezone to respect when handling timestamp values - safecheck : bool - If True, conversion from Arrow to Pandas checks for overflow/truncation - int_to_decimal_coercion_enabled : bool - If True, applies additional coercions in Python before converting to Arrow. - This has performance penalties. - prefers_large_types : bool - If True, prefer large Arrow types (e.g., large_string instead of string). - struct_in_pandas : str, optional - How to represent struct in pandas ("dict", "row", etc.). Default is "dict". - ndarray_as_list : bool, optional - Whether to convert ndarray as list. Default is False. - prefer_int_ext_dtype : bool, optional - Whether to convert integers to Pandas ExtensionDType. Default is False. - df_for_struct : bool, optional - If True, convert struct columns to DataFrame instead of Series. Default is False. - """ - - def __init__( - self, - *, - timezone, - safecheck, - int_to_decimal_coercion_enabled: bool = False, - prefers_large_types: bool = False, - struct_in_pandas: str = "dict", - ndarray_as_list: bool = False, - prefer_int_ext_dtype: bool = False, - df_for_struct: bool = False, - input_type: Optional["StructType"] = None, - arrow_cast: bool = False, - ): - super().__init__() - self._timezone = timezone - self._safecheck = safecheck - self._int_to_decimal_coercion_enabled = int_to_decimal_coercion_enabled - self._prefers_large_types = prefers_large_types - self._struct_in_pandas = struct_in_pandas - self._ndarray_as_list = ndarray_as_list - self._prefer_int_ext_dtype = prefer_int_ext_dtype - self._df_for_struct = df_for_struct - if input_type is not None: - assert isinstance(input_type, StructType) - self._input_type = input_type - self._arrow_cast = arrow_cast - - def dump_stream(self, iterator, stream): - """ - Make ArrowRecordBatches from Pandas Series and serialize. - Each element in iterator is: - - For batched UDFs: tuple of (series, spark_type) tuples: ((s1, t1), (s2, t2), ...) - - For iterator UDFs: single (series, spark_type) tuple directly - """ - - def create_batch( - series_tuples: Tuple[Tuple["pd.Series", DataType], ...], - ) -> "pa.RecordBatch": - series_data = [s for s, _ in series_tuples] - types = [t for _, t in series_tuples] - schema = StructType([StructField(f"_{i}", t) for i, t in enumerate(types)]) - return PandasToArrowConversion.convert( - series_data, - schema, - timezone=self._timezone, - safecheck=self._safecheck, - prefers_large_types=self._prefers_large_types, - int_to_decimal_coercion_enabled=self._int_to_decimal_coercion_enabled, - ) - - super().dump_stream( - (create_batch(_normalize_packed(packed)) for packed in iterator), stream - ) - - def load_stream(self, stream): - """ - Deserialize ArrowRecordBatches to an Arrow table and return as a list of pandas.Series. - """ - yield from map( - lambda batch: ArrowBatchTransformer.to_pandas( - batch, - timezone=self._timezone, - schema=self._input_type, - struct_in_pandas=self._struct_in_pandas, - ndarray_as_list=self._ndarray_as_list, - prefer_int_ext_dtype=self._prefer_int_ext_dtype, - df_for_struct=self._df_for_struct, - ), - super().load_stream(stream), - ) - - def __repr__(self): - return "ArrowStreamPandasSerializer" - - -class ArrowStreamPandasUDFSerializer(ArrowStreamPandasSerializer): - """ - Serializer used by Python worker to evaluate Pandas UDFs - """ - - def __init__( - self, - *, - timezone, - safecheck, - assign_cols_by_name, - df_for_struct: bool = False, - struct_in_pandas: str = "dict", - ndarray_as_list: bool = False, - prefer_int_ext_dtype: bool = False, - arrow_cast: bool = False, - input_type: Optional[StructType] = None, - int_to_decimal_coercion_enabled: bool = False, - prefers_large_types: bool = False, - ignore_unexpected_complex_type_values: bool = False, - is_legacy: bool = False, - ): - super().__init__( - timezone=timezone, - safecheck=safecheck, - int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, - prefers_large_types=prefers_large_types, - struct_in_pandas=struct_in_pandas, - ndarray_as_list=ndarray_as_list, - prefer_int_ext_dtype=prefer_int_ext_dtype, - df_for_struct=df_for_struct, - input_type=input_type, - arrow_cast=arrow_cast, - ) - self._assign_cols_by_name = assign_cols_by_name - self._ignore_unexpected_complex_type_values = ignore_unexpected_complex_type_values - self._is_legacy = is_legacy - - def dump_stream(self, iterator, stream): - """ - Override because Pandas UDFs require a START_ARROW_STREAM before the Arrow stream is sent. - This should be sent after creating the first record batch so in case of an error, it can - be sent back to the JVM before the Arrow stream starts. - - Each element in iterator is: - - For batched UDFs: tuple of (series, spark_type) tuples: ((s1, t1), (s2, t2), ...) - - For iterator UDFs: single (series, spark_type) tuple directly - """ - import pandas as pd - - def create_batch( - series_tuples: Tuple[Tuple["pd.Series", DataType], ...], - ) -> "pa.RecordBatch": - # When struct_in_pandas="dict", UDF must return DataFrame for struct types - if self._struct_in_pandas == "dict": - for s, spark_type in series_tuples: - if isinstance(spark_type, StructType) and not isinstance(s, pd.DataFrame): - raise PySparkValueError( - "Invalid return type. Please make sure that the UDF returns a " - "pandas.DataFrame when the specified return type is StructType." - ) - - series_data = [s for s, _ in series_tuples] - types = [t for _, t in series_tuples] - schema = StructType([StructField(f"_{i}", t) for i, t in enumerate(types)]) - return PandasToArrowConversion.convert( - series_data, - schema, - timezone=self._timezone, - safecheck=self._safecheck, - arrow_cast=self._arrow_cast, - prefers_large_types=self._prefers_large_types, - assign_cols_by_name=self._assign_cols_by_name, - int_to_decimal_coercion_enabled=self._int_to_decimal_coercion_enabled, - ignore_unexpected_complex_type_values=self._ignore_unexpected_complex_type_values, - is_legacy=self._is_legacy, - ) - - batches = self._write_stream_start( - (create_batch(_normalize_packed(packed)) for packed in iterator), - stream, - ) - return ArrowStreamSerializer.dump_stream(self, batches, stream) - - def __repr__(self): - return "ArrowStreamPandasUDFSerializer" - - -class ArrowStreamPandasUDTFSerializer(ArrowStreamPandasUDFSerializer): - """ - Serializer used by Python worker to evaluate Arrow-optimized Python UDTFs. - """ - - def __init__( - self, - *, - timezone, - safecheck, - input_type, - prefer_int_ext_dtype, - int_to_decimal_coercion_enabled, - ): - super().__init__( - timezone=timezone, - safecheck=safecheck, - # The output pandas DataFrame's columns are unnamed. - assign_cols_by_name=False, - # Set to 'False' to avoid converting struct type inputs into a pandas DataFrame. - df_for_struct=False, - # Defines how struct type inputs are converted. If set to "row", struct type inputs - # are converted into Rows. Without this setting, a struct type input would be treated - # as a dictionary. For example, for named_struct('name', 'Alice', 'age', 1), - # if struct_in_pandas="dict", it becomes {"name": "Alice", "age": 1} - # if struct_in_pandas="row", it becomes Row(name="Alice", age=1) - struct_in_pandas="row", - # When dealing with array type inputs, Arrow converts them into numpy.ndarrays. - # To ensure consistency across regular and arrow-optimized UDTFs, we further - # convert these numpy.ndarrays into Python lists. - ndarray_as_list=True, - prefer_int_ext_dtype=prefer_int_ext_dtype, - # Enables explicit casting for mismatched return types of Arrow Python UDTFs. - arrow_cast=True, - input_type=input_type, - # Enable additional coercions for UDTF serialization - int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, - # UDTF-specific: ignore unexpected complex type values in converter - ignore_unexpected_complex_type_values=True, - # Legacy UDTF pandas conversion: enables broader Arrow exception - # handling to allow more implicit type coercions - is_legacy=True, - ) - - def __repr__(self): - return "ArrowStreamPandasUDTFSerializer" - - -# Serializer for SQL_GROUPED_AGG_PANDAS_UDF, SQL_WINDOW_AGG_PANDAS_UDF, -# and SQL_GROUPED_AGG_PANDAS_ITER_UDF -class ArrowStreamAggPandasUDFSerializer(ArrowStreamPandasUDFSerializer): - def __init__( - self, - *, - timezone, - safecheck, - assign_cols_by_name, - prefer_int_ext_dtype, - int_to_decimal_coercion_enabled, - ): - super().__init__( - timezone=timezone, - safecheck=safecheck, - assign_cols_by_name=assign_cols_by_name, - df_for_struct=False, - struct_in_pandas="dict", - ndarray_as_list=False, - prefer_int_ext_dtype=prefer_int_ext_dtype, - arrow_cast=True, - input_type=None, - int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, - ) - - def load_stream(self, stream): - """ - Yield an iterator that produces one tuple of pandas.Series per batch. - Each group yields Iterator[Tuple[pd.Series, ...]], allowing UDF to - process batches one by one without consuming all batches upfront. - """ - for batches in ArrowStreamGroupSerializer.load_stream(self, stream): - # Lazily read and convert Arrow batches to pandas Series one at a time - # from the stream. This avoids loading all batches into memory for the group - series_iter = map( - lambda batch: tuple( - ArrowBatchTransformer.to_pandas( - batch, - timezone=self._timezone, - schema=self._input_type, - struct_in_pandas=self._struct_in_pandas, - ndarray_as_list=self._ndarray_as_list, - prefer_int_ext_dtype=self._prefer_int_ext_dtype, - df_for_struct=self._df_for_struct, - ) - ), - batches, - ) - yield series_iter - # Make sure the batches are fully iterated before getting the next group - for _ in series_iter: - pass - - def __repr__(self): - return "ArrowStreamAggPandasUDFSerializer" - - -# Serializer for SQL_GROUPED_MAP_PANDAS_UDF, SQL_GROUPED_MAP_PANDAS_ITER_UDF -class GroupPandasUDFSerializer(ArrowStreamPandasUDFSerializer): - def __init__( - self, - *, - timezone, - safecheck, - assign_cols_by_name, - prefer_int_ext_dtype, - int_to_decimal_coercion_enabled, - ): - super().__init__( - timezone=timezone, - safecheck=safecheck, - assign_cols_by_name=assign_cols_by_name, - df_for_struct=False, - struct_in_pandas="dict", - ndarray_as_list=False, - prefer_int_ext_dtype=prefer_int_ext_dtype, - arrow_cast=True, - input_type=None, - int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, - ) - - def load_stream(self, stream): - """ - Deserialize Grouped ArrowRecordBatches and yield raw Iterator[pa.RecordBatch]. - Each outer iterator element represents a group. - """ - for batches in ArrowStreamGroupSerializer.load_stream(self, stream): - yield batches - # Make sure the batches are fully iterated before getting the next group - for _ in batches: - pass - - def dump_stream(self, iterator, stream): - """ - Flatten the grouped iterator structure. - """ - # Flatten: Iterator[Iterator[[(df, spark_type)]]] -> Iterator[[(df, spark_type)]] - flattened_iter = (batch for generator in iterator for batch in generator) - super().dump_stream(flattened_iter, stream) - - def __repr__(self): - return "GroupPandasUDFSerializer" - - -class CogroupArrowUDFSerializer(ArrowStreamGroupUDFSerializer): - """ - Serializes pyarrow.RecordBatch data with Arrow streaming format. - - Loads Arrow record batches as `[([pa.RecordBatch], [pa.RecordBatch])]` (one tuple per group) - and serializes `[([pa.RecordBatch], arrow_type)]`. - - Parameters - ---------- - assign_cols_by_name : bool - If True, then DataFrames will get columns by name - """ - - def load_stream(self, stream): - """ - Deserialize Cogrouped ArrowRecordBatches and yield as two `pyarrow.RecordBatch`es. - """ - yield from ArrowStreamCoGroupSerializer.load_stream(self, stream) - - -class CogroupPandasUDFSerializer(ArrowStreamPandasUDFSerializer): - def load_stream(self, stream): - """ - Deserialize Cogrouped ArrowRecordBatches to a tuple of Arrow tables and yield as two - lists of pandas.Series. - """ - import pyarrow as pa - - for left_batches, right_batches in ArrowStreamCoGroupSerializer.load_stream(self, stream): - left_table = pa.Table.from_batches(left_batches) - right_table = pa.Table.from_batches(right_batches) - yield ( - ArrowBatchTransformer.to_pandas( - left_table, - timezone=self._timezone, - schema=from_arrow_schema(left_table.schema), - struct_in_pandas=self._struct_in_pandas, - ndarray_as_list=self._ndarray_as_list, - prefer_int_ext_dtype=self._prefer_int_ext_dtype, - df_for_struct=self._df_for_struct, - ), - ArrowBatchTransformer.to_pandas( - right_table, - timezone=self._timezone, - schema=from_arrow_schema(right_table.schema), - struct_in_pandas=self._struct_in_pandas, - ndarray_as_list=self._ndarray_as_list, - prefer_int_ext_dtype=self._prefer_int_ext_dtype, - df_for_struct=self._df_for_struct, - ), - ) - - -class ApplyInPandasWithStateSerializer(ArrowStreamPandasUDFSerializer): - """ - Serializer used by Python worker to evaluate UDF for applyInPandasWithState. - - Parameters - ---------- - timezone : str - A timezone to respect when handling timestamp values - safecheck : bool - If True, conversion from Arrow to Pandas checks for overflow/truncation - assign_cols_by_name : bool - If True, then Pandas DataFrames will get columns by name - state_object_schema : StructType - The type of state object represented as Spark SQL type - arrow_max_records_per_batch : int - Limit of the number of records that can be written to a single ArrowRecordBatch in memory. - """ - - def __init__( - self, - *, - timezone, - safecheck, - assign_cols_by_name, - prefer_int_ext_dtype, - state_object_schema, - arrow_max_records_per_batch, - prefers_large_var_types, - int_to_decimal_coercion_enabled, - ): - super().__init__( - timezone=timezone, - safecheck=safecheck, - assign_cols_by_name=assign_cols_by_name, - df_for_struct=False, - struct_in_pandas="dict", - ndarray_as_list=False, - prefer_int_ext_dtype=prefer_int_ext_dtype, - arrow_cast=True, - input_type=None, - int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, - prefers_large_types=prefers_large_var_types, - ) - self.pickleSer = CPickleSerializer() - self.utf8_deserializer = UTF8Deserializer() - self.state_object_schema = state_object_schema - - self.result_count_df_type = StructType( - [ - StructField("dataCount", IntegerType()), - StructField("stateCount", IntegerType()), - ] - ) - - self.result_count_pdf_arrow_type = to_arrow_type( - self.result_count_df_type, timezone="UTC", prefers_large_types=prefers_large_var_types - ) - - self.result_state_df_type = StructType( - [ - StructField("properties", StringType()), - StructField("keyRowAsUnsafe", BinaryType()), - StructField("object", BinaryType()), - StructField("oldTimeoutTimestamp", LongType()), - ] - ) - - self.result_state_pdf_arrow_type = to_arrow_type( - self.result_state_df_type, timezone="UTC", prefers_large_types=prefers_large_var_types - ) - self.arrow_max_records_per_batch = ( - arrow_max_records_per_batch if arrow_max_records_per_batch > 0 else 2**31 - 1 - ) - - def load_stream(self, stream): - """ - Read ArrowRecordBatches from stream, deserialize them to populate a list of pair - (data chunk, state), and convert the data into a list of pandas.Series. - - Please refer the doc of inner function `gen_data_and_state` for more details how - this function works in overall. - - In addition, this function further groups the return of `gen_data_and_state` by the state - instance (same semantic as grouping by grouping key) and produces an iterator of data - chunks for each group, so that the caller can lazily materialize the data chunk. - """ - - import pyarrow as pa - import json - from itertools import groupby - from pyspark.sql.streaming.state import GroupState - - def construct_state(state_info_col): - """ - Construct state instance from the value of state information column. - """ - - state_info_col_properties = state_info_col["properties"] - state_info_col_key_row = state_info_col["keyRowAsUnsafe"] - state_info_col_object = state_info_col["object"] - - state_properties = json.loads(state_info_col_properties) - if state_info_col_object: - state_object = self.pickleSer.loads(state_info_col_object) - else: - state_object = None - state_properties["optionalValue"] = state_object - - return GroupState( - keyAsUnsafe=state_info_col_key_row, - valueSchema=self.state_object_schema, - **state_properties, - ) - - def gen_data_and_state(batches): - """ - Deserialize ArrowRecordBatches and return a generator of - `(a list of pandas.Series, state)`. - - The logic on deserialization is following: - - 1. Read the entire data part from Arrow RecordBatch. - 2. Read the entire state information part from Arrow RecordBatch. - 3. Loop through each state information: - 3.A. Extract the data out from entire data via the information of data range. - 3.B. Construct a new state instance if the state information is the first occurrence - for the current grouping key. - 3.C. Leverage the existing state instance if it is already available for the current - grouping key. (Meaning it's not the first occurrence.) - 3.D. Remove the cache of state instance if the state information denotes the data is - the last chunk for current grouping key. - - This deserialization logic assumes that Arrow RecordBatches contain the data with the - ordering that data chunks for same grouping key will appear sequentially. - - This function must avoid materializing multiple Arrow RecordBatches into memory at the - same time. And data chunks from the same grouping key should appear sequentially, to - further group them based on state instance (same state instance will be produced for - same grouping key). - """ - - state_for_current_group = None - - for batch in batches: - batch_schema = batch.schema - data_schema = pa.schema([batch_schema[i] for i in range(0, len(batch_schema) - 1)]) - state_schema = pa.schema( - [ - batch_schema[-1], - ] - ) - - batch_columns = batch.columns - data_columns = batch_columns[0:-1] - state_column = batch_columns[-1] - - data_batch = pa.RecordBatch.from_arrays(data_columns, schema=data_schema) - state_batch = pa.RecordBatch.from_arrays( - [ - state_column, - ], - schema=state_schema, - ) - - state_pandas = ArrowBatchTransformer.to_pandas( - state_batch, - timezone=self._timezone, - schema=None, - struct_in_pandas=self._struct_in_pandas, - ndarray_as_list=self._ndarray_as_list, - prefer_int_ext_dtype=self._prefer_int_ext_dtype, - df_for_struct=self._df_for_struct, - )[0] - - for state_idx in range(0, len(state_pandas)): - state_info_col = state_pandas.iloc[state_idx] - - if not state_info_col: - # no more data with grouping key + state - break - - data_start_offset = state_info_col["startOffset"] - num_data_rows = state_info_col["numRows"] - is_last_chunk = state_info_col["isLastChunk"] - - if state_for_current_group: - # use the state, we already have state for same group and there should be - # some data in same group being processed earlier - state = state_for_current_group - else: - # there is no state being stored for same group, construct one - state = construct_state(state_info_col) - - if is_last_chunk: - # discard the state being cached for same group - state_for_current_group = None - elif not state_for_current_group: - # there's no cached state but expected to have additional data in same group - # cache the current state - state_for_current_group = state - - data_batch_for_group = data_batch.slice(data_start_offset, num_data_rows) - data_pandas = ArrowBatchTransformer.to_pandas( - data_batch_for_group, - timezone=self._timezone, - schema=None, - struct_in_pandas=self._struct_in_pandas, - ndarray_as_list=self._ndarray_as_list, - prefer_int_ext_dtype=self._prefer_int_ext_dtype, - df_for_struct=self._df_for_struct, - ) - - # state info - yield ( - data_pandas, - state, - ) - - _batches = super(ArrowStreamPandasSerializer, self).load_stream(stream) - - data_state_generator = gen_data_and_state(_batches) - - # state will be same object for same grouping key - for _state, _data in groupby(data_state_generator, key=lambda x: x[1]): - yield ( - _data, - _state, - ) - - def dump_stream(self, iterator, stream): - """ - Read through an iterator of (iterator of pandas DataFrame, state), serialize them to Arrow - RecordBatches, and write batches to stream. - """ - - import pandas as pd - - def construct_state_pdf(state): - """ - Construct a pandas DataFrame from the state instance. - """ - - state_properties = state.json().encode("utf-8") - state_key_row_as_binary = state._keyAsUnsafe - if state.exists: - state_object = self.pickleSer.dumps(state._value_schema.toInternal(state._value)) - else: - state_object = None - state_old_timeout_timestamp = state.oldTimeoutTimestamp - - state_dict = { - "properties": [ - state_properties, - ], - "keyRowAsUnsafe": [ - state_key_row_as_binary, - ], - "object": [ - state_object, - ], - "oldTimeoutTimestamp": [ - state_old_timeout_timestamp, - ], - } - - return pd.DataFrame.from_dict(state_dict) - - def construct_record_batch(pdfs, pdf_data_cnt, pdf_schema, state_pdfs, state_data_cnt): - """ - Construct a new Arrow RecordBatch based on output pandas DataFrames and states. Each - one matches to the single struct field for Arrow schema. We also need an extra one to - indicate array length for data and state, so the return value of Arrow RecordBatch will - have schema with three fields, in `count`, `data`, `state` order. - (Readers are expected to access the field via position rather than the name. We do - not guarantee the name of the field.) - - Note that Arrow RecordBatch requires all columns to have all same number of rows, - hence this function inserts empty data for count/state/data with less elements to - compensate. - """ - - max_data_cnt = max(1, max(pdf_data_cnt, state_data_cnt)) - - # We only use the first row in the count column, and fill other rows to be the same - # value, hoping it is more friendly for compression, in case it is needed. - count_dict = { - "dataCount": [pdf_data_cnt] * max_data_cnt, - "stateCount": [state_data_cnt] * max_data_cnt, - } - count_pdf = pd.DataFrame.from_dict(count_dict) - - empty_row_cnt_in_data = max_data_cnt - pdf_data_cnt - empty_row_cnt_in_state = max_data_cnt - state_data_cnt - - empty_rows_pdf = pd.DataFrame( - dict.fromkeys(pdf_schema.names), - index=[x for x in range(0, empty_row_cnt_in_data)], - ) - empty_rows_state = pd.DataFrame( - columns=["properties", "keyRowAsUnsafe", "object", "oldTimeoutTimestamp"], - index=[x for x in range(0, empty_row_cnt_in_state)], - ) - - pdfs.append(empty_rows_pdf) - state_pdfs.append(empty_rows_state) - - merged_pdf = pd.concat(pdfs, ignore_index=True) - merged_state_pdf = pd.concat(state_pdfs, ignore_index=True) - - # Create batch from list of DataFrames, each wrapped as a StructArray. - # Schema fields map to: _0=count, _1=output data, _2=state data - # (types defined in __init__: result_count_df_type, pdf_schema, result_state_df_type) - data = [count_pdf, merged_pdf, merged_state_pdf] - schema = StructType( - [ - StructField("_0", self.result_count_df_type), - StructField("_1", pdf_schema), - StructField("_2", self.result_state_df_type), - ] - ) - return PandasToArrowConversion.convert( - data, - schema, - timezone=self._timezone, - safecheck=self._safecheck, - arrow_cast=self._arrow_cast, - assign_cols_by_name=self._assign_cols_by_name, - int_to_decimal_coercion_enabled=self._int_to_decimal_coercion_enabled, - ) - - def serialize_batches(): - """ - Read through an iterator of (iterator of pandas DataFrame, state), and serialize them - to Arrow RecordBatches. - - This function does batching on constructing the Arrow RecordBatch; a batch will be - serialized to the Arrow RecordBatch when the total number of records exceeds the - configured threshold. - """ - # a set of variables for the state of current batch which will be converted to Arrow - # RecordBatch. - pdfs = [] - state_pdfs = [] - pdf_data_cnt = 0 - state_data_cnt = 0 - - return_schema = None - - for data in iterator: - # data represents the result of each call of user function - packaged_result = data[0] - - # There are two results from the call of user function: - # 1) iterator of pandas DataFrame (output) - # 2) updated state instance - pdf_iter = packaged_result[0][0] - state = packaged_result[0][1] - - # This is static and won't change across batches. - return_schema = packaged_result[1] - - for pdf in pdf_iter: - # We ignore empty pandas DataFrame. - if len(pdf) > 0: - pdf_data_cnt += len(pdf) - pdfs.append(pdf) - - # If the total number of records in current batch exceeds the configured - # threshold, time to construct the Arrow RecordBatch from the batch. - if pdf_data_cnt > self.arrow_max_records_per_batch: - batch = construct_record_batch( - pdfs, pdf_data_cnt, return_schema, state_pdfs, state_data_cnt - ) - - # Reset the variables to start with new batch for further data. - pdfs = [] - state_pdfs = [] - pdf_data_cnt = 0 - state_data_cnt = 0 - - yield batch - - # This has to be performed 'after' evaluating all elements in iterator, so that - # the user function has been completed and the state is guaranteed to be updated. - state_pdf = construct_state_pdf(state) - - state_pdfs.append(state_pdf) - state_data_cnt += 1 - - # processed all output, but current batch may not be flushed yet. - if pdf_data_cnt > 0 or state_data_cnt > 0: - batch = construct_record_batch( - pdfs, pdf_data_cnt, return_schema, state_pdfs, state_data_cnt - ) - - yield batch - - batches = self._write_stream_start(serialize_batches(), stream) - return ArrowStreamSerializer.dump_stream(self, batches, stream) - - -class TransformWithStateInPandasSerializer(ArrowStreamPandasUDFSerializer): - """ - Serializer used by Python worker to evaluate UDF for - :meth:`pyspark.sql.GroupedData.transformWithStateInPandasSerializer`. - - Parameters - ---------- - timezone : str - A timezone to respect when handling timestamp values - safecheck : bool - If True, conversion from Arrow to Pandas checks for overflow/truncation - assign_cols_by_name : bool - If True, then Pandas DataFrames will get columns by name - arrow_max_records_per_batch : int - Limit of the number of records that can be written to a single ArrowRecordBatch in memory. - """ - - def __init__( - self, - *, - timezone, - safecheck, - assign_cols_by_name, - prefer_int_ext_dtype, - arrow_max_records_per_batch, - arrow_max_bytes_per_batch, - int_to_decimal_coercion_enabled, - ): - super().__init__( - timezone=timezone, - safecheck=safecheck, - assign_cols_by_name=assign_cols_by_name, - df_for_struct=False, - struct_in_pandas="dict", - ndarray_as_list=False, - prefer_int_ext_dtype=prefer_int_ext_dtype, - arrow_cast=True, - input_type=None, - int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, - ) - self.arrow_max_records_per_batch = ( - arrow_max_records_per_batch if arrow_max_records_per_batch > 0 else 2**31 - 1 - ) - self.arrow_max_bytes_per_batch = arrow_max_bytes_per_batch - self.key_offsets = None - self.average_arrow_row_size = 0 - self.total_bytes = 0 - self.total_rows = 0 - - def _update_batch_size_stats(self, batch): - """ - Update batch size statistics for adaptive batching. - """ - # Short circuit batch size calculation if the batch size is - # unlimited as computing batch size is computationally expensive. - if self.arrow_max_bytes_per_batch != 2**31 - 1 and batch.num_rows > 0: - batch_bytes = sum( - buf.size for col in batch.columns for buf in col.buffers() if buf is not None - ) - self.total_bytes += batch_bytes - self.total_rows += batch.num_rows - self.average_arrow_row_size = self.total_bytes / self.total_rows - - def load_stream(self, stream): - """ - Read ArrowRecordBatches from stream, deserialize them to populate a list of data chunk, and - convert the data into Rows. - - Please refer the doc of inner function `generate_data_batches` for more details how - this function works in overall. - """ - import pandas as pd - from pyspark.sql.streaming.stateful_processor_util import ( - TransformWithStateInPandasFuncMode, - ) - - def generate_data_batches(batches): - """ - Deserialize ArrowRecordBatches and return a generator of Rows. - - The deserialization logic assumes that Arrow RecordBatches contain the data with the - ordering that data chunks for same grouping key will appear sequentially. - - This function must avoid materializing multiple Arrow RecordBatches into memory at the - same time. And data chunks from the same grouping key should appear sequentially. - """ - - def row_stream(): - for batch in batches: - self._update_batch_size_stats(batch) - data_pandas = ArrowBatchTransformer.to_pandas( - batch, - timezone=self._timezone, - schema=self._input_type, - struct_in_pandas=self._struct_in_pandas, - ndarray_as_list=self._ndarray_as_list, - prefer_int_ext_dtype=self._prefer_int_ext_dtype, - df_for_struct=self._df_for_struct, - ) - for row in pd.concat(data_pandas, axis=1).itertuples(index=False): - batch_key = tuple(row[s] for s in self.key_offsets) - yield (batch_key, row) - - for batch_key, group_rows in groupby(row_stream(), key=lambda x: x[0]): - rows = [] - for _, row in group_rows: - rows.append(row) - if ( - len(rows) >= self.arrow_max_records_per_batch - or len(rows) * self.average_arrow_row_size >= self.arrow_max_bytes_per_batch - ): - yield (batch_key, pd.DataFrame(rows)) - rows = [] - if rows: - yield (batch_key, pd.DataFrame(rows)) - - _batches = super(ArrowStreamPandasSerializer, self).load_stream(stream) - data_batches = generate_data_batches(_batches) - - for k, g in groupby(data_batches, key=lambda x: x[0]): - yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, g) - - yield (TransformWithStateInPandasFuncMode.PROCESS_TIMER, None, None) - - yield (TransformWithStateInPandasFuncMode.COMPLETE, None, None) - - def dump_stream(self, iterator, stream): - """ - Read through an iterator of (iterator of pandas DataFrame), serialize them to Arrow - RecordBatches, and write batches to stream. - """ - - def flatten_iterator(): - # iterator: iter[list[(iter[pandas.DataFrame], pdf_type)]] - for packed in iterator: - iter_pdf_with_type = packed[0] - iter_pdf = iter_pdf_with_type[0] - pdf_type = iter_pdf_with_type[1] - for pdf in iter_pdf: - yield [(pdf, pdf_type)] - - super().dump_stream(flatten_iterator(), stream) - - -class TransformWithStateInPandasInitStateSerializer(TransformWithStateInPandasSerializer): - """ - Serializer used by Python worker to evaluate UDF for - :meth:`pyspark.sql.GroupedData.transformWithStateInPandasInitStateSerializer`. - Parameters - ---------- - Same as input parameters in TransformWithStateInPandasSerializer. - """ - - def __init__( - self, - *, - timezone, - safecheck, - assign_cols_by_name, - prefer_int_ext_dtype, - arrow_max_records_per_batch, - arrow_max_bytes_per_batch, - int_to_decimal_coercion_enabled, - ): - super().__init__( - timezone=timezone, - safecheck=safecheck, - assign_cols_by_name=assign_cols_by_name, - prefer_int_ext_dtype=prefer_int_ext_dtype, - arrow_max_records_per_batch=arrow_max_records_per_batch, - arrow_max_bytes_per_batch=arrow_max_bytes_per_batch, - int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled, - ) - self.init_key_offsets = None - - def load_stream(self, stream): - import pyarrow as pa - import pandas as pd - from pyspark.sql.streaming.stateful_processor_util import ( - TransformWithStateInPandasFuncMode, - ) - - def generate_data_batches(batches): - """ - Deserialize ArrowRecordBatches and return a generator of pandas.Series list. - - The deserialization logic assumes that Arrow RecordBatches contain the data with the - ordering that data chunks for same grouping key will appear sequentially. - See `TransformWithStateInPandasPythonInitialStateRunner` for arrow batch schema sent - from JVM. - This function flatten the columns of input rows and initial state rows and feed them - into the data generator. - """ - - def flatten_columns(cur_batch, col_name): - state_column = cur_batch.column(cur_batch.schema.get_field_index(col_name)) - - # Check if the entire column is null - if state_column.null_count == len(state_column): - # Return empty table with no columns - return pa.Table.from_arrays([], names=[]) - - state_field_names = [ - state_column.type[i].name for i in range(state_column.type.num_fields) - ] - state_field_arrays = [ - state_column.field(i) for i in range(state_column.type.num_fields) - ] - table_from_fields = pa.Table.from_arrays( - state_field_arrays, names=state_field_names - ) - return table_from_fields - - """ - The arrow batch is written in the schema: - schema: StructType = new StructType() - .add("inputData", dataSchema) - .add("initState", initStateSchema) - We'll parse batch into Tuples of (key, inputData, initState) and pass into the Python - data generator. Rows in the same batch may have different grouping keys, - but each batch will have either init_data or input_data, not mix. - """ - - def to_pandas(table): - return ArrowBatchTransformer.to_pandas( - table, - timezone=self._timezone, - schema=self._input_type, - struct_in_pandas=self._struct_in_pandas, - ndarray_as_list=self._ndarray_as_list, - prefer_int_ext_dtype=self._prefer_int_ext_dtype, - df_for_struct=self._df_for_struct, - ) - - def row_stream(): - for batch in batches: - self._update_batch_size_stats(batch) - - data_table = flatten_columns(batch, "inputData") - init_table = flatten_columns(batch, "initState") - - # Check column count - empty table has no columns - has_data = data_table.num_columns > 0 - has_init = init_table.num_columns > 0 - - assert not (has_data and has_init) - - if has_data: - for row in pd.concat(to_pandas(data_table), axis=1).itertuples(index=False): - batch_key = tuple(row[s] for s in self.key_offsets) - yield (batch_key, row, None) - elif has_init: - for row in pd.concat(to_pandas(init_table), axis=1).itertuples(index=False): - batch_key = tuple(row[s] for s in self.init_key_offsets) - yield (batch_key, None, row) - - EMPTY_DATAFRAME = pd.DataFrame() - for batch_key, group_rows in groupby(row_stream(), key=lambda x: x[0]): - rows = [] - init_state_rows = [] - for _, row, init_state_row in group_rows: - if row is not None: - rows.append(row) - if init_state_row is not None: - init_state_rows.append(init_state_row) - - total_len = len(rows) + len(init_state_rows) - if ( - total_len >= self.arrow_max_records_per_batch - or total_len * self.average_arrow_row_size >= self.arrow_max_bytes_per_batch - ): - yield ( - batch_key, - pd.DataFrame(rows) if len(rows) > 0 else EMPTY_DATAFRAME.copy(), - ( - pd.DataFrame(init_state_rows) - if len(init_state_rows) > 0 - else EMPTY_DATAFRAME.copy() - ), - ) - rows = [] - init_state_rows = [] - if rows or init_state_rows: - yield ( - batch_key, - pd.DataFrame(rows) if len(rows) > 0 else EMPTY_DATAFRAME.copy(), - ( - pd.DataFrame(init_state_rows) - if len(init_state_rows) > 0 - else EMPTY_DATAFRAME.copy() - ), - ) - - _batches = super(ArrowStreamPandasSerializer, self).load_stream(stream) - data_batches = generate_data_batches(_batches) - - for k, g in groupby(data_batches, key=lambda x: x[0]): - yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, g) - - yield (TransformWithStateInPandasFuncMode.PROCESS_TIMER, None, None) - - yield (TransformWithStateInPandasFuncMode.COMPLETE, None, None) - - -class TransformWithStateInPySparkRowSerializer(ArrowStreamUDFSerializer): - """ - Serializer used by Python worker to evaluate UDF for - :meth:`pyspark.sql.GroupedData.transformWithState`. - - Parameters - ---------- - arrow_max_records_per_batch : int - Limit of the number of records that can be written to a single ArrowRecordBatch in memory. - """ - - def __init__(self, *, arrow_max_records_per_batch): - super().__init__() - self.arrow_max_records_per_batch = ( - arrow_max_records_per_batch if arrow_max_records_per_batch > 0 else 2**31 - 1 - ) - self.key_offsets = None - - def load_stream(self, stream): - """ - Read ArrowRecordBatches from stream, deserialize them to populate a list of data chunks, - and convert the data into a list of pandas.Series. - - Please refer the doc of inner function `generate_data_batches` for more details how - this function works in overall. - """ - from pyspark.sql.streaming.stateful_processor_util import ( - TransformWithStateInPandasFuncMode, - ) - import itertools - - def generate_data_batches(batches): - """ - Deserialize ArrowRecordBatches and return a generator of Row. - - The deserialization logic assumes that Arrow RecordBatches contain the data with the - ordering that data chunks for same grouping key will appear sequentially. - - This function must avoid materializing multiple Arrow RecordBatches into memory at the - same time. And data chunks from the same grouping key should appear sequentially. - """ - for batch in batches: - DataRow = Row(*batch.schema.names) - - # Iterate row by row without converting the whole batch - num_cols = batch.num_columns - for row_idx in range(batch.num_rows): - # build the key for this row - row_key = tuple(batch[o][row_idx].as_py() for o in self.key_offsets) - row = DataRow(*(batch.column(i)[row_idx].as_py() for i in range(num_cols))) - yield row_key, row - - _batches = super(ArrowStreamUDFSerializer, self).load_stream(stream) - data_batches = generate_data_batches(_batches) - - for k, g in groupby(data_batches, key=lambda x: x[0]): - chained = itertools.chain(g) - chained_values = map(lambda x: x[1], chained) - yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, chained_values) - - yield (TransformWithStateInPandasFuncMode.PROCESS_TIMER, None, None) - - yield (TransformWithStateInPandasFuncMode.COMPLETE, None, None) - - def dump_stream(self, iterator, stream): - """ - Read through an iterator of (iterator of Row), serialize them to Arrow - RecordBatches, and write batches to stream. - """ - import pyarrow as pa - - from pyspark.sql.pandas.types import to_arrow_type - - def flatten_iterator(): - # iterator: iter[list[(iter[Row], spark_type)]] - for packed in iterator: - iter_row_with_type = packed[0] - iter_row = iter_row_with_type[0] - spark_type = iter_row_with_type[1] - - # Convert spark type to arrow type - # TODO: WE need to make this configurable, currently using default values. - arrow_type = to_arrow_type( - spark_type, - timezone="UTC", - prefers_large_types=False, - ) - - rows_as_dict = [] - for row in iter_row: - row_as_dict = row.asDict(True) - rows_as_dict.append(row_as_dict) - - pdf_schema = pa.schema(list(arrow_type)) - record_batch = pa.RecordBatch.from_pylist(rows_as_dict, schema=pdf_schema) - - yield (record_batch, arrow_type) - - return ArrowStreamUDFSerializer.dump_stream(self, flatten_iterator(), stream) - - -class TransformWithStateInPySparkRowInitStateSerializer(TransformWithStateInPySparkRowSerializer): - """ - Serializer used by Python worker to evaluate UDF for - :meth:`pyspark.sql.GroupedData.transformWithStateInPySparkRowInitStateSerializer`. - Parameters - ---------- - Same as input parameters in TransformWithStateInPySparkRowSerializer. - """ - - def __init__(self, *, arrow_max_records_per_batch): - super().__init__(arrow_max_records_per_batch=arrow_max_records_per_batch) - self.init_key_offsets = None - - def load_stream(self, stream): - import pyarrow as pa - from pyspark.sql.streaming.stateful_processor_util import ( - TransformWithStateInPandasFuncMode, - ) - - def generate_data_batches(batches) -> Iterator[Tuple[Any, Optional[Any], Optional[Any]]]: - """ - Deserialize ArrowRecordBatches and return a generator of Row. - The deserialization logic assumes that Arrow RecordBatches contain the data with the - ordering that data chunks for same grouping key will appear sequentially. - See `TransformWithStateInPySparkPythonInitialStateRunner` for arrow batch schema sent - from JVM. - This function flattens the columns of input rows and initial state rows and feed them - into the data generator. - """ - - def extract_rows( - cur_batch, col_name, key_offsets - ) -> Optional[Iterator[Tuple[Any, Any]]]: - data_column = cur_batch.column(cur_batch.schema.get_field_index(col_name)) - - # Check if the entire column is null - if data_column.null_count == len(data_column): - return None - - data_field_names = [ - data_column.type[i].name for i in range(data_column.type.num_fields) - ] - data_field_arrays = [ - data_column.field(i) for i in range(data_column.type.num_fields) - ] - - DataRow = Row(*data_field_names) - - table = pa.Table.from_arrays(data_field_arrays, names=data_field_names) - - if table.num_rows == 0: - return None - - def row_iterator(): - for row_idx in range(table.num_rows): - key = tuple(table.column(o)[row_idx].as_py() for o in key_offsets) - row = DataRow( - *(table.column(i)[row_idx].as_py() for i in range(table.num_columns)) - ) - yield (key, row) - - return row_iterator() - - """ - The arrow batch is written in the schema: - schema: StructType = new StructType() - .add("inputData", dataSchema) - .add("initState", initStateSchema) - We'll parse batch into Tuples of (key, inputData, initState) and pass into the Python - data generator. Each batch will have either init_data or input_data, not mix. - """ - for batch in batches: - # Detect which column has data - each batch contains only one type - input_result = extract_rows(batch, "inputData", self.key_offsets) - init_result = extract_rows(batch, "initState", self.init_key_offsets) - - assert not (input_result is not None and init_result is not None) - - if input_result is not None: - for key, input_data_row in input_result: - yield (key, input_data_row, None) - elif init_result is not None: - for key, init_state_row in init_result: - yield (key, None, init_state_row) - - _batches = super(ArrowStreamUDFSerializer, self).load_stream(stream) - data_batches = generate_data_batches(_batches) - - for k, g in groupby(data_batches, key=lambda x: x[0]): - input_rows = [] - init_rows = [] - - for batch_key, input_row, init_row in g: - if input_row is not None: - input_rows.append(input_row) - if init_row is not None: - init_rows.append(init_row) - - total_len = len(input_rows) + len(init_rows) - if total_len >= self.arrow_max_records_per_batch: - ret_tuple = (iter(input_rows), iter(init_rows)) - yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, ret_tuple) - input_rows = [] - init_rows = [] - - if input_rows or init_rows: - ret_tuple = (iter(input_rows), iter(init_rows)) - yield (TransformWithStateInPandasFuncMode.PROCESS_DATA, k, ret_tuple) - - yield (TransformWithStateInPandasFuncMode.PROCESS_TIMER, None, None) - - yield (TransformWithStateInPandasFuncMode.COMPLETE, None, None) diff --git a/python/pyspark/sql/pandas/types.py b/python/pyspark/sql/pandas/types.py index 4c1ae1a0..1a5c7a33 100644 --- a/python/pyspark/sql/pandas/types.py +++ b/python/pyspark/sql/pandas/types.py @@ -47,6 +47,7 @@ TimestampType, TimestampNTZType, DayTimeIntervalType, + YearMonthIntervalType, ArrayType, MapType, StructType, @@ -410,6 +411,13 @@ def from_arrow_type( spark_type = TimestampType() elif types.is_duration(at): spark_type = DayTimeIntervalType() + elif at.id == 21: # Arrow Type.INTERVAL_MONTHS + # The JVM serializes Spark's YearMonthIntervalType to an Arrow YEAR_MONTH interval + # (an integer number of months); see ArrowUtils.scala / ArrowWriter.scala. Unlike + # DayTimeIntervalType (sent as an Arrow Duration), PyArrow exposes no factory or + # is_*() helper for this type -- only MONTH_DAY_NANO is in pyarrow.types -- so match + # on the stable Arrow type id (Type::INTERVAL_MONTHS == 21). + spark_type = YearMonthIntervalType() elif types.is_list(at): spark_type = ArrayType( elementType=from_arrow_type(at.value_type, prefer_timestamp_ntz), diff --git a/python/pyspark/sql/pandas/utils.py b/python/pyspark/sql/pandas/utils.py index 0f5a73ce..875c77de 100644 --- a/python/pyspark/sql/pandas/utils.py +++ b/python/pyspark/sql/pandas/utils.py @@ -108,7 +108,7 @@ def require_minimum_pyarrow_version() -> None: def require_minimum_numpy_version() -> None: """Raise ImportError if minimum version of NumPy is not installed""" - minimum_numpy_version = "1.22" + minimum_numpy_version = "1.23.2" try: import numpy diff --git a/python/pyspark/sql/profiler.py b/python/pyspark/sql/profiler.py index 60f68c61..eae9b44f 100644 --- a/python/pyspark/sql/profiler.py +++ b/python/pyspark/sql/profiler.py @@ -394,7 +394,7 @@ def dump(id: Union[int, str]) -> None: os.makedirs(path, exist_ok=True) p = os.path.join(path, f"udf_{id}_memory.txt") - with open(p, "w+") as f: + with open(p, "w+", encoding="utf-8") as f: MemoryProfiler._show_results(cm, stream=f) if id is not None: diff --git a/python/pyspark/sql/types.py b/python/pyspark/sql/types.py index d33a91f3..ce52e5b5 100644 --- a/python/pyspark/sql/types.py +++ b/python/pyspark/sql/types.py @@ -73,7 +73,8 @@ class BinaryType(DataType): class DateType(DataType): """Date data type.""" - pass + + EPOCH_ORDINAL = 719163 class TimestampType(DataType): """Timestamp data type.""" @@ -264,6 +265,11 @@ def fromJson(cls, json: dict) -> "UserDefinedType": "YearMonthIntervalType", "DayTimeIntervalType", "VariantType", + "VariantVal", + "GeometryType", + "GeographyType", + "Geometry", + "Geography", "_drop_metadata", "_parse_datatype_json_string", "_parse_datatype_json_value", diff --git a/python/pyspark/sql/variant_utils.py b/python/pyspark/sql/variant_utils.py index 80b7efea..0b044e63 100644 --- a/python/pyspark/sql/variant_utils.py +++ b/python/pyspark/sql/variant_utils.py @@ -563,14 +563,20 @@ def build(self, json_str: str) -> Tuple[bytes, bytes]: # calculation in case of pathological data. max_size = max(dictionary_string_size, num_keys) if max_size > self.size_limit: - raise PySparkValueError(errorClass="VARIANT_SIZE_LIMIT_EXCEEDED", messageParameters={}) + raise PySparkValueError( + errorClass="VARIANT_SIZE_LIMIT", + messageParameters={"sizeLimit": str(self.size_limit)}, + ) offset_size = self._get_integer_size(max_size) offset_start = 1 + offset_size string_start = offset_start + (num_keys + 1) * offset_size metadata_size = string_start + dictionary_string_size if metadata_size > self.size_limit: - raise PySparkValueError(errorClass="VARIANT_SIZE_LIMIT_EXCEEDED", messageParameters={}) + raise PySparkValueError( + errorClass="VARIANT_SIZE_LIMIT", + messageParameters={"sizeLimit": str(self.size_limit)}, + ) metadata = bytearray() header_byte = VariantUtils.VERSION | ((offset_size - 1) << 6) @@ -631,7 +637,10 @@ def _get_integer_size(self, value: int) -> int: def _check_capacity(self, additional: int) -> None: required = len(self.value) + additional if required > self.size_limit: - raise PySparkValueError(errorClass="VARIANT_SIZE_LIMIT_EXCEEDED", messageParameters={}) + raise PySparkValueError( + errorClass="VARIANT_SIZE_LIMIT", + messageParameters={"sizeLimit": str(self.size_limit)}, + ) def _primitive_header(self, type: int) -> bytes: return bytes([(type << 2) | VariantUtils.PRIMITIVE]) diff --git a/python/pyspark/sql/worker/write_into_data_source.py b/python/pyspark/sql/worker/write_into_data_source.py index 1c39eae6..ca97090f 100644 --- a/python/pyspark/sql/worker/write_into_data_source.py +++ b/python/pyspark/sql/worker/write_into_data_source.py @@ -187,7 +187,9 @@ def _main(infile: IO, outfile: IO) -> None: def data_source_write_func(iterator: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: def batch_to_rows() -> Iterator[Row]: for batch in iterator: - columns = [column.to_pylist() for column in batch.columns] + columns = [ + ArrowTableToRowsConversion._to_pylist(column) for column in batch.columns + ] for row in range(0, batch.num_rows): values = [ converters[col](columns[col][row]) for col in range(batch.num_columns) diff --git a/python/pyspark/testing/goldenutils.py b/python/pyspark/testing/goldenutils.py index b1e4af59..d0a191e9 100644 --- a/python/pyspark/testing/goldenutils.py +++ b/python/pyspark/testing/goldenutils.py @@ -15,7 +15,8 @@ # limitations under the License. # -from typing import Any, Optional +from typing import Any, Callable, List, Optional +import inspect import os import time @@ -194,6 +195,85 @@ def save_golden(df: "pd.DataFrame", golden_csv: str, golden_md: Optional[str] = "Install 'tabulate' package to generate markdown files." ) + def compare_or_generate_golden_matrix( + self, + row_names: List[str], + col_names: List[str], + compute_cell: Callable[[str, str], str], + golden_file_prefix: str, + index_name: str = "source \\ target", + overrides: Optional[dict[tuple[str, str], str]] = None, + ) -> None: + """ + Run a matrix of computations and compare against (or generate) a golden file. + + 1. If SPARK_GENERATE_GOLDEN_FILES=1, compute every cell, build a + DataFrame, and save it as the new golden CSV / Markdown file. + 2. Otherwise, load the existing golden file and assert that every cell + matches the freshly computed value. + + Parameters + ---------- + row_names : list[str] + Ordered row labels (becomes the DataFrame index). + col_names : list[str] + Ordered column labels. + compute_cell : (row_name, col_name) -> str + Function that computes the string result for one cell. + golden_file_prefix : str + Prefix for the golden CSV/MD files (without extension). + Files are placed in the same directory as the concrete test file. + index_name : str, default "source \\ target" + Name for the index column in the golden file. + overrides : dict[(row, col) -> str], optional + Version-specific expected values that take precedence over the golden + file. Use this to document known behavioral differences across + library versions (e.g. PyArrow 18 vs 22) directly in the test code, + so that the same golden file works for multiple versions. + """ + generating = self.is_generating_golden() + + test_dir = os.path.dirname(inspect.getfile(type(self))) + golden_csv = os.path.join(test_dir, f"{golden_file_prefix}.csv") + golden_md = os.path.join(test_dir, f"{golden_file_prefix}.md") + + golden = None + if not generating: + golden = self.load_golden_csv(golden_csv) + + errors = [] + results = {} + + for row_name in row_names: + for col_name in col_names: + result = compute_cell(row_name, col_name) + results[(row_name, col_name)] = result + + if not generating: + if overrides and (row_name, col_name) in overrides: + expected = overrides[(row_name, col_name)] + else: + expected = golden.loc[row_name, col_name] + if expected != result: + errors.append( + f"{row_name} -> {col_name}: expected '{expected}', got '{result}'" + ) + + if generating: + import pandas as pd + + index = pd.Index(row_names, name=index_name) + df = pd.DataFrame(index=index) + for col_name in col_names: + df[col_name] = [results[(row, col_name)] for row in row_names] + self.save_golden(df, golden_csv, golden_md) + else: + self.assertEqual( + len(errors), + 0, + f"\n{len(errors)} golden file mismatches:\n" + "\n".join(errors), + ) + @staticmethod def repr_type(t: Any) -> str: """ diff --git a/python/pyspark/version.py b/python/pyspark/version.py index 1d7b8dbf..9cb6f505 100644 --- a/python/pyspark/version.py +++ b/python/pyspark/version.py @@ -16,4 +16,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__: str = "4.2.0" +__version__: str = "4.3.0" diff --git a/python/pyspark/worker_util.py b/python/pyspark/worker_util.py index 8c52bde6..13e08449 100644 --- a/python/pyspark/worker_util.py +++ b/python/pyspark/worker_util.py @@ -24,9 +24,11 @@ from inspect import currentframe, getframeinfo import os import sys -from typing import Any, Generator, IO, Optional +from typing import Any, Generator, IO, Optional, Union, overload import warnings +from pyspark.messages import ZeroCopyByteStream + if "SPARK_TESTING" in os.environ: assert os.environ.get("SPARK_PYTHON_RUNTIME") == "PYTHON_WORKER", ( "This module can only be imported in python woker" @@ -44,7 +46,6 @@ from pyspark.errors import PySparkRuntimeError from pyspark.util import local_connect_and_auth from pyspark.serializers import ( - read_bool, read_int, read_long, write_int, @@ -66,21 +67,27 @@ def add_path(path: str) -> bool: return False -def read_command(serializer: FramedSerializer, file: IO) -> Any: +def read_command(serializer: FramedSerializer, file: Union[IO, bytes, memoryview]) -> Any: if not is_remote_only(): from pyspark.core.broadcast import Broadcast - command = serializer._read_with_length(file) + if isinstance(file, (bytes, memoryview)): + command = serializer.loads(file) + else: + command = serializer._read_with_length(file) if not is_remote_only() and isinstance(command, Broadcast): command = serializer.loads(command.value) return command -def check_python_version(infile: IO) -> None: +def check_python_version(infile_or_version: Union[IO, str]) -> None: """ Check the Python version between the running process and the one used to serialize the command. """ - version = utf8_deserializer.loads(infile) + if isinstance(infile_or_version, str): + version = infile_or_version + else: + version = utf8_deserializer.loads(infile_or_version) worker_version = "%d.%d" % sys.version_info[:2] if version != worker_version: raise PySparkRuntimeError( @@ -130,12 +137,20 @@ def setup_memory_limits(memory_limit_mb: int) -> None: ) -def setup_spark_files(infile: IO) -> None: +@overload +def setup_spark_files(infile_or_spark_files_dir: IO) -> None: ... +@overload +def setup_spark_files(infile_or_spark_files_dir: str, python_includes: list[str]) -> None: ... +def setup_spark_files( + infile_or_spark_files_dir: Union[IO, str], python_includes: Optional[list[str]] = None +) -> None: """ Set up Spark files, archives, and pyfiles. """ - # fetch name of workdir - spark_files_dir = utf8_deserializer.loads(infile) + if isinstance(infile_or_spark_files_dir, str): + spark_files_dir = infile_or_spark_files_dir + else: + spark_files_dir = utf8_deserializer.loads(infile_or_spark_files_dir) if not is_remote_only(): from pyspark.core.files import SparkFiles @@ -145,51 +160,78 @@ def setup_spark_files(infile: IO) -> None: # fetch names of includes (*.zip and *.egg files) and construct PYTHONPATH path_changed = add_path(spark_files_dir) # *.py files that were added will be copied here - num_python_includes = read_int(infile) - for _ in range(num_python_includes): - filename = utf8_deserializer.loads(infile) + if not isinstance(infile_or_spark_files_dir, str): + python_includes = [ + utf8_deserializer.loads(infile_or_spark_files_dir) + for _ in range(read_int(infile_or_spark_files_dir)) + ] + assert python_includes is not None + + for filename in python_includes: path_changed = add_path(os.path.join(spark_files_dir, filename)) or path_changed if path_changed: importlib.invalidate_caches() -def setup_broadcasts(infile: IO) -> None: +@overload +def setup_broadcasts(infile_or_variables: IO[Any]) -> None: ... +@overload +def setup_broadcasts(infile_or_variables: ZeroCopyByteStream) -> None: ... +@overload +def setup_broadcasts( + infile_or_variables: list[tuple[int, Union[str, None]]], conn_info: str, auth_secret: None +) -> None: ... +@overload +def setup_broadcasts( + infile_or_variables: list[tuple[int, Union[str, None]]], conn_info: int, auth_secret: str +) -> None: ... +@overload +def setup_broadcasts( + infile_or_variables: list[tuple[int, Union[str, None]]], + conn_info: Optional[Union[str, int]], + auth_secret: Optional[str], +) -> None: ... +def setup_broadcasts( + infile_or_variables: Union[ZeroCopyByteStream, IO[Any], list[tuple[int, Union[str, None]]]], + conn_info: Optional[Union[str, int]] = None, + auth_secret: Optional[str] = None, +) -> None: """ Set up broadcasted variables. """ if not is_remote_only(): from pyspark.core.broadcast import Broadcast, _broadcastRegistry - # fetch names and values of broadcast variables - needs_broadcast_decryption_server = read_bool(infile) - num_broadcast_variables = read_int(infile) + if isinstance(infile_or_variables, list): + variables = infile_or_variables + else: + from pyspark.worker_message import BroadcastInfo + + broadcast_info = BroadcastInfo.from_stream(infile_or_variables) + conn_info = broadcast_info.conn_info + auth_secret = broadcast_info.auth_secret + variables = broadcast_info.variables + + needs_broadcast_decryption_server = conn_info is not None or auth_secret is not None + if needs_broadcast_decryption_server: - # read the decrypted data from a server in the jvm - conn_info = read_int(infile) - auth_secret = None - if conn_info == -1: - conn_info = utf8_deserializer.loads(infile) - else: - auth_secret = utf8_deserializer.loads(infile) broadcast_sock_file, _ = local_connect_and_auth(conn_info, auth_secret) + else: + broadcast_sock_file = None - for _ in range(num_broadcast_variables): - bid = read_long(infile) + for bid, path in variables: if bid >= 0: - if needs_broadcast_decryption_server: + if path is None: read_bid = read_long(broadcast_sock_file) assert read_bid == bid _broadcastRegistry[bid] = Broadcast(sock_file=broadcast_sock_file) else: - path = utf8_deserializer.loads(infile) _broadcastRegistry[bid] = Broadcast(path=path) - else: - bid = -bid - 1 - _broadcastRegistry.pop(bid) + _broadcastRegistry.pop(-bid - 1) - if needs_broadcast_decryption_server: + if broadcast_sock_file is not None: broadcast_sock_file.write(b"1") broadcast_sock_file.close() @@ -223,29 +265,32 @@ def send_accumulator_updates(outfile: IO) -> None: class Conf: - def __init__(self, infile: Optional[IO] = None) -> None: + def __init__(self, infile_or_dict: Optional[Union[dict[str, str], IO]] = None) -> None: self._conf: dict[str, Any] = {} - if infile is not None: - self.load(infile) - - def load(self, infile: IO) -> None: - num_conf = read_int(infile) - # We do a sanity check here to reduce the possibility to stuck indefinitely - # due to an invalid messsage. If the numer of configurations is obviously - # wrong, we just raise an error directly. - # We hand-pick the configurations to send to the worker so the number should - # be very small (less than 100). - if num_conf < 0 or num_conf > 10000: - raise PySparkRuntimeError( - errorClass="PROTOCOL_ERROR", - messageParameters={ - "failure": f"Invalid number of configurations: {num_conf}", - }, - ) - for _ in range(num_conf): - k = utf8_deserializer.loads(infile) - v = utf8_deserializer.loads(infile) - self._conf[k] = v + if infile_or_dict is not None: + self.load(infile_or_dict) + + def load(self, infile_or_dict: Union[dict[str, str], IO]) -> None: + if isinstance(infile_or_dict, dict): + self._conf = infile_or_dict + else: + num_conf = read_int(infile_or_dict) + # We do a sanity check here to reduce the possibility to stuck indefinitely + # due to an invalid messsage. If the numer of configurations is obviously + # wrong, we just raise an error directly. + # We hand-pick the configurations to send to the worker so the number should + # be very small (less than 100). + if num_conf < 0 or num_conf > 10000: + raise PySparkRuntimeError( + errorClass="PROTOCOL_ERROR", + messageParameters={ + "failure": f"Invalid number of configurations: {num_conf}", + }, + ) + for _ in range(num_conf): + k = utf8_deserializer.loads(infile_or_dict) + v = utf8_deserializer.loads(infile_or_dict) + self._conf[k] = v def get(self, key: str, default: Any = "", *, lower_str: bool = True) -> Any: val = self._conf.get(key, default) diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 6c747627..9b026ab8 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -224,6 +224,48 @@ def is_col(x): assert repr(F._UNSET) == "" +def test_dispatch_keeps_literal_string_args_literal(): + # Guards _FN_COL_ARGS: a bare-str arg at a LITERAL position (e.g. a variant path) + # must reach the core UNCHANGED, not be resolved to a column via _to_col. If it were + # coerced, `variant_delete(v, "$.b")` would send an unresolved attribute col("$.b") + # and fail analysis on the server. Asserting only that a Column is built (as elsewhere) + # can't catch this, so we capture the exact positional args handed to the core: a + # value column is coerced away from str, while a literal path stays a str. + from pyspark.sql import functions as F + + captured = [] + + orig = F._call_function + F._call_function = lambda fname, *args: captured.append((fname, args)) or orig(fname, *args) + def is_col(x): + return type(x).__name__ == "Column" + + try: + # (function call, positions that MUST stay raw literals — a str path or a bool + # flag — and positions that MUST be coerced to a column) + cases = [ + (lambda: F.variant_get("v", "$.a", "int"), {1, 2}, {0}), + (lambda: F.try_variant_get("v", "$.a", "int"), {1, 2}, {0}), + (lambda: F.variant_delete("v", "$.a", "$.b"), {1, 2}, {0}), + (lambda: F.variant_insert("v", "$.b", "x"), {1}, {0, 2}), + (lambda: F.try_variant_insert("v", "$.b", "x"), {1}, {0, 2}), + (lambda: F.variant_set("v", "$.b", "x", True), {1, 3}, {0, 2}), + (lambda: F.try_variant_set("v", "$.b", "x", True), {1, 3}, {0, 2}), + (lambda: F.variant_array_append("v", "$", "x"), {1}, {0, 2}), + (lambda: F.try_variant_array_append("v", "$", "x"), {1}, {0, 2}), + ] + for call, literal_pos, col_pos in cases: + captured.clear() + call() + fname, args = captured[-1] + for i in literal_pos: + assert not is_col(args[i]), f"{fname} arg[{i}] should stay a raw literal, got a Column" + for i in col_pos: + assert is_col(args[i]), f"{fname} arg[{i}] should be coerced to a Column, got {type(args[i]).__name__}" + finally: + F._call_function = orig + + def test_column_alias_metadata_offline(): from pyspark.sql import functions as F @@ -2467,7 +2509,7 @@ def test_drop_metadata_and_create_row(): def test_dropin_version_is_spark_version(): - assert pyspark.__version__ == "4.2.0" + assert pyspark.__version__ == "4.3.0" @pytest.mark.parametrize("cls,name", [(_SS, "version"), (_DF, "dtypes"), (_DF, "executionInfo")]) @@ -2729,3 +2771,91 @@ def test_structtype_tree_string_max_depth(): def test_variant_type_to_internal_accepts_variant_kwarg(): # Reference names the VariantType.toInternal parameter `variant`, not `obj`. assert T.VariantType().toInternal(variant=None) is None + + +# --------------------------------------------- 4.3.0 API-surface additions + + +def test_new_4_3_0_sql_functions_build_columns(): + from pyspark.sql import functions as F + + c = F.col("v") + for fn in (F.collect_union, F.to_base32, F.from_base32): + assert isinstance(fn(c), _Col) + assert isinstance(F.variant_strip_nulls(c), _Col) + assert isinstance(F.variant_strip_nulls(c, include_arrays=False), _Col) + assert isinstance(F.variant_delete(c, F.lit("$.a"), "$.b"), _Col) + + +def test_new_4_3_0_functions_are_exported(): + from pyspark.sql import functions as F + + for name in ( + "variant_delete", + "collect_union", + "to_base32", + "from_base32", + "variant_strip_nulls", + "PandasUDFType", + "ArrowUDFType", + "UserDefinedFunction", + "UserDefinedTableFunction", + "AnalyzeArgument", + "AnalyzeResult", + "SkipRestOfInputTableException", + ): + assert name in F.__all__, name + assert hasattr(F, name), name + + +def test_arrow_udf_type_enum_values(): + from pyspark.sql.pandas.functions import ArrowUDFType, PandasUDFType + + assert (ArrowUDFType.SCALAR, ArrowUDFType.SCALAR_ITER) == (250, 251) + assert (ArrowUDFType.GROUPED_AGG, ArrowUDFType.GROUPED_AGG_ITER) == (252, 254) + assert PandasUDFType.GROUPED_AGG_ITER == 217 + + +def test_datatype_class_constants(): + assert (T.DayTimeIntervalType.DAY, T.DayTimeIntervalType.HOUR) == (0, 1) + assert (T.DayTimeIntervalType.MINUTE, T.DayTimeIntervalType.SECOND) == (2, 3) + assert (T.YearMonthIntervalType.YEAR, T.YearMonthIntervalType.MONTH) == (0, 1) + assert T.DateType.EPOCH_ORDINAL == 719163 + assert T.GeometryType.DEFAULT_CRS == "OGC:CRS84" + assert T.GeometryType.DEFAULT_SRID == 4326 + assert T.GeographyType.DEFAULT_ALG == "SPHERICAL" + assert (T.SpatialType.MIXED_CRS, T.SpatialType.MIXED_SRID) == ("SRID:ANY", -1) + + +def test_types_all_exports_spatial_and_variant(): + for name in ("Geometry", "GeometryType", "Geography", "GeographyType", "VariantVal"): + assert name in T.__all__, name + + +def test_merge_into_writer_nested_class_refs(): + from pyspark.sql.merge import MergeIntoWriter + + assert isinstance(MergeIntoWriter.WhenMatched, type) + assert isinstance(MergeIntoWriter.WhenNotMatched, type) + assert isinstance(MergeIntoWriter.WhenNotMatchedBySource, type) + + +def test_session_hook_is_present_and_passthrough(): + hook = _SS.Hook() + sentinel = object() + assert hook.on_execute_plan(sentinel) is sentinel + + +def test_ml_connect_base_classes_and_functions(): + from pyspark.ml.connect import Estimator, Transformer, Evaluator, Model, PipelineModel + + for cls in (Estimator, Transformer, Evaluator, Model, PipelineModel): + assert isinstance(cls, type) + import pyspark.ml.functions as mf + + assert hasattr(mf, "vector_to_array") and hasattr(mf, "array_to_vector") + + +def test_pyspark_version_exported(): + assert pyspark.__version__ == "4.3.0" + assert "__version__" in pyspark.__all__ diff --git a/scripts/check_vendored_upstream.sh b/scripts/check_vendored_upstream.sh index e96a84e6..51d7b7a8 100644 --- a/scripts/check_vendored_upstream.sh +++ b/scripts/check_vendored_upstream.sh @@ -26,7 +26,7 @@ # fork-adapted files (e.g. pipelines/api.py) — those intentionally differ from upstream. set -euo pipefail -SPARK_TAG="v4.2.0" +SPARK_TAG="v4.3.0-rc3" SPARK_REPO="https://github.com/apache/spark.git" # Directories under python/pyspark vendored verbatim from upstream (whole tree byte-identical). @@ -40,6 +40,8 @@ VENDORED_DIRS=( "sql/worker" "sql/plot" "testing/tests" + # New in 4.3.0: the client-side message-receiver package (pyspark.messages). + "messages" ) # Individual files under python/pyspark vendored verbatim from upstream. Listed per-file @@ -85,6 +87,9 @@ VENDORED_FILES=( "ml/common.py" "ml/dl_util.py" "ml/util.py" + "ml/functions.py" + # ML connect base ABCs (Estimator/Transformer/Evaluator/Model) - pure upstream. + "ml/connect/base.py" # ML: param (param API is unmodified) "ml/param/__init__.py" diff --git a/scripts/gen_fn_dispatch.py b/scripts/gen_fn_dispatch.py index 1bf55396..63f43311 100644 --- a/scripts/gen_fn_dispatch.py +++ b/scripts/gen_fn_dispatch.py @@ -104,6 +104,9 @@ # variant_delete(Column, Vec) -> Column (variadic paths) "broadcast", "variant_delete", + # variant_strip_nulls(Column, bool) -> Column: bool arg, dedicated pyfunc binding + # (the generic Vec dispatch would drop the trailing scalar). + "variant_strip_nulls", # Mixed functions with dedicated pyfunc_* bindings: # sha2(Column, i32) -> Column # window(Column, &str) -> Column diff --git a/scripts/parity_known_failures.txt b/scripts/parity_known_failures.txt index de7d5016..f91e315e 100644 --- a/scripts/parity_known_failures.txt +++ b/scripts/parity_known_failures.txt @@ -6,8 +6,11 @@ # start passing, so the list cannot silently rot. Regenerate with # scripts/gen_parity_skiplist.py on a version bump. # -# NB: many more tests fail against an *older local* Spark dist that lacks the 4.2.0 -# TIME type (UNSUPPORTED_TIME_TYPE); those pass in CI, so they are intentionally NOT +# NB: many more tests fail against an *older local* Spark dist that lacks the TIME +# type (UNSUPPORTED_TIME_TYPE); those pass in CI, so they are intentionally NOT # listed here - the CI drift check confirmed it. Keep this list to the CI reality. +# +# The former arrow-UDF entry (ArrowPythonUDFParityLegacyTests:: +# test_type_coercion_string_to_numeric) was REMOVED for 4.3.0: the drift check reports +# it now passes on the 4.3.0-rc1 CI image, so the gate exercises it for real. pyspark/sql/tests/connect/test_parity_dataframe.py::DataFrameParityTests::test_metadata_column # requires InMemoryCatalog for 'testcat' - a test-only catalog class absent from the Spark distribution -pyspark/sql/tests/connect/arrow/test_parity_arrow_python_udf.py::ArrowPythonUDFParityLegacyTests::test_type_coercion_string_to_numeric # legacy arrow-UDF string->numeric coercion is arrow/pandas-version dependent (fails in the CI image)